diff --git a/packages/core/src/usage-stats/types.ts b/packages/core/src/usage-stats/types.ts index 064dfa11ab..c81e934000 100644 --- a/packages/core/src/usage-stats/types.ts +++ b/packages/core/src/usage-stats/types.ts @@ -75,6 +75,8 @@ export interface UsageLogRow { prefixChangeReason?: PrefixChangeReason; requestShapeHash?: string; requestShapeChangeReason?: PrefixChangeReason; + toolSchemaChangeReason?: ToolSchemaChangeReason; + toolSourceEconomy?: ToolSourceEconomyDiagnostic; promptSegments?: PromptSegmentEstimate[]; contextBudget?: ContextBudgetDiagnostic; } @@ -124,11 +126,15 @@ export interface LlmCallRecord { prefixChangeReason?: PrefixChangeReason; requestShapeHash?: string; requestShapeChangeReason?: PrefixChangeReason; + toolSchemaChangeReason?: ToolSchemaChangeReason; + toolSourceEconomy?: ToolSourceEconomyDiagnostic; cacheMissInputSource?: CacheMissInputSource; promptSegments?: PromptSegmentEstimate[]; contextBudget?: ContextBudgetDiagnostic; } +export type ToolSourceId = string; + export type PrefixChangeReason = | 'first_turn' | 'system_prompt_changed' @@ -139,6 +145,27 @@ export type PrefixChangeReason = | 'stable' | 'unknown'; +export type ToolSchemaChangeReason = + | 'tool_schema_changed' + | 'tool_source_enabled' + | 'tool_source_state_changed'; + +export interface ToolSourceEconomyDiagnostic { + mode: 'full' | 'source_economy'; + enabledSourceIds: ToolSourceId[]; + availableSourceIds?: ToolSourceId[]; + connectorToolName?: string; + coreToolNames?: string[]; + visibleToolNamesBySource?: Record; + visibleToolCount?: number; + fullToolCount?: number; + hiddenToolCount?: number; + visibleToolSchemaChars?: number; + fullToolSchemaChars?: number; + toolSchemaCharReduction?: number; + estimatedToolSchemaTokenReduction?: number; +} + export type CacheMissInputSource = 'explicit' | 'derived'; export type PromptSegmentKind = diff --git a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts index c8865598e5..aca987f28e 100644 --- a/packages/runtime/src/__tests__/ai-sdk-backend.test.ts +++ b/packages/runtime/src/__tests__/ai-sdk-backend.test.ts @@ -27,6 +27,10 @@ import { canonicalizeToolSet, computeRequestShapeDiagnostic, } from '../request-shape.js'; +import { + CONNECT_TOOL_SOURCE_NAME, + ToolSourceEconomyRuntime, +} from '../tool-source-economy.js'; import { ARCHIVED_TOOL_RESULT_PLACEHOLDER_KIND, applyRuntimeEventContextBudget, @@ -1657,6 +1661,228 @@ describe('AiSdkBackend request-shape diagnostics', () => { ); }); + test('classifies strict enabled-source expansion as tool_source_enabled', () => { + const invalid = testTool(INVALID_TOOL_NAME, z.object({ tool: z.string().optional() })); + const initialTools = canonicalizeToolSet([ + testTool('Read', z.object({ path: z.string() })), + testTool(CONNECT_TOOL_SOURCE_NAME, z.object({ source: z.string() })), + ], invalid); + const expandedTools = canonicalizeToolSet([ + testTool('Read', z.object({ path: z.string() })), + testTool('WebFetch', z.object({ url: z.string() })), + testTool(CONNECT_TOOL_SOURCE_NAME, z.object({ source: z.string() })), + ], invalid); + const sourceCatalog = { web: ['WebFetch'] }; + const first = computeRequestShapeDiagnostic({ + connection: connection(), + modelId: 'mock-model-id', + providerTools: initialTools.providerTools, + activeTools: initialTools.activeTools, + priorMessages: [], + toolSourceEconomy: { + mode: 'source_economy', + enabledSourceIds: [], + availableSourceIds: ['web'], + connectorToolName: CONNECT_TOOL_SOURCE_NAME, + coreToolNames: ['Read'], + visibleToolNamesBySource: sourceCatalog, + }, + }, undefined); + const second = computeRequestShapeDiagnostic({ + connection: connection(), + modelId: 'mock-model-id', + providerTools: expandedTools.providerTools, + activeTools: expandedTools.activeTools, + priorMessages: [], + toolSourceEconomy: { + mode: 'source_economy', + enabledSourceIds: ['web'], + availableSourceIds: [], + connectorToolName: CONNECT_TOOL_SOURCE_NAME, + coreToolNames: ['Read'], + visibleToolNamesBySource: sourceCatalog, + }, + }, first); + + assert.equal(second.prefixChangeReason, 'tool_schema_changed'); + assert.equal(second.requestShapeChangeReason, 'tool_schema_changed'); + assert.equal(second.toolSchemaChangeReason, 'tool_source_enabled'); + assert.notEqual(second.prefixHash, first.prefixHash); + }); + + test('tool source economy starts small and connector enables sources for later selections', async () => { + const runtime = new ToolSourceEconomyRuntime([ + { + ...testTool('Write', z.object({ path: z.string(), content: z.string() })), + toolSource: { id: 'files.write', label: 'File writing' }, + }, + { + ...testTool('Read', z.object({ path: z.string() })), + toolSource: { id: 'core' }, + }, + { + ...testTool('WebFetch', z.object({ url: z.string() })), + toolSource: { id: 'web', label: 'Web' }, + }, + ], { mode: 'source_economy' }); + const initial = canonicalizeToolSet( + runtime.selectTools().tools, + testTool(INVALID_TOOL_NAME, z.object({ tool: z.string().optional() })), + ); + + assert.deepEqual( + initial.providerTools.map((tool) => tool.name), + ['Read', CONNECT_TOOL_SOURCE_NAME, INVALID_TOOL_NAME].sort((a, b) => { + if (a === INVALID_TOOL_NAME) return 1; + if (b === INVALID_TOOL_NAME) return -1; + return a.localeCompare(b); + }), + ); + assert.deepEqual(initial.activeTools, ['Read', CONNECT_TOOL_SOURCE_NAME].sort((a, b) => a.localeCompare(b))); + + const connector = initial.providerTools.find((tool) => tool.name === CONNECT_TOOL_SOURCE_NAME); + assert.ok(connector); + const result = await connector.impl({ source: 'web' }, { + sessionId: 'session-1', + turnId: 'turn-1', + cwd: '/tmp/maka', + toolCallId: 'tool-1', + abortSignal: new AbortController().signal, + emitOutput: () => {}, + }); + assert.deepEqual(result, { + ok: true, + source: 'web', + newlyEnabled: true, + enabledSources: ['web'], + availableSources: [{ id: 'files.write', label: 'File writing', toolCount: 1 }], + availableNextRequest: true, + tools: ['WebFetch'], + }); + + const expanded = canonicalizeToolSet( + runtime.selectTools().tools, + testTool(INVALID_TOOL_NAME, z.object({ tool: z.string().optional() })), + ); + assert.equal(expanded.activeTools.includes('WebFetch'), true); + assert.equal(expanded.activeTools.includes('Write'), false); + + const otherRuntime = new ToolSourceEconomyRuntime([ + testTool('Read', z.object({ path: z.string() })), + { ...testTool('WebFetch', z.object({ url: z.string() })), toolSource: { id: 'web' } }, + ], { mode: 'source_economy' }); + assert.equal( + canonicalizeToolSet( + otherRuntime.selectTools().tools, + testTool(INVALID_TOOL_NAME, z.object({ tool: z.string().optional() })), + ).activeTools.includes('WebFetch'), + false, + ); + }); + + test('backend full mode keeps the complete tool surface and omits source connector', async () => { + const model = completionModel(); + const llmRecords: LlmCallRecord[] = []; + const backend = new AiSdkBackend({ + sessionId: 'session-1', + header: header(), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + permissionEngine: new PermissionEngine({ newId: () => 'permission-id', now: () => 1 }), + modelFactory: () => model, + tools: [ + { ...testTool('Read', z.object({ path: z.string() })), toolSource: { id: 'core' } }, + { ...testTool('WebFetch', z.object({ url: z.string() })), toolSource: { id: 'web' } }, + ], + newId: idGenerator(), + now: monotonicClock(), + recordLlmCall: (record) => { + llmRecords.push(record); + }, + }); + + await drain(backend.send({ turnId: 'turn-1', text: 'hi', context: [] })); + + assert.deepEqual(modelToolNames(model), sortedModelToolNames(['Read', 'WebFetch'])); + assert.equal(modelToolNames(model).includes(CONNECT_TOOL_SOURCE_NAME), false); + assert.equal(toolSchemaPromptSegment(llmRecords[0])?.toolCount, 3); + }); + + test('backend connector enables sources for later requests only and preserves prefix reason compatibility', async () => { + const models: MockLanguageModelV3[] = []; + const llmRecords: LlmCallRecord[] = []; + const backend = new AiSdkBackend({ + sessionId: 'session-1', + header: header(), + appendMessage: async () => {}, + connection: connection(), + apiKey: 'sk-test', + modelId: 'mock-model-id', + permissionEngine: new PermissionEngine({ newId: () => 'permission-id', now: () => 1 }), + modelFactory: () => { + const model = completionModel(); + models.push(model); + return model; + }, + tools: [ + { ...testTool('Read', z.object({ path: z.string() })), toolSource: { id: 'core' } }, + { ...testTool('WebFetch', z.object({ url: z.string() })), toolSource: { id: 'web', label: 'Web' } }, + ], + toolSourceEconomy: { mode: 'source_economy' }, + newId: idGenerator(), + now: monotonicClock(), + recordLlmCall: (record) => { + llmRecords.push(record); + }, + }); + + await drain(backend.send({ turnId: 'turn-1', text: 'hi', context: [] })); + + assert.deepEqual( + modelToolNames(models[0]!), + sortedModelToolNames(['Read', CONNECT_TOOL_SOURCE_NAME]), + ); + assert.equal(modelToolNames(models[0]!).includes('WebFetch'), false); + assert.equal(toolSchemaPromptSegment(llmRecords[0])?.toolCount, 3); + const economyRuntime = (backend as unknown as { + toolSourceEconomyRuntime: ToolSourceEconomyRuntime; + }).toolSourceEconomyRuntime; + const connector = economyRuntime.selectTools().tools.find((tool) => tool.name === CONNECT_TOOL_SOURCE_NAME); + assert.ok(connector); + const connectResult = await connector.impl({ source: 'web' }, { + sessionId: 'session-1', + turnId: 'turn-connect', + cwd: '/tmp/maka', + toolCallId: 'tool-1', + abortSignal: new AbortController().signal, + emitOutput: () => {}, + }); + assert.deepEqual(connectResult, { + ok: true, + source: 'web', + newlyEnabled: true, + enabledSources: ['web'], + availableSources: [], + availableNextRequest: true, + tools: ['WebFetch'], + }); + assert.equal(modelToolNames(models[0]!).includes('WebFetch'), false); + + await drain(backend.send({ turnId: 'turn-2', text: 'hi again', context: [] })); + + assert.deepEqual( + modelToolNames(models[1]!), + sortedModelToolNames(['Read', CONNECT_TOOL_SOURCE_NAME, 'WebFetch']), + ); + assert.equal(llmRecords[1]?.prefixChangeReason, 'tool_schema_changed'); + assert.equal(llmRecords[1]?.requestShapeChangeReason, 'tool_schema_changed'); + assert.equal(llmRecords[1]?.toolSchemaChangeReason, 'tool_source_enabled'); + assert.deepEqual(llmRecords[1]?.toolSourceEconomy?.enabledSourceIds, ['web']); + assert.equal(toolSchemaPromptSegment(llmRecords[1])?.toolCount, 4); + }); + test('volatile turn-tail facts do not churn the durable prefix hash', async () => { const events: SessionEvent[] = []; const llmRecords: LlmCallRecord[] = []; @@ -3022,6 +3248,45 @@ function modelCallSettings(model: MockLanguageModelV3): unknown { return rest; } +function modelToolNames(model: MockLanguageModelV3): string[] { + return sortedModelToolNames(Object.keys(modelTools(model))); +} + +function modelTools(model: MockLanguageModelV3): Record { + const call = model.doStreamCalls[0] as unknown as Record | undefined; + const tools = call?.tools; + if (!tools) return {}; + if (Array.isArray(tools)) { + const out: Record = {}; + for (const tool of tools) { + if (tool && typeof tool === 'object') { + const record = tool as Record; + const name = typeof record.name === 'string' + ? record.name + : typeof record.toolName === 'string' + ? record.toolName + : undefined; + if (name) out[name] = tool; + } + } + return out; + } + if (typeof tools === 'object') return tools as Record; + return {}; +} + +function sortedModelToolNames(toolNames: readonly string[]): string[] { + return [...toolNames].sort((a, b) => { + if (a === INVALID_TOOL_NAME) return 1; + if (b === INVALID_TOOL_NAME) return -1; + return a.localeCompare(b); + }); +} + +function toolSchemaPromptSegment(record: LlmCallRecord | undefined): { toolCount?: number } | undefined { + return record?.promptSegments?.find((segment) => segment.kind === 'tool_schema'); +} + function sha256(text: string): string { return createHash('sha256').update(text).digest('hex'); } diff --git a/packages/runtime/src/ai-sdk-backend.ts b/packages/runtime/src/ai-sdk-backend.ts index 2110c4c96b..0f4345c956 100644 --- a/packages/runtime/src/ai-sdk-backend.ts +++ b/packages/runtime/src/ai-sdk-backend.ts @@ -61,6 +61,7 @@ import type { LlmCallRecord, ToolInvocationRecord } from '@maka/core/usage-stats import type { ContextBudgetDiagnostic, PromptSegmentEstimate, + ToolSourceEconomyDiagnostic, } from '@maka/core/usage-stats/types'; import type { JSONValue, ModelMessage } from 'ai'; import { z } from 'zod'; @@ -100,11 +101,16 @@ import { toolSchemaCharsForDiagnostics, type RequestShapeDiagnostic, } from './request-shape.js'; +import { + ToolSourceEconomyRuntime, + type ToolSourceEconomyConfig, +} from './tool-source-economy.js'; import { ARCHIVED_TOOL_RESULT_REWRITE_VERSION, applyRuntimeEventContextBudget, buildPromptSegmentEstimates, collectStaleToolResultArchiveCandidates, + estimateTokens, estimateRuntimeEventsTokens, rawEvidenceRequestReason, retrieveArchivedToolResultsForReplay, @@ -128,6 +134,12 @@ export { export type { MakaTool, MakaToolContext } from './tool-runtime.js'; export { normalizeAiSdkUsage } from './model-adapter.js'; export type { ModelFactory, ModelFactoryInput, RepairableAiSdkToolCall } from './model-adapter.js'; +export type { + ConnectToolSourceResult, + ToolSourceDefinition, + ToolSourceEconomyConfig, + ToolSourceEconomySelection, +} from './tool-source-economy.js'; export type { RunTraceEvent, RunTraceRecorder } from './run-trace.js'; type AiSdkToolResultOutput = @@ -231,6 +243,8 @@ export interface AiSdkBackendInput { /** Canonical-named tools available this session. Backend wraps each with * permission gating before passing to ai-sdk. */ tools: MakaTool[]; + /** Optional opt-in tool source economy mode. Omitted/full mode preserves the full tool surface. */ + toolSourceEconomy?: ToolSourceEconomyConfig; // ── Optional knobs (defaults shown) ──────────────────────────────────── /** ID generator; default `crypto.randomUUID()`. */ @@ -303,6 +317,7 @@ export class AiSdkBackend implements AgentBackend { private readonly maxSteps: number; private readonly toolRuntime: ToolRuntime; private readonly modelAdapter: ModelAdapter; + private readonly toolSourceEconomyRuntime: ToolSourceEconomyRuntime; private aborted = false; private abortController: AbortController | null = null; @@ -320,6 +335,7 @@ export class AiSdkBackend implements AgentBackend { this.newId = input.newId ?? (() => crypto.randomUUID()); this.now = input.now ?? (() => Date.now()); this.maxSteps = input.maxSteps ?? 50; + this.toolSourceEconomyRuntime = new ToolSourceEconomyRuntime(input.tools, input.toolSourceEconomy); this.modelAdapter = new ModelAdapter({ connection: input.connection, apiKey: input.apiKey, @@ -407,7 +423,9 @@ export class AiSdkBackend implements AgentBackend { } // --- Build ai-sdk tools dict with permission-wrapped execute --- - const canonicalTools = canonicalizeToolSet(this.input.tools, buildInvalidMakaTool()); + const toolSourceSelection = this.toolSourceEconomyRuntime.selectTools(); + const invalidTool = buildInvalidMakaTool(); + const canonicalTools = canonicalizeToolSet(toolSourceSelection.tools, invalidTool); const aiSdkTools: Record = {}; for (const t of canonicalTools.providerTools) { aiSdkTools[t.name] = { @@ -450,9 +468,13 @@ export class AiSdkBackend implements AgentBackend { content: this.appendTurnTailPrompt(currentUserContent, turnTailPrompt), }, ]; + const toolSchemaChars = toolSchemaCharsForDiagnostics(canonicalTools.providerTools, activeTools); + const toolSourceDiagnostic = toolSourceSelection.diagnostic !== undefined + ? this.enrichToolSourceDiagnostic(toolSourceSelection.diagnostic, canonicalTools, toolSchemaChars) + : undefined; const promptSegments = buildPromptSegmentEstimates({ systemPrompt, - toolSchemaChars: toolSchemaCharsForDiagnostics(canonicalTools.providerTools, activeTools), + toolSchemaChars, toolCount: canonicalTools.providerTools.length, priorMessages: priorReplay.messages, priorRuntimeEventCount: priorReplay.runtimeEventCount, @@ -469,6 +491,9 @@ export class AiSdkBackend implements AgentBackend { providerTools: canonicalTools.providerTools, activeTools, priorMessages: priorReplay.messages, + ...(toolSourceDiagnostic !== undefined + ? { toolSourceEconomy: toolSourceDiagnostic } + : {}), }, this.priorRequestShape); if (priorReplay.contextBudget?.highWaterReason) { priorReplay.contextBudget.highWaterRequestShapeHashBefore = this.priorRequestShape?.requestShapeHash; @@ -481,6 +506,12 @@ export class AiSdkBackend implements AgentBackend { prefixChangeReason: requestShape.prefixChangeReason, requestShapeHash: requestShape.requestShapeHash, requestShapeChangeReason: requestShape.requestShapeChangeReason, + ...(requestShape.toolSchemaChangeReason !== undefined + ? { toolSchemaChangeReason: requestShape.toolSchemaChangeReason } + : {}), + ...(requestShape.toolSourceEconomy !== undefined + ? { toolSourceEconomy: requestShape.toolSourceEconomy } + : {}), promptSegments, ...(priorReplay.contextBudget ? { contextBudget: priorReplay.contextBudget } : {}), }); @@ -576,6 +607,12 @@ export class AiSdkBackend implements AgentBackend { prefixChangeReason: requestShape.prefixChangeReason, requestShapeHash: requestShape.requestShapeHash, requestShapeChangeReason: requestShape.requestShapeChangeReason, + ...(requestShape.toolSchemaChangeReason !== undefined + ? { toolSchemaChangeReason: requestShape.toolSchemaChangeReason } + : {}), + ...(requestShape.toolSourceEconomy !== undefined + ? { toolSourceEconomy: requestShape.toolSourceEconomy } + : {}), }); const tu: TokenUsageMessage = { type: 'token_usage', @@ -701,6 +738,12 @@ export class AiSdkBackend implements AgentBackend { prefixChangeReason: requestShapeForTelemetry.prefixChangeReason, requestShapeHash: requestShapeForTelemetry.requestShapeHash, requestShapeChangeReason: requestShapeForTelemetry.requestShapeChangeReason, + ...(requestShapeForTelemetry.toolSchemaChangeReason !== undefined + ? { toolSchemaChangeReason: requestShapeForTelemetry.toolSchemaChangeReason } + : {}), + ...(requestShapeForTelemetry.toolSourceEconomy !== undefined + ? { toolSourceEconomy: requestShapeForTelemetry.toolSourceEconomy } + : {}), } : {}), ...(promptSegmentsForTelemetry.length > 0 ? { promptSegments: promptSegmentsForTelemetry } : {}), ...(contextBudgetForTelemetry !== undefined ? { contextBudget: contextBudgetForTelemetry } : {}), @@ -733,6 +776,28 @@ export class AiSdkBackend implements AgentBackend { // Helpers // -------------------------------------------------------------------------- + private enrichToolSourceDiagnostic( + diagnostic: ToolSourceEconomyDiagnostic, + canonicalTools: { providerTools: MakaTool[]; activeTools: string[] }, + visibleToolSchemaChars: number, + ): ToolSourceEconomyDiagnostic { + const fullTools = canonicalizeToolSet(this.input.tools, buildInvalidMakaTool()); + const fullToolSchemaChars = toolSchemaCharsForDiagnostics(fullTools.providerTools, fullTools.activeTools); + const visibleToolNamesExcludingConnector = canonicalTools.activeTools + .filter((toolName) => toolName !== diagnostic.connectorToolName); + const toolSchemaCharReduction = Math.max(0, fullToolSchemaChars - visibleToolSchemaChars); + return { + ...diagnostic, + visibleToolCount: canonicalTools.activeTools.length, + fullToolCount: fullTools.activeTools.length, + hiddenToolCount: Math.max(0, fullTools.activeTools.length - visibleToolNamesExcludingConnector.length), + visibleToolSchemaChars, + fullToolSchemaChars, + toolSchemaCharReduction, + estimatedToolSchemaTokenReduction: estimateTokens(toolSchemaCharReduction), + }; + } + async stop(_reason: 'user_stop' | 'redirect'): Promise { this.aborted = true; this.abortController?.abort(); diff --git a/packages/runtime/src/builtin-tools.ts b/packages/runtime/src/builtin-tools.ts index 472acc1dcb..a1872a4518 100644 --- a/packages/runtime/src/builtin-tools.ts +++ b/packages/runtime/src/builtin-tools.ts @@ -31,6 +31,11 @@ export function buildBuiltinTools(): MakaTool[] { timeout_ms: z.number().int().positive().max(600_000).optional(), }), permissionRequired: true, + toolSource: { + id: 'shell', + label: 'Shell', + description: 'Shell command execution tools.', + }, impl: async ({ command, timeout_ms }, { cwd, abortSignal, emitOutput }) => { const result = await runStreamingShell(command, { cwd, @@ -57,6 +62,11 @@ export function buildBuiltinTools(): MakaTool[] { limit: z.number().int().positive().optional(), }), permissionRequired: false, + toolSource: { + id: 'core', + label: 'Core', + description: 'Read-only file inspection tools.', + }, impl: async ({ path, offset, limit }, { cwd }) => { const abs = await resolveExistingInsideCwd(cwd, path, 'Read'); const content = await fs.readFile(abs, 'utf8'); @@ -72,6 +82,11 @@ export function buildBuiltinTools(): MakaTool[] { description: 'Write content to a file (creates or overwrites). Subject to permission policy.', parameters: z.object({ path: z.string(), content: z.string() }), permissionRequired: true, + toolSource: { + id: 'files.write', + label: 'File writing', + description: 'Workspace file creation and modification tools.', + }, impl: async ({ path, content }, { cwd }) => { const abs = await resolveWritableInsideCwd(cwd, path, 'Write'); await fs.writeFile(abs, content, 'utf8'); @@ -88,6 +103,11 @@ export function buildBuiltinTools(): MakaTool[] { new_string: z.string(), }), permissionRequired: true, + toolSource: { + id: 'files.write', + label: 'File writing', + description: 'Workspace file creation and modification tools.', + }, impl: async ({ path, old_string, new_string }, { cwd }) => { const abs = await resolveExistingInsideCwd(cwd, path, 'Edit'); const current = await fs.readFile(abs, 'utf8'); @@ -110,6 +130,11 @@ export function buildBuiltinTools(): MakaTool[] { cwd: z.string().optional(), }), permissionRequired: false, + toolSource: { + id: 'core', + label: 'Core', + description: 'Read-only file inspection tools.', + }, impl: async ({ pattern, cwd: relCwd }, { cwd }) => { assertRelativeGlobPattern(pattern); const base = relCwd ? await resolveExistingInsideCwd(cwd, relCwd, 'Glob cwd') : await fs.realpath(cwd); @@ -130,6 +155,11 @@ export function buildBuiltinTools(): MakaTool[] { glob: z.string().optional(), }), permissionRequired: false, + toolSource: { + id: 'core', + label: 'Core', + description: 'Read-only file inspection tools.', + }, impl: async ({ pattern, path, glob }, { cwd }) => { const args = ['-n', '--no-heading', '--max-count=50']; if (glob) args.push('--glob', glob); diff --git a/packages/runtime/src/index.ts b/packages/runtime/src/index.ts index 8a960b1be3..6f79b090f1 100644 --- a/packages/runtime/src/index.ts +++ b/packages/runtime/src/index.ts @@ -35,6 +35,8 @@ export type { ModelFactoryInput, RunTraceEvent, RunTraceRecorder, + ToolSourceDefinition, + ToolSourceEconomyConfig, SynthesisCacheLoader, SynthesisCacheLoadInput, SynthesisCacheLoadResult, diff --git a/packages/runtime/src/request-shape.ts b/packages/runtime/src/request-shape.ts index 1dd2a6c7fe..283abdebec 100644 --- a/packages/runtime/src/request-shape.ts +++ b/packages/runtime/src/request-shape.ts @@ -1,7 +1,11 @@ import { Buffer } from 'node:buffer'; import { createHash } from 'node:crypto'; import type { LlmConnection } from '@maka/core/llm-connections'; -import type { PrefixChangeReason } from '@maka/core/usage-stats/types'; +import type { + PrefixChangeReason, + ToolSchemaChangeReason, + ToolSourceEconomyDiagnostic, +} from '@maka/core/usage-stats/types'; import type { ModelMessage } from 'ai'; import { toJSONSchema } from 'zod'; @@ -20,6 +24,7 @@ export interface RequestShapeInput { providerTools: readonly MakaTool[]; activeTools: readonly string[]; priorMessages: readonly ModelMessage[]; + toolSourceEconomy?: ToolSourceEconomyDiagnostic; } export interface RequestShapeComponents { @@ -40,6 +45,8 @@ export interface RequestShapeDiagnostic { requestShapeHash: string; requestShapeChangeReason: PrefixChangeReason; componentHashes: RequestShapeComponents; + toolSchemaChangeReason?: ToolSchemaChangeReason; + toolSourceEconomy?: ToolSourceEconomyDiagnostic; } export function canonicalizeToolSet( @@ -77,6 +84,12 @@ export function computeRequestShapeDiagnostic( const durablePrefixComponents = durableComponents(componentHashes); const prefixHash = stableHash(durablePrefixComponents); const requestShapeHash = stableHash(componentHashes); + const toolSchemaChangeReason = classifyToolSchemaChange( + componentHashes, + prior?.componentHashes, + input.toolSourceEconomy, + prior?.toolSourceEconomy, + ); return { prefixHash, prefixChangeReason: classifyDurablePrefixChange( @@ -84,8 +97,13 @@ export function computeRequestShapeDiagnostic( prior ? durableComponents(prior.componentHashes) : undefined, ), requestShapeHash, - requestShapeChangeReason: classifyRequestShapeChange(componentHashes, prior?.componentHashes), + requestShapeChangeReason: classifyRequestShapeChange( + componentHashes, + prior?.componentHashes, + ), componentHashes, + ...(toolSchemaChangeReason !== undefined ? { toolSchemaChangeReason } : {}), + ...(input.toolSourceEconomy !== undefined ? { toolSourceEconomy: input.toolSourceEconomy } : {}), }; } @@ -132,6 +150,60 @@ function classifyRequestShapeChange( return 'stable'; } +function classifyToolSchemaChange( + current: RequestShapeComponents, + prior: RequestShapeComponents | undefined, + currentEconomy: ToolSourceEconomyDiagnostic | undefined, + priorEconomy: ToolSourceEconomyDiagnostic | undefined, +): ToolSchemaChangeReason | undefined { + if (!prior || current.toolSchemaHash === prior.toolSchemaHash) return undefined; + if (isEnabledSourceStrictSuperset(currentEconomy, priorEconomy) && sourceCatalogStable(currentEconomy, priorEconomy)) { + return 'tool_source_enabled'; + } + if (sourceStateChanged(currentEconomy, priorEconomy)) { + return 'tool_source_state_changed'; + } + return 'tool_schema_changed'; +} + +function isEnabledSourceStrictSuperset( + current: ToolSourceEconomyDiagnostic | undefined, + prior: ToolSourceEconomyDiagnostic | undefined, +): boolean { + if (current?.mode !== 'source_economy' || prior?.mode !== 'source_economy') return false; + const currentIds = new Set(current.enabledSourceIds); + const priorIds = new Set(prior.enabledSourceIds); + if (currentIds.size <= priorIds.size) return false; + for (const sourceId of priorIds) { + if (!currentIds.has(sourceId)) return false; + } + return true; +} + +function sourceCatalogStable( + current: ToolSourceEconomyDiagnostic | undefined, + prior: ToolSourceEconomyDiagnostic | undefined, +): boolean { + if (!current || !prior) return false; + return stableStringify(sourceCatalogShape(current)) === stableStringify(sourceCatalogShape(prior)); +} + +function sourceStateChanged( + current: ToolSourceEconomyDiagnostic | undefined, + prior: ToolSourceEconomyDiagnostic | undefined, +): boolean { + return stableStringify(current ?? null) !== stableStringify(prior ?? null); +} + +function sourceCatalogShape(diagnostic: ToolSourceEconomyDiagnostic): unknown { + return { + mode: diagnostic.mode, + connectorToolName: diagnostic.connectorToolName, + coreToolNames: diagnostic.coreToolNames ?? [], + visibleToolNamesBySource: diagnostic.visibleToolNamesBySource ?? {}, + }; +} + function durableComponents(components: RequestShapeComponents): DurablePrefixComponents { return { modelProviderHash: components.modelProviderHash, diff --git a/packages/runtime/src/run-trace.ts b/packages/runtime/src/run-trace.ts index f5db5a394e..a2ed68b18f 100644 --- a/packages/runtime/src/run-trace.ts +++ b/packages/runtime/src/run-trace.ts @@ -4,6 +4,8 @@ import type { ContextBudgetDiagnostic, PrefixChangeReason, PromptSegmentEstimate, + ToolSchemaChangeReason, + ToolSourceEconomyDiagnostic, } from '@maka/core/usage-stats/types'; export type RunTracePhase = 'turn' | 'model' | 'tool' | 'permission' | 'abort' | 'usage'; @@ -103,6 +105,8 @@ export class RunTrace { prefixChangeReason: PrefixChangeReason; requestShapeHash?: string; requestShapeChangeReason?: PrefixChangeReason; + toolSchemaChangeReason?: ToolSchemaChangeReason; + toolSourceEconomy?: ToolSourceEconomyDiagnostic; promptSegments?: PromptSegmentEstimate[]; contextBudget?: ContextBudgetDiagnostic; }, @@ -141,6 +145,8 @@ export class RunTrace { prefixChangeReason?: PrefixChangeReason; requestShapeHash?: string; requestShapeChangeReason?: PrefixChangeReason; + toolSchemaChangeReason?: ToolSchemaChangeReason; + toolSourceEconomy?: ToolSourceEconomyDiagnostic; }): void { this.emit('usage', 'usage_recorded', 'Token usage recorded', { inputTokens: usage.inputTokens, @@ -159,6 +165,8 @@ export class RunTrace { ...(usage.requestShapeChangeReason !== undefined ? { requestShapeChangeReason: usage.requestShapeChangeReason } : {}), + ...(usage.toolSchemaChangeReason !== undefined ? { toolSchemaChangeReason: usage.toolSchemaChangeReason } : {}), + ...(usage.toolSourceEconomy !== undefined ? { toolSourceEconomy: usage.toolSourceEconomy } : {}), }); } diff --git a/packages/runtime/src/tool-runtime.ts b/packages/runtime/src/tool-runtime.ts index 3eb5ae71a6..acc59c18b0 100644 --- a/packages/runtime/src/tool-runtime.ts +++ b/packages/runtime/src/tool-runtime.ts @@ -14,7 +14,7 @@ import type { PermissionDecision } from '@maka/core/backend-types'; import type { ToolCategory } from '@maka/core/permission'; import type { LlmConnection } from '@maka/core/llm-connections'; import type { SessionHeader } from '@maka/core/session'; -import type { ToolInvocationRecord } from '@maka/core/usage-stats/types'; +import type { ToolInvocationRecord, ToolSourceId } from '@maka/core/usage-stats/types'; import { redactSecrets } from '@maka/core/redaction'; import type { PermissionEngine } from './permission-engine.js'; @@ -42,6 +42,12 @@ export interface MakaTool

{ displayName?: string; /** Optional trusted category override for custom tools. */ categoryHint?: ToolCategory; + /** Optional source grouping used by opt-in tool source economy mode. */ + toolSource?: { + id: ToolSourceId; + label?: string; + description?: string; + }; /** Real tool implementation. Called only after permission allows. */ impl: (args: P, ctx: MakaToolContext) => Promise | R; } diff --git a/packages/runtime/src/tool-source-economy.ts b/packages/runtime/src/tool-source-economy.ts new file mode 100644 index 0000000000..b717a761cd --- /dev/null +++ b/packages/runtime/src/tool-source-economy.ts @@ -0,0 +1,268 @@ +import type { + ToolSourceEconomyDiagnostic, + ToolSourceId, +} from '@maka/core/usage-stats/types'; +import { z } from 'zod'; + +import type { MakaTool } from './tool-runtime.js'; + +export const CONNECT_TOOL_SOURCE_NAME = 'connect_tool_source'; +export const CORE_TOOL_SOURCE_ID = 'core'; + +const DEFAULT_CORE_TOOL_NAMES = ['Read', 'Glob', 'Grep'] as const; + +export interface ToolSourceDefinition { + id: ToolSourceId; + toolNames: readonly string[]; + label?: string; + description?: string; +} + +export type ToolSourceEconomyConfig = + | { mode?: 'full' } + | { + mode: 'source_economy'; + coreToolNames?: readonly string[]; + initialSourceIds?: readonly ToolSourceId[]; + connectorToolName?: string; + sourceDefinitions?: readonly ToolSourceDefinition[]; + }; + +export interface ToolSourceEconomySelection { + tools: MakaTool[]; + diagnostic?: ToolSourceEconomyDiagnostic; +} + +interface SourceInfo { + id: ToolSourceId; + label?: string; + description?: string; + toolNames: string[]; +} + +interface SourceCatalog { + coreToolNames: Set; + sourceByToolName: Map; + sources: Map; +} + +export class ToolSourceEconomyRuntime { + private readonly enabledSourceIds = new Set(); + private readonly mode: 'full' | 'source_economy'; + private readonly connectorToolName: string; + private readonly catalog: SourceCatalog; + + constructor( + private readonly tools: readonly MakaTool[], + config: ToolSourceEconomyConfig | undefined, + ) { + const economyConfig = config?.mode === 'source_economy' ? config : undefined; + this.mode = economyConfig ? 'source_economy' : 'full'; + this.connectorToolName = economyConfig + ? economyConfig.connectorToolName ?? CONNECT_TOOL_SOURCE_NAME + : CONNECT_TOOL_SOURCE_NAME; + this.catalog = buildSourceCatalog(tools, economyConfig); + + if (economyConfig) { + for (const sourceId of economyConfig.initialSourceIds ?? []) { + if (this.catalog.sources.has(sourceId)) { + this.enabledSourceIds.add(sourceId); + } + } + } + } + + selectTools(): ToolSourceEconomySelection { + if (this.mode === 'full') { + return { tools: [...this.tools] }; + } + + const visibleTools = this.tools.filter((tool) => this.isToolVisible(tool)); + const connectorTool = this.buildConnectorTool(); + return { + tools: [...visibleTools, connectorTool], + diagnostic: this.buildDiagnostic(), + }; + } + + private isToolVisible(tool: MakaTool): boolean { + if (this.catalog.coreToolNames.has(tool.name)) return true; + const sourceId = this.catalog.sourceByToolName.get(tool.name); + if (sourceId === undefined) return true; + return this.enabledSourceIds.has(sourceId); + } + + private buildConnectorTool(): MakaTool<{ source: string }, ConnectToolSourceResult> { + return { + name: this.connectorToolName, + description: 'Enable a named tool source for later requests in this backend instance.', + parameters: z.object({ + source: z.string().min(1).describe('Source id to enable for later model requests.'), + }), + permissionRequired: false, + impl: ({ source }) => this.connectSource(source), + }; + } + + private connectSource(source: string): ConnectToolSourceResult { + const sourceInfo = this.catalog.sources.get(source); + if (!sourceInfo) { + return { + ok: false, + source, + error: 'unknown_source', + enabledSources: this.sortedEnabledSourceIds(), + availableSources: this.availableSources(), + }; + } + + const toolNames = sourceInfo.toolNames + .filter((toolName) => !this.catalog.coreToolNames.has(toolName)) + .sort((a, b) => a.localeCompare(b)); + if (toolNames.length === 0) { + return { + ok: false, + source, + error: 'source_has_no_tools', + enabledSources: this.sortedEnabledSourceIds(), + availableSources: this.availableSources(), + }; + } + + const wasEnabled = this.enabledSourceIds.has(source); + this.enabledSourceIds.add(source); + return { + ok: true, + source, + newlyEnabled: !wasEnabled, + enabledSources: this.sortedEnabledSourceIds(), + availableSources: this.availableSources(), + availableNextRequest: true, + tools: toolNames, + }; + } + + private buildDiagnostic(): ToolSourceEconomyDiagnostic { + const availableSourceIds = this.availableSources().map((source) => source.id); + return { + mode: 'source_economy', + enabledSourceIds: this.sortedEnabledSourceIds(), + availableSourceIds, + connectorToolName: this.connectorToolName, + coreToolNames: [...this.catalog.coreToolNames].sort((a, b) => a.localeCompare(b)), + visibleToolNamesBySource: sourceToolNamesById(this.catalog.sources), + }; + } + + private availableSources(): ConnectToolSourceAvailableSource[] { + return [...this.catalog.sources.values()] + .filter((source) => source.toolNames.length > 0 && !this.enabledSourceIds.has(source.id)) + .map((source) => ({ + id: source.id, + ...(source.label !== undefined ? { label: source.label } : {}), + ...(source.description !== undefined ? { description: source.description } : {}), + toolCount: source.toolNames.length, + })) + .sort((a, b) => a.id.localeCompare(b.id)); + } + + private sortedEnabledSourceIds(): ToolSourceId[] { + return [...this.enabledSourceIds].sort((a, b) => a.localeCompare(b)); + } +} + +export type ConnectToolSourceResult = + | { + ok: true; + source: string; + newlyEnabled: boolean; + enabledSources: string[]; + availableSources: ConnectToolSourceAvailableSource[]; + availableNextRequest: true; + tools: string[]; + } + | { + ok: false; + source: string; + error: 'unknown_source' | 'source_has_no_tools'; + enabledSources: string[]; + availableSources: ConnectToolSourceAvailableSource[]; + }; + +export interface ConnectToolSourceAvailableSource { + id: string; + label?: string; + description?: string; + toolCount: number; +} + +function buildSourceCatalog( + tools: readonly MakaTool[], + config: Extract | undefined, +): SourceCatalog { + const coreToolNames = new Set(config?.coreToolNames ?? DEFAULT_CORE_TOOL_NAMES); + const sourceByToolName = new Map(); + const sources = new Map(); + + for (const definition of config?.sourceDefinitions ?? []) { + if (!definition.id || definition.id === CORE_TOOL_SOURCE_ID) continue; + const info = ensureSourceInfo(sources, definition.id, definition.label, definition.description); + for (const toolName of definition.toolNames) { + if (!toolName || coreToolNames.has(toolName) || sourceByToolName.has(toolName)) continue; + sourceByToolName.set(toolName, definition.id); + info.toolNames.push(toolName); + } + } + + for (const tool of tools) { + const source = tool.toolSource; + if (source?.id === CORE_TOOL_SOURCE_ID) { + coreToolNames.add(tool.name); + continue; + } + if (coreToolNames.has(tool.name) || sourceByToolName.has(tool.name) || !source?.id) continue; + const info = ensureSourceInfo(sources, source.id, source.label, source.description); + sourceByToolName.set(tool.name, source.id); + info.toolNames.push(tool.name); + } + + const knownToolNames = new Set(tools.map((tool) => tool.name)); + for (const [sourceId, info] of sources) { + info.toolNames = [...new Set(info.toolNames)] + .filter((toolName) => knownToolNames.has(toolName) && !coreToolNames.has(toolName)) + .sort((a, b) => a.localeCompare(b)); + if (info.toolNames.length === 0) { + sources.delete(sourceId); + } + } + + return { coreToolNames, sourceByToolName, sources }; +} + +function ensureSourceInfo( + sources: Map, + id: ToolSourceId, + label: string | undefined, + description: string | undefined, +): SourceInfo { + const existing = sources.get(id); + if (existing) { + return existing; + } + const next: SourceInfo = { + id, + ...(label !== undefined ? { label } : {}), + ...(description !== undefined ? { description } : {}), + toolNames: [], + }; + sources.set(id, next); + return next; +} + +function sourceToolNamesById(sources: Map): Record { + const out: Record = {}; + for (const source of [...sources.values()].sort((a, b) => a.id.localeCompare(b.id))) { + out[source.id] = [...source.toolNames].sort((a, b) => a.localeCompare(b)); + } + return out; +}