unsloth/studio/frontend/tests/code-tool-placement.test.ts
Michael Han 3cb8ecce78
Studio: add an edit_file tool so agents stop rewriting whole files (#8753)
* Studio: add an edit_file tool so agents stop rewriting whole files

The tool loop had no way to change a file. ALL_TOOLS was web_search, python,
terminal, render_html and search_knowledge_base, so the only way to edit
anything was a whole-file `cat > f <<'EOF'` through terminal or an
open(...).write(...) through python. Both re-send the entire file to change one
line, and both lose whatever the model failed to reproduce verbatim.

Measured on a 520-line source file: a one-line change costs 7,750 output tokens
to rewrite versus 45 to patch. Over five edits with re-reads that is 79,390
tokens against 8,676, which is why tasks that should fit in 64-94K die past
100K.

edit_file replaces an exact string. Not a unified diff: models corrupt @@ hunk
headers far more often than they mis-copy a literal snippet, and a bad hunk
header patches the wrong place instead of failing. A missing or non-unique
old_string is a hard error naming the match count and writes nothing, so the
retry is "add context" rather than "recover a mangled file".

- Preserves CRLF line endings, UTF-8 BOM and file mode. old_string is matched
  against normalized text, so a snippet with plain newlines still matches a
  Windows-authored file instead of failing invisibly.
- Atomic write via temp file and rename, so an interrupted write cannot leave a
  source file half-replaced.
- Contained to the session workdir, checked on the realpath so a planted
  symlink cannot reach out. /mnt/data-style habit paths remap exactly as the
  python shim does.
- Under Full access absolute paths resolve, and the schema says so. Otherwise
  the model assumes it cannot reach a real checkout and falls back to the
  rewrite precisely where files are largest.
- Still prompts in auto mode: python's open(..., "w") already does, so the
  cheaper tool must not become the quiet way around that.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: harden edit_file against the cases raised in review

Six real defects, each reproduced before the fix and covered by a test.

- Receipt was bounded by diff LINES, which bounds nothing when one line is the
  whole file. A 200KB minified source returned a 400KB receipt, twice what the
  tool exists to avoid. Characters are now capped per line and over the receipt:
  the same edit returns 481 chars.
- bool("false") is True, and models emit the JSON string, so replace_all as a
  string turned the multi-match guard off and rewrote every occurrence. The two
  spellings models actually produce are mapped, anything else is refused.
- A FIFO or character device reported st_size 0 and then read forever. This path
  carries no timeout or cancel event, so the turn could not be recovered. Only
  regular files are accepted now.
- An absolute path inside a workdir that itself sits under a habit prefix
  (/workspace/repo) had its own prefix stripped and rejoined onto itself,
  resolving to /workspace/repo/repo/a.py. Paths already inside the workdir skip
  the remap; habit paths outside it still remap as before.
- Two chats sharing a project workspace could both read, both write, and the
  later rename discarded the earlier edit silently. The bytes the edit was
  computed from are compared again before the rename.
- Containment was checked once at resolve time, leaving the whole read and diff
  as a window in which a parent could be swapped for a symlink. It is rechecked
  immediately before the rename.

Left as is: an empty old_string still writes a zero-byte file. Refusing every
existing target would strand the model, since no other old_string can match an
empty file, so nothing could ever write to it. Nothing is lost with no contents,
and the mode is carried over by the write.

* Studio: give edit_file the terminal glyph, not the globe

status_for_tool reports "Editing: name" for edit_file, and toolStatusKind only
treats a "Running" prefix as local, so a file edit on this machine showed the
globe, the same badge a web search gets. It is as local as python and terminal,
so it takes the same glyph.

* Studio: bound the edit_file receipt and make creation atomic

Second review pass. Four findings, each reproduced before the fix.

- The receipt was capped on output but not on what produced it: difflib was fed
  the whole file and its generator drained into a list. replace_all on a file at
  the 16MB cap allocated ~500MB and took 1.3s to return 200 characters. difflib
  now sees only a window around the first change and the generator is consumed
  lazily. Measured on the same 16MB file: 501MB -> 48MB, 1.3s -> 0.05s; a 600KB
  file goes 48MB -> 1MB. Hunk headers are shifted back to real file lines, since
  a receipt pointing at line 3 of a 9000-line file is worse than none.
- Creation checked lexists() and then wrote, so two chats sharing a project
  workspace could both pass the check and the later write drop the earlier file.
  The absent case is now created with O_EXCL, and filling a zero-byte file goes
  through the guarded write rather than clobbering blindly.
- New files came out 0600: mkstemp makes the temp file private and copymode had
  no source to copy from. O_EXCL creation takes the usual umask-derived mode, so
  a group that reads generated files still can.
- enabled_tools in the public request schema still listed only web_search,
  python, terminal and render_html, leaving the new built-in undiscoverable to
  clients reading the OpenAPI schema, and bypass_permissions described only the
  python/terminal sandbox. Both now describe edit_file, including that Full
  access lifts its containment.

Eight new tests, 44 in the file.

* Stop the edit receipt inventing deletions, refuse non-regular targets and unpaired surrogates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove the file a failed create just published

O_EXCL publishes the name before the first byte and the payload goes out
a buffer at a time, so ENOSPC or a quota partway through leaves the
bytes that fit. Reproduced with a real kernel write failure: a 117780
byte create left 4096 bytes cut mid-token, and the retry the error
message asks for is refused for ever, because an empty old_string
refuses a non-empty target and no other old_string exists for a file the
model never saw. close() can report a failure for data written earlier,
so the error can arrive after most of the file is on disk. Unlinking the
inode this call created puts the retry back on the create path, and
keeps O_EXCL rather than mkstemp, whose 0600 would ignore the umask.

* Let edit_file create an empty file

Both strings empty is the documented creation call for __init__.py, py.typed
and .gitkeep, but the identical-strings no-op check ran first and refused it,
so there was no way to write a zero-byte file. Decide creation before that
check, and stop reporting one line for a file with none.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim the review commentary

---------

Co-authored-by: Unsloth <michaelhan@Michaels-MacBook-Pro.local>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: danielhanchen <danielhanchen@users.noreply.github.com>
2026-08-19 07:35:11 -07:00

221 lines
8 KiB
TypeScript

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Where the Code pill runs code.
//
// Before Studio's tool loop reached the general external providers, only
// openai_codex carried studio_tools, so `codeToolsEnabled` on an OpenAI,
// Anthropic or Gemini connection fell through to the hosted branch and sent
// `code_execution` -- the model's code ran in the PROVIDER's sandbox. Now that
// those providers take the Studio branch, the same stored pill would send
// ["python", "terminal"] and run the model's code on the USER's machine. The
// toggle is persisted (unsloth_chat_code_tools_enabled), so nobody re-consents:
// the trust boundary moves during an update, with nothing in the composer or
// the stream saying so.
//
// The rule this file pins: a connection that has its own sandbox keeps it.
// Studio's local python/terminal are for connections that have none.
import assert from "node:assert/strict";
import { readFileSync } from "node:fs";
import { fileURLToPath } from "node:url";
import test from "node:test";
import {
codeToolCanRun,
selectCodeToolNames,
} from "../src/features/chat/api/code-tool-placement.ts";
const SOURCE = readFileSync(
fileURLToPath(new URL("../src/features/chat/api/chat-adapter.ts", import.meta.url)),
"utf8",
);
// ── the rule itself ────────────────────────────────────────────────
test("a provider with its own sandbox keeps running the code there", () => {
assert.deepEqual(
selectCodeToolNames({
codeToolsEnabled: true,
hostedCodeExecutionForThisTurn: true,
providerHostsCodeExecution: true,
}),
{ local: [], hosted: ["code_execution"] },
);
});
test("a provider with a sandbox its MODEL cannot use runs nothing, not local code", () => {
// e.g. an OpenAI connection on a model outside the code-execution family.
// Pre-loop this sent no code tool at all; falling back to python/terminal
// would relocate execution rather than preserve it.
assert.deepEqual(
selectCodeToolNames({
codeToolsEnabled: true,
hostedCodeExecutionForThisTurn: false,
providerHostsCodeExecution: true,
}),
{ local: [], hosted: [] },
);
});
test("a provider with no sandbox uses Studio's own tools", () => {
// llama.cpp / vLLM / Ollama / custom, and the cloud providers that ship no
// code sandbox. Local execution is the only meaning the pill can have there,
// it is what openai_codex has always done, and it paints tool cards under the
// permission gate rather than happening invisibly.
assert.deepEqual(
selectCodeToolNames({
codeToolsEnabled: true,
hostedCodeExecutionForThisTurn: false,
providerHostsCodeExecution: false,
}),
{ local: ["python", "terminal", "edit_file"], hosted: [] },
);
});
test("edit_file is local-only, never a stand-in for a hosted sandbox", () => {
// It must not creep into the hosted branch just because the Code pill is
// what turns it on.
for (const hosted of [true, false]) {
const names = selectCodeToolNames({
codeToolsEnabled: true,
hostedCodeExecutionForThisTurn: hosted,
providerHostsCodeExecution: true,
});
assert.ok(!names.local.includes("edit_file"));
assert.ok(!names.hosted.includes("edit_file"));
}
});
test("the pill being off asks for nothing on either side", () => {
for (const providerHostsCodeExecution of [true, false]) {
assert.deepEqual(
selectCodeToolNames({
codeToolsEnabled: false,
hostedCodeExecutionForThisTurn: providerHostsCodeExecution,
providerHostsCodeExecution,
}),
{ local: [], hosted: [] },
);
}
});
// ── the adapter has to actually use it ─────────────────────────────
// Same technique as hosted-image-tool-with-studio-tools.test.ts: the body is
// built inside a run closure that needs a live runtime, provider store and
// encryption key, so the structural property is read out of the source.
function studioToolsBranch(): string {
const start = SOURCE.indexOf("...(ragEnabled || projectRagEnabled\n");
assert.ok(start > 0, "the Studio-tools enabled_tools list moved");
const end = SOURCE.indexOf("mcp_enabled:", start);
assert.ok(end > start, "the Studio-tools branch moved");
return SOURCE.slice(start, end);
}
test("the Studio branch never hardcodes local code tools", () => {
const branch = studioToolsBranch();
assert.doesNotMatch(
branch,
/codeToolsEnabled \? \["python", "terminal"\]/,
"the Code pill must not send local execution regardless of provider",
);
// Both sides come from the one helper above, so the local and hosted names
// cannot drift apart or both be sent for a single pill.
assert.match(branch, /\.\.\.studioLocalCodeTools/);
assert.match(branch, /\.\.\.hostedCodeToolsForThisTurn/);
});
test("the branch is only taken when a tool Studio itself can run is on", () => {
// Code alone on a hosted-sandbox provider is a hosted request: it must reach
// the hosted branch, which sends no permission_mode. Sending the Studio body
// for it would ask the backend to confirm tool calls on a passthrough request,
// which routes/inference.py answers with a 400.
const gate = SOURCE.slice(
SOURCE.indexOf("...(supportsStudioToolsForThisTurn &&"),
SOURCE.indexOf("enable_tools: true", SOURCE.indexOf("...(supportsStudioToolsForThisTurn &&")),
);
assert.ok(gate.length > 0, "the Studio-tools gate moved");
assert.doesNotMatch(
gate,
/^\s*codeToolsEnabled \|\|$/m,
"a bare codeToolsEnabled sends the Studio body for a hosted-only turn",
);
assert.match(gate, /studioLocalCodeTools\.length > 0/);
});
// ── Whether the pill is offered at all ─────────────────────────────
// Until Studio's loop reached the general external providers, the composer
// keyed the Code pill on the hosted flag alone, so a model without the hosted
// sandbox simply did not offer it. Keying it on the Studio-tools flag instead
// offered it everywhere, including where the rule above deliberately runs
// nothing, and the user got a lit toggle that sent enable_tools: false.
test("a model with its provider's sandbox can run code", () => {
assert.equal(
codeToolCanRun({
hostedCodeExecutionForThisTurn: true,
providerHostsCodeExecution: true,
supportsStudioTools: true,
}),
true,
);
});
test("a model that cannot use its provider's sandbox offers nothing", () => {
assert.equal(
codeToolCanRun({
hostedCodeExecutionForThisTurn: false,
providerHostsCodeExecution: true,
supportsStudioTools: true,
}),
false,
);
});
test("a connection with no sandbox of its own runs Studio's tools", () => {
assert.equal(
codeToolCanRun({
hostedCodeExecutionForThisTurn: false,
providerHostsCodeExecution: false,
supportsStudioTools: true,
}),
true,
);
});
test("and not when the loop cannot run them either", () => {
assert.equal(
codeToolCanRun({
hostedCodeExecutionForThisTurn: false,
providerHostsCodeExecution: false,
supportsStudioTools: false,
}),
false,
);
});
test("the pill is offered exactly when the placement sends something", () => {
for (const hostedCodeExecutionForThisTurn of [true, false]) {
for (const providerHostsCodeExecution of [true, false]) {
const names = selectCodeToolNames({
codeToolsEnabled: true,
hostedCodeExecutionForThisTurn,
providerHostsCodeExecution,
});
const sendsSomething = names.hosted.length > 0 || names.local.length > 0;
assert.equal(
codeToolCanRun({
hostedCodeExecutionForThisTurn,
providerHostsCodeExecution,
supportsStudioTools: true,
}),
sendsSomething,
`hosted=${hostedCodeExecutionForThisTurn} sandbox=${providerHostsCodeExecution}`,
);
}
}
});