From 49627b522aa4ddce7326563cc256a96ddc1993c5 Mon Sep 17 00:00:00 2001 From: Yann Cabral Date: Mon, 24 Aug 2026 21:19:20 +0000 Subject: [PATCH] perf: reuse instruction sources during stream startup --- src/node/services/aiService.test.ts | 54 ++++++- src/node/services/aiService.ts | 28 ++-- .../services/streamContextBuilder.test.ts | 45 ++++++ src/node/services/streamContextBuilder.ts | 28 +++- src/node/services/systemMessage.ts | 139 ++++++++++++------ 5 files changed, 226 insertions(+), 68 deletions(-) diff --git a/src/node/services/aiService.test.ts b/src/node/services/aiService.test.ts index 4a9806bc6e..7c4191ef31 100644 --- a/src/node/services/aiService.test.ts +++ b/src/node/services/aiService.test.ts @@ -304,11 +304,17 @@ function stubCommonStreamMessageDependencies(args: { useRequestedModelString?: boolean; onPlanPayloadMessageIds?: (messageIds: string[]) => void; onBuildStreamSystemContext?: ( - args: Parameters[0] + args: Parameters[0], + instructionSources: Awaited< + ReturnType + >["instructionSources"] ) => void; onPrepareMessagesForProvider?: ( args: Parameters[0] ) => void; + onExtractToolInstructions?: ( + sources: Parameters[0] + ) => void; }): ReturnType> { spyOn(agentResolution, "resolveAgentForStream").mockResolvedValue( resolvedAgentResultFor(args.metadata) @@ -322,8 +328,10 @@ function stubCommonStreamMessageDependencies(args: { }); }); spyOn(streamContextBuilder, "buildStreamSystemContext").mockImplementation((contextArgs) => { - args.onBuildStreamSystemContext?.(contextArgs); + const instructionSources = { global: [], context: [] }; + args.onBuildStreamSystemContext?.(contextArgs, instructionSources); return Promise.resolve({ + instructionSources, agentSystemPromptSections: ["test-agent-prompt"], systemMessage: "test-system-message", systemMessageTokens: 1, @@ -343,7 +351,10 @@ function stubCommonStreamMessageDependencies(args: { const getToolsForModelSpy = spyOn(toolsModule, "getToolsForModel").mockResolvedValue( args.allTools ?? {} ); - spyOn(systemMessageModule, "readToolInstructions").mockResolvedValue({}); + spyOn(systemMessageModule, "extractToolInstructionsFromSources").mockImplementation((sources) => { + args.onExtractToolInstructions?.(sources); + return {}; + }); const providerModelFactory = Reflect.get(args.service, "providerModelFactory") as | ProviderModelFactory @@ -1442,6 +1453,43 @@ describe("AIService.streamMessage compaction boundary slicing", () => { mock.restore(); }); + it("reuses the stream context instruction snapshot for tool extraction", async () => { + using xumHome = new DisposableTempDir("ai-service-instruction-snapshot"); + const projectPath = path.join(xumHome.path, "project"); + await fs.mkdir(projectPath, { recursive: true }); + + const workspaceId = "workspace-instruction-snapshot"; + const metadata = createLocalWorkspaceMetadata(workspaceId, projectPath); + const { config, historyService, initStateManager, service } = createBasicAIService( + xumHome.path + ); + let builtSources: unknown; + let extractedSources: unknown; + stubCommonStreamMessageDependencies({ + service, + config, + historyService, + initStateManager, + metadata, + onBuildStreamSystemContext: (_args, instructionSources) => { + builtSources = instructionSources; + }, + onExtractToolInstructions: (instructionSources) => { + extractedSources = instructionSources; + }, + }); + + const result = await service.streamMessage({ + messages: [createMuxMessage("latest-user", "user", "hello")], + workspaceId, + modelString: "openai:gpt-5.2", + thinkingLevel: "off", + }); + + expect(result.success).toBe(true); + expect(extractedSources).toBe(builtSources); + }); + it("keeps set_goal disabled for one-shot streams that do not opt into agent-created goals", async () => { using xumHome = new DisposableTempDir("ai-service-set-goal-disabled"); const projectPath = path.join(xumHome.path, "project"); diff --git a/src/node/services/aiService.ts b/src/node/services/aiService.ts index 5a1e543d85..d900beec6d 100644 --- a/src/node/services/aiService.ts +++ b/src/node/services/aiService.ts @@ -87,7 +87,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 { extractToolInstructionsFromSources } from "./systemMessage"; import { effectiveAdditionalSystemContext, mergeAdditionalSystemInstructions, @@ -176,7 +176,11 @@ import { DEVTOOLS_RUN_METADATA_ID_HEADER } from "./devToolsHeaderCapture"; import { ProviderModelFactory, modelCostsIncluded } from "./providerModelFactory"; import { prepareMessagesForProvider } from "./messagePipeline"; import { getLegacyModeForAgentMetadata, resolveAgentForStream } from "./agentResolution"; -import { buildPlanInstructions, buildStreamSystemContext } from "./streamContextBuilder"; +import { + buildPlanInstructions, + buildStreamSystemContext, + type StreamSystemContextResult, +} from "./streamContextBuilder"; import { getTokenizerForModel } from "@/node/utils/main/tokenizer"; import { normalizeUsageModelKey, @@ -2142,12 +2146,16 @@ export class AIService extends EventEmitter { // service); tool policy may still strip the tool, which forces a rebuild // below so the prompt never advertises an absent tool. const memoryToolEligible = memoryExperimentEnabled && this.memoryService !== undefined; + const startupInstructionSources: { + current?: StreamSystemContextResult["instructionSources"]; + } = {}; const buildStreamSystemContextForToolset = ( toolset: { advisorToolAvailable: boolean; memoryToolAvailable: boolean }, modelStringForSystem: string = modelString, contextForModel: MemorySessionContext | undefined = memoryContext ) => buildStreamSystemContext({ + instructionSources: startupInstructionSources.current, runtime, metadata, workspacePath, @@ -2180,8 +2188,10 @@ export class AIService extends EventEmitter { advisorToolAvailable: advisorToolEligible, memoryToolAvailable: memoryToolEligible, }); + startupInstructionSources.current = prePolicyStreamSystemContext.instructionSources; recordStartupPhaseTiming("buildStreamSystemContextMs", buildStreamSystemContextStartedAt); const { + instructionSources, agentSystemPromptSections, agentDefinitions, availableSkills, @@ -2250,16 +2260,14 @@ export class AIService extends EventEmitter { const runtimeTempDir = await this.streamManager.createTempDirForStream(streamToken, runtime); recordStartupPhaseTiming("createTempDirForStreamMs", createTempDirForStreamStartedAt); - // Extract tool-specific instructions from AGENTS.md files and agent definition + // Extract tool-specific instructions from the same source snapshot used + // for the system message so instruction files are not read twice. const readToolInstructionsStartedAt = Date.now(); - const toolInstructions = await readToolInstructions( - metadata, - runtime, - workspacePath, + const toolInstructions = extractToolInstructionsFromSources( + instructionSources, capabilityModelString, - agentSystemPromptSections, - cfg.projects, - claudeSkillsCompatExperimentEnabled + metadata, + agentSystemPromptSections ); recordStartupPhaseTiming("readToolInstructionsMs", readToolInstructionsStartedAt); diff --git a/src/node/services/streamContextBuilder.test.ts b/src/node/services/streamContextBuilder.test.ts index a85668b0b7..557adc874c 100644 --- a/src/node/services/streamContextBuilder.test.ts +++ b/src/node/services/streamContextBuilder.test.ts @@ -13,6 +13,7 @@ import { getPlanFilePath } from "@/common/utils/planStorage"; import { LocalRuntime } from "@/node/runtime/LocalRuntime"; import { DisposableTempDir } from "@/node/services/tempDir"; +import { extractToolInstructionsFromSources } from "./systemMessage"; import { buildPlanInstructions, buildStreamSystemContext } from "./streamContextBuilder"; class TestRuntime extends LocalRuntime { @@ -262,6 +263,50 @@ class RestrictedTestRuntime extends TestRuntime { } describe("buildStreamSystemContext", () => { + test("returns the instruction source snapshot used to build the system message", async () => { + using tempRoot = new DisposableTempDir("stream-system-context-instruction-snapshot"); + + const projectPath = path.join(tempRoot.path, "project"); + const xumHome = path.join(tempRoot.path, "xum-home"); + await fs.mkdir(path.join(projectPath, ".xum"), { recursive: true }); + await fs.mkdir(xumHome, { recursive: true }); + await fs.writeFile( + path.join(projectPath, ".xum", "AGENTS.md"), + ["Project prompt guidance.", "", "## Tool: bash", "Use the project bash workflow.", ""].join( + "\n" + ) + ); + + const metadata = createWorkspaceMetadata({ + id: "instruction-snapshot-ws", + name: "instruction-snapshot-workspace", + projectName: "project", + projectPath, + }); + const cfg = createProjectsConfig({ + projectPath, + workspaces: [{ id: metadata.id, name: metadata.name }], + }); + + const result = await buildSystemContextForTest({ + runtime: new TestRuntime(projectPath, xumHome), + metadata, + workspacePath: projectPath, + cfg, + isSubagentWorkspace: false, + }); + + expect(result.systemMessage).toContain("Project prompt guidance."); + expect( + extractToolInstructionsFromSources( + result.instructionSources, + "openai:gpt-5.2", + metadata, + result.agentSystemPromptSections + ).bash + ).toContain("Use the project bash workflow."); + }); + test("includes proactive memory guidance only when the memory tool is available", async () => { using tempRoot = new DisposableTempDir("stream-system-context-memory-guidance"); diff --git a/src/node/services/streamContextBuilder.ts b/src/node/services/streamContextBuilder.ts index 0fe0a61e4a..5b006c098a 100644 --- a/src/node/services/streamContextBuilder.ts +++ b/src/node/services/streamContextBuilder.ts @@ -22,6 +22,7 @@ import type { DesktopCapability } from "@/common/types/desktop"; import type { ProjectsConfig } from "@/common/types/project"; import type { XumToolScope } from "@/common/types/toolScope"; import type { AgentDefinitionScope } from "@/common/types/agentDefinition"; +import type { InstructionSources } from "@/common/types/instructions"; import type { WorkspaceMetadata } from "@/common/types/workspace"; import type { ProvidersConfigMap } from "@/common/orpc/types"; import type { TaskSettings } from "@/common/types/tasks"; @@ -43,7 +44,7 @@ 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 { buildSystemMessageFromSources, loadWorkspaceInstructionSources } from "./systemMessage"; import { getTokenizerForModel } from "@/node/utils/main/tokenizer"; import { resolveModelForMetadata } from "@/common/utils/providers/modelEntries"; import { log } from "./log"; @@ -227,6 +228,8 @@ export async function buildPlanInstructions( /** Options for building the system message context. */ export interface BuildStreamSystemContextOptions { + /** Reuse a source snapshot across policy/model-driven context rebuilds in one startup. */ + instructionSources?: InstructionSources; runtime: Runtime; metadata: WorkspaceMetadata; workspacePath: string; @@ -249,7 +252,7 @@ export interface BuildStreamSystemContextOptions { modelString: string; cfg: ProjectsConfig; providersConfig?: ProvidersConfigMap | null; - mcpServers: Parameters[5]; + mcpServers: Parameters[5]; xumScope?: XumToolScope; loadDesktopCapability?: () => Promise; /** Whether the advisor tool is available for the current agent */ @@ -284,6 +287,8 @@ export interface StreamSystemContextResult { * trailing scoped heading in one section swallow the next section's text. */ agentSystemPromptSections: string[]; + /** Instruction sources loaded once for both prompt and tool-scoped extraction. */ + instructionSources: InstructionSources; /** Full system message string. */ systemMessage: string; /** Token count of the system message. */ @@ -626,10 +631,22 @@ export async function buildStreamSystemContext( effectiveAdditionalInstructions ); + // Load once so prompt assembly and tool-scoped extraction observe the same + // instruction snapshot without reading workspace files twice per startup. + const instructionSources = + opts.instructionSources ?? + (await loadWorkspaceInstructionSources( + metadata, + runtime, + workspacePath, + cfg.projects, + opts.claudeSkillsCompatEnabled + )); + // Build system message from workspace metadata - let systemMessage = await buildSystemMessage( + let systemMessage = buildSystemMessageFromSources( metadata, - runtime, + instructionSources, workspacePath, mergedAdditionalInstructions, modelString, @@ -642,8 +659,6 @@ export async function buildStreamSystemContext( { agentSystemPromptSections, modes: [effectiveMode, agentDefinition.id], - projectConfigs: cfg.projects, - claudeSkillsCompatEnabled: opts.claudeSkillsCompatEnabled, } ); @@ -662,6 +677,7 @@ export async function buildStreamSystemContext( return { agentSystemPromptSections, + instructionSources, systemMessage, systemMessageTokens, agentDefinitions, diff --git a/src/node/services/systemMessage.ts b/src/node/services/systemMessage.ts index 5e3aa5de7d..957dd3d8e2 100644 --- a/src/node/services/systemMessage.ts +++ b/src/node/services/systemMessage.ts @@ -305,6 +305,28 @@ 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 extractToolInstructionsFromSources( + sources: InstructionSources, + modelString: string, + metadata: WorkspaceMetadata, + agentInstructions?: readonly string[] +): Record { + // Tool extraction joins sources highest-precedence first (agent → context → + // global), the opposite of prompt order. `sources.global` is compat-first + // for the prompt, so reverse it here to keep native guidance ahead of the + // ~/.claude/CLAUDE.md compatibility source. + const globalContents = collectInstructionContents([...sources.global].reverse()); + const contextContents = collectInstructionContents(sources.context); + + return extractToolInstructions(globalContents, contextContents, modelString, { + ...getToolAvailabilityOptions({ + workspaceId: metadata.id, + parentWorkspaceId: metadata.parentWorkspaceId, + }), + agentInstructions, + }); +} + export async function readToolInstructions( metadata: WorkspaceMetadata, runtime: Runtime, @@ -325,20 +347,7 @@ export async function readToolInstructions( projectConfigs, claudeSkillsCompatEnabled ); - // Tool extraction joins sources highest-precedence first (agent → context → - // global), the opposite of prompt order. `sources.global` is compat-first - // for the prompt, so reverse it here to keep native guidance ahead of the - // ~/.claude/CLAUDE.md compatibility source. - const globalContents = collectInstructionContents([...sources.global].reverse()); - const contextContents = collectInstructionContents(sources.context); - - return extractToolInstructions(globalContents, contextContents, modelString, { - ...getToolAvailabilityOptions({ - workspaceId: metadata.id, - parentWorkspaceId: metadata.parentWorkspaceId, - }), - agentInstructions, - }); + return extractToolInstructionsFromSources(sources, modelString, metadata, agentInstructions); } /** @@ -481,6 +490,22 @@ function deriveSubProjectRelativePath(projectPath: string, subProjectPath: strin * @param claudeSkillsCompatEnabled - Whether to include ~/.claude/CLAUDE.md before native globals * @returns Structured instruction sources (ordered global and context entries) */ +export async function loadWorkspaceInstructionSources( + metadata: WorkspaceMetadata, + runtime: Runtime, + workspacePath: string, + projectConfigs?: Map, + claudeSkillsCompatEnabled = false +): Promise { + return loadInstructionSources( + metadata, + runtime, + subProjectAwareWorkspaceRoot(metadata, runtime, workspacePath), + projectConfigs, + claudeSkillsCompatEnabled + ); +} + export async function loadInstructionSources( metadata: WorkspaceMetadata, runtime: Runtime, @@ -589,6 +614,25 @@ function buildProjectSettingsInstructionSets( * @param mcpServers - Optional MCP server configuration (name -> command) * @throws Error if metadata or workspacePath invalid */ +export interface BuildSystemMessageFromSourcesOptions { + /** + * Resolved agent prompt as independently-authored sections (agent body, + * subagent append_prompt, advisor guidance, …). Per-section so a trailing + * scoped heading in one section cannot swallow the next section's text. + */ + agentSystemPromptSections?: readonly string[]; + /** + * Active mode identifiers used to extract "Mode: " sections from + * Xum-dedicated instruction sources. The first entry names the injected tag. + */ + modes?: readonly string[]; +} + +interface BuildSystemMessageOptions extends BuildSystemMessageFromSourcesOptions { + projectConfigs?: Map; + claudeSkillsCompatEnabled?: boolean; +} + export async function buildSystemMessage( metadata: WorkspaceMetadata, runtime: Runtime, @@ -596,33 +640,42 @@ export async function buildSystemMessage( additionalSystemInstructions?: string, modelString?: string, mcpServers?: MCPServerMap, - options?: { - /** - * Resolved agent prompt as independently-authored sections (agent body, - * subagent append_prompt, advisor guidance, …). Per-section so a trailing - * scoped heading in one section cannot swallow the next section's text. - */ - agentSystemPromptSections?: readonly string[]; - /** - * Active mode identifiers used to extract "Mode: " sections from - * Xum-dedicated instruction sources: the effective mode (plan/exec/compact) - * plus the agent id, so "Mode: plan" covers custom plan-like agents and - * "Mode: " covers per-agent sections. The first entry names the - * injected tag. Duplicates are ignored. - */ - modes?: readonly string[]; - /** - * Project configs from ~/.mux/config.json, used to append per-project - * `customInstructions` (Settings → Instructions) to the prompt. - */ - projectConfigs?: Map; - /** Read ~/.claude/CLAUDE.md as a lowest-precedence global compatibility source. */ - claudeSkillsCompatEnabled?: boolean; - } + options?: BuildSystemMessageOptions ): Promise { if (!metadata) throw new Error("Invalid workspace metadata: metadata is required"); if (!workspacePath) throw new Error("Invalid workspace path: workspacePath is required"); + const workspaceRootPath = subProjectAwareWorkspaceRoot(metadata, runtime, workspacePath); + const instructionSources = await loadInstructionSources( + metadata, + runtime, + workspaceRootPath, + options?.projectConfigs, + options?.claudeSkillsCompatEnabled + ); + return buildSystemMessageFromSources( + metadata, + instructionSources, + workspacePath, + additionalSystemInstructions, + modelString, + mcpServers, + options + ); +} + +export function buildSystemMessageFromSources( + metadata: WorkspaceMetadata, + instructionSources: InstructionSources, + workspacePath: string, + additionalSystemInstructions?: string, + modelString?: string, + mcpServers?: MCPServerMap, + options?: BuildSystemMessageFromSourcesOptions +): string { + if (!metadata) throw new Error("Invalid workspace metadata: metadata is required"); + if (!workspacePath) throw new Error("Invalid workspace path: workspacePath is required"); + // Read instruction sets // Get runtime type from metadata (defaults to "local" for legacy workspaces without runtimeConfig) const runtimeType = metadata.runtimeConfig?.type ?? "local"; @@ -648,18 +701,6 @@ export async function buildSystemMessage( // tool descriptions (agent_skill_read, task) for better model attention per Anthropic // best practices. See tools.ts ToolConfiguration.availableSkills/availableSubagents. - // Read instruction sets - // Sub-project workspaces pass the execution path (root + subProject); fall - // 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, - options?.projectConfigs, - options?.claudeSkillsCompatEnabled - ); // Xum-dedicated per-file contents (/.xum/AGENTS.md context files, then // native ~/.xum global files). Claude compatibility instructions are shared. // Scoped Model:/Mode: directives are honored ONLY in Xum-dedicated sources