diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 464e4763..4028bbaa 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -165,7 +165,10 @@ backslashes, symlink escapes, unexpected fields, and host roots are rejected. Workspace mutation remains disabled unless the operator starts the worker with `--allow-workspace-writes` (or `LIBRECHAT_CODE_ALLOW_WORKSPACE_WRITES=true`). That adds bounded `write_file` -and exact-match `edit_file` operations. Writes are limited to 1 MiB of UTF-8 +and `edit_file` operations. `edit_file` matches exactly by default; the +negotiated `tolerant_match` and `replace_all` edit features add +whitespace-tolerant matching and multi-location replacement (see +`packages/code/README.md`). Writes are limited to 1 MiB of UTF-8 text, require an existing in-workspace parent directory, reject symlinks, and commit atomically. The worker capability is an enforcement boundary; LibreChat should still route every mutation through its configurable tool-approval hooks. diff --git a/packages/code/README.md b/packages/code/README.md index 95d86485..67a2167b 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -590,6 +590,46 @@ installed as one atomic mutation. Code API dispatches the batch form only after the worker and server negotiate `batch` in `editFileModes`. Revision-fenced edits likewise require the negotiated `expected_base_sha256` entry in `editFileFeatures`. + +Edits apply in order, each to the text the earlier ones produced. When any edit +fails, the worker still checks the rest and rejects the whole batch with one +`EDIT_CONFLICT` whose message lists every failing edit by position. A +missing edit names the nearest candidate line and flags elided (`...`) or +line-numbered `oldText`, a whitespace-only difference, or CRLF line endings. +An ambiguous edit gives its match count and line numbers. Overlapping +occurrences count as separate locations. Detailed source-line excerpts require +`read_file` or `preview_edit` on the same workspace; edit-only workers return a +generic failure and its error code without revealing file contents. A preview +itself exposes the resulting file text, so it is read-capable. + +Two optional features change matching, each negotiated in `editFileFeatures` +before Code API dispatches it: + +- `tolerant_match`: a request-level `matching: 'tolerant'` falls back from an + exact match to, in order, `line-trimmed` (ignores trailing whitespace and + CRLF), `indentation-flexible` (a uniformly shifted block, with `newText` + moved to the file's indentation) and `whitespace-normalized` (any whitespace + run between complete whitespace-delimited tokens, never a prefix or suffix + of another token). Without `replaceAll`, a match must still be unique; + replacements use the matched line's ending even in mixed-ending files. The + whitespace-normalized tier peels the complete shared newline-and-indentation + wrapper from `newText` even when CRLF/LF or nearby spaces differ, without + removing intentional extra line breaks or duplicating the source line ending. + Boundary whitespace claimed by `oldText`, including the number of line breaks, + must exist beside the matched tokens in the source; an attempt to remove it + with a token-only fallback fails rather than silently preserving it. Exact and + line-window matches can still replace terminators. + Excessively repetitive indentation candidates fail closed with a request + for more context rather than scanning every long window. Dense files reuse a + compact newline index for matching and diagnostics; overlapping exact matches + are counted without restarting a scan at each offset. +- `replace_all`: a batch edit's `replaceAll: true` replaces every + non-overlapping match instead of requiring exactly one, and still fails when + nothing matches. + +A request that sets `matching` or any `replaceAll` receives `matches`, one +`{ strategy, occurrences }` entry per edit. Requests that set neither receive +exactly the legacy result. Only IDs, names, protocol version, supported operations, and negotiated write modes appear in worker capabilities; absolute host paths remain local to the worker process. @@ -655,8 +695,10 @@ non-regular files, and commit through an owner-only temporary file followed by an atomic rename. The worker syncs the containing directory and verifies that the installed inode still contains the requested bytes before reporting success. Edits replace text only when the requested old text occurs exactly -once and reject if the file changes before commit. These operations do not -create directories or execute commands. +once (unless negotiated `replaceAll` selects every non-overlapping match) and +reject if the file changes before commit. Intermediate replacements are bounded +before construction, including `replaceAll`; these operations do not create +directories or execute commands. Register one directory already present on the worker machine with the worker-directory option: diff --git a/packages/code/src/edits.test.ts b/packages/code/src/edits.test.ts new file mode 100644 index 00000000..0151d793 --- /dev/null +++ b/packages/code/src/edits.test.ts @@ -0,0 +1,602 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { + applyTextEdits, + EDIT_DIAGNOSTIC_MAX_CHARS, + WorkspaceEditMatchError, + WorkspaceEditOutputLimitError, +} from './edits.js'; + +function rejection(run: () => unknown): WorkspaceEditMatchError { + try { + run(); + } catch (error) { + assert.ok(error instanceof WorkspaceEditMatchError); + return error; + } + assert.fail('expected the edits to be rejected'); +} + +test('exact matching keeps its single-location contract', () => { + const applied = applyTextEdits('alpha\nbeta\n', [{ oldText: 'beta', newText: 'gamma' }]); + assert.deepEqual(applied, { + text: 'alpha\ngamma\n', + matches: [{ strategy: 'exact', occurrences: 1 }], + }); +}); + +test('overlapping occurrences make an exact edit ambiguous', () => { + const error = rejection(() => applyTextEdits('aaa', [{ oldText: 'aa', newText: 'b' }])); + assert.match(error.message, /matched 2 locations/); + const applied = applyTextEdits('aaaa', [{ oldText: 'aa', newText: 'b', replaceAll: true }]); + assert.deepEqual(applied, { text: 'bb', matches: [{ strategy: 'exact', occurrences: 2 }] }); +}); + +test('later edits see the result of earlier ones', () => { + const applied = applyTextEdits('one', [ + { oldText: 'one', newText: 'two' }, + { oldText: 'two', newText: 'three' }, + ]); + assert.equal(applied.text, 'three'); +}); + +test('an ambiguous edit names every location instead of a bare conflict', () => { + const text = 'const a = 1;\nreturn a;\nconst b = 2;\nreturn a;\n'; + const error = rejection(() => + applyTextEdits(text, [{ oldText: 'return a;', newText: 'return b;' }]), + ); + assert.match(error.message, /matched 2 locations at lines 2, 4/); + assert.match(error.message, /include more surrounding lines/); + assert.match(error.message, /nothing was written/); +}); + +test('a batch reports every failing edit with its position in one rejection', () => { + const text = 'a\nb\nc\nb\n'; + const error = rejection(() => + applyTextEdits(text, [ + { oldText: 'a', newText: 'A' }, + { oldText: 'missing', newText: 'x' }, + { oldText: 'c', newText: 'C' }, + { oldText: 'b', newText: 'B' }, + ]), + ); + assert.deepEqual( + error.failures.map((failure) => failure.index), + [1, 3], + ); + assert.match(error.message, /^2 of 4 workspace edits did not apply/); + assert.match(error.message, /\nEdit 2: old_text was not found/); + assert.match(error.message, /\nEdit 4: old_text matched 2 locations at lines 2, 4/); +}); + +test('a missing edit points at the line it most likely meant', () => { + const text = 'function load(user) {\n return fetchUser(user.id);\n}\n'; + const error = rejection(() => + applyTextEdits(text, [ + { oldText: 'function load(user) {\n return fetchUser(user);\n}', newText: 'x' }, + ]), + ); + assert.match(error.message, /its first line appears at line 1, but the lines after it differ/); +}); + +test('a missing edit explains elided and line-numbered old_text', () => { + const text = 'start\nmiddle\nend\n'; + const elided = rejection(() => + applyTextEdits(text, [{ oldText: 'start\n// ...\nend', newText: 'x' }]), + ); + assert.match(elided.message, /elision placeholder/); + const numbered = rejection(() => + applyTextEdits(text, [{ oldText: '1 | start\n2 | middle', newText: 'x' }]), + ); + assert.match(numbered.message, /line-number prefixes/); +}); + +test('exact mode says when only the whitespace differs', () => { + const shifted = rejection(() => + applyTextEdits('class A {\n a();\n b();\n}\n', [{ oldText: 'a();\nb();', newText: 'x' }]), + ); + assert.match(shifted.message, /exists at line 2 with different whitespace \(indentation-flexible\)/); + const reflowed = rejection(() => + applyTextEdits('if (ready) {\n run();\n}\n', [ + { oldText: 'if (ready) {\n run();\n}', newText: 'x' }, + ]), + ); + assert.match(reflowed.message, /exists at line 1 with different whitespace \(whitespace-normalized\)/); +}); + +test('tolerant matching ignores trailing whitespace and keeps CRLF line endings', () => { + const text = 'first \r\nsecond\t\r\nthird\r\n'; + const applied = applyTextEdits( + text, + [{ oldText: 'first\nsecond', newText: 'one\ntwo' }], + 'tolerant', + ); + assert.equal(applied.text, 'one\r\ntwo\r\nthird\r\n'); + assert.deepEqual(applied.matches, [{ strategy: 'line-trimmed', occurrences: 1 }]); +}); + +test('a trailing newline in old_text matches through the line terminator', () => { + const text = 'keep\ndrop\nkeep too\n'; + const applied = applyTextEdits(text, [{ oldText: 'drop \n', newText: '' }], 'tolerant'); + assert.equal(applied.text, 'keep\nkeep too\n'); +}); + +test('indentation-flexible matches move new_text to the file indentation', () => { + const text = 'class A {\n method() {\n return 1;\n }\n}\n'; + const applied = applyTextEdits( + text, + [ + { + oldText: 'method() {\n return 1;\n}', + newText: 'method() {\n const value = 2;\n return value;\n}', + }, + ], + 'tolerant', + ); + assert.equal( + applied.text, + 'class A {\n method() {\n const value = 2;\n return value;\n }\n}\n', + ); + assert.deepEqual(applied.matches, [{ strategy: 'indentation-flexible', occurrences: 1 }]); +}); + +test('whitespace-normalized matches do not splice prefixes or suffixes of tokens', () => { + for (const [source, oldText] of [ + ['prereturn value;\n', 'return value;'], + ['return valueSuffix\n', 'return value'], + ['return value;postfix\n', 'return value;'], + ]) { + const error = rejection(() => applyTextEdits(source, [{ oldText, newText: 'changed' }], 'tolerant')); + assert.match(error.message, /old_text was not found/); + } + const valid = applyTextEdits('return value;\n', [ + { oldText: 'return value;', newText: 'return changed;' }, + ], 'tolerant'); + assert.equal(valid.text, 'return changed;\n'); + assert.deepEqual(valid.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('large whitespace-normalized edits match without compiling a request-sized regular expression', () => { + const oldText = `head ${'part '.repeat(16_000)}tail`; + const source = `before ${oldText.replace(/ /g, '\t')} after`; + const applied = applyTextEdits(source, [{ oldText, newText: 'result' }], 'tolerant'); + assert.equal(applied.text, 'before result after'); + assert.deepEqual(applied.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('whitespace-normalized matches count overlapping token sequences but replaceAll does not overlap', () => { + const source = 'a b a\tb a'; + const edit = { oldText: 'a b a', newText: 'x' }; + const error = rejection(() => applyTextEdits(source, [edit], 'tolerant')); + assert.match(error.message, /matched 2 locations/); + const replaced = applyTextEdits(source, [{ ...edit, replaceAll: true }], 'tolerant'); + assert.equal(replaced.text, 'x\tb a'); + assert.deepEqual(replaced.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +for (const fileEnding of ['\n', '\r\n'] as const) { + for (const [oldEnding, newEnding] of [['\r\n', '\n'], ['\n', '\r\n']] as const) { + test(`whitespace-normalized matching peels ${JSON.stringify(oldEnding)} to ${JSON.stringify(newEnding)} in ${JSON.stringify(fileEnding)} files`, () => { + const source = `foo bar${fileEnding}next`; + const result = applyTextEdits(source, [{ + oldText: `foo bar${oldEnding}`, + newText: `baz qux${newEnding}`, + }], 'tolerant'); + assert.equal(result.text, `baz qux${fileEnding}next`); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); + }); + } +} + +test('whitespace-normalized matching peels leading and trailing CRLF/LF wrappers in a batch', () => { + const result = applyTextEdits('header\r\n \tfoo bar\r\nnext\r\n', [ + { oldText: '\n foo bar\r\n', newText: '\r\n baz qux\n' }, + { oldText: 'next', newText: 'done' }, + ], 'tolerant'); + assert.equal(result.text, 'header\r\n \tbaz qux\r\ndone\r\n'); + assert.deepEqual(result.matches, [ + { strategy: 'whitespace-normalized', occurrences: 1 }, + { strategy: 'exact', occurrences: 1 }, + ]); +}); + +test('whitespace-normalized matching peels a trailing newline even if adjacent spaces differ', () => { + const result = applyTextEdits('foo bar \r\nnext', [{ + oldText: 'foo bar \r\n', newText: 'baz qux\n', + }], 'tolerant'); + assert.equal(result.text, 'baz qux \r\nnext'); +}); + +test('whitespace-normalized matching does not prepend new indentation beside preserved source indentation', () => { + const source = 'header\n foo bar\n'; + const result = applyTextEdits(source, [{ + oldText: '\n foo bar', newText: '\n\tbaz qux', + }], 'tolerant'); + assert.equal(result.text, 'header\n baz qux\n'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('whitespace-normalized matching preserves extra leading blank lines when indentation changes', () => { + const source = 'header\n foo bar\n'; + const result = applyTextEdits(source, [{ + oldText: '\n foo bar', newText: '\n\t\n\tbaz qux', + }], 'tolerant'); + assert.equal(result.text, 'header\n \n\tbaz qux\n'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('whitespace-normalized replaceAll preserves extra leading blank lines independently', () => { + const source = 'one\n foo bar\ntwo\n foo bar\n'; + const result = applyTextEdits(source, [{ + oldText: '\n foo bar', newText: '\n\t\n\tbaz qux', replaceAll: true, + }], 'tolerant'); + assert.equal(result.text, 'one\n \n\tbaz qux\ntwo\n \n\tbaz qux\n'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 2 }]); +}); + +test('whitespace-normalized matching preserves extra lines after an identical leading wrapper', () => { + const result = applyTextEdits('header\r\n foo bar', [{ + oldText: '\r\n foo bar', newText: '\r\n \r\n\tbaz qux', + }], 'tolerant'); + assert.equal(result.text, 'header\r\n \r\n\tbaz qux'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('whitespace-normalized matching preserves multiple extra leading blank lines', () => { + const source = 'header\r\n \r\n foo bar\r\n'; + const result = applyTextEdits(source, [{ + oldText: '\n \n foo bar', newText: '\r\n\t\r\n\t\r\n\t\r\n\tbaz qux', + }], 'tolerant'); + assert.equal(result.text, 'header\r\n \r\n \r\n\t\r\n\tbaz qux\r\n'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('whitespace-normalized matching preserves extra trailing blank lines when boundary indentation changes', () => { + const source = 'foo bar\r\n next'; + const result = applyTextEdits(source, [{ + oldText: 'foo bar\n ', newText: 'baz qux\n\t\n\t', + }], 'tolerant'); + assert.equal(result.text, 'baz qux\r\n\t\r\n next'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('whitespace-normalized matching leaves intentional indentation on an extra blank line', () => { + const result = applyTextEdits('foo bar\r\nnext', [{ + oldText: 'foo bar\n', newText: 'baz qux\n\t\n', + }], 'tolerant'); + assert.equal(result.text, 'baz qux\r\n\t\r\nnext'); +}); + +test('whitespace-normalized matching does not duplicate differing trailing spaces', () => { + const result = applyTextEdits('foo bar \nnext', [{ + oldText: 'foo bar \n', newText: 'baz qux\t\n', + }], 'tolerant'); + assert.equal(result.text, 'baz qux \nnext'); +}); + +test('whitespace-normalized matching peels a leading newline even if adjacent spaces differ', () => { + const result = applyTextEdits('header\r\n foo bar\r\nnext', [{ + oldText: '\n foo bar', newText: '\nchanged', + }], 'tolerant'); + assert.equal(result.text, 'header\r\n changed\r\nnext'); +}); + +test('whitespace-normalized matching rejects absent required boundary whitespace', () => { + const cases = [ + { source: 'foo bar', oldText: 'foo bar\r\n', newText: 'baz qux\n' }, + { source: 'foo bar next', oldText: 'foo bar\n', newText: 'baz qux\n' }, + { source: 'previous foo bar', oldText: '\nfoo bar', newText: '\nbaz qux' }, + { source: 'foo bar', oldText: ' foo bar', newText: ' baz qux' }, + { source: 'foo bar', oldText: 'foo bar ', newText: 'baz qux ' }, + { source: 'foo bar\nnext', oldText: 'foo bar\n\n', newText: 'baz qux\n\n' }, + { source: 'header\nfoo bar\nnext', oldText: '\n\nfoo bar', newText: '\n\nbaz qux' }, + { source: 'header\n \n foo bar', oldText: '\n \n foo bar', newText: '\n\tbaz qux' }, + ]; + for (const edit of cases) { + const { source, ...request } = edit; + const error = rejection(() => applyTextEdits(source, [request], 'tolerant')); + assert.match(error.message, /old_text was not found/, `unmatched boundary: ${JSON.stringify(edit)}`); + } +}); + +test('whitespace-normalized scanning skips boundary-invalid matches and keeps later valid ones', () => { + const source = 'header foo bar header\nfoo bar\nend'; + const result = applyTextEdits(source, [{ + oldText: '\nfoo bar\n', newText: '\nbaz qux\n', replaceAll: true, + }], 'tolerant'); + assert.equal(result.text, 'header foo bar header\nbaz qux\nend'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('whitespace-normalized matching rejects boundary removal outside its token range', () => { + const cases = [ + { source: 'foo bar\r\nnext', oldText: 'foo bar\r\n', newText: 'baz qux' }, + { source: 'header\r\n foo bar', oldText: '\n foo bar', newText: 'baz qux' }, + { source: ' foo bar', oldText: ' foo bar', newText: 'baz qux' }, + { source: 'foo bar ', oldText: 'foo bar ', newText: 'baz qux' }, + { source: 'header\nfoo bar\nnext', oldText: '\nfoo bar\n', newText: '\n' }, + { source: 'foo bar\n\nnext', oldText: 'foo bar\n\n', newText: 'baz qux\n' }, + ]; + for (const edit of cases) { + const { source, ...request } = edit; + const error = rejection(() => applyTextEdits(source, [request], 'tolerant')); + assert.match(error.message, /old_text was not found/, `unremovable boundary: ${JSON.stringify(edit)}`); + } + const exact = applyTextEdits('foo bar\r\nnext', [{ + oldText: 'foo bar\r\n', newText: 'baz qux', + }], 'tolerant'); + assert.equal(exact.text, 'baz quxnext'); + assert.deepEqual(exact.matches, [{ strategy: 'exact', occurrences: 1 }]); +}); + +test('whitespace-normalized matching preserves two present boundary line breaks', () => { + const result = applyTextEdits('foo bar\n\nnext', [{ + oldText: 'foo bar\r\n\r\n', newText: 'baz qux\n\n', + }], 'tolerant'); + assert.equal(result.text, 'baz qux\n\nnext'); +}); + +test('whitespace-normalized matching preserves an intentional extra line break', () => { + const result = applyTextEdits('foo bar\r\nnext', [{ + oldText: 'foo bar\r\n', newText: 'baz qux\n\n', + }], 'tolerant'); + assert.equal(result.text, 'baz qux\r\n\r\nnext'); +}); + +test('line-trimmed replacements adopt the matched line ending in mixed files', () => { + const source = 'header\r\nfoo \nbar \nnext\n'; + const result = applyTextEdits(source, [{ oldText: 'foo\nbar', newText: 'baz\nqux' }], 'tolerant'); + assert.equal(result.text, 'header\r\nbaz\nqux\nnext\n'); + assert.deepEqual(result.matches, [{ strategy: 'line-trimmed', occurrences: 1 }]); +}); + +test('indentation-flexible replacements keep the selected line ending in mixed files', () => { + const source = 'header\r\n method() {\n return 1;\n }\nend\n'; + const result = applyTextEdits(source, [{ + oldText: 'method() {\n return 1;\n}', + newText: 'method() {\r\n return 2;\r\n}', + }], 'tolerant'); + assert.equal(result.text, 'header\r\n method() {\n return 2;\n }\nend\n'); + assert.deepEqual(result.matches, [{ strategy: 'indentation-flexible', occurrences: 1 }]); +}); + +test('line-trimmed replaceAll keeps each local ending in mixed files', () => { + const source = 'foo \r\nbar \r\nfoo \nbar \nend'; + const result = applyTextEdits(source, [{ + oldText: 'foo\nbar', newText: 'baz\nqux', replaceAll: true, + }], 'tolerant'); + assert.equal(result.text, 'baz\r\nqux\r\nbaz\nqux\nend'); + assert.deepEqual(result.matches, [{ strategy: 'line-trimmed', occurrences: 2 }]); +}); + +test('exact replacements preserve caller line endings rather than normalizing them', () => { + const source = 'foo bar\r\nnext\r\n'; + const result = applyTextEdits(source, [{ oldText: 'foo bar\r\n', newText: 'baz qux\n' }], 'tolerant'); + assert.equal(result.text, 'baz qux\nnext\r\n'); + assert.deepEqual(result.matches, [{ strategy: 'exact', occurrences: 1 }]); +}); + +test('whitespace-normalized replacement adopts the matched line ending in mixed files', () => { + const original = 'header\r\nfoo bar\nnext\n'; + const result = applyTextEdits(original, [{ + oldText: 'foo bar\n', newText: 'baz\nqux\n', + }], 'tolerant'); + assert.equal(result.text, 'header\r\nbaz\nqux\nnext\n'); +}); + +test('whitespace-normalized replaceAll peels boundary newlines on every mixed-ending match', () => { + const source = 'foo bar\r\nfoo bar\nend'; + const result = applyTextEdits(source, [{ + oldText: 'foo bar\r\n', newText: 'baz qux\n', replaceAll: true, + }], 'tolerant'); + assert.equal(result.text, 'baz qux\r\nbaz qux\nend'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 2 }]); +}); + +test('whitespace-normalized replaceAll uses each matched line ending in mixed files', () => { + const source = 'foo bar\r\nfoo bar\nend'; + const result = applyTextEdits(source, [{ + oldText: 'foo bar', newText: 'baz\nqux', replaceAll: true, + }], 'tolerant'); + assert.equal(result.text, 'baz\r\nqux\r\nbaz\nqux\nend'); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 2 }]); +}); + +test('whitespace-normalized replaceAll finds many same-line matches without rescanning newline tails', () => { + const source = 'foo bar '.repeat(16_000); + const start = performance.now(); + const result = applyTextEdits(source, [{ + oldText: 'foo bar', newText: 'baz\nqux', replaceAll: true, + }], 'tolerant'); + assert.equal(result.text, 'baz\nqux '.repeat(16_000)); + assert.deepEqual(result.matches, [{ strategy: 'whitespace-normalized', occurrences: 16_000 }]); + assert.ok(performance.now() - start < 2_000, 'line endings must not be searched from each match'); +}); + +test('whitespace-normalized matches do not indent new_text twice', () => { + const text = ' total = price *\n quantity;\n'; + const applied = applyTextEdits( + text, + [{ oldText: ' total = price * quantity;', newText: ' total = price * quantity * rate;' }], + 'tolerant', + ); + assert.equal(applied.text, ' total = price * quantity * rate;\n'); + assert.deepEqual(applied.matches, [{ strategy: 'whitespace-normalized', occurrences: 1 }]); +}); + +test('tolerant matching still refuses an ambiguous edit', () => { + const error = rejection(() => + applyTextEdits('x = 1 \ny = 2\nx = 1\n', [{ oldText: 'x = 1', newText: 'x = 3' }], 'tolerant'), + ); + assert.match(error.message, /matched 2 locations at lines 1, 3/); +}); + +test('exact overlap counts UTF-16 offsets while replaceAll consumes whole characters', () => { + const error = rejection(() => applyTextEdits('😀😀😀', [ + { oldText: '😀😀', newText: 'x' }, + ])); + assert.match(error.message, /matched 2 locations/); + const result = applyTextEdits('😀😀😀', [{ oldText: '😀😀', newText: 'x', replaceAll: true }]); + assert.equal(result.text, 'x😀'); + assert.deepEqual(result.matches, [{ strategy: 'exact', occurrences: 1 }]); +}); + +test('long overlapping exact matches are counted without repeatedly rescanning the input', () => { + const text = 'a'.repeat(450_000); + const oldText = 'a'.repeat(220_000); + const started = performance.now(); + const error = rejection(() => applyTextEdits(text, [{ oldText, newText: 'b' }])); + assert.match(error.message, /matched 230001 locations/); + assert.ok(performance.now() - started < 3_000, 'overlapping ambiguity must be counted in linear time'); +}); + +test('highly repeated exact matches report the count with bounded line samples', () => { + const error = rejection(() => applyTextEdits('z'.repeat(200_000), [ + { oldText: 'z', newText: 'y' }, + ])); + assert.match(error.message, /matched 200000 locations at lines 1, 1, 1, 1, 1 and 199995 more/); +}); + +test('replaceAll replaces every location and reports the count', () => { + const applied = applyTextEdits('foo(); bar(); foo();', [ + { oldText: 'foo()', newText: 'baz()', replaceAll: true }, + ]); + assert.equal(applied.text, 'baz(); bar(); baz();'); + assert.deepEqual(applied.matches, [{ strategy: 'exact', occurrences: 2 }]); +}); + +test('replaceAll over line windows never overlaps its own matches', () => { + const applied = applyTextEdits( + 'a\na\na\na\n', + [{ oldText: 'a\na', newText: 'b', replaceAll: true }], + 'tolerant', + ); + assert.equal(applied.text, 'b\nb\n'); + assert.deepEqual(applied.matches, [{ strategy: 'exact', occurrences: 2 }]); +}); + +test('replaceAll bounds memory for a million identical and changed matches', () => { + const source = 'x'.repeat(1024 * 1024); + const noop = applyTextEdits(source, [ + { oldText: 'x', newText: 'x', replaceAll: true }, + { oldText: 'x', newText: 'y', replaceAll: true }, + ]); + assert.equal(noop.text, 'y'.repeat(source.length)); + assert.deepEqual(noop.matches, [ + { strategy: 'exact', occurrences: source.length }, + { strategy: 'exact', occurrences: source.length }, + ]); +}); + +test('replaceAll preserves non-overlapping literal edits when identical hits precede changed hits', () => { + for (const [source, oldText, newText] of [ + ['aaaaa', 'aa', 'a'], + ['foo foo foo', 'foo', 'bar'], + ['😀😀😀', '😀😀', 'x'], + ['x\nx\nx', 'x', ''], + ['aba aba', 'aba', 'aba'], + ]) { + const result = applyTextEdits(source, [{ oldText, newText, replaceAll: true }]); + assert.equal(result.text, source.split(oldText).join(newText)); + assert.deepEqual(result.matches, [{ strategy: 'exact', occurrences: source.split(oldText).length - 1 }]); + } + const source = 'foo\tbar foo bar'; + const tolerant = applyTextEdits(source, [{ + oldText: 'foo bar', newText: 'foo\tbar', replaceAll: true, + }], 'tolerant'); + assert.equal(tolerant.text, 'foo\tbar foo\tbar'); + assert.deepEqual(tolerant.matches, [{ strategy: 'whitespace-normalized', occurrences: 2 }]); +}); + +test('replaceAll handles repeated tolerant and non-overlapping line-window matches', () => { + const lines = 'foo \r\nbar \r\n'.repeat(4000); + const result = applyTextEdits(lines, [ + { oldText: 'foo\nbar', newText: 'baz\nqux', replaceAll: true }, + { oldText: 'baz qux', newText: 'updated', replaceAll: true }, + ], 'tolerant'); + assert.equal(result.text, 'updated\r\n'.repeat(4000)); + assert.deepEqual(result.matches, [ + { strategy: 'line-trimmed', occurrences: 4000 }, + { strategy: 'whitespace-normalized', occurrences: 4000 }, + ]); +}); + +test('replaceAll rejects oversized intermediate output before constructing it', () => { + assert.throws( + () => applyTextEdits('x'.repeat(1_000_000), [ + { oldText: 'x', newText: 'y'.repeat(100_000), replaceAll: true }, + ]), + WorkspaceEditOutputLimitError, + ); +}); + +test('replaceAll still fails when nothing matches', () => { + const error = rejection(() => + applyTextEdits('abc', [{ oldText: 'xyz', newText: '', replaceAll: true }]), + ); + assert.match(error.message, /old_text was not found/); +}); + +test('line diagnostic index follows earlier successful batch edits', () => { + const error = rejection(() => applyTextEdits('alpha\nbeta\n', [ + { oldText: 'not present', newText: 'skip' }, + { oldText: 'alpha', newText: 'first\nsecond' }, + { oldText: 'beta\nmissing', newText: 'nope' }, + ])); + assert.match(error.message, /Edit 3: old_text was not found; its first line appears at line 3/); +}); + +test('newline-dense files do not materialize millions of line objects across failed batch edits', () => { + const source = 'a\n'.repeat(450_000); + const edits = Array.from({ length: 80 }, () => ({ oldText: 'missing\ntext', newText: 'other' })); + const started = performance.now(); + const error = rejection(() => applyTextEdits(source, edits)); + assert.equal(error.failures.length, edits.length); + assert.ok(performance.now() - started < 10_000, 'a bounded file must not re-index every failed diagnostic'); +}); + +test('first-line hints report bounded samples even when the line repeats throughout a file', () => { + const error = rejection(() => applyTextEdits('a\n'.repeat(100_000), [ + { oldText: 'a\nmissing', newText: 'replacement' }, + ])); + assert.match(error.message, /its first line appears at lines 1, 2, 3, 4, 5 and 99995 more/); +}); + +test('repetitive indentation candidates fail closed after a bounded comparison budget', () => { + const error = rejection(() => applyTextEdits(' a\n'.repeat(1_200), [ + { oldText: `${'a\n'.repeat(199)}a`, newText: 'changed' }, + ], 'tolerant')); + assert.match(error.message, /too many repetitive line-window candidates/); +}); + +test('repetitive long line windows return a missing-edit diagnosis without quadratic scans', () => { + const text = 'line\n'.repeat(24_000); + const oldText = `${'line\n'.repeat(6_000)}missing`; + const started = performance.now(); + const error = rejection(() => applyTextEdits(text, [{ oldText, newText: 'replacement' }])); + assert.match(error.message, /did not apply and nothing was written/); + assert.ok(performance.now() - started < 2_000, 'a bounded edit must not compare every long window'); +}); + +test('a hundred failing multi-line edits on a large file are diagnosed quickly', () => { + const text = Array.from({ length: 30_000 }, (_, index) => ` line ${index} = value;`).join('\n'); + const edits = Array.from({ length: 100 }, (_, index) => ({ + oldText: `line ${index} = value;\n other ${index};\n more ${index};`, + newText: 'y', + })); + const started = performance.now(); + rejection(() => applyTextEdits(text, edits)); + assert.ok(performance.now() - started < 5_000); +}); + +test('diagnostics for a hundred failing edits stay within the settlement bound', () => { + const edits = Array.from({ length: 100 }, (_, index) => ({ + oldText: `missing ${'x'.repeat(200)} ${index}`, + newText: 'y', + })); + const error = rejection(() => applyTextEdits('content\n'.repeat(1000), edits)); + assert.ok(error.message.length <= EDIT_DIAGNOSTIC_MAX_CHARS); + assert.match(error.message, /^100 of 100 workspace edits did not apply/); + assert.match(error.message, /more failing edits not shown/); +}); diff --git a/packages/code/src/edits.ts b/packages/code/src/edits.ts new file mode 100644 index 00000000..0a6ff108 --- /dev/null +++ b/packages/code/src/edits.ts @@ -0,0 +1,700 @@ +import { BRIDGE_WORKSPACE_WRITE_MAX_BYTES } from './protocol.js'; + +import type { + WorkspaceEditMatch, + WorkspaceEditMatching, + WorkspaceEditMatchStrategy, + WorkspaceTextEdit, +} from './protocol.js'; + +/** + * The Code API rejects a settlement whose error exceeds 4096 characters, and + * callers prefix their own context, so diagnostics stay well below that. + */ +export const EDIT_DIAGNOSTIC_MAX_CHARS = 3000; +const MAX_REPORTED_LINES = 5; +const MAX_SNIPPET_CHARS = 120; +/** A highly repetitive indentation candidate must not monopolize the worker. */ +const MAX_LINE_WINDOW_VERIFICATIONS = 100_000; +const MAX_REPLACEMENT_CHUNK_CHARS = 16 * 1024; + +interface MatchedRange { + start: number; + end: number; + /** Replacement for this range, already adapted to its indentation and line endings. */ + replacement: string; +} + +type MatchOutcome = + | { status: 'matched'; strategy: WorkspaceEditMatchStrategy; ranges: MatchedRange[]; occurrences: number; updated?: string } + | { status: 'ambiguous'; strategy: WorkspaceEditMatchStrategy; count: number; starts: number[] } + | { status: 'limit' } + | { status: 'none' }; + +interface CollectedMatches { + count: number; + source: string; + projectedLength: number; + /** Only the first range is needed for a unique edit or an exact-mode hint. */ + ranges: MatchedRange[]; + starts: number[]; + output?: { chunks: string[]; pending: string; cursor: number }; +} + +export interface EditFailure { + /** Zero-based position of the edit in the request. */ + index: number; + reason: string; +} + +export class WorkspaceEditMatchError extends Error { + constructor( + readonly failures: EditFailure[], + readonly editCount: number, + ) { + super(formatEditFailures(failures, editCount)); + this.name = 'WorkspaceEditMatchError'; + } +} + +export class WorkspaceEditOutputLimitError extends Error { + constructor() { + super('Workspace file exceeds write limit'); + this.name = 'WorkspaceEditOutputLimitError'; + } +} + +export interface AppliedEdits { + text: string; + matches: WorkspaceEditMatch[]; +} + +/** + * Applies ordered replacements to `text`. Each edit sees the result of the + * edits before it. Every edit is attempted even after a failure, so a caller + * learns about all unmatched or ambiguous edits from one rejection; nothing is + * returned unless every edit applied. + */ +export function applyTextEdits( + text: string, + edits: readonly WorkspaceTextEdit[], + matching: WorkspaceEditMatching = 'exact', +): AppliedEdits { + let working = text; + let lineIndex: LineIndex | undefined; + const lines = (): LineIndex => (lineIndex ??= new LineIndex(working)); + const matches: WorkspaceEditMatch[] = []; + const failures: EditFailure[] = []; + edits.forEach((edit, index) => { + const outcome = findEditMatch(working, edit, matching, lines); + if (outcome.status === 'matched') { + working = outcome.updated ?? replaceRanges(working, outcome.ranges); + lineIndex = undefined; + matches.push({ strategy: outcome.strategy, occurrences: outcome.occurrences }); + return; + } + failures.push({ + index, + reason: + outcome.status === 'ambiguous' + ? describeAmbiguous(working, outcome.strategy, outcome.count, outcome.starts) + : outcome.status === 'limit' + ? 'old_text has too many repetitive line-window candidates; include more surrounding lines or use an exact match' + : describeMissing(working, edit.oldText, matching, lines), + }); + }); + if (failures.length > 0) { + throw new WorkspaceEditMatchError(failures, edits.length); + } + return { text: working, matches }; +} + +function findEditMatch( + text: string, + edit: WorkspaceTextEdit, + matching: WorkspaceEditMatching, + lines: () => LineIndex, +): MatchOutcome { + const strategies = matching === 'tolerant' ? TOLERANT_STRATEGIES : EXACT_STRATEGIES; + for (const find of strategies) { + const outcome = find(text, edit, lines); + if (outcome.status !== 'none') return outcome; + } + return { status: 'none' }; +} + +function collectedMatches(text: string, replaceAll?: boolean): CollectedMatches { + return { + count: 0, source: text, projectedLength: text.length, starts: [], ranges: [], + ...(replaceAll === true ? { output: { chunks: [], pending: '', cursor: 0 } } : {}), + }; +} + +function appendReplacement(output: NonNullable, part: string): void { + if (!part) return; + output.pending += part; + if (output.pending.length >= MAX_REPLACEMENT_CHUNK_CHARS) { + output.chunks.push(output.pending); + output.pending = ''; + } +} + +function finishReplacement(matches: CollectedMatches): string { + const output = matches.output!; + if (output.cursor === 0) return matches.source; + appendReplacement(output, matches.source.slice(output.cursor)); + return output.chunks.join('') + output.pending; +} + +function collectMatch( + matches: CollectedMatches, + start: number, + end: number, + replacement: string, + replaceAll?: boolean, +): void { + if (replaceAll === true) { + matches.projectedLength += replacement.length - (end - start); + // Even deleting every remaining source character cannot bring this + // intermediate below the 1 MiB limit. Do not retain more matches. + if (matches.projectedLength - (matches.source.length - end) > BRIDGE_WORKSPACE_WRITE_MAX_BYTES) { + throw new WorkspaceEditOutputLimitError(); + } + // Leave unchanged ranges in the source suffix rather than constructing a + // million identical output pieces for a no-op replaceAll edit. + if (replacement !== matches.source.slice(start, end)) { + const output = matches.output!; + appendReplacement(output, matches.source.slice(output.cursor, start)); + appendReplacement(output, replacement); + output.cursor = end; + } + } + matches.count++; + if (matches.starts.length < MAX_REPORTED_LINES) matches.starts.push(start); + if (matches.count === 1) matches.ranges.push({ start, end, replacement }); +} + +function resolve( + strategy: WorkspaceEditMatchStrategy, + matches: CollectedMatches, + replaceAll: boolean | undefined, +): MatchOutcome { + if (matches.count === 0) return { status: 'none' }; + if (matches.count === 1 || replaceAll === true) { + return { + status: 'matched', strategy, ranges: matches.ranges, occurrences: matches.count, + ...(replaceAll === true ? { updated: finishReplacement(matches) } : {}), + }; + } + return { status: 'ambiguous', strategy, count: matches.count, starts: matches.starts }; +} + +/** + * Overlapping occurrences count toward ambiguity (`aa` in `aaa` is two + * locations), while `replaceAll` replaces non-overlapping occurrences. + */ +function findExact(text: string, edit: WorkspaceTextEdit): MatchOutcome { + if (edit.oldText.length === 0) return { status: 'none' }; + const matches = collectedMatches(text, edit.replaceAll); + const prefix = prefixTable(edit.oldText); + let matched = 0; + for (let index = 0; index < text.length; index++) { + while (matched > 0 && text[index] !== edit.oldText[matched]) matched = prefix[matched - 1]; + if (text[index] === edit.oldText[matched]) matched++; + if (matched !== edit.oldText.length) continue; + collectMatch(matches, index - matched + 1, index + 1, edit.newText, edit.replaceAll); + // Ambiguity counts overlaps, but replaceAll must consume disjoint ranges. + matched = edit.replaceAll === true ? 0 : prefix[matched - 1]; + } + return resolve('exact', matches, edit.replaceAll); +} + +interface Line { + /** Offset of the first character of the line. */ + start: number; + /** Offset just past the line's content, excluding `\r\n` or `\n`. */ + end: number; + /** Offset just past the line terminator, or `undefined` for an unterminated last line. */ + next: number | undefined; + text: string; +} + +/** One bounded index per source revision, shared by matching and diagnostics. */ +class LineIndex { + private readonly newlines: Uint32Array; + readonly length: number; + + constructor(private readonly source: string) { + this.newlines = new Uint32Array(newlineCount(source)); + let index = 0; + for (let offset = source.indexOf('\n'); offset >= 0; offset = source.indexOf('\n', offset + 1)) { + this.newlines[index++] = offset; + } + // A trailing terminator leaves an empty final line, just as before. + this.length = this.newlines.length + 1; + } + + start(index: number): number { + return index === 0 ? 0 : this.newlines[index - 1] + 1; + } + + next(index: number): number | undefined { + return index < this.newlines.length ? this.newlines[index] + 1 : undefined; + } + + end(index: number): number { + const start = this.start(index); + const end = index < this.newlines.length ? this.newlines[index] : this.source.length; + return end > start && this.source[end - 1] === '\r' ? end - 1 : end; + } + + text(index: number): string { + return this.source.slice(this.start(index), this.end(index)); + } + + at(index: number): Line { + return { start: this.start(index), end: this.end(index), next: this.next(index), text: this.text(index) }; + } +} + +/** + * Lines a line-window strategy must match. A trailing line break means "through + * the end of the last line", not "followed by an empty line". + */ +function neededLines(oldText: string): { lines: string[]; throughTerminator: boolean } { + const normalized = oldText.replace(/\r\n/g, '\n'); + const throughTerminator = normalized.endsWith('\n'); + return { + lines: (throughTerminator ? normalized.slice(0, -1) : normalized).split('\n'), + throughTerminator, + }; +} + +/** Fallback when neither the matched line nor its neighbors have a terminator. */ +function fileLineEnding(text: string): '\r\n' | '\n' { + return text.includes('\r\n') ? '\r\n' : '\n'; +} + +function lineEndingAt(text: string, newline: number): '\r\n' | '\n' { + return newline > 0 && text[newline - 1] === '\r' ? '\r\n' : '\n'; +} + +function leadingWhitespace(line: string): string { + return /^[ \t]*/.exec(line)?.[0] ?? ''; +} + +function commonIndent(lines: readonly string[]): string { + let common: string | undefined; + for (const line of lines) { + if (line.trim().length === 0) continue; + const indent = leadingWhitespace(line); + if (common === undefined) { + common = indent; + continue; + } + let shared = 0; + while (shared < common.length && shared < indent.length && common[shared] === indent[shared]) { + shared++; + } + common = common.slice(0, shared); + } + return common ?? ''; +} + +function withLineEnding(value: string, ending: '\r\n' | '\n'): string { + const normalized = value.replace(/\r\n/g, '\n'); + return ending === '\r\n' ? normalized.replace(/\n/g, '\r\n') : normalized; +} + +function prefixTable(values: ArrayLike): Uint32Array { + const prefix = new Uint32Array(values.length); + for (let index = 1, matched = 0; index < values.length; index++) { + while (matched > 0 && values[index] !== values[matched]) matched = prefix[matched - 1]; + if (values[index] === values[matched]) matched++; + prefix[index] = matched; + } + return prefix; +} + +/** + * Line-window strategies compare whole lines, so a match always spans from the + * start of its first line to the end of its last line's content. The file's + * own line terminators are kept, and the replacement adopts them. + */ +function findLineWindows( + text: string, + edit: WorkspaceTextEdit, + strategy: 'line-trimmed' | 'indentation-flexible', + getLines: () => LineIndex, +): MatchOutcome { + const { lines: needle, throughTerminator } = neededLines(edit.oldText); + if (needle.every((line) => line.trim().length === 0)) return { status: 'none' }; + const lines = getLines(); + const ending = fileLineEnding(text); + const needleIndent = strategy === 'indentation-flexible' ? commonIndent(needle) : ''; + const normalizedNeedle = needle.map((line) => + (strategy === 'indentation-flexible' ? stripIndent(line, needleIndent) : line).trimEnd(), + ); + if (needle.length > lines.length) return { status: 'none' }; + // For line-trimmed matching, the normalized lines are the complete comparison. + // For indentation-flexible matching, whole-line content is a linear-time + // prefilter; only complete candidates need their relative indent verified. + const sought = strategy === 'line-trimmed' + ? normalizedNeedle + : normalizedNeedle.map((line) => line.trimStart()); + const prefix = prefixTable(sought); + const collected = collectedMatches(text, edit.replaceAll); + let matched = 0; + let verifications = 0; + for (let index = 0; index < lines.length; index++) { + const value = strategy === 'line-trimmed' ? lines.text(index).trimEnd() : lines.text(index).trim(); + while (matched > 0 && value !== sought[matched]) matched = prefix[matched - 1]; + if (value === sought[matched]) matched++; + if (matched !== sought.length) continue; + + const first = index - sought.length + 1; + const end = throughTerminator ? lines.next(index) : lines.end(index); + let windowIndent = ''; + let valid = end !== undefined; + if (valid && strategy === 'indentation-flexible') { + if (verifications + needle.length > MAX_LINE_WINDOW_VERIFICATIONS) return { status: 'limit' }; + const window = Array.from({ length: needle.length }, (_, offset) => lines.at(first + offset)); + verifications += needle.length; + windowIndent = commonIndent(window.map((line) => line.text)); + valid = window.every((line, offset) => + stripIndent(line.text, windowIndent).trimEnd() === normalizedNeedle[offset], + ); + } + if (valid) { + const replacement = strategy === 'indentation-flexible' + ? reindent(edit.newText, needleIndent, windowIndent) + : edit.newText; + const lineFeed = lines.next(first) ?? (first > 0 ? lines.next(first - 1) : undefined); + const localEnding = lineFeed === undefined ? ending : lineEndingAt(text, lineFeed - 1); + collectMatch(collected, lines.start(first), end!, withLineEnding(replacement, localEnding), edit.replaceAll); + } + // A replacement cannot consume overlapping lines; ambiguity still counts them. + matched = valid && edit.replaceAll === true ? 0 : prefix[matched - 1]; + } + return resolve(strategy, collected, edit.replaceAll); +} + +function stripIndent(line: string, indent: string): string { + return line.startsWith(indent) ? line.slice(indent.length) : line.trimStart(); +} + +/** + * Moves `value` from the indentation the caller wrote to the file's own. Lines + * the caller indented less than its `old_text` keep their indentation as written. + */ +function reindent(value: string, from: string, to: string): string { + if (from === to) return value; + return value + .replace(/\r\n/g, '\n') + .split('\n') + .map((line) => + line.trim().length > 0 && line.startsWith(from) ? to + line.slice(from.length) : line, + ) + .join('\n'); +} + +function findLineTrimmed(text: string, edit: WorkspaceTextEdit, lines: () => LineIndex): MatchOutcome { + return findLineWindows(text, edit, 'line-trimmed', lines); +} + +function findIndentationFlexible(text: string, edit: WorkspaceTextEdit, lines: () => LineIndex): MatchOutcome { + return findLineWindows(text, edit, 'indentation-flexible', lines); +} + +function newlineCount(text: string): number { + let count = 0; + for (let index = text.indexOf('\n'); index >= 0; index = text.indexOf('\n', index + 1)) count++; + return count; +} + +/** A boundary supplied by oldText must exist beside the candidate tokens. */ +function hasBoundaryWhitespace( + text: string, + position: number, + direction: -1 | 1, + requiredNewlines: number, +): boolean { + for (let index = position; index >= 0 && index < text.length; index += direction) { + const char = text[index]; + if (!/\s/.test(char)) break; + if (char === '\n') requiredNewlines--; + if (requiredNewlines <= 0) return true; + } + return false; +} + +/** A token-only match leaves the source's outer whitespace untouched. */ +function peelBoundaryWhitespace( + text: string, + boundary: string, + side: 'leading' | 'trailing', +): string | undefined { + if (!boundary) return text; + if (side === 'leading' && text.startsWith(boundary)) return text.slice(boundary.length); + if (side === 'trailing' && text.endsWith(boundary)) return text.slice(0, -boundary.length); + + const requiredNewlines = newlineCount(boundary); + if (!requiredNewlines) return undefined; + const whitespace = (side === 'leading' ? /^\s*/ : /\s*$/).exec(text)?.[0] ?? ''; + if (newlineCount(whitespace) < requiredNewlines) return undefined; + + if (side === 'leading') { + let end = 0; + for (let count = 0; count < requiredNewlines; count++) { + end = whitespace.indexOf('\n', end) + 1; + } + // Consume the equivalent indentation, stopping before the first extra + // newline. Extra blank lines belong to the replacement, not the wrapper. + while (end < whitespace.length && whitespace[end] !== '\n') end++; + return text.slice(end); + } + + if (newlineCount(whitespace) === requiredNewlines) { + return text.slice(0, text.length - whitespace.length); + } + let start = whitespace.length; + for (let count = 0; count < requiredNewlines; count++) { + start = whitespace.lastIndexOf('\n', start - 1); + } + // Indentation before this newline belongs to the extra blank line. + return text.slice(0, text.length - (whitespace.length - start)); +} + +/** + * Tolerates any run of whitespace, including line breaks, between tokens. The + * match starts and ends on a token, so whitespace the caller wrapped around + * `old_text` is also peeled off `new_text` rather than inserted twice. + */ +function findWhitespaceNormalized(text: string, edit: WorkspaceTextEdit): MatchOutcome { + const oldText = edit.oldText.replace(/\r\n/g, '\n'); + const tokens = oldText.trim().split(/\s+/).filter(Boolean); + if (tokens.length < 2) return { status: 'none' }; + // The matched range contains tokens only; leave the file's boundary whitespace + // outside it. Normalize the caller's line endings *before* peeling equivalent + // wrappers, otherwise CRLF/LF differences insert a second line break. + const leading = /^\s*/.exec(oldText)?.[0] ?? ''; + const trailing = /\s*$/.exec(oldText)?.[0] ?? ''; + const leadingNewlines = newlineCount(leading); + const trailingNewlines = newlineCount(trailing); + const normalizedNewText = edit.newText.replace(/\r\n/g, '\n'); + const withoutLeading = peelBoundaryWhitespace(normalizedNewText, leading, 'leading'); + if (withoutLeading === undefined) return { status: 'none' }; + const newText = peelBoundaryWhitespace(withoutLeading, trailing, 'trailing'); + if (newText === undefined) return { status: 'none' }; + const lfReplacement = newText; + const crlfReplacement = withLineEnding(newText, '\r\n'); + // Match entire whitespace-delimited tokens, never an identifier prefix or + // suffix. A fixed-size regex tokenizes the file; KMP keeps repetitive input + // linear without compiling user-provided text as a regular expression. + const prefix = prefixTable(tokens); + const tokenStarts = new Uint32Array(tokens.length); + const collected = collectedMatches(text, edit.replaceAll); + const words = /\S+/g; + let matched = 0; + let tokenIndex = 0; + // Advance the newline cursor only forwards. replaceAll may encounter many + // matches on one long line, so searching from each match would be quadratic. + let nextNewline = text.indexOf('\n'); + let previousNewline = -1; + for (let word = words.exec(text); word != null; word = words.exec(text)) { + tokenStarts[tokenIndex % tokens.length] = word.index; + while (matched > 0 && word[0] !== tokens[matched]) matched = prefix[matched - 1]; + if (word[0] === tokens[matched]) matched++; + if (matched === tokens.length) { + const start = tokenStarts[(tokenIndex + 1) % tokens.length]; + const end = word.index + word[0].length; + const boundariesMatch = + (leading.length === 0 || hasBoundaryWhitespace(text, start - 1, -1, leadingNewlines)) && + (trailing.length === 0 || hasBoundaryWhitespace(text, end, 1, trailingNewlines)); + if (boundariesMatch) { + while (nextNewline >= 0 && nextNewline < start) { + previousNewline = nextNewline; + nextNewline = text.indexOf('\n', nextNewline + 1); + } + const nearestNewline = nextNewline >= 0 ? nextNewline : previousNewline; + const replacement = nearestNewline >= 0 && lineEndingAt(text, nearestNewline) === '\r\n' + ? crlfReplacement + : lfReplacement; + collectMatch(collected, start, end, replacement, edit.replaceAll); + } + matched = boundariesMatch && edit.replaceAll === true ? 0 : prefix[matched - 1]; + } + tokenIndex++; + } + return resolve('whitespace-normalized', collected, edit.replaceAll); +} + +type Strategy = (text: string, edit: WorkspaceTextEdit, lines: () => LineIndex) => MatchOutcome; + +/** + * Loosest last. Line-window strategies run before whitespace normalization + * because they can carry `new_text` over to the file's indentation; a + * whitespace-normalized match would insert it exactly as written. + */ +const RELAXED_STRATEGIES: readonly Strategy[] = [ + findLineTrimmed, + findIndentationFlexible, + findWhitespaceNormalized, +]; +const EXACT_STRATEGIES: readonly Strategy[] = [findExact]; +const TOLERANT_STRATEGIES: readonly Strategy[] = [findExact, ...RELAXED_STRATEGIES]; + +function replaceRanges(text: string, ranges: readonly MatchedRange[]): string { + const length = ranges.reduce( + (total, range) => total + range.replacement.length - (range.end - range.start), + text.length, + ); + // Every valid UTF-8 string has at least this many encoded bytes. Reject a + // pathological replaceAll before allocating an unbounded intermediate string. + if (length > BRIDGE_WORKSPACE_WRITE_MAX_BYTES) throw new WorkspaceEditOutputLimitError(); + let result = ''; + let cursor = 0; + for (const range of ranges) { + result += text.slice(cursor, range.start) + range.replacement; + cursor = range.end; + } + return result + text.slice(cursor); +} + +function lineNumberAt(text: string, offset: number): number { + let line = 1; + for (let index = text.indexOf('\n'); index >= 0 && index < offset; index = text.indexOf('\n', index + 1)) { + line++; + } + return line; +} + +function formatLineList(text: string, starts: readonly number[], count = starts.length): string { + const shown = starts.slice(0, MAX_REPORTED_LINES).map((start) => lineNumberAt(text, start)); + const more = count - shown.length; + return `line${shown.length === 1 ? '' : 's'} ${shown.join(', ')}${more > 0 ? ` and ${more} more` : ''}`; +} + +function describeAmbiguous( + text: string, + strategy: WorkspaceEditMatchStrategy, + count: number, + starts: readonly number[], +): string { + const how = strategy === 'exact' ? '' : ` (${strategy})`; + return `old_text matched ${count} locations${how} at ${formatLineList(text, starts, count)}; include more surrounding lines so it matches exactly one`; +} + +function snippet(line: string): string { + const trimmed = line.trim(); + return JSON.stringify( + trimmed.length > MAX_SNIPPET_CHARS ? `${trimmed.slice(0, MAX_SNIPPET_CHARS)}…` : trimmed, + ); +} + +const ELISION_LINE = /^\s*(?:(?:\/\/|#|--|\/\*|\*|))?\s*$/; +const LINE_NUMBER_PREFIX = /^\s*\d+\s*(?:\||:|\t)/; + +function describeMissing( + text: string, + oldText: string, + matching: WorkspaceEditMatching, + lines: () => LineIndex, +): string { + const hints: string[] = []; + const nonBlank = neededLines(oldText).lines.filter((line) => line.trim().length > 0); + if (nonBlank.some((line) => ELISION_LINE.test(line))) { + hints.push('it contains an elision placeholder ("..."); copy the exact lines instead of abbreviating'); + } + if (nonBlank.length > 0 && nonBlank.every((line) => LINE_NUMBER_PREFIX.test(line))) { + hints.push('it appears to include line-number prefixes from read_file output; remove them'); + } + if (matching === 'exact') { + let tolerant: MatchOutcome | undefined; + for (const find of RELAXED_STRATEGIES) { + tolerant = find(text, { oldText, newText: '' }, lines); + if (tolerant.status !== 'none') break; + } + if (tolerant?.status === 'matched') { + hints.push( + `the same text exists at ${formatLineList(text, [tolerant.ranges[0].start])} with different whitespace (${tolerant.strategy}); copy that whitespace exactly`, + ); + } else if (text.includes('\r\n') && !oldText.includes('\r\n') && oldText.includes('\n')) { + hints.push('the file uses CRLF line endings'); + } + } + const nearest = nearestLine(nonBlank[0], lines); + if (nearest != null && hints.length === 0) { + hints.push( + nearest.exact + ? `its first line appears at ${formatLineList(text, nearest.starts, nearest.count)}, but the lines after it differ` + : `the closest line is ${formatLineList(text, nearest.starts)}: ${snippet(nearest.text)}`, + ); + } + return `old_text was not found${hints.length > 0 ? `; ${hints.join('; ')}` : ''}`; +} + +/** Finds where the first line of a failed edit most likely belongs. */ +function nearestLine( + firstLine: string | undefined, + getLines: () => LineIndex, +): { exact: boolean; count: number; starts: number[]; text: string } | undefined { + const target = firstLine?.trim(); + if (!target) return undefined; + const lines = getLines(); + const starts: number[] = []; + let count = 0; + let firstMatch = ''; + for (let index = 0; index < lines.length; index++) { + const content = lines.text(index); + if (content.trim() !== target) continue; + if (count++ === 0) firstMatch = content; + if (starts.length < MAX_REPORTED_LINES) starts.push(lines.start(index)); + } + if (count > 0) { + return { exact: true, count, starts, text: firstMatch }; + } + const tokens = new Set(target.split(/\W+/).filter((token) => token.length > 1)); + if (tokens.size < 2) return undefined; + let best: Line | undefined; + let bestScore = 0; + for (let index = 0; index < lines.length; index++) { + const content = lines.text(index); + let score = 0; + for (const token of new Set(content.split(/\W+/))) { + if (tokens.has(token)) score++; + } + if (score > bestScore) { + best = { start: lines.start(index), end: lines.end(index), next: lines.next(index), text: content }; + bestScore = score; + } + } + return best != null && bestScore / tokens.size >= 0.5 + ? { exact: false, count: 1, starts: [best.start], text: best.text } + : undefined; +} + +function formatEditFailures(failures: readonly EditFailure[], editCount: number): string { + if (editCount === 1) { + return `Workspace edit did not apply and nothing was written: ${failures[0]?.reason ?? 'no match'}.`.slice( + 0, + EDIT_DIAGNOSTIC_MAX_CHARS, + ); + } + let message = `${failures.length} of ${editCount} workspace edits did not apply, so nothing was written. Every other edit matched.`; + let shown = 0; + for (const failure of failures) { + const line = `\nEdit ${failure.index + 1}: ${failure.reason}.`; + if (message.length + line.length > EDIT_DIAGNOSTIC_MAX_CHARS - 120) break; + message += line; + shown++; + } + const hidden = failures.length - shown; + if (hidden > 0) { + message += `\n${hidden} more failing edit${hidden === 1 ? '' : 's'} not shown.`; + } + if (failures.some((failure) => failure.index > 0)) { + message += '\nLine numbers account for the earlier edits in this batch.'; + } + return message.slice(0, EDIT_DIAGNOSTIC_MAX_CHARS); +} diff --git a/packages/code/src/protocol.test.ts b/packages/code/src/protocol.test.ts index 5c26c462..abb378bb 100644 --- a/packages/code/src/protocol.test.ts +++ b/packages/code/src/protocol.test.ts @@ -543,6 +543,164 @@ test('workspace mutations accept bounded UTF-8 requests and exact result shapes' ); }); +test('edit matching and replaceAll are opt-in and reported only when requested', () => { + const tolerant = { + protocolVersion: 1 as const, + operation: 'edit_file' as const, + workspaceId: 'primary', + path: 'notes.txt', + matching: 'tolerant' as const, + edits: [ + { oldText: 'hello', newText: 'goodbye' }, + { oldText: 'world', newText: 'BYOM', replaceAll: true }, + ], + }; + const result = { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + replacements: 2, + bytesWritten: 12, + }; + const matches = [ + { strategy: 'line-trimmed', occurrences: 1 }, + { strategy: 'exact', occurrences: 3 }, + ]; + assert.equal(isWorkspaceToolRequest(tolerant), true); + assert.equal(isWorkspaceToolRequest({ ...tolerant, operation: 'preview_edit' }), true); + assert.equal(isWorkspaceToolRequest({ ...tolerant, matching: 'fuzzy' }), false); + assert.equal( + isWorkspaceToolRequest({ + ...tolerant, + edits: [{ oldText: 'a', newText: 'b', replaceAll: 'yes' }], + }), + false, + ); + assert.equal( + isWorkspaceToolRequest({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + oldText: 'a', + newText: 'b', + replaceAll: true, + }), + false, + 'replaceAll is only accepted per batch edit', + ); + + assert.equal(isWorkspaceToolResult(tolerant, { ...result, matches }), true); + assert.equal(isWorkspaceToolResult(tolerant, result), false, 'an opted-in request must report matches'); + assert.equal( + isWorkspaceToolResult(tolerant, { ...result, matches: matches.slice(0, 1) }), + false, + ); + assert.equal( + isWorkspaceToolResult(tolerant, { + ...result, + matches: [matches[0], { strategy: 'exact', occurrences: 0 }], + }), + false, + ); + assert.equal( + isWorkspaceToolResult(tolerant, { + ...result, + matches: [{ strategy: 'exact', occurrences: 2 }, matches[1]], + }), + false, + 'only replaceAll edits may replace more than one location', + ); + + const exactReplaceAll = { + ...tolerant, + matching: undefined, + edits: [{ oldText: 'a', newText: 'b', replaceAll: true }], + }; + delete exactReplaceAll.matching; + assert.equal( + isWorkspaceToolResult(exactReplaceAll, { + ...result, + replacements: 1, + matches: [{ strategy: 'line-trimmed', occurrences: 2 }], + }), + false, + 'without tolerant matching every edit matches exactly', + ); + + const legacy = { ...tolerant, edits: [{ oldText: 'a', newText: 'b' }] }; + delete (legacy as { matching?: string }).matching; + assert.equal( + isWorkspaceToolResult(legacy, { + ...result, + replacements: 1, + matches: [{ strategy: 'exact', occurrences: 1 }], + }), + false, + 'legacy requests never receive a matches field', + ); +}); + +test('edit features advertise any unique subset of the known features', () => { + const valid = { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file', 'edit_file'], + workspaces: [{ id: 'primary' }], + editFileModes: ['single', 'batch'], + }, + }; + const withFeatures = (editFileFeatures: unknown) => ({ + ...valid, + workspaceTools: { ...valid.workspaceTools, editFileFeatures }, + }); + for (const features of [ + ['expected_base_sha256'], + ['tolerant_match'], + ['expected_base_sha256', 'tolerant_match', 'replace_all'], + ]) { + assert.equal(isValidBridgeWorkerCapabilities(withFeatures(features)), true, features.join(',')); + } + for (const features of [[], ['fuzzy'], ['replace_all', 'replace_all']]) { + assert.equal(isValidBridgeWorkerCapabilities(withFeatures(features)), false, features.join(',')); + } +}); + +test('preview-only workers may advertise tolerant matching and replace-all, but not edit-only hashes', () => { + const previewOnly = { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file', 'preview_edit'], + workspaces: [{ id: 'primary' }], + editFileModes: ['single', 'batch'], + }, + }; + for (const features of [ + ['tolerant_match'], + ['replace_all'], + ['tolerant_match', 'replace_all'], + ]) { + assert.equal(isValidBridgeWorkerCapabilities({ + ...previewOnly, + workspaceTools: { ...previewOnly.workspaceTools, editFileFeatures: features }, + }), true, features.join(',')); + } + assert.equal(isValidBridgeWorkerCapabilities({ + ...previewOnly, + workspaceTools: { + ...previewOnly.workspaceTools, + editFileFeatures: ['expected_base_sha256', 'tolerant_match'], + }, + }), false); +}); + test('workspace commands require bounded sandbox inputs and outputs', () => { const request = { protocolVersion: 1 as const, diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 7657fc9c..bc8482c9 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -280,7 +280,29 @@ export type BridgeWorkspaceToolOperation = export type WorkspaceWriteFileMode = 'replace' | 'create'; export type WorkspaceEditFileMode = 'single' | 'batch'; -export type WorkspaceEditFileFeature = 'expected_base_sha256'; +export type WorkspaceEditFileFeature = + | 'expected_base_sha256' + | 'tolerant_match' + | 'replace_all'; +/** Every edit feature this protocol version defines, for capability validation. */ +export const WORKSPACE_EDIT_FILE_FEATURES: readonly WorkspaceEditFileFeature[] = [ + 'expected_base_sha256', + 'tolerant_match', + 'replace_all', +]; +/** `tolerant` falls back from exact matching to whitespace-tolerant strategies. */ +export type WorkspaceEditMatching = 'exact' | 'tolerant'; +export type WorkspaceEditMatchStrategy = + | 'exact' + | 'line-trimmed' + | 'whitespace-normalized' + | 'indentation-flexible'; +const WORKSPACE_EDIT_MATCH_STRATEGIES = new Set([ + 'exact', + 'line-trimmed', + 'whitespace-normalized', + 'indentation-flexible', +]); export type WorkspaceListFileFeature = 'after_path'; export type WorkspaceProgrammaticLanguage = 'bash'; @@ -433,6 +455,8 @@ interface WorkspaceEditFileRequestBase { path: string; /** Refuses the mutation unless current file bytes match this preview revision. */ expectedBaseSha256?: string; + /** Requires the `tolerant_match` edit feature. Omitted means `exact`. */ + matching?: WorkspaceEditMatching; } export interface WorkspaceSingleEditFileRequest @@ -459,6 +483,15 @@ export type WorkspaceEditFileRequest = export interface WorkspaceTextEdit { oldText: string; newText: string; + /** Replaces every match instead of requiring exactly one. Requires `replace_all`. */ + replaceAll?: boolean; +} + +/** How one edit matched. Present only when the request set `matching` or `replaceAll`. */ +export interface WorkspaceEditMatch { + strategy: WorkspaceEditMatchStrategy; + /** Locations replaced; always 1 unless the edit set `replaceAll`. */ + occurrences: number; } export interface WorkspaceEditFileResult { @@ -466,8 +499,10 @@ export interface WorkspaceEditFileResult { operation: 'edit_file'; workspaceId: string; path: string; + /** Number of edits applied, which is always the number requested. */ replacements: number; bytesWritten: number; + matches?: WorkspaceEditMatch[]; } interface WorkspacePreviewEditRequestBase { @@ -476,6 +511,8 @@ interface WorkspacePreviewEditRequestBase { workspaceId: string; workspaceInstanceId?: string; path: string; + /** Requires the `tolerant_match` edit feature. Omitted means `exact`. */ + matching?: WorkspaceEditMatching; } export interface WorkspaceSinglePreviewEditRequest @@ -506,6 +543,7 @@ export interface WorkspacePreviewEditResult { baseSha256: string; replacements: number; bytesWritten: number; + matches?: WorkspaceEditMatch[]; } export interface WorkspaceExecuteCommandRequest { @@ -599,6 +637,7 @@ const WORKSPACE_EDIT_REQUEST_KEYS = new Set([ 'newText', 'edits', 'expectedBaseSha256', + 'matching', ]); const WORKSPACE_PREVIEW_EDIT_REQUEST_KEYS = new Set([ 'protocolVersion', @@ -609,8 +648,10 @@ const WORKSPACE_PREVIEW_EDIT_REQUEST_KEYS = new Set([ 'oldText', 'newText', 'edits', + 'matching', ]); -const WORKSPACE_TEXT_EDIT_KEYS = new Set(['oldText', 'newText']); +const WORKSPACE_TEXT_EDIT_KEYS = new Set(['oldText', 'newText', 'replaceAll']); +const WORKSPACE_EDIT_MATCH_KEYS = new Set(['strategy', 'occurrences']); const WORKSPACE_COMMAND_REQUEST_KEYS = new Set([ 'environmentAction', 'protocolVersion', @@ -663,6 +704,7 @@ const WORKSPACE_EDIT_RESULT_KEYS = new Set([ 'path', 'replacements', 'bytesWritten', + 'matches', ]); const WORKSPACE_PREVIEW_EDIT_RESULT_KEYS = new Set([ 'protocolVersion', @@ -674,6 +716,7 @@ const WORKSPACE_PREVIEW_EDIT_RESULT_KEYS = new Set([ 'baseSha256', 'replacements', 'bytesWritten', + 'matches', ]); const WORKSPACE_COMMAND_RESULT_KEYS = new Set([ 'protocolVersion', @@ -1145,6 +1188,44 @@ function isWithinRequestedPath(candidate: string, requested?: string): boolean { ); } +/** Whether a request asked for per-edit match reporting (and so must receive it). */ +export function workspaceEditRequestReportsMatches( + request: WorkspaceEditFileRequest | WorkspacePreviewEditRequest, +): boolean { + return ( + request.matching !== undefined || + (request.edits?.some((edit) => edit.replaceAll !== undefined) ?? false) + ); +} + +function isValidWorkspaceEditMatches( + request: WorkspaceEditFileRequest | WorkspacePreviewEditRequest, + matches: unknown, +): boolean { + if (!workspaceEditRequestReportsMatches(request)) return matches === undefined; + const edits: WorkspaceTextEdit[] = request.edits ?? [ + { oldText: request.oldText ?? '', newText: request.newText ?? '' }, + ]; + return ( + Array.isArray(matches) && + matches.length === edits.length && + matches.every((match: unknown, index) => { + if (typeof match !== 'object' || match === null) return false; + const candidate = match as Record; + return ( + hasOnlyKeys(candidate, WORKSPACE_EDIT_MATCH_KEYS) && + WORKSPACE_EDIT_MATCH_STRATEGIES.has( + candidate.strategy as WorkspaceEditMatchStrategy, + ) && + (request.matching === 'tolerant' || candidate.strategy === 'exact') && + Number.isSafeInteger(candidate.occurrences) && + Number(candidate.occurrences) >= 1 && + (edits[index]?.replaceAll === true || candidate.occurrences === 1) + ); + }) + ); +} + function isValidWorkspaceEditRequest( request: Record, ): boolean { @@ -1155,6 +1236,13 @@ function isValidWorkspaceEditRequest( ) { return false; } + if ( + request.matching !== undefined && + request.matching !== 'exact' && + request.matching !== 'tolerant' + ) { + return false; + } const edits = hasBatch ? request.edits : [{ oldText: request.oldText, newText: request.newText }]; @@ -1185,7 +1273,9 @@ function isValidWorkspaceEditRequest( candidate.oldText || typeof candidate.newText !== 'string' || Buffer.from(candidate.newText).toString('utf8') !== - candidate.newText + candidate.newText || + (candidate.replaceAll !== undefined && + typeof candidate.replaceAll !== 'boolean') ) { return false; } @@ -1493,7 +1583,8 @@ export function isWorkspaceToolResult( result.replacements === replacements && Number.isSafeInteger(result.bytesWritten) && Number(result.bytesWritten) >= 0 && - Number(result.bytesWritten) <= BRIDGE_WORKSPACE_WRITE_MAX_BYTES + Number(result.bytesWritten) <= BRIDGE_WORKSPACE_WRITE_MAX_BYTES && + isValidWorkspaceEditMatches(request, result.matches) ); } @@ -1514,7 +1605,8 @@ export function isWorkspaceToolResult( Number(result.bytesWritten) === new TextEncoder().encode(content).byteLength + (result.hasUtf8Bom ? 3 : 0) && - Number(result.bytesWritten) <= BRIDGE_WORKSPACE_WRITE_MAX_BYTES + Number(result.bytesWritten) <= BRIDGE_WORKSPACE_WRITE_MAX_BYTES && + isValidWorkspaceEditMatches(request, result.matches) ); } @@ -1638,9 +1730,20 @@ export function isValidBridgeWorkspaceToolCapabilities( if ( capabilities.editFileFeatures !== undefined && (!Array.isArray(capabilities.editFileFeatures) || - capabilities.editFileFeatures.length !== 1 || - !capabilities.operations.includes('edit_file') || - capabilities.editFileFeatures[0] !== 'expected_base_sha256') + capabilities.editFileFeatures.length < 1 || + capabilities.editFileFeatures.length > + WORKSPACE_EDIT_FILE_FEATURES.length || + (!capabilities.operations.includes('edit_file') && + !capabilities.operations.includes('preview_edit')) || + !capabilities.editFileFeatures.every((feature: unknown) => + WORKSPACE_EDIT_FILE_FEATURES.includes( + feature as WorkspaceEditFileFeature, + ) && + (feature !== 'expected_base_sha256' || + (capabilities.operations as string[]).includes('edit_file')), + ) || + new Set(capabilities.editFileFeatures).size !== + capabilities.editFileFeatures.length) ) { return false; } diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 2a5896fd..4d1471a7 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -161,6 +161,19 @@ function errorCode(value: object): string | undefined { return undefined; } +/** Edit conflicts can include source excerpts, so disclosure needs read or full-preview access. */ +function workspaceCanReadSource( + capabilities: BridgeWorkerCapabilities['workspaceTools'], + workspaceId: string, +): boolean { + const workspace = capabilities?.workspaces.find((entry) => entry.id === workspaceId); + if (!workspace || !capabilities) return false; + return (['read_file', 'preview_edit'] as const).some( + (operation) => capabilities.operations.includes(operation) && + (workspace.operations == null || workspace.operations.includes(operation)), + ); +} + function workspaceCapabilitiesMatch( advertised: NonNullable, executor: NonNullable, @@ -363,7 +376,8 @@ function supportedWorkspaceCapabilities( }); if (workspaces.length === 0) return undefined; const editFileFeatures = desired.editFileFeatures?.filter((feature) => - registration.supportedWorkspaceEditFileFeatures?.includes(feature), + registration.supportedWorkspaceEditFileFeatures?.includes(feature) && + (feature !== 'expected_base_sha256' || operations.includes('edit_file')), ); const listFileFeatures = desired.listFileFeatures?.filter((feature) => registration.supportedWorkspaceListFileFeatures?.includes(feature), @@ -392,7 +406,7 @@ function supportedWorkspaceCapabilities( ...(supportsEditRequests && editFileModes?.length ? { editFileModes } : {}), - ...(operations.includes('edit_file') && editFileFeatures?.length + ...(supportsEditRequests && editFileFeatures?.length ? { editFileFeatures } : {}), ...(operations.includes('list_files') && listFileFeatures?.length @@ -1652,9 +1666,13 @@ export class BridgeWorker { ); } if ( - workspaceRequest.operation === 'edit_file' && - workspaceRequest.expectedBaseSha256 !== undefined && - !advertised.editFileFeatures?.includes('expected_base_sha256') + (workspaceRequest.operation === 'edit_file' && + workspaceRequest.expectedBaseSha256 !== undefined && + !advertised.editFileFeatures?.includes('expected_base_sha256')) || + (workspaceRequest.matching !== undefined && + !advertised.editFileFeatures?.includes('tolerant_match')) || + (workspaceRequest.edits?.some((edit) => edit.replaceAll !== undefined) === true && + !advertised.editFileFeatures?.includes('replace_all')) ) { throw new BridgeProtocolError( 'Workspace edit feature is not advertised', @@ -1926,9 +1944,14 @@ export class BridgeWorker { error instanceof WorkspaceToolError ? { errorCode: error.code } : {}), - error: (error instanceof Error - ? error.message - : 'Sandbox execution failed' + error: (assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) && + assignment.request.operation === 'edit_file' && + !workspaceCanReadSource(this.activeCapabilities.workspaceTools, assignment.request.workspaceId) + ? 'Workspace edit failed; source diagnostics require read access' + : error instanceof Error + ? error.message + : 'Sandbox execution failed' ).slice(0, MAX_SETTLEMENT_ERROR_LENGTH), }; } diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index 96a5d1f0..71ef2cd3 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -2,6 +2,7 @@ import assert from 'node:assert/strict'; import test from 'node:test'; import { BridgeProtocolError } from './protocol.js'; +import type { WorkspaceToolRequest } from './protocol.js'; import { BridgeWorker, BridgeWorkspaceQuarantinedError } from './worker.js'; import { SandboxWorkspaceTools, WorkspaceToolError } from './workspace.js'; @@ -659,6 +660,49 @@ test('worker advertises only edit modes and features negotiated by Code API', as ]); }); +test('worker negotiates preview-only tolerant features and drops edit-only hashes', async () => { + const registrations: Array> = []; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' as const, 'preview_edit' as const, 'edit_file' as const], + workspaces: [{ id: 'primary' }], + editFileModes: ['single' as const, 'batch' as const], + editFileFeatures: ['expected_base_sha256' as const, 'tolerant_match' as const, 'replace_all' as const], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + async execute() { throw new Error('must not execute'); }, + }, + workspaceMutationQuarantine: mutationQuarantine(), + fetchImpl: async (_input, init) => { + const capabilities = JSON.parse(String(init?.body)).capabilities.workspaceTools; + registrations.push(capabilities); + return Response.json({ + protocolVersion: 1, workerId: 'vm-1', incarnationId, + registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, + supportedWorkspaceToolOperations: ['read_file', 'preview_edit'], + supportedWorkspaceEditFileModes: ['single', 'batch'], + supportedWorkspaceEditFileFeatures: ['tolerant_match', 'replace_all', 'expected_base_sha256'], + }); + }, + }); + await worker.register(); + assert.deepEqual(registrations.at(-1)?.operations, ['read_file', 'preview_edit']); + assert.deepEqual(registrations.at(-1)?.editFileFeatures, ['tolerant_match', 'replace_all']); +}); + test('worker drops file operations when no request mode is compatible', async () => { const registrations: Array> = []; const workspaceCapabilities = { @@ -3165,6 +3209,98 @@ test('worker rejects workspace operations outside its advertised capability', as assert.match(String(settlement?.error), /operation is not advertised/i); }); +for (const [operation, extras] of [ + ['preview_edit', { matching: 'tolerant', oldText: 'a', newText: 'b' }], + ['preview_edit', { matching: 'exact', oldText: 'a', newText: 'b' }], + ['preview_edit', { edits: [{ oldText: 'a', newText: 'b', replaceAll: false }] }], + ['edit_file', { edits: [{ oldText: 'a', newText: 'b', replaceAll: true }] }], +] as const) { + test(`worker refuses unnegotiated ${operation} feature before dispatch`, async () => { + let executions = 0; + let armed = 0; + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['preview_edit' as const, 'edit_file' as const], + editFileModes: ['single' as const, 'batch' as const], + workspaces: [{ id: 'primary' }], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { statefulWorkspace: true, sandboxProfile: 'native-srt', runtimes: ['bash'], + workspaceTools: workspaceCapabilities }, + workspaceTools: { capabilities: workspaceCapabilities, async execute() { + executions++; + throw new Error('must not execute'); + } }, + workspaceMutationQuarantine: mutationQuarantine(undefined, () => armed++), + fetchImpl: async (_input, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: `assignment-${operation}-${String('matching' in extras ? extras.matching : 'replace')}`, + workerId: 'vm-1', incarnationId, generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_tool', + request: { protocolVersion: 1, operation, workspaceId: 'primary', path: 'notes.txt', ...extras } as WorkspaceToolRequest, + }); + assert.equal(executions, 0); + assert.equal(armed, 0); + assert.equal(settlement?.status, 'rejected'); + assert.match(String(settlement?.error), /edit feature is not advertised/i); + }); +} + +for (const [label, operations, allowed, mayRead] of [ + ['edit-only worker', ['edit_file'], ['edit_file'], false], + ['read denied by workspace', ['read_file', 'edit_file'], ['edit_file'], false], + ['read permitted by workspace', ['read_file', 'edit_file'], ['read_file', 'edit_file'], true], + ['preview permits full file reads', ['preview_edit', 'edit_file'], ['preview_edit', 'edit_file'], true], +] as const) { + test(`edit diagnostics respect ${label} authorization`, async () => { + const source = 'confidential-source-line'; + let settlement: Record | undefined; + const capabilities = { + protocolVersion: 1 as const, + operations: [...operations], + editFileModes: ['single' as const], + workspaces: [{ id: 'primary', operations: [...allowed] }], + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', runtimes: ['bash'], + workspaceTools: capabilities }, + workspaceTools: { capabilities, mutationFailuresAreAtomic: true, async execute() { + throw new WorkspaceToolError(`old_text was not found; closest line: "${source}"`, 'EDIT_CONFLICT'); + } }, + workspaceMutationQuarantine: mutationQuarantine(), + fetchImpl: async (_url, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'edit-source-authorization', + workerId: 'vm-1', incarnationId, generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_tool', + request: { protocolVersion: 1, operation: 'edit_file', workspaceId: 'primary', + path: 'notes.txt', oldText: 'missing', newText: 'replacement' }, + }); + assert.equal(settlement?.status, 'rejected'); + assert.equal(settlement?.errorCode, 'EDIT_CONFLICT'); + assert.equal(String(settlement?.error).includes(source), mayRead); + }); +} + test('worker rejects legacy replacement writes outside its advertised mode', async () => { let executions = 0; let settlement: Record | undefined; diff --git a/packages/code/src/workspace.test.ts b/packages/code/src/workspace.test.ts index 3dd7d2f4..9b805ab0 100644 --- a/packages/code/src/workspace.test.ts +++ b/packages/code/src/workspace.test.ts @@ -939,7 +939,7 @@ test('writable workspaces create, replace, and exactly edit files', async (t) => ], writeFileModes: ['replace', 'create'], editFileModes: ['single', 'batch'], - editFileFeatures: ['expected_base_sha256'], + editFileFeatures: ['expected_base_sha256', 'tolerant_match', 'replace_all'], listFileFeatures: ['after_path'], }); await tools.execute({ @@ -1351,6 +1351,303 @@ test('exact edits reject missing or repeated text without changing the file', as assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), 'aaa'); }); +test('a failed edit batch reports every failing edit and writes nothing', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = 'alpha\nbeta\nbeta\ngamma\n'; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + + await assert.rejects( + tools.execute({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + edits: [ + { oldText: 'alpha', newText: 'ALPHA' }, + { oldText: 'beta', newText: 'BETA' }, + { oldText: 'delta', newText: 'DELTA' }, + ], + }), + (error: unknown) => + error instanceof WorkspaceToolError && + error.code === 'EDIT_CONFLICT' && + /^2 of 3 workspace edits did not apply/.test(error.message) && + /Edit 2: old_text matched 2 locations at lines 2, 3/.test(error.message) && + /Edit 3: old_text was not found/.test(error.message), + ); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); +}); + +test('tolerant edits and replaceAll report how each edit matched', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); + t.after(() => rm(root, { recursive: true, force: true })); + await writeFile(join(root, 'app.ts'), 'if (ok) { \r\n run();\r\n}\r\nrun();\r\n'); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + + const preview = await tools.execute({ + protocolVersion: 1, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'app.ts', + matching: 'tolerant', + edits: [{ oldText: 'if (ok) {\n run();\n}', newText: 'if (ok) {\n go();\n}' }], + }); + assert.equal(preview.operation, 'preview_edit'); + assert.deepEqual(preview.operation === 'preview_edit' && preview.matches, [ + { strategy: 'line-trimmed', occurrences: 1 }, + ]); + + const edit = await tools.execute({ + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'app.ts', + matching: 'tolerant', + edits: [ + { oldText: 'if (ok) {\n run();\n}', newText: 'if (ok) {\n go();\n}' }, + { oldText: 'run();', newText: 'stop();', replaceAll: true }, + ], + }); + assert.deepEqual(edit, { + protocolVersion: 1, + operation: 'edit_file', + workspaceId: 'primary', + path: 'app.ts', + replacements: 2, + bytesWritten: 32, + matches: [ + { strategy: 'line-trimmed', occurrences: 1 }, + { strategy: 'exact', occurrences: 1 }, + ], + }); + assert.equal( + await readFile(join(root, 'app.ts'), 'utf8'), + 'if (ok) {\r\n go();\r\n}\r\nstop();\r\n', + ); +}); + +test('preview and edit peel equivalent newline wrappers without doubling line breaks or losing BOM', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-newlines-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = '\uFEFFheader\r\n \tfoo bar\r\nnext\r\n'; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + const request = { + protocolVersion: 1 as const, + operation: 'preview_edit' as const, + workspaceId: 'primary', + path: 'notes.txt', + matching: 'tolerant' as const, + edits: [ + { oldText: '\n foo bar\r\n', newText: '\r\n baz qux\n' }, + { oldText: 'next', newText: 'done' }, + ], + }; + const preview = await tools.execute(request); + assert.equal(preview.operation, 'preview_edit'); + if (preview.operation !== 'preview_edit') assert.fail('expected preview result'); + assert.equal(preview.content, 'header\r\n \tbaz qux\r\ndone\r\n'); + assert.equal(preview.hasUtf8Bom, true); + assert.equal(preview.bytesWritten, Buffer.byteLength(`\uFEFF${preview.content}`)); + assert.equal(isWorkspaceToolResult(request, preview), true); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + const editRequest = { + ...request, + operation: 'edit_file' as const, + expectedBaseSha256: preview.baseSha256, + }; + const edited = await tools.execute(editRequest); + assert.equal(edited.operation, 'edit_file'); + if (edited.operation !== 'edit_file') assert.fail('expected edit result'); + assert.deepEqual(edited.matches, [ + { strategy: 'whitespace-normalized', occurrences: 1 }, + { strategy: 'exact', occurrences: 1 }, + ]); + assert.equal(isWorkspaceToolResult(editRequest, edited), true); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), `\uFEFF${preview.content}`); +}); + +test('tolerant previews and edits keep source indentation when newText uses a different indent', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-indent-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = 'header\n foo bar\n'; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + const edit = { + protocolVersion: 1 as const, workspaceId: 'primary', path: 'notes.txt', + matching: 'tolerant' as const, oldText: '\n foo bar', newText: '\n\tbaz qux', + }; + const preview = await tools.execute({ ...edit, operation: 'preview_edit' }); + assert.equal(preview.operation === 'preview_edit' && preview.content, 'header\n baz qux\n'); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + const result = await tools.execute({ ...edit, operation: 'edit_file' }); + assert.equal(result.operation === 'edit_file' && result.matches?.[0]?.strategy, 'whitespace-normalized'); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), 'header\n baz qux\n'); +}); + +test('tolerant previews and edits retain an extra leading blank line across a fenced batch', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-extra-line-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = '\uFEFFheader\r\n foo bar\r\nend\r\n'; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + const request = { + protocolVersion: 1 as const, workspaceId: 'primary', path: 'notes.txt', + matching: 'tolerant' as const, + edits: [ + { oldText: '\n foo bar', newText: '\r\n\t\r\n\tbaz qux' }, + { oldText: 'end', newText: 'done' }, + ], + }; + const previewRequest = { ...request, operation: 'preview_edit' as const }; + const preview = await tools.execute(previewRequest); + if (preview.operation !== 'preview_edit') assert.fail('expected preview result'); + assert.equal(preview.content, 'header\r\n \r\n\tbaz qux\r\ndone\r\n'); + assert.equal(preview.hasUtf8Bom, true); + assert.equal(isWorkspaceToolResult(previewRequest, preview), true); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + const editRequest = { ...request, operation: 'edit_file' as const, + expectedBaseSha256: preview.baseSha256 }; + const edited = await tools.execute(editRequest); + if (edited.operation !== 'edit_file') assert.fail('expected edit result'); + assert.deepEqual(edited.matches, [ + { strategy: 'whitespace-normalized', occurrences: 1 }, + { strategy: 'exact', occurrences: 1 }, + ]); + assert.equal(isWorkspaceToolResult(editRequest, edited), true); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), `\uFEFF${preview.content}`); +}); + +test('tolerant previews and edits reject missing source boundary newlines without changing the file', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-missing-newline-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = 'foo bar'; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + for (const operation of ['preview_edit', 'edit_file'] as const) { + await assert.rejects(tools.execute({ + protocolVersion: 1, operation, workspaceId: 'primary', path: 'notes.txt', + matching: 'tolerant', oldText: 'foo bar\r\n', newText: 'baz qux\n', + }), (error: unknown) => error instanceof WorkspaceToolError && error.code === 'EDIT_CONFLICT'); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + } +}); + +test('tolerant previews and edits reject boundary removal outside the token match', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-remove-newline-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = 'foo bar\r\nnext'; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + for (const operation of ['preview_edit', 'edit_file'] as const) { + await assert.rejects(tools.execute({ + protocolVersion: 1, operation, workspaceId: 'primary', path: 'notes.txt', + matching: 'tolerant', oldText: 'foo bar\r\n', newText: 'baz qux', + }), (error: unknown) => error instanceof WorkspaceToolError && error.code === 'EDIT_CONFLICT'); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + } +}); + +test('tolerant previews and edits refuse partial tokens without changing the file', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-boundary-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = 'prereturn value;\n'; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + for (const operation of ['preview_edit', 'edit_file'] as const) { + await assert.rejects(tools.execute({ + protocolVersion: 1, operation, workspaceId: 'primary', path: 'notes.txt', + matching: 'tolerant', oldText: 'return value;', newText: 'return changed;', + }), (error: unknown) => error instanceof WorkspaceToolError && error.code === 'EDIT_CONFLICT'); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + } +}); + +test('large whitespace-only differences work for preview and edit_file', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-large-')); + t.after(() => rm(root, { recursive: true, force: true })); + const oldText = `head ${'part '.repeat(16_000)}tail`; + const original = `before ${oldText.replace(/ /g, '\t')} after`; + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + const base = { protocolVersion: 1 as const, workspaceId: 'primary', path: 'notes.txt', + matching: 'tolerant' as const, oldText, newText: 'result' }; + const preview = await tools.execute({ ...base, operation: 'preview_edit' }); + assert.equal(preview.operation === 'preview_edit' && preview.content, 'before result after'); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + const edited = await tools.execute({ ...base, operation: 'edit_file' }); + assert.equal(edited.operation === 'edit_file' && edited.matches?.[0]?.strategy, 'whitespace-normalized'); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), 'before result after'); +}); + +test('large replaceAll previews and edits return bounded content and accurate match counts', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-stream-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = 'x'.repeat(256 * 1024); + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + const request = { protocolVersion: 1 as const, workspaceId: 'primary', path: 'notes.txt', + edits: [{ oldText: 'x', newText: 'y', replaceAll: true }] }; + const previewRequest = { ...request, operation: 'preview_edit' as const }; + const preview = await tools.execute(previewRequest); + if (preview.operation !== 'preview_edit') assert.fail('expected preview result'); + assert.equal(preview.content, 'y'.repeat(original.length)); + assert.equal(preview.bytesWritten, original.length); + assert.deepEqual(preview.matches, [{ strategy: 'exact', occurrences: original.length }]); + assert.equal(isWorkspaceToolResult(previewRequest, preview), true); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + const editRequest = { ...request, operation: 'edit_file' as const, + expectedBaseSha256: preview.baseSha256 }; + const edited = await tools.execute(editRequest); + assert.equal(edited.operation === 'edit_file' && edited.bytesWritten, original.length); + assert.deepEqual(edited.operation === 'edit_file' && edited.matches, preview.matches); + assert.equal(isWorkspaceToolResult(editRequest, edited), true); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), preview.content); +}); + +test('oversized replaceAll previews and edits fail before writing the source file', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'librechat-code-edit-limit-')); + t.after(() => rm(root, { recursive: true, force: true })); + const original = 'x'.repeat(10_000); + await writeFile(join(root, 'notes.txt'), original); + const tools = await LocalWorkspaceTools.create({ + workspaces: [{ id: 'primary', root, writable: true }], + }); + for (const operation of ['preview_edit', 'edit_file'] as const) { + await assert.rejects(tools.execute({ + protocolVersion: 1, + operation, + workspaceId: 'primary', + path: 'notes.txt', + edits: [{ oldText: 'x', newText: 'y'.repeat(100_000), replaceAll: true }], + }), (error: unknown) => error instanceof WorkspaceToolError && + error.code === 'WRITE_LIMIT_EXCEEDED' && !error.mutationMayHaveCommitted); + assert.equal(await readFile(join(root, 'notes.txt'), 'utf8'), original); + } +}); + test('writes reject symlink targets and missing parent directories', async (t) => { const parent = await mkdtemp(join(tmpdir(), 'librechat-code-workspace-')); t.after(() => rm(parent, { recursive: true, force: true })); @@ -1930,6 +2227,8 @@ test('composes sandboxed commands without exposing them on unconfigured workspac assert.deepEqual(tools.capabilities.editFileModes, ['single', 'batch']); assert.deepEqual(tools.capabilities.editFileFeatures, [ 'expected_base_sha256', + 'tolerant_match', + 'replace_all', ]); assert.deepEqual(tools.capabilities.listFileFeatures, ['after_path']); assert.deepEqual( diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index 65757d39..0b9328d0 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -20,10 +20,19 @@ import { isValidBridgeWorkspaceToolCapabilities, isWorkspaceToolRequest, isWorkspaceToolResult, + WORKSPACE_EDIT_FILE_FEATURES, + workspaceEditRequestReportsMatches, } from './protocol.js'; +import { + applyTextEdits, + WorkspaceEditMatchError, + WorkspaceEditOutputLimitError, +} from './edits.js'; +import type { AppliedEdits } from './edits.js'; import type { BridgeWorkspaceDescriptor, + WorkspaceEditMatch, BridgeWorkspaceToolCapabilities, WorkspaceReadFileRequest, WorkspaceReadFileResult, @@ -722,7 +731,10 @@ async function editWorkspaceFile( 'EDIT_CONFLICT', ); } - const { updated, replacements } = applyWorkspaceEdits(original, request); + const { updated, replacements, matches } = applyWorkspaceEdits( + original, + request, + ); await atomicWriteConfinedFile( root, request.path, @@ -741,6 +753,7 @@ async function editWorkspaceFile( path: request.path, replacements, bytesWritten: updated.byteLength, + ...(matches ? { matches } : {}), }; } catch (error) { if (error instanceof WorkspaceToolError) throw error; @@ -753,7 +766,7 @@ async function editWorkspaceFile( function applyWorkspaceEdits( original: Buffer, request: WorkspaceEditFileRequest | WorkspacePreviewEditRequest, -): { updated: Buffer; replacements: number } { +): { updated: Buffer; replacements: number; matches?: WorkspaceEditMatch[] } { const hasBom = original[0] === 0xef && original[1] === 0xbb && original[2] === 0xbf; const body = hasBom ? original.subarray(3) : original; @@ -767,20 +780,19 @@ function applyWorkspaceEdits( const edits = request.edits ?? [ { oldText: request.oldText ?? '', newText: request.newText ?? '' }, ]; - let updatedText = text; - for (const edit of edits) { - const first = updatedText.indexOf(edit.oldText); - if (first < 0 || updatedText.indexOf(edit.oldText, first + 1) >= 0) { - throw new WorkspaceToolError( - 'Workspace edit must match exactly once', - 'EDIT_CONFLICT', - ); + let applied: AppliedEdits; + try { + applied = applyTextEdits(text, edits, request.matching ?? 'exact'); + } catch (error) { + if (error instanceof WorkspaceEditMatchError) { + throw new WorkspaceToolError(error.message, 'EDIT_CONFLICT'); + } + if (error instanceof WorkspaceEditOutputLimitError) { + throw new WorkspaceToolError(error.message, 'WRITE_LIMIT_EXCEEDED'); } - updatedText = - updatedText.slice(0, first) + - edit.newText + - updatedText.slice(first + edit.oldText.length); + throw error; } + const updatedText = applied.text; const updatedBody = Buffer.from(updatedText, 'utf8'); const updated = hasBom ? Buffer.concat([Buffer.from([0xef, 0xbb, 0xbf]), updatedBody]) @@ -791,7 +803,13 @@ function applyWorkspaceEdits( 'WRITE_LIMIT_EXCEEDED', ); } - return { updated, replacements: edits.length }; + return { + updated, + replacements: edits.length, + ...(workspaceEditRequestReportsMatches(request) + ? { matches: applied.matches } + : {}), + }; } async function previewWorkspaceEdit( @@ -806,7 +824,10 @@ async function previewWorkspaceEdit( 'EXECUTION_ABORTED', ); } - const { updated, replacements } = applyWorkspaceEdits(original, request); + const { updated, replacements, matches } = applyWorkspaceEdits( + original, + request, + ); if (signal?.aborted) { throw new WorkspaceToolError( 'Workspace tool execution aborted', @@ -825,6 +846,7 @@ async function previewWorkspaceEdit( baseSha256: createHash('sha256').update(original).digest('hex'), replacements, bytesWritten: updated.byteLength, + ...(matches ? { matches } : {}), }; } @@ -1497,7 +1519,7 @@ export class LocalWorkspaceTools implements WorkspaceToolExecutor { ...(anyWritable ? { writeFileModes: ['replace', 'create'] } : {}), ...(anyWritable ? { editFileModes: ['single', 'batch'] } : {}), ...(anyWritable - ? { editFileFeatures: ['expected_base_sha256'] } + ? { editFileFeatures: [...WORKSPACE_EDIT_FILE_FEATURES] } : {}), listFileFeatures: ['after_path'], }; diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 63640111..28035847 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -432,7 +432,11 @@ describe('paired bridge HTTP API', () => { ], supportedWorkspaceWriteFileModes: ['replace', 'create'], supportedWorkspaceEditFileModes: ['single', 'batch'], - supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], + supportedWorkspaceEditFileFeatures: [ + 'expected_base_sha256', + 'tolerant_match', + 'replace_all', + ], supportedWorkspaceListFileFeatures: ['after_path'], }); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index b3d549b1..e3946abc 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -12,6 +12,7 @@ import { isValidBridgeWorkerCapabilities, isValidBridgeWorkerId, isWorkspaceToolErrorCode, + WORKSPACE_EDIT_FILE_FEATURES, } from '../../../packages/code/src/protocol'; import { BridgePairingError, RedisBridgePairingStore } from './pairing'; import { BridgeStoreError, RedisBridgeStore } from './store'; @@ -535,7 +536,7 @@ router.post( ], supportedWorkspaceWriteFileModes: ['replace', 'create'], supportedWorkspaceEditFileModes: ['single', 'batch'], - supportedWorkspaceEditFileFeatures: ['expected_base_sha256'], + supportedWorkspaceEditFileFeatures: [...WORKSPACE_EDIT_FILE_FEATURES], supportedWorkspaceListFileFeatures: ['after_path'], supportedWorkspaceProgrammaticLanguages: ['bash'], supportedWorkspaceInstanceTypes: ['git_worktree'], diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 0ca75566..ab2616f8 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -168,12 +168,21 @@ function supportsWorkspaceTool( const mode = request.edits === undefined ? 'single' : 'batch'; const modes = capabilities?.editFileModes; const supportsMode = modes == null ? mode === 'single' : modes.includes(mode); - if (request.operation === 'preview_edit') return supportsMode; + const features = capabilities?.editFileFeatures ?? []; + const supportsMatching = + request.matching === undefined || features.includes('tolerant_match'); + const supportsReplaceAll = + request.edits?.some((edit) => edit.replaceAll !== undefined) !== true || + features.includes('replace_all'); + if (request.operation === 'preview_edit') { + return supportsMode && supportsMatching && supportsReplaceAll; + } return ( supportsMode && + supportsMatching && + supportsReplaceAll && (request.expectedBaseSha256 === undefined || - capabilities?.editFileFeatures?.includes('expected_base_sha256') === - true) + features.includes('expected_base_sha256')) ); } return true; diff --git a/service/src/bridge/workspace-store.test.ts b/service/src/bridge/workspace-store.test.ts index f248aa40..90f68da3 100644 --- a/service/src/bridge/workspace-store.test.ts +++ b/service/src/bridge/workspace-store.test.ts @@ -2,8 +2,9 @@ import { afterEach, expect, test } from 'bun:test'; import RedisMock from 'ioredis-mock'; import type Redis from 'ioredis'; +import type { WorkspaceToolRequest } from '../../../packages/code/src/protocol'; -import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { BRIDGE_PROTOCOL_VERSION, isValidBridgeWorkerCapabilities } from '../../../packages/code/src/protocol'; import { RedisBridgeStore } from './store'; const redis = new RedisMock() as unknown as Redis; @@ -643,6 +644,131 @@ test('rejects fenced edits from workers without the negotiated feature', async ( expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); }); +const featureGatedEdits: Array<[string, WorkspaceToolRequest]> = [ + [ + 'tolerant matching', + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + matching: 'tolerant', + edits: [{ oldText: 'before', newText: 'after' }], + }, + ], + [ + 'tolerant previews', + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'preview_edit', + workspaceId: 'primary', + path: 'notes.txt', + matching: 'tolerant', + edits: [{ oldText: 'before', newText: 'after' }], + }, + ], + [ + 'replaceAll edits', + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'edit_file', + workspaceId: 'primary', + path: 'notes.txt', + edits: [{ oldText: 'before', newText: 'after', replaceAll: true }], + }, + ], +]; + +test.each(featureGatedEdits)( + 'rejects %s from workers without the negotiated feature', + async (_label, request) => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['preview_edit', 'edit_file'], + editFileModes: ['single', 'batch'], + editFileFeatures: ['expected_base_sha256'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + + await expect( + store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await redis.keys('codeapi:bridge:v1:assignment:*')).toHaveLength(0); + }, +); + +test('preview-only workers can dispatch negotiated tolerant and replace-all previews', async () => { + expect(isValidBridgeWorkerCapabilities({ + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['preview_edit'], + editFileModes: ['single', 'batch'], + editFileFeatures: ['tolerant_match', 'replace_all'], + workspaces: [{ id: 'primary' }], + }, + })).toBe(true); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['preview_edit'], + editFileModes: ['single', 'batch'], + editFileFeatures: ['tolerant_match', 'replace_all'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const request = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'preview_edit' as const, + workspaceId: 'primary', + path: 'notes.txt', + matching: 'tolerant' as const, + edits: [{ oldText: 'before', newText: 'after', replaceAll: true }], + }; + const completion = store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request, + deadlineAtMs: Date.now() + 5_000, + signal: new AbortController().signal, + }); + const assignment = await store.lease('workspace-worker', incarnationId, 1_000); + expect(assignment).toMatchObject({ executionKind: 'workspace_tool', request }); + await store.settle('workspace-worker', assignment!.assignmentId, { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment!.generation, + leaseToken: assignment!.leaseToken, + incarnationId, + status: 'rejected', + error: 'preview tested', + }); + await expect(completion).resolves.toMatchObject({ status: 'rejected', error: 'preview tested' }); +}); + test('rejects batch previews from workers without the negotiated mode', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION,