build-tools.ts

import path from "path"

import type OpenAI from "openai"

import type { ProviderSettings, ModeConfig, ModelInfo, ToolGroup } from "@shofer/types"
import { toolGroupNameSchema, getHost } from "@shofer/types"
import { customToolRegistry } from "../custom-tools/custom-tool-registry.js"
import { toolGroupRegistry } from "../tool-groups/category-registry.js"
import { toolsLog } from "../logging/subsystems.js"
import { formatNative } from "../custom-tools/format-native.js"
import { pluginRegistry } from "../plugins/plugin-registry.js"
import { setPrivateToolInvokeMap } from "../tools/private-tool-registry.js"

import { type TaskProviderLike } from "../task-provider/index.js"
import { getRooDirectoriesForCwd } from "../services/shofer-config/index.js"

import { getNativeTools } from "../prompts/tools/native-tools/index.js"
import { getMcpServerTools } from "../prompts/tools/native-tools/mcp_server.js"
import { filterNativeToolsForMode, filterMcpToolsForMode } from "../prompts/tools/filter-tools-for-mode.js"
import { applyToolSchemaTiers } from "../prompts/tools/tool-stubs.js"
import {
	DESCRIBE_TOOLS_TOOL_NAME,
	defaultModeSlug,
	getGroupName,
	getModeBySlug,
	getToolsForMode,
	modeStubsToolSchemas,
	resolveModeConfig,
} from "@shofer/types"
import { resolveToolAlias } from "../tools/tool-aliases.js"
import { recordToolSchemas } from "../tools/tool-schema-registry.js"

interface BuildToolsOptions {
	provider: TaskProviderLike
	cwd: string
	mode: string | undefined
	customModes: ModeConfig[] | undefined
	experiments: Record<string, boolean> | undefined
	apiConfiguration: ProviderSettings | undefined
	disabledTools?: string[]
	modelInfo?: ModelInfo
	/**
	 * If true, returns all tools without mode filtering, but also includes
	 * the list of allowed tool names for use with allowedFunctionNames.
	 * This enables providers that support function call restrictions (e.g., Gemini)
	 * to pass all tool definitions while restricting callable tools.
	 */
	includeAllToolsWithRestrictions?: boolean
	/**
	 * Per-task JSON Schema override for the `attempt_completion` tool's
	 * `result` parameter. Threaded from {@link Task.completionSchema}.
	 * When set, the generic `result: string` is replaced with a structured
	 * object schema.
	 */
	completionSchema?: Record<string, unknown>
	/**
	 * When true, this task's title was locked by its spawning parent (via
	 * `new_task`'s `title`), so the `set_task_title` tool is omitted from the
	 * tool list entirely — the agent can't even attempt a rename. Threaded from
	 * {@link Task.nameLocked}.
	 */
	titleLocked?: boolean
	/**
	 * Per-task tool-group allow-list.
	 * Threaded from {@link Task.agentToolGroups}. When set, the final tool list
	 * is intersected with these groups — a tool survives only if its group is
	 * declared (native tools resolved via `TOOL_GROUPS`, MCP tools via the
	 * `mcp` group, private/custom tools via their own group), EXCEPT
	 * `ALWAYS_AVAILABLE_TOOLS`, which are always retained so the agent can still
	 * complete/coordinate. This is a restriction layered on top of mode
	 * filtering — it can only remove tools, never add them. Unknown group names
	 * are ignored.
	 */
	agentToolGroups?: string[]
}

interface BuildToolsResult {
	/**
	 * The tools to pass to the model.
	 * If includeAllToolsWithRestrictions is true, this includes ALL tools.
	 * Otherwise, it includes only mode-filtered tools.
	 */
	tools: OpenAI.Chat.ChatCompletionTool[]
	/**
	 * The names of tools that are allowed to be called based on mode restrictions.
	 * Only populated when includeAllToolsWithRestrictions is true.
	 * Use this with allowedFunctionNames in providers that support it.
	 */
	allowedFunctionNames?: string[]
}

/**
 * Extracts the function name from a tool definition.
 */
function getToolName(tool: OpenAI.Chat.ChatCompletionTool): string {
	return (tool as OpenAI.Chat.ChatCompletionFunctionTool).function.name
}

// ──────────────────────────────────────────────
// Private Tool Provider system
// ──────────────────────────────────────────────
//
// Extensions register tools via the `shofer.privateToolProviders` VS Code
// configuration, not via `vscode.lm.tools` (which is Copilot's interface).
// Each provider exposes two commands:
//
//   <getDefinitionsCommand>  → returns Array<{name, description, inputSchema, group?}>
//   <invokeToolCommand>      → takes (name, input), returns {content, is_error?}
//
// See docs/tool-registration-interface.md for the full contract.

/** Configuration shape for a single private tool provider. */
interface PrivateToolProviderConfig {
	/** VS Code command ID that returns all tool definitions. */
	getDefinitionsCommand: string
	/** VS Code command ID that invokes a tool by name. */
	invokeToolCommand: string
}

/** A tool definition returned by a provider's getDefinitions command. */
interface PrivateToolDef {
	name: string
	description: string
	inputSchema: object
	/** Optional tool group override. Falls back to provider config. */
	group?: string
}

/** Metadata for a single tool discovered from a private provider. */
interface PrivateToolMeta {
	tool: OpenAI.Chat.ChatCompletionFunctionTool
	group: ToolGroup
	/** The VS Code command to invoke this tool at execution time. */
	invokeCommand: string
}

/**
 * Read all registered private tool providers from config and discover
 * their tools. Returns a combined list with group assignments and
 * invocation commands.
 *
 * Config key: `shofer.privateToolProviders`
 */
async function getPrivateLmToolMeta(): Promise<PrivateToolMeta[]> {
	const providers = getHost().config.get<Record<string, PrivateToolProviderConfig>>(
		"shofer",
		"privateToolProviders",
		{},
	)

	const allMeta: PrivateToolMeta[] = []

	for (const [providerId, providerCfg] of Object.entries(providers)) {
		if (!providerCfg?.getDefinitionsCommand || !providerCfg?.invokeToolCommand) continue

		try {
			const definitions = await getHost().workspace.executeCommand<PrivateToolDef[] | undefined>(
				providerCfg.getDefinitionsCommand,
			)

			if (!definitions || !Array.isArray(definitions)) continue

			for (const def of definitions) {
				const group = resolvePrivateToolGroup(providerId, def)
				allMeta.push({
					tool: {
						type: "function" as const,
						function: {
							name: def.name,
							description: def.description || def.name,
							parameters: (def.inputSchema || {
								type: "object",
								properties: {},
							}) as OpenAI.FunctionParameters,
						},
					},
					group,
					invokeCommand: providerCfg.invokeToolCommand,
				})
			}
		} catch {
			// Provider extension not installed or not activated — skip.
		}
	}

	return allMeta
}

/**
 * Whether `name` is usable as a category name.
 *
 * Builtins need no separate arm: every builtin name is itself a valid slug, so
 * "builtin or slug" collapses to "slug" (`toolGroupNameSchema`'s contract).
 */
function isValidToolGroupName(name: string): boolean {
	return toolGroupNameSchema.safeParse(name).success
}

/**
 * Accept a private tool's DECLARED group: register the category (minting it when
 * the name is new) and record the tool → group mapping so the approval path
 * resolves the same group `filterPrivateToolsForMode` uses for visibility.
 *
 * Without the mapping the two paths disagree — visibility reads the declaration
 * while approval infers from the name prefix — and a `salesforce` tool is shown
 * as salesforce yet gated as `uncategorized`: its toggle on, and the tool still
 * asking.
 */
function acceptPrivateToolGroup(toolName: string, group: string): string {
	toolGroupRegistry.registerToolMapping(toolName, group)
	return group
}

/**
 * Resolve the ToolGroup for a private tool:
 *  1. If the tool definition has an explicit `group`, validate and use it.
 *  2. Fall back to the provider's `shofer.<providerId>.toolGroups` config.
 *  3. Default to "uncategorized".
 *
 * A name that is a valid slug is accepted whether or not anything has used it
 * before — that is how a provider mints a category. Only a malformed name falls
 * through to `uncategorized`, which is the bucket for tools that declared
 * nothing usable.
 */
function resolvePrivateToolGroup(providerId: string, def: PrivateToolDef): ToolGroup {
	// 1. Explicit group in the definition
	if (def.group && isValidToolGroupName(def.group)) {
		return acceptPrivateToolGroup(def.name, def.group)
	}

	// 2. Provider-level config
	try {
		const toolGroups = getHost().config.get<Record<string, string> | undefined>(
			`shofer.${providerId}`,
			"toolGroups",
			undefined,
		)
		if (toolGroups && typeof toolGroups[def.name] === "string") {
			const declared = toolGroups[def.name]!
			if (isValidToolGroupName(declared)) {
				return acceptPrivateToolGroup(def.name, declared)
			}
		}
	} catch {
		// Config read failed.
	}

	return "uncategorized"
}

/**
 * Filter private tools by mode, using each tool's assigned group.
 *
 * @param privateMeta - Private tool metadata with group assignments
 * @param mode - Current mode slug
 * @param customModes - Custom mode configurations
 * @returns Filtered private tool definitions
 */
function filterPrivateToolsForMode(
	privateMeta: PrivateToolMeta[],
	mode: string | undefined,
	customModes: ModeConfig[] | undefined,
): OpenAI.Chat.ChatCompletionFunctionTool[] {
	const modeSlug = mode ?? defaultModeSlug
	const modeConfig = getModeBySlug(modeSlug, customModes)

	if (!modeConfig) {
		return privateMeta.map((m) => m.tool)
	}

	const allowedGroups = new Set<string>((modeConfig.tools ?? []).map((g) => getGroupName(g)))

	if (allowedGroups.size === 0) {
		return privateMeta.map((m) => m.tool)
	}

	return privateMeta.filter((meta) => allowedGroups.has(meta.group)).map((meta) => meta.tool)
}

/**
 * Builds the complete tools array for native protocol requests.
 * Combines native tools and MCP tools, filtered by mode restrictions.
 *
 * @param options - Configuration options for building the tools
 * @returns Array of filtered native and MCP tools
 */
export async function buildNativeToolsArray(options: BuildToolsOptions): Promise<OpenAI.Chat.ChatCompletionTool[]> {
	const result = await buildNativeToolsArrayWithRestrictions(options)
	return result.tools
}

/** The already-mode-filtered tool categories, plus the private-tool metadata
 *  needed to resolve each private tool's group. */
export interface ToolCategories {
	native: OpenAI.Chat.ChatCompletionTool[]
	mcp: OpenAI.Chat.ChatCompletionTool[]
	custom: OpenAI.Chat.ChatCompletionFunctionTool[]
	private: OpenAI.Chat.ChatCompletionFunctionTool[]
	privateMeta: PrivateToolMeta[]
}

/**
 * Apply a task's declared tool-group restriction to the
 * already-mode-filtered tool set. Pure intersection — only removes tools,
 * never adds them.
 *
 * Semantics:
 * - `agentToolGroups === undefined` → no restriction (returns categories as-is).
 * - declared (incl. `[]`) → keep a tool only if its group is declared. Native
 *   tools are matched via `getToolsForMode` (which also re-adds
 *   `ALWAYS_AVAILABLE_TOOLS`, so `attempt_completion` etc. always survive — a
 *   restricted agent can still complete stakes). MCP tools belong to the `mcp`
 *   group; native custom tools to `write`; private tools carry their own group.
 * - Malformed group names are dropped (fail-closed): `tools: [Bad_Name]`
 *   restricts to always-available only. A valid SLUG is kept even when nothing
 *   has registered it yet — a task restricted to a category whose server has not
 *   connected is legitimate; its tools arrive when it does.
 */
export function restrictToolsToDeclaredGroups(
	agentToolGroups: string[] | undefined,
	categories: ToolCategories,
): Omit<ToolCategories, "privateMeta"> {
	const { native, mcp, custom, private: priv, privateMeta } = categories
	if (agentToolGroups === undefined) {
		return { native, mcp, custom, private: priv }
	}
	const declaredGroups = agentToolGroups.filter((g): g is ToolGroup => isValidToolGroupName(g))
	const declaredSet = new Set<ToolGroup>(declaredGroups)
	const allowedNativeNames = new Set(getToolsForMode(declaredGroups))
	const privateGroupByName = new Map(privateMeta.map((m) => [getToolName(m.tool), m.group]))
	return {
		native: native.filter((t) => allowedNativeNames.has(resolveToolAlias(getToolName(t)))),
		mcp: declaredSet.has("mcp") ? mcp : [],
		custom: declaredSet.has("write") ? custom : [],
		private: priv.filter((t) => {
			const g = privateGroupByName.get(getToolName(t))
			return g !== undefined && declaredSet.has(g)
		}),
	}
}

/**
 * Builds the complete tools array for native protocol requests with optional mode restrictions.
 */
export async function buildNativeToolsArrayWithRestrictions(options: BuildToolsOptions): Promise<BuildToolsResult> {
	const {
		provider,
		cwd,
		mode,
		customModes,
		experiments,
		apiConfiguration,
		disabledTools,
		modelInfo,
		includeAllToolsWithRestrictions,
	} = options

	// A conversational turn (`toolCallingEnabled === false`) has no tool plane at
	// all. Gating here rather than at each metadata literal means every consumer
	// — the main request, condensation, context-window recovery — sees the same
	// empty list, and no tool definition is ever assembled (nor any plugin tool
	// collected) for a turn that cannot call one.
	if (apiConfiguration?.toolCallingEnabled === false) {
		return { tools: [] }
	}

	const mcpHub = provider.getMcpHub()

	// Category II managers are reached through their host registries (Chunk B): the
	// concrete singletons live in VS Code `src`; a headless host leaves the registry
	// unset and the features are simply off.

	const filterSettings = {
		todoListEnabled: apiConfiguration?.todoListEnabled ?? true,
		disabledTools,
		modelInfo,
	}

	const supportsImages = modelInfo?.supportsImages ?? false

	const nativeTools = getNativeTools({
		supportsImages,
		completionSchema: options.completionSchema,
		titleLocked: options.titleLocked,
	})

	const filteredNativeTools = filterNativeToolsForMode(
		nativeTools,
		mode,
		customModes,
		experiments,
		filterSettings,
		mcpHub,
	)

	const mcpTools = getMcpServerTools(mcpHub)
	const mcpToolMeta = mcpHub?.getMcpToolMetadata() ?? []
	const filteredMcpTools = filterMcpToolsForMode(mcpTools, mcpToolMeta, mode, customModes, experiments)

	let nativeCustomTools: OpenAI.Chat.ChatCompletionFunctionTool[] = []
	// §10: register plugin-contributed tools so they are assembled AND executable
	// through the unified custom-tool path (independent of the customTools experiment).
	const pluginTools = await pluginRegistry.collectTools({ workspacePath: cwd, cwd, mode })
	for (const def of pluginTools) {
		customToolRegistry.register(def, "plugin")
	}
	// Diagnostic: which plugin tools made it into this build (output channel). If a plugin
	// tool a user expects (e.g. `ask_live_memory`) is missing here, the plugin either
	// isn't loaded yet or its `registerTools` failed (see the `Plugin:<name>` log category).
	toolsLog.debug(
		`[plugin-tools] ${pluginTools.length} contributed (registry rev ${pluginRegistry.revision}): ` +
			(pluginTools.map((t) => t.name).join(", ") || "(none)"),
	)
	if (experiments?.customTools) {
		const toolDirs = getRooDirectoriesForCwd(cwd).map((dir) => path.join(dir, "tools"))
		await customToolRegistry.loadFromDirectoriesIfStale(toolDirs)
	}
	const serializedCustomTools = customToolRegistry.getAllSerialized()
	if (serializedCustomTools.length > 0) {
		nativeCustomTools = serializedCustomTools.map(formatNative)
	}

	// Discover all tools from private providers (extensions using the
	// shofer.privateToolProviders config convention).
	const privateMeta = await getPrivateLmToolMeta()
	const allPrivateTools = privateMeta.map((m) => m.tool)
	const filteredPrivateTools = filterPrivateToolsForMode(privateMeta, mode, customModes)

	// Build the invoke-command lookup map for the execution layer (core registry).
	setPrivateToolInvokeMap(privateMeta.map((m) => [getToolName(m.tool), m.invokeCommand]))

	// Per-task tool-group restriction (no-op when the field is absent).
	const {
		native: restrictedNative,
		mcp: restrictedMcp,
		custom: restrictedCustom,
		private: restrictedPrivate,
	} = restrictToolsToDeclaredGroups(options.agentToolGroups, {
		native: filteredNativeTools,
		mcp: filteredMcpTools,
		custom: nativeCustomTools,
		private: filteredPrivateTools,
		privateMeta,
	})

	const filteredTools = [...restrictedNative, ...restrictedMcp, ...restrictedCustom, ...restrictedPrivate]

	// The FULL contracts of everything this mode admits, recorded before the stub
	// tier reduces them — this is what `describe_tools` hands back, so the
	// described contract is by construction the one the call is validated
	// against. Recorded even when the mode declares no tiering: the tool is not
	// offered there, and a registry that only exists in one configuration is a
	// registry nobody can test.
	const modeConfig = resolveModeConfig(mode ?? defaultModeSlug, customModes)
	recordToolSchemas(modeConfig.slug, filteredTools)

	// Presentation tier (`ModeConfig.tools_full_schema`): a mode that declares one
	// keeps its always-on tools whole and reduces the rest to stubs. A mode that
	// declares none is returned exactly what it was returned before.
	const { tools: tieredTools, stubbed } = applyToolSchemaTiers(filteredTools, modeConfig)
	if (stubbed.length > 0) {
		toolsLog.debug(
			`[tool-stubs] mode ${modeConfig.slug}: ${tieredTools.length - stubbed.length} full schema(s), ` +
				`${stubbed.length} stub(s) recoverable via describe_tools`,
		)
	}

	if (includeAllToolsWithRestrictions) {
		// Providers that take every definition and restrict callability by name
		// (Gemini's `allowedFunctionNames`) get the same two tiers. `describe_tools`
		// is dropped here for a mode that declares no tiering, for the same reason
		// `computeToolAccess` drops it: a mode with no stubs must see exactly the
		// tools it saw before.
		const everything = [...nativeTools, ...mcpTools, ...nativeCustomTools, ...allPrivateTools].filter(
			(tool) => modeStubsToolSchemas(modeConfig) || getToolName(tool) !== DESCRIBE_TOOLS_TOOL_NAME,
		)
		const allTools = applyToolSchemaTiers(everything, modeConfig).tools
		const allowedFunctionNames = tieredTools.map((tool) => resolveToolAlias(getToolName(tool)))
		return { tools: allTools, allowedFunctionNames }
	}

	return { tools: tieredTools }
}