diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index 70587bc046a..a1782b2322b 100644 --- a/src/node/services/aiService.test.ts +++ b/src/node/services/aiService.test.ts @@ -330,6 +330,7 @@ function stubCommonStreamMessageDependencies(args: { agentDefinitions: undefined, availableSkills: undefined, ancestorPlanFilePaths: [], + instructionSources: { global: null, context: [] }, }); }); spyOn(messagePipeline, "prepareMessagesForProvider").mockImplementation((pipelineArgs) => { @@ -343,6 +344,7 @@ function stubCommonStreamMessageDependencies(args: { const getToolsForModelSpy = spyOn(toolsModule, "getToolsForModel").mockResolvedValue( args.allTools ?? {} ); + spyOn(systemMessageModule, "toolInstructionsFromSources").mockReturnValue({}); spyOn(systemMessageModule, "readToolInstructions").mockResolvedValue({}); const providerModelFactory = Reflect.get(args.service, "providerModelFactory") as diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 7cecc11e4c0..181e04b3ae4 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -72,7 +72,7 @@ import { sumUsageHistory, getTotalCost } from "@/common/utils/tokens/usageAggreg import { createDisplayUsage } from "@/common/utils/tokens/displayUsage"; import { normalizeToCanonical } from "@/common/utils/ai/models"; import { extractChunkDeltaText } from "@/common/utils/ai/streamChunks"; -import { readToolInstructions } from "./systemMessage"; +import { toolInstructionsFromSources } from "./systemMessage"; import { effectiveAdditionalSystemContext, mergeAdditionalSystemInstructions, @@ -1664,6 +1664,7 @@ export class AIService extends EventEmitter { agentDefinitions, availableSkills, ancestorPlanFilePaths, + instructionSources, } = prePolicyStreamSystemContext; let systemMessageTokens = prePolicyStreamSystemContext.systemMessageTokens; let systemMessage = prePolicyStreamSystemContext.systemMessage; @@ -1676,32 +1677,68 @@ export class AIService extends EventEmitter { // Generate stream token and create temp directory for tools const streamToken = this.streamManager.generateStreamToken(); - let mcpTools: Record | undefined; - let mcpStats: MCPWorkspaceStats | undefined; let mcpSetupDurationMs = 0; - if (this.mcpServerManager) { - const mcpToolSetupStartedAt = Date.now(); - try { - const result = await this.mcpServerManager.getToolsForWorkspace({ - workspaceId, - projectPath: metadata.projectPath, - runtime, - workspacePath, - trusted: projectTrusted, - overrides: mcpOverrides, - projectSecrets: await secretsToRecord(projectSecrets, this.opResolver), - }); + const mcpToolSetupStartedAt = Date.now(); + const createTempDirForStreamStartedAt = Date.now(); + const readToolInstructionsStartedAt = Date.now(); + const loadSessionUsageStartedAt = Date.now(); + const toolInstructions = toolInstructionsFromSources( + instructionSources, + metadata, + capabilityModelString, + agentSystemPromptSections + ); + recordStartupPhaseTiming("readToolInstructionsMs", readToolInstructionsStartedAt); + const [mcpSetupResult, runtimeTempDir, sessionCostsUsd] = await Promise.all([ + this.mcpServerManager + ? (async (): Promise<{ + tools: Record | undefined; + stats: MCPWorkspaceStats | undefined; + }> => { + try { + const result = await this.mcpServerManager!.getToolsForWorkspace({ + workspaceId, + projectPath: metadata.projectPath, + runtime, + workspacePath, + trusted: projectTrusted, + overrides: mcpOverrides, + projectSecrets: await secretsToRecord(projectSecrets, this.opResolver), + }); + return { tools: result.tools, stats: result.stats }; + } catch (error) { + workspaceLog.error("Failed to start MCP servers", { error }); + return { tools: undefined, stats: undefined }; + } finally { + mcpSetupDurationMs = Date.now() - mcpToolSetupStartedAt; + startupPhaseTimingsMs.mcpToolSetupMs = mcpSetupDurationMs; + } + })() + : Promise.resolve({ tools: undefined, stats: undefined }), + this.streamManager.createTempDirForStream(streamToken, runtime).then((tempDir) => { + recordStartupPhaseTiming("createTempDirForStreamMs", createTempDirForStreamStartedAt); + return tempDir; + }), + (async (): Promise => { + try { + if (!this.sessionUsageService) { + return undefined; + } + const sessionUsage = await this.sessionUsageService.getSessionUsage(workspaceId); + if (!sessionUsage) { + return undefined; + } + const allUsage = sumUsageHistory(Object.values(sessionUsage.byModel)); + return getTotalCost(allUsage); + } finally { + recordStartupPhaseTiming("loadSessionUsageMs", loadSessionUsageStartedAt); + } + })(), + ]); - mcpTools = result.tools; - mcpStats = result.stats; - } catch (error) { - workspaceLog.error("Failed to start MCP servers", { error }); - } finally { - mcpSetupDurationMs = Date.now() - mcpToolSetupStartedAt; - startupPhaseTimingsMs.mcpToolSetupMs = mcpSetupDurationMs; - } - } + const mcpTools = mcpSetupResult.tools; + const mcpStats = mcpSetupResult.stats; // Tool search (tool-search experiment): assembly-time gate. The runtime // holder makes getToolsForModel create the tool_catalog_search tool; its `state` @@ -1711,33 +1748,6 @@ export class AIService extends EventEmitter { const toolSearchRuntime: ToolSearchRuntime | undefined = toolSearchExperimentEnabled && Object.keys(mcpTools ?? {}).length > 0 ? {} : undefined; - const createTempDirForStreamStartedAt = Date.now(); - const runtimeTempDir = await this.streamManager.createTempDirForStream(streamToken, runtime); - recordStartupPhaseTiming("createTempDirForStreamMs", createTempDirForStreamStartedAt); - - // Extract tool-specific instructions from AGENTS.md files and agent definition - const readToolInstructionsStartedAt = Date.now(); - const toolInstructions = await readToolInstructions( - metadata, - runtime, - workspacePath, - capabilityModelString, - agentSystemPromptSections - ); - recordStartupPhaseTiming("readToolInstructionsMs", readToolInstructionsStartedAt); - - // Calculate cumulative session costs for MUX_COSTS_USD env var - let sessionCostsUsd: number | undefined; - const loadSessionUsageStartedAt = Date.now(); - if (this.sessionUsageService) { - const sessionUsage = await this.sessionUsageService.getSessionUsage(workspaceId); - if (sessionUsage) { - const allUsage = sumUsageHistory(Object.values(sessionUsage.byModel)); - sessionCostsUsd = getTotalCost(allUsage); - } - } - recordStartupPhaseTiming("loadSessionUsageMs", loadSessionUsageStartedAt); - // Get model-specific tools with workspace path (correct for local or remote) emitStartupBreadcrumb("loading_tools"); const getToolsForModelStartedAt = Date.now(); diff --git a/src/node/services/streamContextBuilder.ts b/src/node/services/streamContextBuilder.ts index 583363b5c32..f675d70aa1e 100644 --- a/src/node/services/streamContextBuilder.ts +++ b/src/node/services/streamContextBuilder.ts @@ -43,7 +43,9 @@ import { isAgentEffectivelyDisabled } from "@/node/services/agentDefinitions/age import { resolveAgentInheritanceChain } from "@/node/services/agentDefinitions/resolveAgentInheritanceChain"; import { discoverAgentSkills } from "@/node/services/agentSkills/agentSkillsService"; import { resolveSkillStorageContext } from "@/node/services/agentSkills/skillStorageContext"; -import { buildSystemMessage } from "./systemMessage"; +import { buildSystemMessage, loadInstructionSources } from "./systemMessage"; +import { resolveWorkspaceRootPath } from "@/node/runtime/runtimeHelpers"; +import type { InstructionSources } from "@/common/types/instructions"; import { getTokenizerForModel } from "@/node/utils/main/tokenizer"; import { resolveModelForMetadata } from "@/common/utils/providers/modelEntries"; import { log } from "./log"; @@ -217,7 +219,6 @@ export async function buildPlanInstructions( } } } - return { effectiveAdditionalInstructions, planFilePath, planContentForTransition }; } @@ -292,6 +293,7 @@ export interface StreamSystemContextResult { availableSkills: Awaited> | undefined; /** Exact ancestor plan files surfaced in the prompt and forwarded through tool configuration. */ ancestorPlanFilePaths: string[]; + instructionSources: InstructionSources; } const MAX_ANCESTOR_PLAN_PATH_HOPS = 32; @@ -441,7 +443,6 @@ function resolveAncestorPlanContext(args: { }); ancestorPlanFilePaths.push(normalizedPlanFilePath); } - return { entries: filteredEntries, ancestorPlanFilePaths, @@ -526,85 +527,79 @@ export async function buildStreamSystemContext( const workspaceLog = log.withFields({ workspaceId, workspaceName: metadata.name }); - // Resolve the body with inheritance (prompt.append merges with base). - // Use agentDefinition.id (may have fallen back to exec) instead of effectiveAgentId. - const resolvedBody = await resolveAgentBody( - agentDiscoveryRuntime, - agentDiscoveryPath, - agentDefinition.id, - { - skipScopesAbove: getSkipScopesAboveForKnownScope(agentDefinition.scope), - } - ); - - let subagentAppendPrompt: string | undefined; - if (isSubagentWorkspace) { - try { - const resolvedFrontmatter = await resolveAgentFrontmatter( - agentDiscoveryRuntime, - agentDiscoveryPath, - agentDefinition.id, - { - skipScopesAbove: getSkipScopesAboveForKnownScope(agentDefinition.scope), - } - ); - subagentAppendPrompt = resolvedFrontmatter.subagent?.append_prompt; - } catch (error: unknown) { - workspaceLog.debug("Failed to resolve agent frontmatter for subagent append_prompt", { - agentId: agentDefinition.id, - error: getErrorMessage(error), - }); - } - } + const skillCtx = resolveSkillStorageContext({ + runtime, + workspacePath, + muxScope, + includeClaudeSkills: opts.claudeSkillsCompatEnabled, + }); + const workspaceRootPath = metadata.subProjectPath?.trim() + ? resolveWorkspaceRootPath(metadata, runtime) + : workspacePath; + const [ + resolvedBody, + subagentAppendPrompt, + agentDefinitions, + availableSkills, + instructionSources, + ] = await Promise.all([ + resolveAgentBody(agentDiscoveryRuntime, agentDiscoveryPath, agentDefinition.id, { + skipScopesAbove: getSkipScopesAboveForKnownScope(agentDefinition.scope), + }), + isSubagentWorkspace + ? (async (): Promise => { + try { + const resolvedFrontmatter = await resolveAgentFrontmatter( + agentDiscoveryRuntime, + agentDiscoveryPath, + agentDefinition.id, + { + skipScopesAbove: getSkipScopesAboveForKnownScope(agentDefinition.scope), + } + ); + return resolvedFrontmatter.subagent?.append_prompt; + } catch (error: unknown) { + workspaceLog.debug("Failed to resolve agent frontmatter for subagent append_prompt", { + agentId: agentDefinition.id, + error: getErrorMessage(error), + }); + return undefined; + } + })() + : Promise.resolve(undefined), + !isSubagentWorkspace + ? discoverAvailableSubagentsForToolContext({ + runtime: agentDiscoveryRuntime, + workspacePath: agentDiscoveryPath, + cfg, + loadDesktopCapability, + }) + : Promise.resolve(undefined), + (async () => { + try { + return await discoverAgentSkills(skillCtx.runtime, skillCtx.workspacePath, { + roots: skillCtx.roots, + containment: skillCtx.containment, + includeClaudeSkills: opts.claudeSkillsCompatEnabled, + }); + } catch (error) { + workspaceLog.warn("Failed to discover agent skills for tool description", { error }); + return undefined; + } + })(), + loadInstructionSources(metadata, runtime, workspaceRootPath), + ]); const agentSystemPromptSections = [resolvedBody]; if (isSubagentWorkspace && subagentAppendPrompt) { agentSystemPromptSections.push(subagentAppendPrompt); } if (advisorToolAvailable) { - // Keep prompt guidance in lockstep with actual tool availability for the agent. agentSystemPromptSections.push(buildAdvisorGuidanceSection()); } if (opts.memoryToolAvailable) { - // Same lockstep rule: the post-policy system-context rebuild strips this - // section when tool policy removes the memory tool. agentSystemPromptSections.push(buildMemoryGuidanceSection()); } - - // Discover available agent definitions for sub-agent context (only for top-level workspaces). - // - // NOTE: discoverAgentDefinitions returns disabled agents too, so Settings can surface them. - // For tool descriptions (task tool), filter to agents that are effectively enabled. - let agentDefinitions: Awaited> | undefined; - if (!isSubagentWorkspace) { - agentDefinitions = await discoverAvailableSubagentsForToolContext({ - runtime: agentDiscoveryRuntime, - workspacePath: agentDiscoveryPath, - cfg, - loadDesktopCapability, - }); - } - - // Discover available skills for tool description context - const skillCtx = resolveSkillStorageContext({ - runtime, - workspacePath, - muxScope, - includeClaudeSkills: opts.claudeSkillsCompatEnabled, - }); - - let availableSkills: Awaited> | undefined; - try { - availableSkills = await discoverAgentSkills(skillCtx.runtime, skillCtx.workspacePath, { - roots: skillCtx.roots, - containment: skillCtx.containment, - // Used only for the project-runtime default-roots fallback (skillCtx.roots undefined). - includeClaudeSkills: opts.claudeSkillsCompatEnabled, - }); - } catch (error) { - workspaceLog.warn("Failed to discover agent skills for tool description", { error }); - } - const ancestorPlanContext = resolveAncestorPlanContext({ metadata, workspaceId, @@ -618,8 +613,6 @@ export async function buildStreamSystemContext( formatAncestorPlanPathInstructions(ancestorPlanContext.entries), effectiveAdditionalInstructions ); - - // Build system message from workspace metadata let systemMessage = await buildSystemMessage( metadata, runtime, @@ -627,27 +620,18 @@ export async function buildStreamSystemContext( mergedAdditionalInstructions, modelString, mcpServers, - // "Mode: " sections in Mux-dedicated instruction sources match the - // effective mode (so "Mode: plan" also covers custom plan-like agents) - // and the agent id (so per-agent sections work). The effective mode names - // the injected tag; agentDefinition.id (may have fallen back - // to exec) is the prompt actually in effect. - { agentSystemPromptSections, modes: [effectiveMode, agentDefinition.id] } + { + agentSystemPromptSections, + modes: [effectiveMode, agentDefinition.id], + instructionSources, + } ); - - // Append the hot-memories block (memory-hot-set sub-experiment). Placed at - // the end of the system message so the most recent stable prompt prefix - // stays byte-identical for provider prompt caching. The memory index lives - // in the memory tool description (same disclosure mechanic as skills). if (opts.memoryToolAvailable && opts.hotMemoriesBlock) { systemMessage = `${systemMessage}\n\n${opts.hotMemoriesBlock}`; } - - // Count system message tokens for cost tracking const metadataModel = resolveModelForMetadata(modelString, providersConfig ?? null); const tokenizer = await getTokenizerForModel(modelString, metadataModel); const systemMessageTokens = await tokenizer.countTokens(systemMessage); - return { agentSystemPromptSections, systemMessage, @@ -655,6 +639,7 @@ export async function buildStreamSystemContext( agentDefinitions, availableSkills, ancestorPlanFilePaths: ancestorPlanContext.ancestorPlanFilePaths, + instructionSources, }; } diff --git a/src/node/services/systemMessage.ts b/src/node/services/systemMessage.ts index 3d5a51cc6a7..c150c92ece1 100644 --- a/src/node/services/systemMessage.ts +++ b/src/node/services/systemMessage.ts @@ -289,6 +289,26 @@ export function extractToolInstructions( * @param agentInstructions - Optional agent definition body (searched first for tool sections) * @returns Map of tool names to their additional instructions */ +export function toolInstructionsFromSources( + sources: InstructionSources, + metadata: WorkspaceMetadata, + modelString: string, + agentInstructions?: readonly string[] +): Record { + return extractToolInstructions( + collectInstructionContents([sources.global]), + collectInstructionContents(sources.context), + modelString, + { + ...getToolAvailabilityOptions({ + workspaceId: metadata.id, + parentWorkspaceId: metadata.parentWorkspaceId, + }), + agentInstructions, + } + ); +} + export async function readToolInstructions( metadata: WorkspaceMetadata, runtime: Runtime, @@ -296,21 +316,9 @@ export async function readToolInstructions( modelString: string, agentInstructions?: readonly string[] ): Promise> { - // Tool instructions read the same `AGENTS.md` files as the system prompt; - // anchor at the workspace root so sub-project workspaces still see parent - // project tool sections (see `loadInstructionSources` doc). const workspaceRootPath = subProjectAwareWorkspaceRoot(metadata, runtime, workspacePath); const sources = await loadInstructionSources(metadata, runtime, workspaceRootPath); - const globalContents = collectInstructionContents([sources.global]); - const contextContents = collectInstructionContents(sources.context); - - return extractToolInstructions(globalContents, contextContents, modelString, { - ...getToolAvailabilityOptions({ - workspaceId: metadata.id, - parentWorkspaceId: metadata.parentWorkspaceId, - }), - agentInstructions, - }); + return toolInstructionsFromSources(sources, metadata, modelString, agentInstructions); } /** @@ -461,11 +469,12 @@ export async function loadInstructionSources( // (root + subProject) for a sub-project workspace would silently lose the // parent project's AGENTS.md, so we require root explicitly. See // `resolveWorkspaceRootPath` in `@/node/runtime/runtimeHelpers`. - const global = await readInstructionSet(getSystemDirectory(), INSTRUCTION_SCOPE.GLOBAL); - const context = isMultiProject(metadata) - ? await readMultiProjectContextInstructions(metadata, runtime, workspaceRootPath) - : await readSingleProjectContextInstructions(metadata, runtime, workspaceRootPath); - + const [global, context] = await Promise.all([ + readInstructionSet(getSystemDirectory(), INSTRUCTION_SCOPE.GLOBAL), + isMultiProject(metadata) + ? readMultiProjectContextInstructions(metadata, runtime, workspaceRootPath) + : readSingleProjectContextInstructions(metadata, runtime, workspaceRootPath), + ]); return { global, context }; } @@ -517,6 +526,7 @@ export async function buildSystemMessage( * injected tag. Duplicates are ignored. */ modes?: readonly string[]; + instructionSources?: InstructionSources; } ): Promise { if (!metadata) throw new Error("Invalid workspace metadata: metadata is required"); @@ -552,7 +562,9 @@ export async function buildSystemMessage( // back to the resolved root so the parent project's AGENTS.md is still read. // For non-sub-project workspaces this is a no-op (root === execution path). const workspaceRootPath = subProjectAwareWorkspaceRoot(metadata, runtime, workspacePath); - const instructionSources = await loadInstructionSources(metadata, runtime, workspaceRootPath); + const instructionSources = + options?.instructionSources ?? + (await loadInstructionSources(metadata, runtime, workspaceRootPath)); // Mux-dedicated per-file contents (/.mux/AGENTS.md context files, then // the global ~/.mux/AGENTS.md set, which is Mux-dedicated by construction). // Scoped Model:/Mode: directives are honored ONLY in Mux-dedicated sources diff --git a/src/node/utils/main/instructionFiles.ts b/src/node/utils/main/instructionFiles.ts index 971db0815b1..2283a69fbe8 100644 --- a/src/node/utils/main/instructionFiles.ts +++ b/src/node/utils/main/instructionFiles.ts @@ -147,51 +147,49 @@ async function readInstructionSetWith( // are honored there and we must not look for a nested ~/.mux/.mux/AGENTS.md. const isGlobalScope = scope === INSTRUCTION_SCOPE.GLOBAL; - const base = await readBaseInstructionFile(reader, directory, scope, projectName, isGlobalScope); - - const local = base.exists - ? await readSingleFile( - reader, - directory, - LOCAL_INSTRUCTION_FILENAME, - scope, - true, - projectName, - isGlobalScope - ) - : ({ exists: false } satisfies ReadInstructionFileResult); - - // Mux-dedicated companion: /.mux/AGENTS.md (+ .local.md). Read - // independently of the shared base file so a repo can provide only - // Mux-specific instructions. Skipped for the global set (see above). - let muxBase: ReadInstructionFileResult = { exists: false }; - let muxLocal: ReadInstructionFileResult = { exists: false }; - if (!isGlobalScope) { - const muxDirectory = path.join(directory, MUX_INSTRUCTION_SUBDIR); - muxBase = await readSingleFile( - reader, - muxDirectory, - MUX_INSTRUCTION_FILENAME, - scope, - false, - projectName, - true - ); - if (muxBase.exists) { - muxLocal = await readSingleFile( - reader, - muxDirectory, - LOCAL_INSTRUCTION_FILENAME, - scope, - true, - projectName, - true - ); - } - } - + const muxDirectory = isGlobalScope ? null : path.join(directory, MUX_INSTRUCTION_SUBDIR); + // Shared base + mux base are independent; parallelize to cut remote round-trips. + const [base, muxBase] = await Promise.all([ + readBaseInstructionFile(reader, directory, scope, projectName, isGlobalScope), + muxDirectory + ? readSingleFile( + reader, + muxDirectory, + MUX_INSTRUCTION_FILENAME, + scope, + false, + projectName, + true + ) + : Promise.resolve({ exists: false } satisfies ReadInstructionFileResult), + ]); if (!base.exists && !muxBase.exists) return null; + const [local, muxLocal] = await Promise.all([ + base.exists + ? readSingleFile( + reader, + directory, + LOCAL_INSTRUCTION_FILENAME, + scope, + true, + projectName, + isGlobalScope + ) + : Promise.resolve({ exists: false } satisfies ReadInstructionFileResult), + muxDirectory && muxBase.exists + ? readSingleFile( + reader, + muxDirectory, + LOCAL_INSTRUCTION_FILENAME, + scope, + true, + projectName, + true + ) + : Promise.resolve({ exists: false } satisfies ReadInstructionFileResult), + ]); + const files: InstructionFile[] = [ base.exists ? base.file : null, local.exists ? local.file : null,