From f7386db5530d118fc648da2f249d34455d84c2a7 Mon Sep 17 00:00:00 2001 From: li-xiu-qi Date: Fri, 30 Oct 2026 17:04:12 +0800 Subject: [PATCH] fix(coding-agent): route permission scans by tool origin --- docs/command-permissions.md | 19 + packages/coding-agent/docs/extensions.md | 2 + packages/coding-agent/docs/sdk.md | 2 +- .../coding-agent/src/core/agent-session.ts | 10 +- .../coding-agent/src/core/extensions/index.ts | 1 + .../src/core/extensions/loader.ts | 15 +- .../coding-agent/src/core/extensions/types.ts | 9 +- packages/coding-agent/src/core/index.ts | 1 + packages/coding-agent/src/core/tool-origin.ts | 15 + .../src/features/step-capabilities.ts | 28 +- packages/coding-agent/src/index.ts | 1 + packages/coding-agent/src/step/mcp.ts | 1 + packages/coding-agent/src/step/permissions.ts | 45 ++- .../coding-agent/src/step/tool-profile.ts | 5 +- packages/coding-agent/test/suite/harness.ts | 4 +- .../test/tool-origin-permissions.test.ts | 334 ++++++++++++++++++ 16 files changed, 469 insertions(+), 23 deletions(-) create mode 100644 packages/coding-agent/src/core/tool-origin.ts create mode 100644 packages/coding-agent/test/tool-origin-permissions.test.ts diff --git a/docs/command-permissions.md b/docs/command-permissions.md index 806229d39..d2b7d242b 100644 --- a/docs/command-permissions.md +++ b/docs/command-permissions.md @@ -53,6 +53,25 @@ command execution must supply the corresponding `shellContext` to the controller Manual `!` input and RPC `bash` use the separate `user_bash` event and are outside this agent-tool policy. +### Tool-origin routing + +Tool registrations carry optional `origin` provenance: `core`, `builtin`, +`extension`, or `mcp`. Core tools and built-in capabilities run command analysis +only for the dedicated `run_command`, `bash`, and `powershell` tools. Their other +structured inputs, including workflow source text, are not reinterpreted as shell +commands. MCP tools, extension tools, and calls without provenance retain the +conservative heuristic scan of command-bearing fields. + +The host assigns this provenance at registration. Disk-loaded extensions and +ordinary SDK `customTools` always use `extension`. Definitions emitted by +`createStepToolProfile` retain `core` through an object-identity marker when the +host injects that profile through `customTools`; cloning a definition loses this +exception. Inline extension factories are host-composed code and may identify +first-party capabilities as `builtin` and MCP adapters as `mcp`. Core provenance +is otherwise reserved for the base tool registry. An `origin` value is not an +authorization grant: explicit tool overrides still run first, and trusted shell +tools retain dangerous-command and incomplete-analysis checks. + ## Syntax, command semantics, and rules `shell-analysis.ts` uses the pinned `unbash` parser and explicitly traverses its diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 396947e00..e9ad6b5c3 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -2045,6 +2045,8 @@ pi.registerTool({ }); ``` +**Tool origin:** `origin?: "core" | "builtin" | "extension" | "mcp"` is optional host provenance used by command-permission routing and propagated to `tool_call` handlers. Ordinary extension tools should omit it. Disk-loaded extensions and ordinary SDK `customTools` are always registered as `extension`; the host-injected `createStepToolProfile` definitions retain `core` through an object-identity marker. Inline factories are host-composed code and may mark first-party capabilities as `builtin` and MCP adapters as `mcp`; `core` is otherwise reserved for the base tool registry. See [Command permissions](../../../docs/command-permissions.md#tool-origin-routing). + **Usage accounting:** If a tool makes nested LLM calls, return their combined `Usage` as `usage`. Step persists it on the tool result and includes it in footer, `/session`, and RPC session totals. `tool_result` handlers can inspect or replace this value. **Signaling errors:** To mark a tool execution as failed (sets `isError: true` on the result and reports it to the LLM), throw an error from `execute`. Returning a value never sets the error flag regardless of what properties you include in the return object. diff --git a/packages/coding-agent/docs/sdk.md b/packages/coding-agent/docs/sdk.md index 916c93f54..363ea211d 100644 --- a/packages/coding-agent/docs/sdk.md +++ b/packages/coding-agent/docs/sdk.md @@ -608,7 +608,7 @@ const { session } = await createAgentSession({ Use `defineTool()` for standalone definitions and arrays like `customTools: [myTool]`. Inline `pi.registerTool({ ... })` already infers parameter types correctly. -Custom tools passed via `customTools` are combined with extension-registered tools. Extensions loaded by the ResourceLoader can also register tools via `pi.registerTool()`. +Custom tools passed via `customTools` are combined with extension-registered tools. Extensions loaded by the ResourceLoader can also register tools via `pi.registerTool()`. Ordinary SDK custom tools always use the `extension` tool origin, so a self-declared `origin` cannot switch them to a trusted command-permission tier. Definitions emitted by `createStepToolProfile` retain their registered `core` provenance when the host injects that profile through `customTools`; cloning a definition removes this exception. If you pass `tools`, include each custom or extension tool name you want enabled, for example `tools: ["read", "bash", "my_tool"]`. diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 54acd866d..99a62a0b8 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -126,6 +126,7 @@ import type { SettingsManager } from "./settings-manager.ts"; import type { SlashCommandInfo } from "./slash-commands.ts"; import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.ts"; import { type BuildSystemPromptOptions, buildSystemPrompt, type SystemPromptProduct } from "./system-prompt.ts"; +import { isMarkedCoreToolDefinition } from "./tool-origin.ts"; import { type BashOperations, createLocalBashOperations } from "./tools/bash.ts"; import { createAllToolDefinitions } from "./tools/index.ts"; import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.ts"; @@ -576,6 +577,7 @@ export class AgentSession { type: "tool_call", toolName: toolCall.name, toolCallId: toolCall.id, + origin: this._toolDefinitions.get(toolCall.name)?.definition.origin, input: args as Record, }); } catch (err) { @@ -1520,6 +1522,7 @@ export class AgentSession { description: definition.description, parameters: definition.parameters, promptGuidelines: definition.promptGuidelines, + origin: definition.origin, sourceInfo, })); } @@ -3432,7 +3435,10 @@ export class AgentSession { const allCustomTools = [ ...registeredTools, ...this._customTools.map((definition) => ({ - definition, + definition: { + ...definition, + origin: isMarkedCoreToolDefinition(definition) ? ("core" as const) : ("extension" as const), + }, sourceInfo: createSyntheticSourceInfo(``, { source: "sdk" }), })), ].filter((tool) => isAllowedTool(tool.definition.name)); @@ -3442,7 +3448,7 @@ export class AgentSession { .map(([name, definition]) => [ name, { - definition, + definition: { ...definition, origin: definition.origin ?? "core" }, sourceInfo: createSyntheticSourceInfo(``, { source: "builtin" }), }, ]), diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index d472cbf01..d9ce1850d 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -161,6 +161,7 @@ export type { ToolExecutionStartEvent, ToolExecutionUpdateEvent, ToolInfo, + ToolOrigin, ToolRenderResultOptions, ToolResultEvent, ToolResultEventResult, diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index 1ce609c85..8233438f7 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -262,6 +262,17 @@ export function createExtensionRuntime(): ExtensionRuntime { return runtime; } +/** + * Disk-loaded extensions cannot self-assign a trusted tier. Inline factories + * come directly from the host and may tag first-party capabilities and MCP + * adapters; core provenance is reserved for the base tool registry. + */ +function withExtensionToolOrigin(tool: ToolDefinition, extension: Extension): ToolDefinition { + const isInline = extension.path.startsWith("<") && extension.path.endsWith(">"); + const origin = isInline && (tool.origin === "builtin" || tool.origin === "mcp") ? tool.origin : "extension"; + return { ...tool, origin }; +} + /** * Create the ExtensionAPI for an extension. * Registration methods write to the extension object. @@ -305,7 +316,7 @@ function createExtensionAPI( registerTool(tool: ToolDefinition): void { assertActive(); extension.tools.set(tool.name, { - definition: tool, + definition: withExtensionToolOrigin(tool, extension), sourceInfo: extension.sourceInfo, }); runtime.refreshTools(); @@ -319,7 +330,7 @@ function createExtensionAPI( } for (const tool of tools) { extension.tools.set(tool.name, { - definition: tool, + definition: withExtensionToolOrigin(tool, extension), sourceInfo: extension.sourceInfo, }); changed = true; diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 7fc046175..1e86ce04b 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -514,6 +514,9 @@ export interface ToolRenderContext { isError: boolean; } +/** Registration provenance used to select the appropriate permission policy. */ +export type ToolOrigin = "core" | "builtin" | "extension" | "mcp"; + /** * Tool definition for registerTool(). */ @@ -524,6 +527,8 @@ export interface ToolDefinition string | undefined; export type GetActiveToolsHandler = () => string[]; /** Tool info with name, description, parameter schema, prompt guidelines, and source metadata. */ -export type ToolInfo = Pick & { +export type ToolInfo = Pick & { sourceInfo: SourceInfo; }; diff --git a/packages/coding-agent/src/core/index.ts b/packages/coding-agent/src/core/index.ts index f2cd0a129..1c728c8d5 100644 --- a/packages/coding-agent/src/core/index.ts +++ b/packages/coding-agent/src/core/index.ts @@ -71,6 +71,7 @@ export { type ToolCallEvent, type ToolCallEventResult, type ToolDefinition, + type ToolOrigin, type ToolRenderResultOptions, type ToolResultEvent, type TurnEndEvent, diff --git a/packages/coding-agent/src/core/tool-origin.ts b/packages/coding-agent/src/core/tool-origin.ts new file mode 100644 index 000000000..e0606d532 --- /dev/null +++ b/packages/coding-agent/src/core/tool-origin.ts @@ -0,0 +1,15 @@ +import type { ToolDefinition } from "./extensions/types.ts"; + +const coreToolDefinitions = new WeakSet(); + +/** Mark a freshly created first-party profile definition without making provenance cloneable. */ +export function markCoreToolDefinition(tool: T): T { + tool.origin = "core"; + coreToolDefinitions.add(tool); + return tool; +} + +/** Object identity prevents arbitrary SDK definitions from opting into the core tier. */ +export function isMarkedCoreToolDefinition(tool: ToolDefinition): boolean { + return coreToolDefinitions.has(tool); +} diff --git a/packages/coding-agent/src/features/step-capabilities.ts b/packages/coding-agent/src/features/step-capabilities.ts index 563719291..4dcb9f9ee 100644 --- a/packages/coding-agent/src/features/step-capabilities.ts +++ b/packages/coding-agent/src/features/step-capabilities.ts @@ -10,6 +10,19 @@ import { registerWorkflowChildAcl } from "./workflow/acl-extension.ts"; import { createStepWorkflowExtension, type StepWorkflowExtensionOptions } from "./workflow/step-workflow.ts"; import { createUltraloopOptInExtension, type UltraloopTurnState } from "./workflow/ultraloop-opt-in.ts"; +/** Tag every tool registered by the first-party capability bundle without naming individual tools in policy code. */ +function withBuiltinToolOrigin(pi: ExtensionAPI): ExtensionAPI { + return { + ...pi, + registerTool: (tool) => pi.registerTool({ ...tool, origin: "builtin" }), + registerTools: (tools, options) => + pi.registerTools( + tools.map((tool) => ({ ...tool, origin: "builtin" })), + options, + ), + }; +} + export interface StepCapabilitiesExtensionOptions { subagent?: StepSubagentExtensionOptions; /** Optional process reporter; capability adapters never own delivery. */ @@ -38,13 +51,14 @@ export function createStepCapabilitiesExtension(options: StepCapabilitiesExtensi turnState: workflowTurnState, }); return (pi: ExtensionAPI): void => { - registerWorkflowChildAcl(pi, options.telemetry); - registerStepClarifyUserExtension(pi, options.telemetry); - plan(pi); - tasks(pi); - subagent(pi); - workflow(pi); - ultraloopOptIn(pi); + const builtinPi = withBuiltinToolOrigin(pi); + registerWorkflowChildAcl(builtinPi, options.telemetry); + registerStepClarifyUserExtension(builtinPi, options.telemetry); + plan(builtinPi); + tasks(builtinPi); + subagent(builtinPi); + workflow(builtinPi); + ultraloopOptIn(builtinPi); }; } diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index e78ad92cb..67344d037 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -178,6 +178,7 @@ export type { ToolExecutionStartEvent, ToolExecutionUpdateEvent, ToolInfo, + ToolOrigin, ToolRenderResultOptions, ToolResultEvent, TurnEndEvent, diff --git a/packages/coding-agent/src/step/mcp.ts b/packages/coding-agent/src/step/mcp.ts index e6680b2d6..690ebd5c1 100644 --- a/packages/coding-agent/src/step/mcp.ts +++ b/packages/coding-agent/src/step/mcp.ts @@ -775,6 +775,7 @@ function createRemoteTool(server: ConnectedServer, remote: McpTool) { name, label: remote.title?.trim() || remote.name, description: remote.description?.trim() || `MCP tool '${remote.name}' from server '${server.name}'.`, + origin: "mcp" as const, // The provider boundary accepts JSON Schema directly. Reconstructing it // as TypeBox types loses constraints, local references and combinators. parameters: remote.inputSchema, diff --git a/packages/coding-agent/src/step/permissions.ts b/packages/coding-agent/src/step/permissions.ts index edfaf6f31..cbfe72eca 100644 --- a/packages/coding-agent/src/step/permissions.ts +++ b/packages/coding-agent/src/step/permissions.ts @@ -13,6 +13,7 @@ import type { ExtensionUIContext, ToolCallEvent, ToolCallEventResult, + ToolOrigin, } from "../core/extensions/types.ts"; import { getShellConfig } from "../utils/shell.ts"; @@ -414,16 +415,24 @@ export function decideStepToolCall( state: StepPermissionState, overrides: Readonly> | undefined = state.toolOverrides, shellContext?: ShellExecutionContext, + origin?: ToolOrigin, ): StepToolDecision { const thenRun = getThenRunCommand(toolName, input); if (thenRun !== undefined) { const { then_run: _thenRun, ...mutationInput } = input; - const mutation = decideStepToolCall(toolName, mutationInput, state, overrides, shellContext); - const verification = decideStepToolCall("run_command", { command: thenRun }, state, overrides, shellContext); + const mutation = decideStepToolCall(toolName, mutationInput, state, overrides, shellContext, origin); + const verification = decideStepToolCall( + "run_command", + { command: thenRun }, + state, + overrides, + shellContext, + "core", + ); return combineThenRunDecisions(mutation, verification, thenRun); } const normalizedName = toolName.trim().toLowerCase(); - const command = extractCommand(input); + const command = extractCommand(normalizedName, input, origin); let analysis: CommandPolicyAnalysis | undefined; if (command !== undefined) { try { @@ -560,9 +569,16 @@ export class StepPermissionController { * differs without a terminal (see unattendedCause). A caller that omits it * asks for the interactive policy. */ - decide(toolName: string, input: Record, hasUI = true): StepToolDecision { + decide(toolName: string, input: Record, hasUI = true, origin?: ToolOrigin): StepToolDecision { const cause = this.unattendedCause(hasUI); - return decideStepToolCall(toolName, input, this.effectiveState(cause), this.toolOverrides, this.shellContext()); + return decideStepToolCall( + toolName, + input, + this.effectiveState(cause), + this.toolOverrides, + this.shellContext(), + origin, + ); } getOverrides(): Record { @@ -618,7 +634,14 @@ export class StepPermissionController { const input = event.input; const cause = this.unattendedCause(context.hasUI); const state = this.effectiveState(cause); - const decision = decideStepToolCall(event.toolName, input, state, this.toolOverrides, this.shellContext()); + const decision = decideStepToolCall( + event.toolName, + input, + state, + this.toolOverrides, + this.shellContext(), + event.origin, + ); if (decision.action === "allow") return undefined; if (decision.action === "deny") { @@ -761,7 +784,15 @@ function findToolOverride( return undefined; } -function extractCommand(input: Record): string | undefined { +function extractCommand( + toolName: string, + input: Record, + origin: ToolOrigin | undefined, +): string | undefined { + if (origin === "core" || origin === "builtin") { + if (toolName !== "run_command" && toolName !== "bash" && toolName !== "powershell") return undefined; + return typeof input.command === "string" ? input.command : undefined; + } for (const key of ["command", "cmd", "script"]) { const value = input[key]; if (typeof value === "string") return value; diff --git a/packages/coding-agent/src/step/tool-profile.ts b/packages/coding-agent/src/step/tool-profile.ts index 4ec53ecd7..51507e41b 100644 --- a/packages/coding-agent/src/step/tool-profile.ts +++ b/packages/coding-agent/src/step/tool-profile.ts @@ -23,6 +23,7 @@ import type { ToolRenderContext, ToolRenderResultOptions, } from "../core/extensions/types.ts"; +import { markCoreToolDefinition } from "../core/tool-origin.ts"; import { ShellToolError } from "../core/tools/bash.ts"; import type { EditToolOptions } from "../core/tools/edit.ts"; import { generateDiffString, generateUnifiedPatch, normalizeToLF } from "../core/tools/edit-diff.ts"; @@ -1771,8 +1772,8 @@ export function createStepToolProfile(cwd: string, options: StepToolProfileOptio editFile, runCommand, ]; - const findTools = createFindToolsDefinition(profileWithoutFindTools); - return [...profileWithoutFindTools, findTools]; + const findTools: AnyToolDefinition = createFindToolsDefinition(profileWithoutFindTools); + return [...profileWithoutFindTools, findTools].map(markCoreToolDefinition); } /** True for the Step model-facing names. */ diff --git a/packages/coding-agent/test/suite/harness.ts b/packages/coding-agent/test/suite/harness.ts index e33f95957..6c6784d27 100644 --- a/packages/coding-agent/test/suite/harness.ts +++ b/packages/coding-agent/test/suite/harness.ts @@ -22,7 +22,7 @@ import { convertToLlm } from "../../src/core/messages.ts"; import { SessionManager } from "../../src/core/session-manager.ts"; import type { Settings } from "../../src/core/settings-manager.ts"; import { SettingsManager } from "../../src/core/settings-manager.ts"; -import type { InlineExtension, ResourceLoader } from "../../src/index.ts"; +import type { InlineExtension, ResourceLoader, ToolDefinition } from "../../src/index.ts"; import { type CreateTestExtensionsResultInput, createTestExtensionsResult, @@ -65,6 +65,7 @@ export interface HarnessOptions { settings?: Partial; systemPrompt?: string; tools?: AgentTool[]; + customTools?: ToolDefinition[]; initialActiveToolNames?: string[]; allowedToolNames?: string[]; excludedToolNames?: string[]; @@ -187,6 +188,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise, + origin: ToolOrigin | undefined, + state: StepPermissionState = stepPermissionStateForPreset("approve-for-me"), +) { + return decideStepToolCall(toolName, input, state, undefined, undefined, origin); +} + +describe("tool-origin permission routing", () => { + it.each(["builtin", "core"] as const)("skips shell heuristics for non-shell %s tools", (origin) => { + for (const script of [hazardousWorkflowScript, ordinaryWorkflowScript]) { + expect(decideWithOrigin("workflow", { script }, origin)).toMatchObject({ + action: "allow", + hazardous: false, + }); + expect(decideWithOrigin("workflow", { script }, origin).analysisIncomplete).toBeUndefined(); + } + }); + + it.each(["builtin", "core"] as const)("keeps command analysis on dedicated shell tools from %s", (origin) => { + for (const toolName of ["run_command", "bash", "powershell"]) { + expect(decideWithOrigin(toolName, { command: "rm -rf ./tmp/cache" }, origin)).toMatchObject({ + action: "confirm", + hazardous: true, + }); + } + }); + + it.each(["mcp", "extension", undefined] as const)( + "retains conservative script analysis for origin %s, even when the tool is named workflow", + (origin) => { + expect(decideWithOrigin("workflow", { script: hazardousWorkflowScript }, origin)).toMatchObject({ + action: "confirm", + hazardous: true, + }); + }, + ); + + it("does not let trusted-origin routing bypass an explicit tool denial", () => { + expect( + decideStepToolCall( + "workflow", + { script: hazardousWorkflowScript }, + stepPermissionStateForPreset("approve-for-me"), + { workflow: "deny" }, + undefined, + "builtin", + ), + ).toMatchObject({ action: "deny" }); + }); +}); + +describe("tool-origin registration", () => { + const directories: string[] = []; + + afterEach(async () => { + vi.unstubAllEnvs(); + await Promise.all(directories.splice(0).map((directory) => rm(directory, { recursive: true, force: true }))); + }); + + it("marks every Step core-profile tool as core", () => { + const tools = createStepToolProfile(process.cwd()); + expect(tools.length).toBeGreaterThan(0); + for (const tool of tools) expect(tool.origin, tool.name).toBe("core"); + }); + + it("marks tools registered by Step capabilities as builtin", async () => { + vi.stubEnv("STEP_DISABLE_WORKFLOW", ""); + const extension = await loadExtensionFromFactory( + createStepCapabilitiesExtension({ + workflow: { vmExecutor: async () => ({ value: null, meta: {} }) }, + }), + process.cwd(), + createEventBus(), + createExtensionRuntime(), + ); + + expect(extension.tools.has("workflow")).toBe(true); + expect(extension.tools.size).toBeGreaterThan(1); + for (const [name, tool] of extension.tools) expect(tool.definition.origin, name).toBe("builtin"); + }); + + it("defaults unmarked inline tools to the conservative extension tier", async () => { + const extension = await loadExtensionFromFactory( + (pi) => { + pi.registerTool({ + name: "inline_custom", + label: "Inline custom", + description: "Custom inline tool", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }), + }); + }, + process.cwd(), + createEventBus(), + createExtensionRuntime(), + ); + + expect(extension.tools.get("inline_custom")?.definition.origin).toBe("extension"); + }); + + it("does not allow an inline factory to declare core provenance", async () => { + const extension = await loadExtensionFromFactory( + (pi) => { + pi.registerTool({ + name: "claimed_core", + label: "Claimed core", + description: "Attempts to self-assign core provenance", + origin: "core", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }), + }); + }, + process.cwd(), + createEventBus(), + createExtensionRuntime(), + ); + + expect(extension.tools.get("claimed_core")?.definition.origin).toBe("extension"); + }); + + it("does not allow a disk extension to declare a trusted origin for itself", async () => { + const directory = await mkdtemp(path.join(tmpdir(), "step-tool-origin-extension-")); + directories.push(directory); + const extensionPath = path.join(directory, "extension.ts"); + await writeFile( + extensionPath, + `export default function (pi) { + pi.registerTool({ + name: "claimed_builtin", + label: "Claimed builtin", + description: "Attempts to self-assign a trusted origin", + origin: "builtin", + parameters: {}, + execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }), + }); + }`, + ); + + const result = await loadExtensions([extensionPath], directory, createEventBus(), createExtensionRuntime()); + expect(result.errors).toEqual([]); + expect(result.extensions[0]?.tools.get("claimed_builtin")?.definition.origin).toBe("extension"); + }); +}); + +describe("tool-origin propagation through the agent loop", () => { + let harness: Harness | undefined; + const directories: string[] = []; + + afterEach(async () => { + await harness?.session.abort(); + await harness?.session.extensionRunner.emit({ type: "session_shutdown", reason: "quit" }); + harness?.cleanup(); + harness = undefined; + vi.unstubAllEnvs(); + await Promise.all(directories.splice(0).map((directory) => rm(directory, { recursive: true, force: true }))); + }); + + it("keeps SDK custom tools in the extension tier even when they claim a trusted origin", async () => { + harness = await createHarness({ + tools: [], + customTools: [ + { + name: "sdk_workflow", + label: "SDK workflow", + description: "SDK custom tool with a self-declared origin", + origin: "builtin", + parameters: Type.Object({ script: Type.String() }), + execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }), + }, + ], + }); + + const tool = harness.session.getAllTools().find((candidate) => candidate.name === "sdk_workflow"); + expect(tool?.origin).toBe("extension"); + expect(decideWithOrigin("sdk_workflow", { script: hazardousWorkflowScript }, tool?.origin)).toMatchObject({ + action: "confirm", + hazardous: true, + }); + }); + + it("preserves only the marked core profile through host customTools injection", async () => { + const coreTool = createStepToolProfile(process.cwd()).find((candidate) => candidate.name === "read_file")!; + harness = await createHarness({ + tools: [], + customTools: [coreTool, { ...coreTool, name: "cloned_core" }], + }); + + expect(harness.session.getAllTools().find((candidate) => candidate.name === "read_file")?.origin).toBe("core"); + const cloned = harness.session.getAllTools().find((candidate) => candidate.name === "cloned_core"); + expect(cloned?.origin).toBe("extension"); + expect(decideWithOrigin("cloned_core", { script: hazardousWorkflowScript }, cloned?.origin)).toMatchObject({ + action: "confirm", + hazardous: true, + }); + }); + + it("propagates the mcp tier through a real stdio catalog registration", async () => { + const directory = await mkdtemp(path.join(tmpdir(), "step-tool-origin-mcp-")); + directories.push(directory); + const serverPath = path.join(directory, "server.mjs"); + await writeFile( + serverPath, + `let buffer = ""; +process.stdin.setEncoding("utf8"); +process.stdin.on("data", (chunk) => { + buffer += chunk; + for (;;) { + const newline = buffer.indexOf("\\n"); + if (newline < 0) break; + const line = buffer.slice(0, newline).trim(); + buffer = buffer.slice(newline + 1); + if (!line) continue; + const message = JSON.parse(line); + if (message.id === undefined) continue; + let result; + if (message.method === "initialize") { + result = { + protocolVersion: message.params?.protocolVersion ?? "2025-03-26", + capabilities: { tools: {} }, + serverInfo: { name: "origin-fixture", version: "1.0.0" }, + }; + } else if (message.method === "tools/list") { + result = { tools: [{ + name: "workflow", + description: "MCP tool whose name collides with the built-in workflow tool", + inputSchema: { type: "object", properties: { script: { type: "string" } } }, + }] }; + } else if (message.method === "tools/call") { + result = { content: [{ type: "text", text: "mcp-ok" }] }; + } else if (message.method === "ping") { + result = {}; + } else { + process.stdout.write(JSON.stringify({ jsonrpc: "2.0", id: message.id, error: { code: -32601, message: "Unknown method" } }) + "\\n"); + continue; + } + process.stdout.write(JSON.stringify({ jsonrpc: "2.0", id: message.id, result }) + "\\n"); + } +});`, + ); + const configDir = path.join(directory, "config"); + await mkdir(configDir, { recursive: true }); + await writeFile( + path.join(configDir, "config.toml"), + [ + "[mcp_servers.steppage__steppage]", + "enabled = false", + "[mcp_servers.origin_fixture]", + `command = ${JSON.stringify(process.execPath)}`, + `args = [${JSON.stringify(serverPath)}]`, + "startup_timeout_sec = 5", + "tool_timeout_sec = 5", + "", + ].join("\n"), + ); + vi.stubEnv("STEP_CODING_AGENT_DIR", path.join(configDir, "agent")); + vi.stubEnv("STEPCODE_STORAGE_ROOT_DIR", path.join(directory, "storage")); + vi.stubEnv("HOME", directory); + vi.stubEnv("USERPROFILE", directory); + + harness = await createHarness({ + tools: [], + extensionFactories: [createStepExtension({ permission: { env: {}, initialPreset: "approve-for-me" } })], + }); + await harness.session.bindExtensions({ mode: "print" }); + + const tool = harness.session.getAllTools().find((candidate) => candidate.name === "origin_fixture__workflow"); + expect(tool?.origin).toBe("mcp"); + expect(decideWithOrigin(tool!.name, { script: hazardousWorkflowScript }, tool?.origin)).toMatchObject({ + action: "confirm", + hazardous: true, + }); + }); + + it("runs the built-in workflow sandbox without shell-analysis approval", async () => { + vi.stubEnv("STEP_DISABLE_WORKFLOW", ""); + const vmExecutor = vi.fn(async () => ({ value: "sandbox-ok", meta: {} })); + harness = await createHarness({ + tools: [], + initialActiveToolNames: ["workflow"], + extensionFactories: [ + createStepExtension({ permission: { env: {}, initialPreset: "approve-for-me" } }), + createStepCapabilitiesExtension({ workflow: { vmExecutor } }), + ], + }); + await harness.session.bindExtensions({ mode: "print" }); + harness.setResponses([ + fauxAssistantMessage( + fauxToolCall( + "workflow", + { script: hazardousWorkflowScript, args: { items: [] } }, + { id: "workflow-origin" }, + ), + { stopReason: "toolUse" }, + ), + fauxAssistantMessage("done"), + ]); + + await harness.session.prompt("Run the sandbox workflow"); + + expect(vmExecutor).toHaveBeenCalledTimes(1); + expect( + harness.session.messages.find( + (message) => message.role === "toolResult" && message.toolCallId === "workflow-origin", + ), + ).toMatchObject({ isError: false }); + }); +});