@tanstack/ai-code-mode/dist/esm/harness.js

Secure TypeScript Code Mode for TanStack AI agents to execute sandboxed tool orchestration programs.

import { createCodeMode } from "./create-code-mode.js";
import { toolDefinition } from "@tanstack/ai";
import { PermissionCallCapability, PermissionDecisionCapability, PermissionResources, PermissionRules, configOption, decidePermission, definePlugin, durableTool } from "@tanstack/ai-harness";
//#region src/harness.ts
/** The permission modes of `permissions()`, for the `mode` setting. */
var MODES = [
	"default",
	"plan",
	"acceptEdits",
	"bypass"
];
/**
* The default timeout of a program that can wait for the user to approve a
* call: 10 minutes.
*/
var APPROVAL_TIMEOUT_MS = 6e5;
/** The error of a call that a stop cut while it ran. */
var CUT_BY_STOP = "This call was cut by a stop and does not run again. The tool may or may not have run. Check before you retry.";
/**
* A name that works as a JavaScript identifier. Code mode turns each tool
* into an `external_<name>` function, and MCP tool names often have `-`.
*/
function identifierOf(name) {
	const safe = name.replace(/[^A-Za-z0-9_$]/g, "_");
	return /^[0-9]/.test(safe) ? `_${safe}` : safe;
}
/** A tool that runs on the server: not a client tool, and it has `execute`. */
function isServerTool(tool) {
	return ("__toolSide" in tool ? tool.__toolSide : "server") === "server" && typeof tool.execute === "function";
}
/**
* A short hash of the arguments of a call. A replay compares it, so a
* program that changed after a stop does not get the result of another
* call.
*/
function hashOf(args) {
	const json = JSON.stringify(args);
	let hash = 2166136261;
	for (const char of json ?? "") hash = Math.imul(hash ^ (char.codePointAt(0) ?? 0), 16777619);
	return (hash >>> 0).toString(16);
}
/**
* A result that the log can keep as a step value. JSON drops a function or
* a symbol, as the isolate does, and a `bigint` is not JSON.
*/
function isKeepable(value) {
	const type = typeof value;
	return type !== "function" && type !== "symbol" && type !== "bigint";
}
/** A `step` of a durable tool call. */
function isToolStep(value) {
	return typeof value === "object" && value !== null && "do" in value && typeof value.do === "function" && "ask" in value && typeof value.ask === "function";
}
/**
* The `step` of the call that runs the program. The harness gives it to
* the durable `execute_typescript`.
*/
function stepOf(context) {
	const step = context && "step" in context ? context.step : void 0;
	return isToolStep(step) ? step : void 0;
}
/**
* A tool-mode agent asks the user by stopping its turn. A call inside the
* isolate cannot stop that turn, so it does not run.
*/
var failClosedInAgent = (error) => {
	if (error instanceof Error && error.name === "MCPInputRequiredError") throw new Error("This call needs the approval of the user, and this agent cannot ask for it, so the call did not run.");
	throw error;
};
/**
* The hook of one program. Each call of a `guarded` tool gets the decision
* of a direct tool call: it runs, it throws, or it waits for the user. The
* log keeps its result in a step of the `execute_typescript` call, so a
* replay of the program after a stop does not run a finished call again.
*/
function guardCalls({ guarded, decide, context, signal }) {
	const step = stepOf(context);
	let count = 0;
	const hook = async (call, run) => {
		const tool = guarded.get(call.name);
		if (!tool) return run();
		if (!step) throw new Error(`Code mode cannot check this call of ${tool.name}, so it did not run.`);
		count += 1;
		const n = count;
		const { name } = tool;
		const hash = hashOf(call.args);
		const ran = (value) => ({
			ok: true,
			value: isKeepable(value) ? value : void 0,
			hash
		});
		const failed = (error) => ({
			ok: false,
			error,
			hash
		});
		const attempt = async () => {
			const checked = await decide?.({
				tool: name,
				args: call.args
			}) ?? { decision: "allow" };
			switch (checked.decision) {
				case "deny": return failed(checked.error);
				case "ask": {
					const reply = await step.ask({
						name: `approve:${n}:${hash}`,
						message: checked.message,
						schema: checked.schema,
						signal
					}).catch(failClosedInAgent);
					const refusal = await checked.settle(reply);
					if (refusal !== void 0) return failed(refusal);
					break;
				}
			}
			if (signal.aborted) throw new Error(`The program ended, so this call of ${name} did not run.`);
			if (tool.replay !== "safe") {
				let isFirstStart = false;
				await step.do(`start:${n}:${name}`, () => {
					isFirstStart = true;
					return true;
				});
				if (!isFirstStart) throw new Error(CUT_BY_STOP);
			}
			try {
				return ran(await run());
			} catch (error) {
				return failed(error instanceof Error ? error.message : String(error));
			}
		};
		const kept = await step.do(`call:${n}:${name}`, attempt);
		if (kept.hash !== hash) throw new Error(`The program changed after a stop: call ${n} of ${name} has other arguments now, so it did not run.`);
		if (!kept.ok) throw new Error(kept.error);
		return kept.value;
	};
	return hook;
}
/**
* `tool` (`execute_typescript`) as a durable tool. Its execute gets `step`,
* so the log keeps the checked calls of each program. After a stop, the
* next host runs the program again, and a finished call returns its kept
* result.
*/
function durable(tool) {
	const { execute } = tool;
	if (!execute) throw new Error("execute_typescript has no execute function.");
	return durableTool(toolDefinition({
		name: tool.name,
		description: tool.description,
		inputSchema: tool.inputSchema,
		outputSchema: tool.outputSchema
	}), (args, context) => execute(args, context));
}
/** The prompt part for the functions that check the permission rules. */
function checksPrompt(identifiers) {
	return `
### Permission checks

These functions check the permission rules before they run, as a direct tool call does: ${[...identifiers].map((identifier) => `\`external_${identifier}\``).join(", ")}. A call can wait for the user to approve it. When the user refuses the call or a rule denies it, the call throws an error that says why.
`;
}
/**
* A harness plugin that gives the model an `execute_typescript` tool. The
* model writes one TypeScript program that calls several tools, and the
* program runs in the isolate of `driver` (any `@tanstack/ai-isolate-*`
* driver). The tools it calls leave the model's tool list, including tools
* that plugins find at run time, such as MCP tools after `/connect`.
*
* By default, only tools that are safe to run without a question move into
* code mode. The other tools stay normal tool calls. A tool with
* `metadata.codeMode: true` moves even when `include` does not pick it, but
* only when it is safe in the same way.
*
* `include` can pick more tools, also tools with permission checks, such as
* the coding tools of `workspaceTools()`. Each call of such a tool in a
* program gets the same decision as a direct tool call from `permissions()`:
* it runs, it throws an error with the reason, or the program waits for the
* user to approve it (a question that `session.answer` answers). The log
* keeps the result of each such call, so after a stop the next host runs
* the program again without running a finished call twice. `include` never
* moves a tool with `needsApproval`, or a tool that the rules deny.
*
* With `lazy: true`, the model gets only the names of the moved tools and
* asks `discover_tools` for the signatures. With `directCalls: 'keep'`, the
* model can still call the moved tools directly.
*
* @example
* ```ts
* import { codeMode } from '@tanstack/ai-code-mode/harness'
* import { createQuickJSIsolateDriver } from '@tanstack/ai-isolate-quickjs'
*
* defineHarness({
*   name: 'acme/agent',
*   adapter,
*   plugins: () => [codeMode({ driver: createQuickJSIsolateDriver() })],
* })
* ```
*/
function codeMode(options) {
	const { include, lazy = false, directCalls = "remove", ...config } = options;
	/** The tools that `include` picked but cannot move, warned once each. */
	const warned = /* @__PURE__ */ new Set();
	const warnKept = (logger, name) => {
		if (warned.has(name)) return;
		warned.add(name);
		logger.warn(`[@tanstack/ai-code-mode] include picked the tool ${JSON.stringify(name)}, but it needs approval, or the rules deny it. Such a tool stays a normal tool call.`);
	};
	return definePlugin({
		name: "tanstack/code-mode",
		optionalRequires: [PermissionDecisionCapability, PermissionCallCapability],
		setup: (ctx) => {
			const rules = ctx.collect(PermissionRules);
			const resources = ctx.collect(PermissionResources);
			let prompt = "";
			const allowCodeMode = config.getSnippetBindings ? [] : ["execute_typescript", "discover_tools"].map((tool) => PermissionRules.item({
				tool,
				decision: "allow",
				kind: "read"
			}));
			return {
				config: { directCalls: configOption.select({
					options: ["remove", "keep"],
					default: directCalls,
					description: "remove: the tools of code mode leave the tool list. keep: the model can also call them directly."
				}) },
				prompts: [{
					id: "tanstack/code-mode",
					text: () => prompt
				}],
				contribute: allowCodeMode,
				prepareTools: ({ tools }) => {
					prompt = "";
					const decide = ctx.getOptional(PermissionDecisionCapability);
					const decisionOf = (tool, mode) => decide?.(tool.name, mode) ?? decidePermission(rules, tool.name, mode);
					const hasChecks = (tool) => resources.some((map) => Object.hasOwn(map, tool.name));
					const safe = (tool) => !tool.needsApproval && !hasChecks(tool) && decisionOf(tool, "plan") === "allow";
					const value = ctx.config.get("mode");
					const mode = MODES.find((each) => each === value) ?? "default";
					const canInclude = (tool) => !tool.needsApproval && decisionOf(tool, mode) !== "deny";
					const moves = (tool) => {
						if (tool.metadata?.codeMode === true && safe(tool)) return true;
						if (!include) return safe(tool);
						if (!include(tool)) return false;
						if (canInclude(tool)) return true;
						warnKept(ctx.logger, tool.name);
						return false;
					};
					const byIdentifier = /* @__PURE__ */ new Map();
					for (const tool of tools.filter(isServerTool)) {
						if (!moves(tool)) continue;
						const identifier = identifierOf(tool.name);
						if (!byIdentifier.has(identifier)) byIdentifier.set(identifier, tool);
					}
					if (byIdentifier.size === 0) return tools;
					const guarded = new Map([...byIdentifier].filter(([, tool]) => hasChecks(tool) || decisionOf(tool, mode) !== "allow"));
					const isGuarded = guarded.size > 0;
					const timeout = config.timeout ?? (isGuarded ? APPROVAL_TIMEOUT_MS : void 0);
					const created = createCodeMode({
						...config,
						timeout,
						tools: [...byIdentifier].map(([identifier, tool]) => identifier === tool.name && !lazy ? tool : {
							...tool,
							name: identifier,
							...lazy ? { lazy } : {}
						}),
						aroundCalls: isGuarded ? (context, signal) => guardCalls({
							guarded,
							decide: ctx.getOptional(PermissionCallCapability),
							context,
							signal
						}) : void 0
					});
					prompt = isGuarded ? `${created.systemPrompt}${checksPrompt(guarded.keys())}` : created.systemPrompt;
					const program = isGuarded ? durable(created.tool) : created.tool;
					const moved = new Set([...byIdentifier.values()].map((tool) => tool.name));
					return [...ctx.config.get("directCalls") === "keep" ? tools : tools.filter((tool) => !moved.has(tool.name)), ...created.tools.map((tool) => tool === created.tool ? program : tool)];
				}
			};
		}
	});
}
//#endregion
export { codeMode };

//# sourceMappingURL=harness.js.map