diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 3e233348bc6..adb65f6cffe 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -61,6 +61,8 @@ const { buildToolApprovalExecutionConfig, collectAttachedCodeEnvironmentAgentIds, collectAttachedCodeEnvironmentPolicySettings, + collectAttachedCodeApprovalPolicies, + collectAttachedCodeRoutePolicies, buildAttachedCodeEnvironmentAdmissionHooks, resolveAttachedCodeApprovalMode, markNativeCodeToolApprovalRequests, @@ -1986,7 +1988,7 @@ class AgentClient extends BaseClient { const topLevelAgents = [this.options.agent, ...(this.agentConfigs?.values() ?? [])]; const codeApprovalMode = resolveAttachedCodeApprovalMode( this.options.req.body.codeApprovalMode, - collectAttachedCodeEnvironmentPolicySettings(topLevelAgents), + collectAttachedCodeApprovalPolicies(topLevelAgents), agentsEConfig?.toolApproval?.enabled !== false, ); const persistedCodeEnvironmentDecision = resolvePersistableCodeEnvironmentDecision({ @@ -4563,7 +4565,7 @@ class AgentClient extends BaseClient { collectAttachedCodeEnvironmentPolicySettings(topLevelAgents); const codeApprovalMode = resolveAttachedCodeApprovalMode( this.options.req.body.codeApprovalMode, - attachedCodeEnvironmentSettings, + collectAttachedCodeApprovalPolicies(topLevelAgents), agentsEConfig?.toolApproval?.enabled !== false, ); const effectiveToolApprovalPolicy = resolveToolApprovalPolicy({ @@ -4584,6 +4586,7 @@ class AgentClient extends BaseClient { attachedCodeEnvironmentAgentIds, attachedCodeEnvironmentSettings, codeApprovalMode, + collectAttachedCodeRoutePolicies(topLevelAgents), ), ]; const askUserQuestionAdminDisabled = isAskUserQuestionAdminDisabled(appConfig); diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index c9ef4f2d33d..333979348b5 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -24,6 +24,7 @@ const { buildAgentContextAttachmentsByAgentId, collectCodeExecutionProfileRoutes, getLazySubagentConfigId, + createRoutedGraphMemberLoader, createViewableSubagentLoader, resolveCodeExecutionContext, resolveCodeExecutionWorkspaceSelections, @@ -32,7 +33,7 @@ const { resolveCodeExecutionWorkspaceContext, resolveSubagentCodeWorkspaceInheritance, resolveSubagentCodeAvailability, - guardUnavailableSubagent, + guardRoutableSubagent, createStatefulCodeEnvironmentPolicyError, buildSubagentThreadTaskConfig, backgroundCompletionWakeupsEnabled, @@ -43,6 +44,7 @@ const { getCodeApiAuthHeaders, getCodeExecutionBaseUrl, getAuthorizedRunFileSnapshot, + isRunFileSharingRequested, encodeAndFormatDocuments, encodeAndFormatAudios, encodeAndFormatVideos, @@ -53,6 +55,7 @@ const { retainScheduleMCPCompletion, getScheduleMCPExecution, getMCPRequestContext, + createSubagentCodeRouting, } = require('@librechat/api'); const { ResourceType, @@ -401,6 +404,10 @@ const initializeClientWithProvider = async ({ * }>} */ const agentToolContexts = new Map(); + /** Per-call subagent machine routing; assigned once the request's code inputs are known. */ + let subagentCodeRouting; + const getRoutedToolContext = (agentId, executionContext) => + subagentCodeRouting?.getToolContext(agentId, executionContext); let runFileBindings; const resolveMcpServerName = (toolName, agentId) => { if (typeof toolName !== 'string' || typeof agentId !== 'string') { @@ -449,6 +456,7 @@ const initializeClientWithProvider = async ({ const soleContext = agentToolContexts.size === 1 ? agentToolContexts.values().next().value : null; const trustedContext = + getRoutedToolContext(executingAgentId, metadata.executionContext) ?? (typeof executingAgentId === 'string' ? agentToolContexts.get(executingAgentId) : null) ?? soleContext; const callbackMetadata = { ...metadata }; @@ -480,7 +488,10 @@ const initializeClientWithProvider = async ({ runSignal, executionContext, ) => { - const ctx = runFileBindings.getContext(agentId, executionContext) ?? {}; + const ctx = + getRoutedToolContext(agentId, executionContext) ?? + runFileBindings.getContext(agentId, executionContext) ?? + {}; logger.debug(`[ON_TOOL_EXECUTE] ctx found: ${!!ctx.userMCPAuthMap}, agent: ${ctx.agent?.id}`); logger.debug(`[ON_TOOL_EXECUTE] toolRegistry size: ${ctx.toolRegistry?.size ?? 'undefined'}`); @@ -576,6 +587,7 @@ const initializeClientWithProvider = async ({ req, agentToolContexts, resolvePrimaryAgentId: () => primaryConfig?.id, + resolveExecutionContext: getRoutedToolContext, }), }; @@ -1138,7 +1150,24 @@ const initializeClientWithProvider = async ({ backgroundToolsAvailable, }); - const toLazySubagentMetadata = async (agent) => { + subagentCodeRouting = createSubagentCodeRouting({ + allowEnvironmentSelection: + appConfig.endpoints?.agents?.statefulCodeSessions?.allowEnvironmentSelection, + persistedSelections: admittedConversation?.codeWorkspaces, + requestedSelections: runtimeRequestBody?.codeWorkspaces, + environments: appConfig?.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments, + userId, + conversationId, + getAppConfig, + getInheritedEnvironments: () => req.codeWorkspaceInheritance, + sharedRunFiles: isRunFileSharingRequested({ + policy: appConfig.endpoints?.agents?.fileSharing, + agent: primaryConfig, + }), + }); + + /** The code flags a lazy subagent runs with in this request. */ + const getSubagentCodeFlags = (agent) => { const configuredCodeEnvironments = appConfig?.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments; const attachedEnvironmentOptOut = optsOutOfAttachedCodeEnvironment( @@ -1166,6 +1195,21 @@ const initializeClientWithProvider = async ({ ) { throw createStatefulCodeEnvironmentPolicyError(statefulCodeEnvironment); } + return { + configuredCodeEnvironments, + lazyCodeEnvAvailable, + statefulCodeSessions, + statefulCodeEnvironment, + }; + }; + + const toLazySubagentMetadata = async (agent) => { + const { + configuredCodeEnvironments, + lazyCodeEnvAvailable, + statefulCodeSessions, + statefulCodeEnvironment, + } = getSubagentCodeFlags(agent); const codeAvailability = await resolveSubagentCodeAvailability({ agentId: agent.id, codeEnvAvailable: lazyCodeEnvAvailable, @@ -1198,16 +1242,29 @@ const initializeClientWithProvider = async ({ }, }); const { codeExecutionContext } = codeAvailability; - const { - alwaysApplySkillPrimes, - historicalToolNames, - historicalMcpServerNames, - skillAuthoringAvailable, - } = await lazyHistoryResolver.resolve({ - agent, - codeExecutionAvailable: lazyCodeEnvAvailable, - memoryAvailable, - }); + const [ + { + alwaysApplySkillPrimes, + historicalToolNames, + historicalMcpServerNames, + skillAuthoringAvailable, + }, + { subagentHostArgs, codeExecutionChoices }, + ] = await Promise.all([ + lazyHistoryResolver.resolve({ + agent, + codeExecutionAvailable: lazyCodeEnvAvailable, + memoryAvailable, + }), + waitForAbort( + subagentCodeRouting.describe( + agent, + { statefulCodeSessions, statefulCodeEnvironment }, + signal, + ), + signal, + ), + ]); return copyToolApprovalAdmissionMetadata( { id: agent.id, @@ -1225,6 +1282,9 @@ const initializeClientWithProvider = async ({ ...codeAvailability, statefulCodeEnvironment, codeSessionKey: codeExecutionContext?.codeSessionKey, + subagentHostArgs, + codeExecutionChoices, + subagentCodeFlags: { statefulCodeSessions, statefulCodeEnvironment }, includeReasoningHistory: getIncludeReasoningHistory(agent), alwaysApplySkillPrimes, historicalToolNames, @@ -1319,6 +1379,8 @@ const initializeClientWithProvider = async ({ configId, context, lazyChildren, + codeFlags, + codeWorkspaceUnavailable, viewAccessChecked = false, }) => { throwIfAborted(context.signal); @@ -1336,6 +1398,17 @@ const initializeClientWithProvider = async ({ if (!validation.isValid) { throw new Error(validation.error?.message ?? `Subagent ${agentId} failed model validation.`); } + const placement = await waitForAbort( + subagentCodeRouting.place({ + agent, + flags: codeFlags ?? {}, + context, + unavailableReason: codeWorkspaceUnavailable, + }), + context.signal, + ); + throwIfAborted(context.signal); + const routedAgent = placement.agent; const scopedSkillIds = resolveAgentScopedSkillIds({ agent, accessibleSkillIds, @@ -1349,77 +1422,80 @@ const initializeClientWithProvider = async ({ ephemeralSkillsToggle, }); const config = await waitForAbort( - initializeAgent( - { - req, - res, - agent, - loadTools: createToolLoader( + subagentCodeRouting.settle( + placement, + initializeAgent( + { req, res, - context.signal, - streamId, - true, - jobCreatedAt, - upstreamTokenProvider, - upstreamTokenProviderResolver, - ), - requestFiles, - authorizedRunFiles: getAuthorizedRunFileSnapshot({ - policy: appConfig.endpoints?.agents?.fileSharing, - agent: primaryConfig, - files: primaryConfig.currentRequestAttachments, - }), - conversationId, - parentMessageId, - requestBody: runtimeRequestBody, - endpointOption: { ...endpointOption, endpoint: EModelEndpoint.agents }, - allowedProviders, - accessibleSkillIds: scopedSkillIds, - skillAuthoringAvailable: canAuthorSkillFiles({ - agent, - scopedEditableSkillIds, - skillCreateAllowed, - skillsCapabilityEnabled, - ephemeralSkillsToggle, - }), - codeEnvAvailable, - fileSearchAvailable, - resolveLinkedInstructions, - recordLinkedPromptUsage, - backgroundToolsAvailable, - toolIntentsAvailable, - statefulSessionsAvailable, - allowedStatefulCodeEnvironments, - memoryAvailable, - skillStates, - defaultActiveOnShare, - signal: context.signal, - }, - { - getProjectFiles: db.getProjectFiles, - getFiles: db.getFiles, - getUserKey: db.getUserKey, - getMessages: db.getMessages, - getConvoFiles: db.getConvoFiles, - getAccessibleMcpServerNames, - updateFilesUsage: db.updateFilesUsage, - getUserKeyValues: db.getUserKeyValues, - getUserCodeFiles: db.getUserCodeFiles, - getDeferredProvisionFiles: db.getDeferredProvisionFiles, - getToolFilesByIds: db.getToolFilesByIds, - getCodeGeneratedFiles: db.getCodeGeneratedFiles, - filterFilesByAgentAccess, - listSkillsByAccess: skillDbMethods.listSkillsByAccess, - listAlwaysApplySkills: skillDbMethods.listAlwaysApplySkills, - getSkillByName: skillDbMethods.getSkillByName, - provisionToCodeEnv, - provisionToVectorDB, - checkSessionsAlive, - loadCodeApiKey, - updateFile: db.updateFile, - getRoleByName: db.getRoleByName, - }, + agent: routedAgent, + loadTools: createToolLoader( + req, + res, + context.signal, + streamId, + true, + jobCreatedAt, + upstreamTokenProvider, + upstreamTokenProviderResolver, + ), + requestFiles, + authorizedRunFiles: getAuthorizedRunFileSnapshot({ + policy: appConfig.endpoints?.agents?.fileSharing, + agent: primaryConfig, + files: primaryConfig.currentRequestAttachments, + }), + conversationId, + parentMessageId, + requestBody: runtimeRequestBody, + endpointOption: { ...endpointOption, endpoint: EModelEndpoint.agents }, + allowedProviders, + accessibleSkillIds: scopedSkillIds, + skillAuthoringAvailable: canAuthorSkillFiles({ + agent, + scopedEditableSkillIds, + skillCreateAllowed, + skillsCapabilityEnabled, + ephemeralSkillsToggle, + }), + codeEnvAvailable, + fileSearchAvailable, + resolveLinkedInstructions, + recordLinkedPromptUsage, + backgroundToolsAvailable, + toolIntentsAvailable, + statefulSessionsAvailable, + allowedStatefulCodeEnvironments, + memoryAvailable, + skillStates, + defaultActiveOnShare, + signal: context.signal, + }, + { + getProjectFiles: db.getProjectFiles, + getFiles: db.getFiles, + getUserKey: db.getUserKey, + getMessages: db.getMessages, + getConvoFiles: db.getConvoFiles, + getAccessibleMcpServerNames, + updateFilesUsage: db.updateFilesUsage, + getUserKeyValues: db.getUserKeyValues, + getUserCodeFiles: db.getUserCodeFiles, + getDeferredProvisionFiles: db.getDeferredProvisionFiles, + getToolFilesByIds: db.getToolFilesByIds, + getCodeGeneratedFiles: db.getCodeGeneratedFiles, + filterFilesByAgentAccess, + listSkillsByAccess: skillDbMethods.listSkillsByAccess, + listAlwaysApplySkills: skillDbMethods.listAlwaysApplySkills, + getSkillByName: skillDbMethods.getSkillByName, + provisionToCodeEnv, + provisionToVectorDB, + checkSessionsAlive, + loadCodeApiKey, + updateFile: db.updateFile, + getRoleByName: db.getRoleByName, + }, + ), ), context.signal, ); @@ -1428,14 +1504,35 @@ const initializeClientWithProvider = async ({ if (config.userMCPAuthMap) { Object.assign(userMCPAuthMap, config.userMCPAuthMap); } - agentToolContexts.set(agentId, buildAgentToolContext({ agent, config })); + subagentCodeRouting.attach(agentToolContexts, { + agentId, + context, + placement, + codeExecutionContext: config.codeExecutionContext, + toolContext: buildAgentToolContext({ agent: routedAgent, config }), + }); endpointTokenConfigByAgentId.set(agentId, config.endpointTokenConfig); return config; }; - const initializeLazySubagent = async ({ agentId, configId, context, lazyChildren }) => { + const initializeLazySubagent = async ({ + agentId, + configId, + context, + lazyChildren, + codeFlags, + codeWorkspaceUnavailable, + }) => { throwIfAborted(context.signal); const agent = await waitForAbort(db.getAgentWithVersionCount({ id: agentId }), context.signal); - return initializeLoadedSubagent({ agent, agentId, configId, context, lazyChildren }); + return initializeLoadedSubagent({ + agent, + agentId, + configId, + context, + lazyChildren, + codeFlags, + codeWorkspaceUnavailable, + }); }; const buildLazySubagentDescriptors = async (agent, depth = 0, ancestors = new Set()) => { @@ -1509,6 +1606,8 @@ const initializeClientWithProvider = async ({ statefulCodeEnvironment: metadata.statefulCodeEnvironment, codeExecutionContext: metadata.codeExecutionContext, codeSessionKey: metadata.codeSessionKey, + subagentHostArgs: metadata.subagentHostArgs, + codeExecutionChoices: metadata.codeExecutionChoices, includeReasoningHistory: metadata.includeReasoningHistory, skillAuthoringAvailable: metadata.skillAuthoringAvailable, alwaysApplySkillPrimes: metadata.alwaysApplySkillPrimes, @@ -1517,19 +1616,34 @@ const initializeClientWithProvider = async ({ lazySubagentConfigs: lazyChildren, subagentAgentConfigs: eagerChildren, subagentGraphMemberMetadata, - ...guardUnavailableSubagent({ + settle: (context, resolveInputs) => + subagentCodeRouting.settleExecution(context, resolveInputs, { + onCommit: () => publishSharedConfig(context), + }), + ...guardRoutableSubagent({ description: metadata.description, codeWorkspaceUnavailable: metadata.codeWorkspaceUnavailable, + subagentHostArgs: metadata.subagentHostArgs, resolve: async (context) => { const config = await initializeLazySubagent({ agentId: metadata.id, configId: metadata.configId, context, lazyChildren, + codeFlags: metadata.subagentCodeFlags, + codeWorkspaceUnavailable: metadata.codeWorkspaceUnavailable, }); config.subagentAgentConfigs = eagerChildren; - graphMemberConfigsById.set(config.id, config); - await resolveGraphSubagentsFor(config, context.signal); + if (!subagentCodeRouting.isRouted(context.executionId)) { + sharedConfigByCall.set(context, config); + } + await resolveGraphSubagentsFor( + config, + context.signal, + subagentCodeRouting.routesChildren(context.executionId) + ? context.executionId + : undefined, + ); return config; }, }), @@ -1576,6 +1690,14 @@ const initializeClientWithProvider = async ({ const graphMemberConfigsById = new Map( rootSubagentConfigs.filter((config) => config?.id).map((config) => [config.id, config]), ); + /** A lazy call's config joins the shared graph cache only once its route commits. */ + const sharedConfigByCall = new WeakMap(); + const publishSharedConfig = (context) => { + const config = sharedConfigByCall.get(context); + if (config != null && !graphMemberConfigsById.has(config.id)) { + graphMemberConfigsById.set(config.id, config); + } + }; const graphMemberLoadsById = new Map(); const initializeGraphMember = createConcurrencyLimiter(SUBAGENT_GRAPH_LOAD_CONCURRENCY); const loadGraphMemberOnce = async (memberId) => { @@ -1597,6 +1719,7 @@ const initializeClientWithProvider = async ({ configId: getLazySubagentConfigId(agent), context: { signal }, lazyChildren: [], + codeFlags: getSubagentCodeFlags(agent), viewAccessChecked: true, }); graphMemberConfigsById.set(memberId, config); @@ -1634,7 +1757,28 @@ const initializeClientWithProvider = async ({ return waitForAbort(pending, graphSignal); }; - async function resolveGraphSubagentsFor(config, graphSignal = signal) { + const loadRoutedGraphMember = createRoutedGraphMemberLoader({ + getShared: (memberId) => graphMemberConfigsById.get(memberId), + isSkipped: (memberId) => skippedAgentIds.has(memberId), + skip: (memberId) => skippedAgentIds.add(memberId), + getAgent: (memberId, graphSignal) => + waitForAbort(db.getAgentWithVersionCount({ id: memberId }), graphSignal), + canView: (agent, memberId, graphSignal) => hasSubagentViewAccess(agent, memberId, graphSignal), + initialize: ({ agent, memberId, context }) => + initializeLoadedSubagent({ + agent, + agentId: memberId, + configId: getLazySubagentConfigId(agent), + context, + lazyChildren: [], + codeFlags: getSubagentCodeFlags(agent), + viewAccessChecked: true, + }), + isFatal: (error, graphSignal) => + isFatalAgentInitializationError(error, { signal: graphSignal }), + }); + + async function resolveGraphSubagentsFor(config, graphSignal = signal, routedParentRunId) { throwIfAborted(graphSignal); const definitions = subagentsAvailableForRun && config.subagents?.enabled === true @@ -1662,7 +1806,13 @@ const initializeClientWithProvider = async ({ } const memberConfigs = await Promise.all( memberIds.map((memberId) => - initializeGraphMember(() => loadGraphMember(memberId, graphSignal)), + memberId === config.id + ? config + : initializeGraphMember(() => + routedParentRunId == null + ? loadGraphMember(memberId, graphSignal) + : loadRoutedGraphMember(memberId, routedParentRunId, graphSignal), + ), ), ); throwIfAborted(graphSignal); diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index 1b847ff1ad7..ac9b3f5fc06 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -2685,6 +2685,10 @@ describe('initializeClient — subagent loading', () => { delete process.env.TEST_INHERIT_WORKSPACE_TOKEN; }); + /** Selects a lazy child the way the run does: the whole resolution inside its settle hook. */ + const select = (descriptor, context) => + descriptor.settle(context, () => descriptor.resolve(context)); + it("routes the reviewer to its parent's machine and workspace, not its own default", async () => { const req = await setup(); const fetchSpy = mockWorkerStatus(); @@ -2720,6 +2724,524 @@ describe('initializeClient — subagent loading', () => { }); }); + it('routes a reviewer to the machine its parent names and keeps it there this turn', async () => { + const req = await setup(); + const fetchSpy = mockWorkerStatus(); + try { + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + const descriptor = agentClientArgs.agent.lazySubagentConfigs[0]; + const routedAgents = []; + mockInitializeAgent.mockImplementation(async (params) => { + routedAgents.push(params.agent); + return { + ...makeSubagentConfig(SUBAGENT_ID), + codeExecutionContext: { + environmentId: params.agent.code_environment_id, + environmentType: 'attached', + }, + }; + }); + + await descriptor.resolve({ + signal: new AbortController().signal, + executionId: 'run-routed', + hostArgs: { machine: SKYNET }, + }); + await descriptor.resolve({ + signal: new AbortController().signal, + executionId: 'run-later', + }); + await expect( + descriptor.resolve({ + signal: new AbortController().signal, + executionId: 'run-unadmitted', + hostArgs: { machine: 'code-unadmitted' }, + }), + ).rejects.toThrow(); + + expect(routedAgents).toHaveLength(2); + expect(routedAgents[0]).toMatchObject({ + id: SUBAGENT_ID, + code_environment_id: SKYNET, + code_environment_ids: [SKYNET], + }); + expect(routedAgents[1]).toMatchObject({ + code_environment_id: SKYNET, + code_environment_ids: [SKYNET], + }); + } finally { + fetchSpy.mockRestore(); + } + }); + + it("routes a routed reviewer's graph members to its machine when they may use it", async () => { + const req = await setup(); + const member = await createAgent({ + id: 'agent_graph_member', + name: 'Graph member', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['execute_code'], + stateful_code_sessions: true, + stateful_code_environment: 'user', + code_environment_id: LIA_RAG, + code_environment_ids: [SKYNET], + }); + await grantView(member); + const fetchSpy = mockWorkerStatus(); + try { + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + const descriptor = agentClientArgs.agent.lazySubagentConfigs[0]; + const initialized = []; + mockInitializeAgent.mockImplementation(async (params) => { + initialized.push(params.agent); + return { + ...makeSubagentConfig(params.agent.id), + ...(params.agent.id === SUBAGENT_ID + ? { + subagents: { + enabled: true, + allowSelf: false, + agent_ids: [], + graphs: [ + { + type: 'review_team', + name: 'Review team', + description: 'Reviews together.', + agent_ids: [member.id], + edges: [], + entryAgentId: member.id, + resultAgentId: member.id, + }, + ], + }, + } + : {}), + codeExecutionContext: { + environmentId: params.agent.code_environment_id, + environmentType: 'attached', + }, + }; + }); + + const config = await descriptor.resolve({ + signal: new AbortController().signal, + executionId: 'run-routed-graph', + hostArgs: { machine: SKYNET }, + }); + + expect(initialized.map((agent) => [agent.id, agent.code_environment_id])).toEqual([ + [SUBAGENT_ID, SKYNET], + [member.id, SKYNET], + ]); + expect(initialized[1].code_environment_ids).toEqual([SKYNET]); + expect(config.subagentGraphConfigs?.[0]?.memberConfigs?.[0]?.id).toBe(member.id); + } finally { + fetchSpy.mockRestore(); + } + }); + + it('gives back the machine of a failed initialization so a retry may name another', async () => { + const req = await setup(); + const fetchSpy = mockWorkerStatus(); + try { + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + const descriptor = agentClientArgs.agent.lazySubagentConfigs[0]; + mockInitializeAgent.mockRejectedValueOnce(new Error('tools failed to load')); + mockInitializeAgent.mockImplementation(async (params) => ({ + ...makeSubagentConfig(SUBAGENT_ID), + codeExecutionContext: { + environmentId: params.agent.code_environment_id, + environmentType: 'attached', + }, + })); + + await expect( + descriptor.resolve({ signal: new AbortController().signal, executionId: 'run-failed' }), + ).rejects.toThrow('tools failed to load'); + const config = await descriptor.resolve({ + signal: new AbortController().signal, + executionId: 'run-retry', + hostArgs: { machine: SKYNET }, + }); + + expect(config.codeExecutionContext.environmentId).toBe(SKYNET); + } finally { + fetchSpy.mockRestore(); + } + }); + + it('gives back a routed reviewer\u2019s machine when its graph resolution is canceled', async () => { + const req = await setup(); + const member = await createAgent({ + id: 'agent_graph_member', + name: 'Graph member', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['execute_code'], + stateful_code_sessions: true, + stateful_code_environment: 'user', + code_environment_id: LIA_RAG, + code_environment_ids: [SKYNET], + }); + await grantView(member); + const fetchSpy = mockWorkerStatus(); + try { + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + const descriptor = agentClientArgs.agent.lazySubagentConfigs[0]; + let memberStarted; + const started = new Promise((resolve) => { + memberStarted = resolve; + }); + let memberLoads = 0; + mockInitializeAgent.mockImplementation(async (params) => { + if (params.agent.id === member.id && memberLoads++ === 0) { + memberStarted(); + await new Promise(() => {}); + } + return { + ...makeSubagentConfig(params.agent.id), + ...(params.agent.id === SUBAGENT_ID + ? { + subagents: { + enabled: true, + allowSelf: false, + agent_ids: [], + graphs: [ + { + type: 'review_team', + name: 'Review team', + description: 'Reviews together.', + agent_ids: [member.id], + edges: [], + entryAgentId: member.id, + resultAgentId: member.id, + }, + ], + }, + } + : {}), + codeExecutionContext: { + environmentId: params.agent.code_environment_id, + environmentType: 'attached', + }, + }; + }); + const controller = new AbortController(); + + const canceled = select(descriptor, { + signal: controller.signal, + executionId: 'run-canceled', + hostArgs: { machine: SKYNET }, + }); + await started; + controller.abort(); + await expect(canceled).rejects.toBeDefined(); + const config = await select(descriptor, { + signal: new AbortController().signal, + executionId: 'run-retry', + hostArgs: { machine: LIA_RAG }, + }); + + expect(config.codeExecutionContext.environmentId).toBe(LIA_RAG); + } finally { + fetchSpy.mockRestore(); + } + }); + + it('does not reuse the config of a resolution whose graph failed', async () => { + const req = await setup(); + const member = await createAgent({ + id: 'agent_graph_member', + name: 'Graph member', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['execute_code'], + stateful_code_sessions: true, + stateful_code_environment: 'user', + code_environment_id: LIA_RAG, + code_environment_ids: [SKYNET], + }); + await grantView(member); + const fetchSpy = mockWorkerStatus(); + try { + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + const descriptor = agentClientArgs.agent.lazySubagentConfigs[0]; + let memberStarted; + const started = new Promise((resolve) => { + memberStarted = resolve; + }); + let memberLoads = 0; + mockInitializeAgent.mockImplementation(async (params) => { + if (params.agent.id === member.id && memberLoads++ === 0) { + memberStarted(); + await new Promise(() => {}); + } + return { + ...makeSubagentConfig(params.agent.id), + ...(params.agent.id === SUBAGENT_ID + ? { + subagents: { + enabled: true, + allowSelf: false, + agent_ids: [], + graphs: [ + { + type: 'review_team', + name: 'Review team', + description: 'Reviews together.', + agent_ids: [SUBAGENT_ID, member.id], + edges: [], + entryAgentId: SUBAGENT_ID, + resultAgentId: member.id, + }, + ], + }, + } + : {}), + codeExecutionContext: { + environmentId: params.agent.code_environment_id, + environmentType: 'attached', + }, + }; + }); + const controller = new AbortController(); + + const canceled = select(descriptor, { + signal: controller.signal, + executionId: 'run-canceled', + }); + await started; + controller.abort(); + await expect(canceled).rejects.toBeDefined(); + const config = await select(descriptor, { + signal: new AbortController().signal, + executionId: 'run-retry', + hostArgs: { machine: LIA_RAG }, + }); + const graphReviewer = config.subagentGraphConfigs?.[0]?.memberConfigs?.find( + (memberConfig) => memberConfig.id === SUBAGENT_ID, + ); + + expect(config.codeExecutionContext.environmentId).toBe(LIA_RAG); + expect(graphReviewer?.codeExecutionContext.environmentId).toBe(LIA_RAG); + } finally { + fetchSpy.mockRestore(); + } + }); + + it('shares a lazy child config with later graphs only once its selection commits', async () => { + const req = await setup(); + const member = await createAgent({ + id: 'agent_graph_member', + name: 'Graph member', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['execute_code'], + stateful_code_sessions: true, + stateful_code_environment: 'user', + code_environment_id: LIA_RAG, + code_environment_ids: [], + }); + await grantView(member); + mockInitializeAgent.mockResolvedValue({ + ...makePrimaryConfig({ + subagents: { enabled: true, allowSelf: false, agent_ids: [SUBAGENT_ID, member.id] }, + }), + statefulCodeSessions: true, + codeExecutionContext: { environmentId: LIA_RAG, environmentType: 'attached' }, + }); + const fetchSpy = mockWorkerStatus(); + try { + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + const descriptorFor = (id) => + agentClientArgs.agent.lazySubagentConfigs.find((descriptor) => descriptor.id === id); + mockInitializeAgent.mockImplementation(async (params) => ({ + ...makeSubagentConfig(params.agent.id), + ...(params.agent.id === SUBAGENT_ID + ? { + subagents: { + enabled: true, + allowSelf: false, + agent_ids: [], + graphs: [ + { + type: 'review_team', + name: 'Review team', + description: 'Reviews together.', + agent_ids: [member.id], + edges: [], + entryAgentId: member.id, + resultAgentId: member.id, + }, + ], + }, + } + : {}), + codeExecutionContext: { + environmentId: params.agent.code_environment_id, + environmentType: 'attached', + }, + })); + const reviewWith = async (executionId) => { + const config = await select(descriptorFor(SUBAGENT_ID), { + signal: new AbortController().signal, + executionId, + hostArgs: { machine: SKYNET }, + }); + return config.subagentGraphConfigs?.[0]?.memberConfigs?.[0]; + }; + + const unsettled = await descriptorFor(member.id).resolve({ + signal: new AbortController().signal, + executionId: 'run-member-unsettled', + }); + expect(await reviewWith('run-review-a')).not.toBe(unsettled); + + const settled = await select(descriptorFor(member.id), { + signal: new AbortController().signal, + executionId: 'run-member-settled', + }); + expect(await reviewWith('run-review-b')).toBe(settled); + } finally { + fetchSpy.mockRestore(); + } + }); + + it("holds a graph member's default machine while it initializes", async () => { + const req = await setup(); + const member = await createAgent({ + id: 'agent_graph_member', + name: 'Graph member', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['execute_code'], + stateful_code_sessions: true, + stateful_code_environment: 'user', + code_environment_id: LIA_RAG, + code_environment_ids: [SKYNET], + }); + await grantView(member); + mockInitializeAgent.mockResolvedValue({ + ...makePrimaryConfig({ + subagents: { enabled: true, allowSelf: false, agent_ids: [SUBAGENT_ID, member.id] }, + }), + statefulCodeSessions: true, + codeExecutionContext: { environmentId: LIA_RAG, environmentType: 'attached' }, + }); + const fetchSpy = mockWorkerStatus(); + try { + await initializeClient({ + req, + res: {}, + signal: new AbortController().signal, + endpointOption: makeEndpointOption(), + }); + const descriptorFor = (id) => + agentClientArgs.agent.lazySubagentConfigs.find((descriptor) => descriptor.id === id); + let memberStarted; + const started = new Promise((resolve) => { + memberStarted = resolve; + }); + let releaseMember; + const memberGate = new Promise((resolve) => { + releaseMember = resolve; + }); + let memberLoads = 0; + mockInitializeAgent.mockImplementation(async (params) => { + if (params.agent.id === member.id && memberLoads++ === 0) { + memberStarted(); + await memberGate; + } + return { + ...makeSubagentConfig(params.agent.id), + ...(params.agent.id === SUBAGENT_ID + ? { + subagents: { + enabled: true, + allowSelf: false, + agent_ids: [], + graphs: [ + { + type: 'review_team', + name: 'Review team', + description: 'Reviews together.', + agent_ids: [member.id], + edges: [], + entryAgentId: member.id, + resultAgentId: member.id, + }, + ], + }, + } + : {}), + codeExecutionContext: { + environmentId: params.agent.code_environment_id, + environmentType: 'attached', + }, + }; + }); + + const reviewing = descriptorFor(SUBAGENT_ID).resolve({ + signal: new AbortController().signal, + executionId: 'run-reviewer', + }); + await started; + const moved = descriptorFor(member.id) + .resolve({ + signal: new AbortController().signal, + executionId: 'run-member', + hostArgs: { machine: SKYNET }, + }) + .then( + () => 'routed', + () => 'refused', + ); + const outcome = await moved; + releaseMember(); + await reviewing; + + expect(outcome).toBe('refused'); + } finally { + fetchSpy.mockRestore(); + } + }); + it('keeps a reviewer that may not use the parent machine on its own default', async () => { const req = await setup({ reviewer: { code_environment_ids: [] } }); const fetchSpy = mockWorkerStatus(); diff --git a/api/server/services/Files/provisionCallback.js b/api/server/services/Files/provisionCallback.js index 328cca2d9be..0d3993b35bf 100644 --- a/api/server/services/Files/provisionCallback.js +++ b/api/server/services/Files/provisionCallback.js @@ -4,11 +4,17 @@ const db = require('~/models'); /* Wiring only: binds this workspace's provisioning service and file model to the * ON_TOOL_EXECUTE callback implemented in packages/api. */ -function createProvisionFilesCallback({ req, agentToolContexts, resolvePrimaryAgentId }) { +function createProvisionFilesCallback({ + req, + agentToolContexts, + resolvePrimaryAgentId, + resolveExecutionContext, +}) { return createCallback({ req, agentToolContexts, resolvePrimaryAgentId, + resolveExecutionContext, provisionToCodeEnv, provisionToVectorDB, updateFile: db.updateFile, diff --git a/packages/api/src/agents/__tests__/run-codeTools.test.ts b/packages/api/src/agents/__tests__/run-codeTools.test.ts index 7602d5c9aa1..286735db7d7 100644 --- a/packages/api/src/agents/__tests__/run-codeTools.test.ts +++ b/packages/api/src/agents/__tests__/run-codeTools.test.ts @@ -66,7 +66,13 @@ jest.mock('~/agents/checkpointer', () => ({ getAgentCheckpointer: jest.fn().mockResolvedValue({}), })); -import { HookRegistry, InMemorySubagentTaskStore, Run } from '@librechat/agents'; +import { + Constants, + HookRegistry, + InMemorySubagentTaskStore, + Run, + executeHooks, +} from '@librechat/agents'; function makeAgent(overrides?: Record) { return { @@ -211,6 +217,62 @@ describe('createRun code-tool eager/session wiring', () => { expect(childInput.additional_instructions?.match(/\/mnt\/data\/data\.csv/g)).toHaveLength(1); }); + it('settles a lazy child only after the run has its inputs, so a late cancel is given back', async () => { + const controller = new AbortController(); + const outcomes: string[] = []; + const settle = jest.fn((_context: unknown, resolveInputs: () => Promise) => + resolveInputs().then( + (inputs) => { + outcomes.push('kept'); + return inputs; + }, + (error: unknown) => { + outcomes.push('given back'); + throw error; + }, + ), + ); + const resolve = jest.fn(async () => { + controller.abort(new Error('canceled after resolving')); + return makeAgent({ id: 'child' }); + }); + const config = await captureRunConfig( + makeAgent({ + subagents: { enabled: true, allowSelf: false, agent_ids: ['child'] }, + lazySubagentConfigs: [{ id: 'child', configId: 'child:1', resolve, settle }], + }), + ); + const [parentInput] = (config.graphConfig as { agents: AgentInputs[] }).agents; + const resolveInputs = parentInput.subagentConfigs?.[0].resolveAgentInputs; + if (!resolveInputs) throw new Error('Missing lazy subagent resolver'); + + await expect( + resolveInputs({ signal: controller.signal } as SubagentResolveContext), + ).rejects.toThrow('canceled after resolving'); + expect(outcomes).toEqual(['given back']); + }); + + it('declares per-call machine choices only on lazy children that have them', async () => { + const resolve = jest.fn().mockResolvedValue(makeAgent({ id: 'child' })); + const subagentHostArgs = { + machine: { description: 'Machine the subagent runs on.', enum: ['laptop'] }, + }; + const config = await captureRunConfig( + makeAgent({ + subagents: { enabled: true, allowSelf: false, agent_ids: ['child', 'plain'] }, + lazySubagentConfigs: [ + { id: 'child', configId: 'child:1', resolve, subagentHostArgs }, + { id: 'plain', configId: 'plain:1', resolve }, + ], + }), + ); + const [parentInput] = (config.graphConfig as { agents: AgentInputs[] }).agents; + const [routed, plain] = parentInput.subagentConfigs ?? []; + + expect(routed).toHaveProperty('hostArgs', subagentHostArgs); + expect(plain).not.toHaveProperty('hostArgs'); + }); + it.each([false, true])( 'reconciles a lazy child’s live setup file, parent already provisioned=%s', async (parentProvisioned) => { @@ -570,3 +632,139 @@ describe('createRun code-tool eager/session wiring', () => { ).resolves.toEqual(expect.objectContaining({ updatedOutput: expect.any(String) })); }); }); + +describe('attached-machine policy for lazily initialized graph members', () => { + beforeEach(() => jest.clearAllMocks()); + + const attachedOn = (environmentId: string, commandExecution: 'allow' | 'deny') => ({ + environmentType: 'attached', + environmentId, + codeEnvironmentConfigSchema: { + permissions: { commandExecution: { allowed: ['allow', 'ask', 'deny'], default: 'ask' } }, + }, + codeEnvironmentSettings: { permissions: { commandExecution } }, + }); + + it('stops applying attached-machine policy to a routable child that resolved on managed code', async () => { + const choice = { ...attachedOn('buildbox', 'deny'), statefulSessions: true }; + const resolve = jest + .fn() + .mockResolvedValue( + makeAgent({ id: 'child', codeExecutionContext: { environmentType: 'managed' } }), + ); + await createRun({ + agents: [ + makeAgent({ + subagents: { enabled: true, allowSelf: false, agent_ids: ['child'] }, + lazySubagentConfigs: [ + { id: 'child', configId: 'child:1', resolve, codeExecutionChoices: [choice] }, + ], + }), + ] as never, + signal: new AbortController().signal, + streaming: true, + streamUsage: true, + hitlCapable: true, + }); + const runConfig = (Run.create as jest.Mock).mock.calls[0][0] as Record; + /** The attached-machine policy's own verdict, apart from the run's other approval hooks. */ + const attachedVerdicts = async () => { + const verdicts: unknown[] = []; + for (const matcher of (runConfig.hooks as HookRegistry).getMatchers('PreToolUse')) { + for (const hook of matcher.hooks) { + const result = await hook( + { + hook_event_name: 'PreToolUse', + runId: 'run-1', + toolName: Constants.BASH_TOOL, + toolInput: { command: 'touch out.txt' }, + toolUseId: 'call-1', + executingAgentId: 'child', + }, + new AbortController().signal, + ); + if (String(result.reason ?? '').includes('attached code environment')) { + verdicts.push(result.decision); + } + } + } + return verdicts; + }; + + expect(await attachedVerdicts()).toEqual(['ask']); + + const [parentInput] = (runConfig.graphConfig as { agents: AgentInputs[] }).agents; + const resolveInputs = parentInput.subagentConfigs?.[0].resolveAgentInputs; + if (!resolveInputs) throw new Error('Missing lazy subagent resolver'); + await resolveInputs({ signal: new AbortController().signal } as SubagentResolveContext); + + expect(await attachedVerdicts()).toEqual([]); + }); + + it('applies the policy of the machine a routed child sent its graph member to', async () => { + const definition = { + type: 'review_team', + name: 'Review team', + description: 'A one-member team', + agent_ids: ['member'], + edges: [], + entry_agent_id: 'member', + result_agent_id: 'member', + }; + const routedMember = makeAgent({ + id: 'member', + codeExecutionContext: attachedOn('buildbox', 'deny'), + }); + const child = makeAgent({ + id: 'child', + subagentGraphConfigs: [{ definition, memberConfigs: [routedMember] }], + }); + const resolve = jest.fn().mockResolvedValue(child); + await createRun({ + agents: [ + makeAgent({ + subagents: { enabled: true, allowSelf: false, agent_ids: ['child'] }, + lazySubagentConfigs: [ + { + id: 'child', + configId: 'child:1', + resolve, + subagentGraphMemberMetadata: [ + { id: 'member', codeExecutionContext: attachedOn('laptop', 'allow') }, + ], + }, + ], + }), + ] as never, + signal: new AbortController().signal, + streaming: true, + streamUsage: true, + hitlCapable: true, + }); + const runConfig = (Run.create as jest.Mock).mock.calls[0][0] as Record; + const decide = async () => + ( + await executeHooks({ + registry: runConfig.hooks as HookRegistry, + input: { + hook_event_name: 'PreToolUse', + runId: 'run-1', + toolName: Constants.BASH_TOOL, + toolInput: { command: 'ls' }, + toolUseId: 'call-1', + executingAgentId: 'member', + }, + matchQuery: Constants.BASH_TOOL, + }) + ).decision; + + expect(await decide()).toBe('allow'); + + const [parentInput] = (runConfig.graphConfig as { agents: AgentInputs[] }).agents; + const resolveInputs = parentInput.subagentConfigs?.[0].resolveAgentInputs; + if (!resolveInputs) throw new Error('Missing lazy subagent resolver'); + await resolveInputs({ signal: new AbortController().signal } as SubagentResolveContext); + + expect(await decide()).toBe('deny'); + }); +}); diff --git a/packages/api/src/agents/codeFilesSession.spec.ts b/packages/api/src/agents/codeFilesSession.spec.ts index 32c2a1f2b78..5d95de3c2df 100644 --- a/packages/api/src/agents/codeFilesSession.spec.ts +++ b/packages/api/src/agents/codeFilesSession.spec.ts @@ -664,6 +664,38 @@ describe('collectCodeExecutionProfileRoutes', () => { expect(routes).toEqual([{ codeExecutionContext: graphContext, codeSessionKeys: [graphKey] }]); }); + it('includes the per-call machines of a graph member a lazy child initializes later', () => { + const memberKey = 'execute_code:stateful:attached:buildbox:member'; + const buildbox = { + baseUrl: 'https://bridge.example.com/v1', + codeSessionKey: memberKey, + executionProfile: 'stateful' as const, + executionRouteKey: 'attached:buildbox', + runtimeSessionHint: 'member', + statefulSessions: true, + environmentId: 'buildbox', + environmentType: 'attached' as const, + }; + + const routes = collectCodeExecutionProfileRoutes([ + { + id: 'parent', + codeEnvAvailable: false, + lazySubagentConfigs: [ + { + id: 'reviewer', + codeEnvAvailable: false, + subagentGraphMemberMetadata: [ + { id: 'member', codeEnvAvailable: false, codeExecutionChoices: [buildbox] }, + ], + }, + ], + }, + ]); + + expect(routes).toEqual([{ codeExecutionContext: buildbox, codeSessionKeys: [memberKey] }]); + }); + it('keeps configured stateful deployments in separate routing namespaces', () => { const context = (executionRouteKey: string, baseUrl: string) => ({ baseUrl, diff --git a/packages/api/src/agents/codeFilesSession.ts b/packages/api/src/agents/codeFilesSession.ts index 7a59af3a69c..8ea8aa753d3 100644 --- a/packages/api/src/agents/codeFilesSession.ts +++ b/packages/api/src/agents/codeFilesSession.ts @@ -20,12 +20,16 @@ export interface CodeFilesAgent { id?: string; codeEnvAvailable?: boolean; codeExecutionContext?: CodeExecutionContext; + /** Routes a parent may place this subagent on per call; each is primed like the default. */ + codeExecutionChoices?: CodeExecutionContext[]; codeSessionKey?: string; primedCodeFiles?: CodeEnvFile[]; statefulCodeSessions?: boolean; statefulCodeEnvironment?: StatefulCodeEnvironment; subagentAgentConfigs?: CodeFilesAgent[]; lazySubagentConfigs?: CodeFilesAgent[]; + /** Graph members a lazy child initializes only when it is selected. */ + subagentGraphMemberMetadata?: CodeFilesAgent[]; subagentGraphConfigs?: Array<{ memberConfigs: CodeFilesAgent[] }>; } @@ -42,6 +46,7 @@ function enqueueCodeFilesChildren( for (const child of [ ...(agent.subagentAgentConfigs ?? []), ...(agent.lazySubagentConfigs ?? []), + ...(agent.subagentGraphMemberMetadata ?? []), ]) { if (child && !visited.has(child)) queue.push(child); } @@ -83,14 +88,20 @@ export function collectCodeExecutionProfileRoutes( conversationId: scope.conversationId, }) : undefined); - if (agent.codeEnvAvailable === true && context) { - const routeKey = getCodeExecutionRouteKey(context); + const addRoute = (routeContext: CodeExecutionContext, codeSessionKey: string): void => { + const routeKey = getCodeExecutionRouteKey(routeContext); const route = routes.get(routeKey) ?? { - codeExecutionContext: context, + codeExecutionContext: routeContext, codeSessionKeys: new Set(), }; - route.codeSessionKeys.add(agent.codeSessionKey ?? context.codeSessionKey); + route.codeSessionKeys.add(codeSessionKey); routes.set(routeKey, route); + }; + if (agent.codeEnvAvailable === true && context) { + addRoute(context, agent.codeSessionKey ?? context.codeSessionKey); + } + for (const choice of agent.codeExecutionChoices ?? []) { + addRoute(choice, choice.codeSessionKey); } enqueueCodeFilesChildren(agent, queue, visited); } diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts index 1567eacf938..60279b7472e 100644 --- a/packages/api/src/agents/execution.ts +++ b/packages/api/src/agents/execution.ts @@ -106,8 +106,39 @@ export function getCodeWorkspaceSelections( type CodeExecutionApprovalAgent = { id?: string | null; codeExecutionContext?: CodeExecutionContext | null; + /** Alternate machines a parent may route this subagent to per call. */ + codeExecutionChoices?: readonly CodeExecutionContext[] | null; }; +function hashCodeExecutionTarget(agentId: string | null, context: CodeExecutionContext): string { + return createHash('sha256') + .update( + JSON.stringify([ + agentId, + context.executionProfile, + context.baseUrl, + context.codeSessionKey, + context.executionRouteKey ?? null, + context.runtimeSessionHint ?? null, + context.environmentId ?? null, + context.environmentType ?? null, + context.bridgeWorkerId ?? null, + context.codeWorkspace == null + ? null + : { + environmentId: context.codeWorkspace.environmentId, + workspaceId: context.codeWorkspace.workspaceId, + workspaceInstanceId: context.codeWorkspace.workspaceInstanceId ?? null, + operations: [...new Set(context.codeWorkspace.operations)].sort(), + ...(context.codeWorkspace.environment + ? { definitionFingerprint: context.codeWorkspace.environment.fingerprint } + : {}), + }, + ]), + ) + .digest('hex'); +} + const CODE_EXECUTION_TARGET_HASH = /^[a-f0-9]{64}$/; const MAX_CODE_EXECUTION_APPROVAL_TARGETS = 128; @@ -120,39 +151,28 @@ export function captureCodeExecutionApprovalBinding( agents: readonly (CodeExecutionApprovalAgent | null | undefined)[], ): Agents.CodeExecutionApprovalBinding | undefined { const targetsByIdentity = new Map(); + const addTarget = (target: Agents.CodeExecutionApprovalTargetBinding): void => { + targetsByIdentity.set(`${target.agentId ?? ''}\u0000${target.targetHash}`, target); + }; for (const agent of agents) { - const context = agent?.codeExecutionContext; - if (context?.statefulSessions !== true) { + const agentId = agent?.id ?? null; + const routeHashes = [agent?.codeExecutionContext, ...(agent?.codeExecutionChoices ?? [])] + .filter((context): context is CodeExecutionContext => context?.statefulSessions === true) + .map((context) => hashCodeExecutionTarget(agentId, context)); + if (routeHashes.length === 0) { continue; } - const targetHash = createHash('sha256') - .update( - JSON.stringify([ - agent?.id ?? null, - context.executionProfile, - context.baseUrl, - context.codeSessionKey, - context.executionRouteKey ?? null, - context.runtimeSessionHint ?? null, - context.environmentId ?? null, - context.environmentType ?? null, - context.bridgeWorkerId ?? null, - context.codeWorkspace == null - ? null - : { - environmentId: context.codeWorkspace.environmentId, - workspaceId: context.codeWorkspace.workspaceId, - workspaceInstanceId: context.codeWorkspace.workspaceInstanceId ?? null, - operations: [...new Set(context.codeWorkspace.operations)].sort(), - ...(context.codeWorkspace.environment - ? { definitionFingerprint: context.codeWorkspace.environment.fingerprint } - : {}), - }, - ]), - ) - .digest('hex'); - const target = { agentId: agent?.id ?? null, targetHash }; - targetsByIdentity.set(`${target.agentId ?? ''}\u0000${target.targetHash}`, target); + /** An agent with per-call choices folds its default and every choice into one + * target, so each agent contributes at most one entry to the bounded binding. */ + addTarget({ + agentId, + targetHash: + (agent?.codeExecutionChoices?.length ?? 0) === 0 + ? routeHashes[0] + : createHash('sha256') + .update(JSON.stringify(['choices', [...new Set(routeHashes)].sort()])) + .digest('hex'), + }); } const targets = [...targetsByIdentity.values()]; if (targets.length === 0) { diff --git a/packages/api/src/agents/files/binding.ts b/packages/api/src/agents/files/binding.ts index 239be439a41..18477d1c1db 100644 --- a/packages/api/src/agents/files/binding.ts +++ b/packages/api/src/agents/files/binding.ts @@ -20,6 +20,7 @@ type ProvisionCallback = ( names: string[], agentId?: string, signal?: AbortSignal, + context?: SubagentExecutionContext, ) => Promise; export interface ChatRunFileBindings extends RunFileHost { @@ -174,7 +175,7 @@ export function createChatRunFileBindings({ signal?: AbortSignal, context?: Parameters[1], ) => { - if (!host.session.isActive()) return fallback(names, agentId, signal); + if (!host.session.isActive()) return fallback(names, agentId, signal, context); if (!agentId || !signal) throw new Error('Shared-file provisioning requires an execution identity.'); return host.provisionPrepared(names, agentId, signal, context); diff --git a/packages/api/src/agents/files/session.ts b/packages/api/src/agents/files/session.ts index b3657b7ba03..eddb4d33294 100644 --- a/packages/api/src/agents/files/session.ts +++ b/packages/api/src/agents/files/session.ts @@ -23,6 +23,21 @@ export interface SharedRunArtifact { source: RunArtifactDescriptor; } +/** Whether a root agent's run shares files with its subagents under the deployment policy. */ +export function isRunFileSharingRequested({ + policy, + agent, +}: { + policy?: TAgentsEndpoint['fileSharing']; + agent: Pick; +}): boolean { + return ( + policy?.enabled === true && + agent.subagents?.enabled === true && + agent.subagents.shareFiles === true + ); +} + export function getAuthorizedRunFileSnapshot({ policy, agent, @@ -32,11 +47,7 @@ export function getAuthorizedRunFileSnapshot({ agent: Pick; files: readonly TFile[]; }): readonly TFile[] | undefined { - return policy?.enabled === true && - agent.subagents?.enabled === true && - agent.subagents.shareFiles === true - ? files - : undefined; + return isRunFileSharingRequested({ policy, agent }) ? files : undefined; } export interface RunFilePreparation { diff --git a/packages/api/src/agents/hitl/byom.spec.ts b/packages/api/src/agents/hitl/byom.spec.ts index c1e21bea171..567dcbac08a 100644 --- a/packages/api/src/agents/hitl/byom.spec.ts +++ b/packages/api/src/agents/hitl/byom.spec.ts @@ -3,6 +3,8 @@ import { assertAttachedCodeEnvironmentApprovalSupported, buildAttachedCodeEnvironmentAdmissionHooks, collectAttachedCodeEnvironmentAgentIds, + collectAttachedCodeApprovalPolicies, + collectAttachedCodeRoutePolicies, collectAttachedCodeEnvironmentPolicySettings, createAttachedCodeEnvironmentPolicyHook, isStatefulCodeEnvironmentToolName, @@ -25,7 +27,7 @@ const fullAccessSettings: AttachedCodeEnvironmentPolicySettings = { describe('full access', () => { test('allows native writes and commands while preserving mandatory skill review', async () => { const settings = new Map([['attached-agent', fullAccessSettings]]); - const mode = resolveAttachedCodeApprovalMode('fullAccess', settings); + const mode = resolveAttachedCodeApprovalMode('fullAccess', settings.values()); const hook = createAttachedCodeEnvironmentPolicyHook(new Set(settings.keys()), settings, mode); for (const toolName of ['create_file', 'edit_file', 'bash_tool', 'execute_code']) { await expect( @@ -70,7 +72,7 @@ describe('full access', () => { expect(() => resolveAttachedCodeApprovalMode( 'fullAccess', - new Map(reverse ? entries.reverse() : entries), + new Map(reverse ? entries.reverse() : entries).values(), ), ).toThrow('not permitted'); }, @@ -82,7 +84,7 @@ describe('full access', () => { const hook = createAttachedCodeEnvironmentPolicyHook( attachedIds, settings, - resolveAttachedCodeApprovalMode('fullAccess', settings), + resolveAttachedCodeApprovalMode('fullAccess', settings.values()), ); settings.set('attached-agent', { ...fullAccessSettings, @@ -96,7 +98,7 @@ describe('full access', () => { await expect( hook({ toolName: 'bash_tool', executingAgentId: 'lazy-agent' } as never, signal), ).resolves.toMatchObject({ decision: 'ask' }); - expect(() => resolveAttachedCodeApprovalMode('fullAccess', settings, false)).toThrow( + expect(() => resolveAttachedCodeApprovalMode('fullAccess', settings.values(), false)).toThrow( 'not permitted', ); }); @@ -255,7 +257,7 @@ describe('createAttachedCodeEnvironmentPolicyHook', () => { }, ], ]); - const mode = resolveAttachedCodeApprovalMode('acceptEdits', settings); + const mode = resolveAttachedCodeApprovalMode('acceptEdits', settings.values()); const hook = createAttachedCodeEnvironmentPolicyHook( new Set(['attached-agent']), settings, @@ -272,19 +274,13 @@ describe('createAttachedCodeEnvironmentPolicyHook', () => { test('rejects accept edits when the attached machine excludes file-write allow', () => { expect(() => - resolveAttachedCodeApprovalMode( - 'acceptEdits', - new Map([ - [ - 'attached-agent', - { - configSchema: { - permissions: { fileWrite: { allowed: ['ask'], default: 'ask' } }, - }, - }, - ], - ]), - ), + resolveAttachedCodeApprovalMode('acceptEdits', [ + { + configSchema: { + permissions: { fileWrite: { allowed: ['ask'], default: 'ask' } }, + }, + }, + ]), ).toThrow('not permitted'); }); @@ -307,7 +303,7 @@ describe('createAttachedCodeEnvironmentPolicyHook', () => { }, ], ]); - const mode = resolveAttachedCodeApprovalMode('acceptEdits', settings); + const mode = resolveAttachedCodeApprovalMode('acceptEdits', settings.values()); const hook = createAttachedCodeEnvironmentPolicyHook( new Set(['permissive-agent', 'restrictive-agent']), settings, @@ -334,7 +330,7 @@ describe('createAttachedCodeEnvironmentPolicyHook', () => { }, ], ]); - const mode = resolveAttachedCodeApprovalMode('acceptEdits', settings); + const mode = resolveAttachedCodeApprovalMode('acceptEdits', settings.values()); const hook = createAttachedCodeEnvironmentPolicyHook(attachedIds, settings, mode); attachedIds.add('lazy-agent'); @@ -350,11 +346,11 @@ describe('createAttachedCodeEnvironmentPolicyHook', () => { }); test('rejects an explicit mode when approvals are disabled by the administrator', () => { - expect(resolveAttachedCodeApprovalMode('ask', new Map(), false)).toBeUndefined(); - expect(() => resolveAttachedCodeApprovalMode('acceptEdits', new Map(), false)).toThrow( + expect(resolveAttachedCodeApprovalMode('ask', [], false)).toBeUndefined(); + expect(() => resolveAttachedCodeApprovalMode('acceptEdits', [], false)).toThrow( 'not permitted', ); - expect(resolveAttachedCodeApprovalMode(undefined, new Map(), false)).toBeUndefined(); + expect(resolveAttachedCodeApprovalMode(undefined, [], false)).toBeUndefined(); }); test.each(['create_file', 'edit_file'])( @@ -601,6 +597,118 @@ describe('collectAttachedCodeEnvironmentAgentIds', () => { }); }); +describe('collectAttachedCodeApprovalPolicies', () => { + const acceptEditsMachine = { + environmentType: 'attached', + codeEnvironmentConfigSchema: fullAccessSettings.configSchema, + codeEnvironmentSettings: { permissions: { fileWrite: 'allow' as const } }, + }; + + test('validates a mode against a fallback-routable child by its per-call machines', () => { + const agents = [ + { + id: 'root', + codeExecutionContext: { environmentType: 'managed' }, + lazySubagentConfigs: [ + { + id: 'reviewer', + codeExecutionChoices: [{ environmentType: 'managed' }, acceptEditsMachine], + }, + ], + }, + ]; + + expect(() => + resolveAttachedCodeApprovalMode( + 'acceptEdits', + collectAttachedCodeEnvironmentPolicySettings(agents).values(), + ), + ).toThrow('not permitted'); + expect( + resolveAttachedCodeApprovalMode('acceptEdits', collectAttachedCodeApprovalPolicies(agents)), + ).toBe('acceptEdits'); + }); + + test('lets an asking per-call machine make a permissive-by-default child pause-capable', () => { + const agents = [ + { + id: 'reviewer', + codeExecutionContext: { + ...acceptEditsMachine, + codeEnvironmentSettings: { + permissions: { fileWrite: 'allow' as const, commandExecution: 'allow' as const }, + }, + }, + codeExecutionChoices: [ + { + environmentType: 'attached', + codeEnvironmentConfigSchema: fullAccessSettings.configSchema, + }, + ], + }, + ]; + const ids = collectAttachedCodeEnvironmentAgentIds(agents); + const settings = collectAttachedCodeEnvironmentPolicySettings(agents); + const pausing = (hooks: ReturnType) => + hooks.flatMap((hook) => hook.toolNames ?? []); + + expect(pausing(buildAttachedCodeEnvironmentAdmissionHooks(ids, settings))).not.toContain( + 'bash_tool', + ); + expect( + pausing( + buildAttachedCodeEnvironmentAdmissionHooks( + ids, + settings, + undefined, + collectAttachedCodeRoutePolicies(agents), + ), + ), + ).toContain('bash_tool'); + }); + + test('does not treat a missing default as asking when every per-call machine decides', () => { + const decides = { + environmentType: 'attached', + codeEnvironmentConfigSchema: fullAccessSettings.configSchema, + codeEnvironmentSettings: { + permissions: { fileWrite: 'allow' as const, commandExecution: 'deny' as const }, + }, + }; + const agents = [{ id: 'reviewer', codeExecutionChoices: [decides, decides] }]; + + expect( + buildAttachedCodeEnvironmentAdmissionHooks( + collectAttachedCodeEnvironmentAgentIds(agents), + collectAttachedCodeEnvironmentPolicySettings(agents), + undefined, + collectAttachedCodeRoutePolicies(agents), + ), + ).toEqual([]); + }); + + test('requires every machine a child may be routed to before granting full access', () => { + const agents = [ + { + id: 'reviewer', + codeExecutionContext: { ...acceptEditsMachine, codeEnvironmentSettings: undefined }, + codeExecutionChoices: [ + { + environmentType: 'attached', + codeEnvironmentConfigSchema: fullAccessSettings.configSchema, + }, + { environmentType: 'attached' }, + ], + }, + ]; + + expect(collectAttachedCodeApprovalPolicies(agents)).toHaveLength(3); + expect(() => + resolveAttachedCodeApprovalMode('fullAccess', collectAttachedCodeApprovalPolicies(agents)), + ).toThrow('not permitted'); + }); +}); + describe('assertAttachedCodeEnvironmentApprovalSupported', () => { test('rejects attached environments on callers without an approval/resume surface', () => { expect(() => diff --git a/packages/api/src/agents/hitl/byom.ts b/packages/api/src/agents/hitl/byom.ts index 4cd0c737abf..ca9fd569b81 100644 --- a/packages/api/src/agents/hitl/byom.ts +++ b/packages/api/src/agents/hitl/byom.ts @@ -65,6 +65,14 @@ type CodeEnvironmentPolicyAgent = { codeEnvironmentConfigSchema?: CodeEnvironmentUserConfigSchema; codeEnvironmentSettings?: CodeEnvironmentUserSettings; }; + /** Machines a parent may route this subagent to per call. */ + codeExecutionChoices?: + | readonly { + environmentType?: string; + codeEnvironmentConfigSchema?: CodeEnvironmentUserConfigSchema; + codeEnvironmentSettings?: CodeEnvironmentUserSettings; + }[] + | null; subagentAgentConfigs?: readonly (CodeEnvironmentPolicyAgent | null | undefined)[]; lazySubagentConfigs?: readonly (CodeEnvironmentPolicyAgent | null | undefined)[]; subagentGraphMemberMetadata?: readonly (CodeEnvironmentPolicyAgent | null | undefined)[]; @@ -189,7 +197,9 @@ export function collectAttachedCodeEnvironmentAgentIds( ): Set { const attachedAgentIds = new Set(); for (const agent of collectCodeEnvironmentPolicyAgents(roots)) { - if (agent.id && agent.codeExecutionContext?.environmentType === 'attached') { + const routable = + agent.codeExecutionChoices?.some((choice) => choice.environmentType === 'attached') === true; + if (agent.id && (agent.codeExecutionContext?.environmentType === 'attached' || routable)) { attachedAgentIds.add(agent.id); } } @@ -212,6 +222,36 @@ export function collectAttachedCodeEnvironmentPolicySettings( return settingsByAgentId; } +/** + * Each agent's attached-machine policies: its default route and every machine a parent may route + * it to per call. Preflight decisions use all of them, since the call picks the machine later. + */ +export function collectAttachedCodeRoutePolicies( + roots: readonly (CodeEnvironmentPolicyAgent | null | undefined)[], +): Map { + const policiesByAgentId = new Map(); + for (const agent of collectCodeEnvironmentPolicyAgents(roots)) { + const policies = [agent.codeExecutionContext, ...(agent.codeExecutionChoices ?? [])] + .filter((context) => context?.environmentType === 'attached') + .map((context) => ({ + configSchema: context?.codeEnvironmentConfigSchema, + settings: context?.codeEnvironmentSettings, + skillAuthoringAvailable: agent.skillAuthoringAvailable === true, + })); + if (agent.id && policies.length > 0) { + policiesByAgentId.set(agent.id, [...(policiesByAgentId.get(agent.id) ?? []), ...policies]); + } + } + return policiesByAgentId; +} + +/** Every attached-machine policy the run may execute under, for validating the requested mode. */ +export function collectAttachedCodeApprovalPolicies( + roots: readonly (CodeEnvironmentPolicyAgent | null | undefined)[], +): AttachedCodeEnvironmentPolicySettings[] { + return [...collectAttachedCodeRoutePolicies(roots).values()].flat(); +} + function permissionDecision( policy: AttachedCodeEnvironmentPolicySettings | undefined, category: PermissionCategory, @@ -236,7 +276,7 @@ function permissionDecision( export function resolveAttachedCodeApprovalMode( requested: unknown, - settingsByAgentId: ReadonlyMap, + policies: Iterable, approvalsEnabled = true, ): CodeApprovalMode | undefined { if (!approvalsEnabled) { @@ -251,7 +291,7 @@ export function resolveAttachedCodeApprovalMode( } let resolved: CodeApprovalMode | undefined; let rejection: Error | undefined; - for (const policy of settingsByAgentId.values()) { + for (const policy of policies) { try { resolved = resolveCodeApprovalMode(requested, { environment: 'attached', @@ -284,6 +324,11 @@ export function buildAttachedCodeEnvironmentAdmissionHooks( attachedAgentIds: ReadonlySet, settingsByAgentId: ReadonlyMap = new Map(), mode?: CodeApprovalMode, + /** Machines a parent may route an agent to per call; any of them that asks can pause. */ + routePoliciesByAgentId: ReadonlyMap< + string, + readonly AttachedCodeEnvironmentPolicySettings[] + > = new Map(), ): ResolvedToolApprovalHook[] { const hook = createAttachedCodeEnvironmentPolicyHook(attachedAgentIds, settingsByAgentId, mode); const hooks: ResolvedToolApprovalHook[] = []; @@ -291,11 +336,15 @@ export function buildAttachedCodeEnvironmentAdmissionHooks( const askCommandAgents = new Set(); const skillAuthoringAgents = new Set(); for (const agentId of attachedAgentIds) { - const policy = settingsByAgentId.get(agentId); - if (permissionDecision(policy, 'fileWrite', mode) === 'ask') askFileAgents.add(agentId); - if (permissionDecision(policy, 'commandExecution', mode) === 'ask') - askCommandAgents.add(agentId); - if (policy?.skillAuthoringAvailable === true) skillAuthoringAgents.add(agentId); + const routePolicies = routePoliciesByAgentId.get(agentId) ?? []; + const policies = routePolicies.length > 0 ? routePolicies : [settingsByAgentId.get(agentId)]; + const asks = (category: PermissionCategory): boolean => + policies.some((policy) => permissionDecision(policy, category, mode) === 'ask'); + if (asks('fileWrite')) askFileAgents.add(agentId); + if (asks('commandExecution')) askCommandAgents.add(agentId); + if (policies.some((policy) => policy?.skillAuthoringAvailable === true)) { + skillAuthoringAgents.add(agentId); + } } if (askFileAgents.size > 0) { hooks.push({ diff --git a/packages/api/src/agents/lazySubagents.spec.ts b/packages/api/src/agents/lazySubagents.spec.ts index 5e30e1f7e6d..8a17d562841 100644 --- a/packages/api/src/agents/lazySubagents.spec.ts +++ b/packages/api/src/agents/lazySubagents.spec.ts @@ -5,6 +5,7 @@ import { guardUnavailableSubagent, getLazySubagentConfigId, createViewableSubagentLoader, + createRoutedGraphMemberLoader, resolveSubagentCodeAvailability, } from './lazySubagents'; import { CodeWorkspaceSelectionError } from '~/code/errors'; @@ -227,3 +228,68 @@ describe('subagent code availability', () => { }); }); }); + +describe('createRoutedGraphMemberLoader', () => { + const signal = new AbortController().signal; + type Deps = Parameters>[0]; + const setup = (overrides: Partial = {}) => { + const skipped = new Set(); + const deps: Deps = { + getShared: jest.fn((_memberId: string): string | undefined => undefined), + isSkipped: jest.fn((memberId: string) => skipped.has(memberId)), + skip: jest.fn((memberId: string) => { + skipped.add(memberId); + }), + getAgent: jest.fn(async (memberId: string) => ({ id: memberId })), + canView: jest.fn(async () => true), + initialize: jest.fn( + async ({ memberId, context }: { memberId: string; context: { executionId: string } }) => + `${memberId}@${context.executionId}`, + ), + isFatal: jest.fn(() => false), + ...overrides, + }; + return { deps, load: createRoutedGraphMemberLoader<{ id: string }, string>(deps) }; + }; + + it('initializes a member as a child of the routed parent execution', async () => { + const { deps, load } = setup(); + + await expect(load('member', 'run-parent', signal)).resolves.toBe( + 'member@run-parent:graph:member', + ); + expect(deps.initialize).toHaveBeenCalledWith({ + agent: { id: 'member' }, + memberId: 'member', + context: { signal, parentRunId: 'run-parent', executionId: 'run-parent:graph:member' }, + }); + }); + + it('reuses a shared member and never initializes a skipped or unviewable one', async () => { + const shared = setup({ getShared: jest.fn(() => 'shared-config') }); + await expect(shared.load('member', 'run-parent', signal)).resolves.toBe('shared-config'); + expect(shared.deps.getAgent).not.toHaveBeenCalled(); + + const hidden = setup({ canView: jest.fn(async () => false) }); + await expect(hidden.load('member', 'run-parent', signal)).resolves.toBeNull(); + await expect(hidden.load('member', 'run-other', signal)).resolves.toBeNull(); + expect(hidden.deps.getAgent).toHaveBeenCalledTimes(1); + expect(hidden.deps.initialize).not.toHaveBeenCalled(); + }); + + it('drops a member that fails to initialize unless the failure is fatal', async () => { + const failure = new Error('tools failed to load'); + const recoverable = setup({ initialize: jest.fn().mockRejectedValue(failure) }); + await expect(recoverable.load('member', 'run-parent', signal)).resolves.toBeNull(); + + const fatal = setup({ + initialize: jest.fn().mockRejectedValue(failure), + isFatal: jest.fn(() => true), + }); + await expect(fatal.load('member', 'run-parent', signal)).rejects.toBe(failure); + + const aborted = new AbortController(); + aborted.abort(new Error('canceled')); + await expect(fatal.load('member', 'run-parent', aborted.signal)).rejects.toThrow('canceled'); + }); +}); diff --git a/packages/api/src/agents/lazySubagents.ts b/packages/api/src/agents/lazySubagents.ts index 9b70e13666d..021774507b5 100644 --- a/packages/api/src/agents/lazySubagents.ts +++ b/packages/api/src/agents/lazySubagents.ts @@ -193,6 +193,71 @@ export function createViewableSubagentLoader({ }; } +/** The SDK-facing context a routed graph member initializes under. */ +export interface RoutedGraphMemberContext { + signal: AbortSignal; + /** The per-call-routed parent execution the member follows. */ + parentRunId: string; + /** Stable per parent and member; routing groups the member's route with the parent's. */ + executionId: string; +} + +/** + * Loads the graph members a per-call-routed parent resolves. A member already shared + * this request (initialized on its own route) is reused; otherwise it initializes as + * a child of the parent's execution so it may follow the parent's machine, and its + * config stays with that execution rather than the request-wide graph cache. + */ +export function createRoutedGraphMemberLoader({ + getShared, + isSkipped, + skip, + getAgent, + canView, + initialize, + isFatal, +}: { + getShared: (memberId: string) => TConfig | undefined; + isSkipped: (memberId: string) => boolean; + skip: (memberId: string) => void; + getAgent: (memberId: string, signal: AbortSignal) => Promise; + canView: (agent: TAgent, memberId: string, signal: AbortSignal) => Promise; + initialize: (input: { + agent: TAgent; + memberId: string; + context: RoutedGraphMemberContext; + }) => Promise; + isFatal: (error: unknown, signal: AbortSignal) => boolean; +}): (memberId: string, parentRunId: string, signal: AbortSignal) => Promise { + return async (memberId, parentRunId, signal) => { + if (signal.aborted) { + throw signal.reason ?? new Error('Subagent resolution was aborted.'); + } + const shared = getShared(memberId); + if (shared != null) { + return shared; + } + if (isSkipped(memberId)) { + return null; + } + const agent = await getAgent(memberId, signal); + if (agent == null || !(await canView(agent, memberId, signal))) { + skip(memberId); + return null; + } + const executionId = `${parentRunId}:graph:${memberId}`; + try { + return await initialize({ agent, memberId, context: { signal, parentRunId, executionId } }); + } catch (error) { + if (isFatal(error, signal)) { + throw error; + } + logger.error(`[lazySubagents] Error initializing routed graph member ${memberId}:`, error); + return null; + } + }; +} + export interface SubagentCodeAvailability { codeEnvAvailable: boolean; statefulCodeSessions: boolean; diff --git a/packages/api/src/agents/run.ts b/packages/api/src/agents/run.ts index 054c3e39ce1..6548723bdcf 100644 --- a/packages/api/src/agents/run.ts +++ b/packages/api/src/agents/run.ts @@ -63,6 +63,7 @@ import type { TerminalSteerHook } from '~/agents/steering/runtime'; import type { LangfuseTraceContext } from '~/langfuse/identity'; import type { ResolvedAlwaysApplySkill } from '~/agents/skills'; import type { CodeExecutionContext } from '~/agents/execution'; +import type { SubagentCodeHostArgSpecs } from '~/code/targets'; import type { MCPToolAlias } from '~/tools/classification'; import type { ReviewedToolApprovals } from './hitl/modes'; import type { SubagentUsageEvent } from '~/agents/usage'; @@ -75,6 +76,7 @@ import { collectNativeEditFileAgentIds, collectAttachedCodeEnvironmentAgentIds, collectAttachedCodeEnvironmentPolicySettings, + collectAttachedCodeApprovalPolicies, createAttachedCodeEnvironmentPolicyHook, resolveAttachedCodeApprovalMode, } from '~/agents/hitl/byom'; @@ -529,11 +531,20 @@ type LazySubagentAgent = Pick< | 'mcpToolAliases' > & { configId: string; + /** Per-call machine choices declared on the child's subagent call. */ + subagentHostArgs?: SubagentCodeHostArgSpecs; + /** The routes behind those choices, for run-wide gates and approval bindings. */ + codeExecutionChoices?: CodeExecutionContext[]; subagentAgentConfigs?: RunAgent[]; lazySubagentConfigs?: LazySubagentAgent[]; /** Lightweight graph-member metadata used only by run-wide capability gates. */ subagentGraphMemberMetadata?: SubagentTreeNode[]; resolve: (context: SubagentResolveContext) => Promise; + /** + * Wraps the whole selection, through the inputs handed to the SDK, so host state the + * call reserved is kept only when the child is actually exposed. + */ + settle?: (context: SubagentResolveContext, resolveInputs: () => Promise) => Promise; }; type SubagentTreeNode = Pick< @@ -551,6 +562,7 @@ type SubagentTreeNode = Pick< | 'includeReasoningHistory' | 'mcpToolAliases' > & { + codeExecutionChoices?: CodeExecutionContext[]; subagentAgentConfigs?: SubagentTreeNode[]; lazySubagentConfigs?: SubagentTreeNode[]; subagentGraphMemberMetadata?: SubagentTreeNode[]; @@ -1619,7 +1631,44 @@ function createLazySubagentConfig( prebuiltGraphInputs?: ReadonlyMap, onResolvedAgent?: (agent: RunAgent) => void, ): SubagentConfig { - return { + const resolveInputs = async (context: SubagentResolveContext): Promise => { + if (context.signal.aborted) { + throw context.signal.reason ?? new Error('Subagent resolution was aborted.'); + } + const resolvedChild = await child.resolve(context); + if (context.signal.aborted) { + throw context.signal.reason ?? new Error('Subagent resolution was aborted.'); + } + onResolvedAgent?.(resolvedChild); + /** Graph members initialized with the child may run on its per-call machine. */ + for (const graph of resolvedChild.subagentGraphConfigs ?? []) { + for (const member of graph.memberConfigs) { + onResolvedAgent?.(member); + } + } + const childInputs = buildIsolatedAgentInputs(resolvedChild, toInput); + const resolutionState: SubagentBuildState = { + configCount: 1, + rootAgentIds: [resolvedChild.id], + }; + const grandchildConfigs = buildSubagentConfigs( + resolvedChild, + childInputs, + toInput, + resolutionState, + agentsEConfig, + ancestors, + depth, + prebuiltGraphInputs, + false, + onResolvedAgent, + ); + if (grandchildConfigs.length > 0) { + childInputs.subagentConfigs = grandchildConfigs; + } + return childInputs; + }; + const config: SubagentConfig = { type: child.id, name: child.name ?? child.id, description: @@ -1628,38 +1677,15 @@ function createLazySubagentConfig( configId: child.configId, allowNested: true, maxTurns: resolveSubagentMaxTurns(agentsEConfig, child), - resolveAgentInputs: async (context) => { - if (context.signal.aborted) { - throw context.signal.reason ?? new Error('Subagent resolution was aborted.'); - } - const resolvedChild = await child.resolve(context); - if (context.signal.aborted) { - throw context.signal.reason ?? new Error('Subagent resolution was aborted.'); - } - onResolvedAgent?.(resolvedChild); - const childInputs = buildIsolatedAgentInputs(resolvedChild, toInput); - const resolutionState: SubagentBuildState = { - configCount: 1, - rootAgentIds: [resolvedChild.id], - }; - const grandchildConfigs = buildSubagentConfigs( - resolvedChild, - childInputs, - toInput, - resolutionState, - agentsEConfig, - ancestors, - depth, - prebuiltGraphInputs, - false, - onResolvedAgent, - ); - if (grandchildConfigs.length > 0) { - childInputs.subagentConfigs = grandchildConfigs; - } - return childInputs; - }, + resolveAgentInputs: (context) => + child.settle == null + ? resolveInputs(context) + : child.settle(context, () => resolveInputs(context)), }; + /** Assigned rather than spread so the field typechecks against SDKs that predate it. */ + return child.subagentHostArgs == null + ? config + : Object.assign(config, { hostArgs: child.subagentHostArgs }); } function enqueueSubagentChildren( @@ -1762,7 +1788,7 @@ function anyAgentHasCodeEnv(agents: RunAgent[]): boolean { continue; } visited.add(agent.id); - if (agent.codeEnvAvailable === true) { + if (agent.codeEnvAvailable === true || (agent.codeExecutionChoices?.length ?? 0) > 0) { return true; } enqueueSubagentChildren(agent, pending, visited); @@ -2637,7 +2663,7 @@ export async function createRun({ const attachedCodeEnvironmentSettings = collectAttachedCodeEnvironmentPolicySettings(agents); const codeApprovalMode = resolveAttachedCodeApprovalMode( requestedCodeApprovalMode, - attachedCodeEnvironmentSettings, + collectAttachedCodeApprovalPolicies(agents), agentsEndpointConfig?.toolApproval?.enabled !== false, ); assertAttachedCodeEnvironmentApprovalSupported({ @@ -2853,6 +2879,11 @@ export async function createRun({ settings: resolvedAgent.codeExecutionContext.codeEnvironmentSettings, skillAuthoringAvailable: resolvedAgent.skillAuthoringAvailable === true, }); + } else { + /** A routable child counted as attached before its call chose a route; it resolved + * off attached machines, and runs on that one route for the whole request. */ + attachedCodeEnvironmentAgentIds.delete(resolvedAgent.id); + attachedCodeEnvironmentSettings.delete(resolvedAgent.id); } const discoveredAliases = collectRunMCPToolAliases([resolvedAgent]).filter( ({ name, aliasName }) => { diff --git a/packages/api/src/code/index.ts b/packages/api/src/code/index.ts index 45ac13a95f1..237e3a04561 100644 --- a/packages/api/src/code/index.ts +++ b/packages/api/src/code/index.ts @@ -10,3 +10,4 @@ export * from './decision'; export * from './inheritance'; export * from './errors'; export * from './instructions'; +export * from './targets'; diff --git a/packages/api/src/code/targets.integration.spec.ts b/packages/api/src/code/targets.integration.spec.ts new file mode 100644 index 00000000000..bb33a4958a0 --- /dev/null +++ b/packages/api/src/code/targets.integration.spec.ts @@ -0,0 +1,212 @@ +import mongoose, { Types } from 'mongoose'; +import { MongoMemoryServer } from 'mongodb-memory-server'; +import { createMethods, createModels, tenantStorage } from '@librechat/data-schemas'; +import type { AppConfig } from '@librechat/data-schemas'; +import type { CodeEnvironmentPrincipalContext } from './environments'; +import type { CodeEnvironmentConfig } from '~/agents/execution'; +import type { SubagentCodeAgent } from './targets'; +import { createCodeEnvironmentRegistry } from './environments'; +import { mergeAccessibleCodeEnvironments } from './config'; +import { createSubagentCodeRouting } from './targets'; + +/** + * Stands in for the SDK's `SubagentHostArgumentError` until LibreChat depends + * on an `@librechat/agents` release that exports it. + */ +jest.mock('@librechat/agents', () => { + const actual = jest.requireActual('@librechat/agents'); + class MockSubagentHostArgumentError extends Error { + constructor(name: string, reason: 'unavailable' | 'not_allowed') { + super('Subagent host argument was rejected.'); + this.name = 'SubagentHostArgumentError'; + Object.assign(this, { argument: name, rejection: reason }); + } + } + return { + ...actual, + SubagentHostArgumentError: actual.SubagentHostArgumentError ?? MockSubagentHostArgumentError, + }; +}); + +const BASE_URL = 'https://bridge.example'; +const TOKEN_ENV = 'TEST_SUBAGENT_TARGETS_INTEGRATION_TOKEN'; + +const controlPlane: CodeEnvironmentConfig = { + id: 'control-plane', + name: 'Control plane', + type: 'attached', + owner: 'deployment', + baseURL: BASE_URL, + pairing: { allowPrincipalWorkers: true, tokenEnv: TOKEN_ENV }, +}; + +const deploymentConfig = { + endpoints: { agents: { statefulCodeSessions: { environments: [controlPlane] } } }, +} as AppConfig; + +const sealed = [{ environmentId: 'danny-vm', workspaceId: 'agents' }]; +const reviewer: SubagentCodeAgent = { + id: 'agent_reviewer', + code_environment_ids: ['danny-vm'], +}; +const flags = { statefulCodeSessions: true, statefulCodeEnvironment: 'conversation' }; + +function actor(userId: Types.ObjectId): CodeEnvironmentPrincipalContext { + return { userId, role: 'USER', idOnTheSource: null }; +} + +function serveReadyWorker(): void { + jest.spyOn(globalThis, 'fetch').mockImplementation(async (input) => { + const workerId = /\/bridge\/workers\/([^/]+)\/status$/.exec(String(input))?.[1] ?? ''; + return new Response( + JSON.stringify({ + protocolVersion: 1, + workerId, + online: true, + ready: true, + leaseExpiresInMs: 45_000, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file', 'execute_command'], + workspaces: [{ id: 'agents' }], + }, + }, + }), + ); + }); +} + +describe('subagent code targets against the principal-scoped machine list', () => { + let mongoServer: MongoMemoryServer; + + beforeAll(async () => { + mongoServer = await MongoMemoryServer.create(); + await mongoose.connect(mongoServer.getUri()); + createModels(mongoose); + await Promise.all(Object.values(mongoose.models).map((model) => model.init())); + }); + + afterAll(async () => { + await mongoose.disconnect(); + await mongoServer.stop(); + }); + + beforeEach(async () => { + await mongoose.connection.dropDatabase(); + await createMethods(mongoose).seedDefaultRoles(); + process.env[TOKEN_ENV] = `token-${Math.random()}`; + serveReadyWorker(); + }); + + afterEach(() => { + jest.restoreAllMocks(); + delete process.env[TOKEN_ENV]; + }); + + async function registerMachine(ownerId: Types.ObjectId): Promise { + await createCodeEnvironmentRegistry(mongoose).register({ + actor: actor(ownerId), + environment: { + id: 'danny-vm', + name: "Danny's VM", + type: 'attached', + baseURL: BASE_URL, + controlPlaneId: 'control-plane', + workerId: 'w-danny', + }, + }); + } + + async function routingFor(userId: Types.ObjectId) { + const appConfig = await mergeAccessibleCodeEnvironments({ + appConfig: deploymentConfig, + deploymentConfig, + actor: actor(userId), + registry: createCodeEnvironmentRegistry(mongoose), + }); + return createSubagentCodeRouting({ + allowEnvironmentSelection: true, + persistedSelections: sealed, + environments: appConfig.endpoints?.agents?.statefulCodeSessions?.environments, + userId: userId.toString(), + conversationId: 'convo-1', + getAppConfig: async () => deploymentConfig, + }); + } + + it('offers and routes to a machine the requesting principal can use', async () => { + const ownerId = new Types.ObjectId(); + await registerMachine(ownerId); + const routing = await routingFor(ownerId); + + const description = await routing.describe(reviewer, flags); + const placed = await routing.place({ + agent: reviewer, + flags, + context: { executionId: 'run', hostArgs: { machine: 'danny-vm' } }, + }); + + expect(description.subagentHostArgs?.machine.enum).toEqual(['danny-vm']); + expect(description.subagentHostArgs?.workspace.enum).toEqual(['agents']); + expect(placed.agent.code_environment_id).toBe('danny-vm'); + expect(placed.target?.context).toMatchObject({ + environmentId: 'danny-vm', + bridgeWorkerId: 'w-danny', + codeWorkspace: { workspaceId: 'agents' }, + }); + }); + + it('refuses another principal the same admitted machine', async () => { + const ownerId = new Types.ObjectId(); + await registerMachine(ownerId); + const routing = await routingFor(new Types.ObjectId()); + + await expect(routing.describe(reviewer, flags)).resolves.toEqual({}); + await expect( + routing.place({ + agent: reviewer, + flags, + context: { executionId: 'run', hostArgs: { machine: 'danny-vm' } }, + }), + ).rejects.toMatchObject({ + argument: 'machine', + rejection: 'not_allowed', + }); + }); + + it('refuses a machine registered in another tenant', async () => { + const ownerId = new Types.ObjectId(); + await tenantStorage.run({ tenantId: 'tenant-a' }, () => registerMachine(ownerId)); + + await tenantStorage.run({ tenantId: 'tenant-b' }, async () => { + const routing = await routingFor(ownerId); + await expect(routing.describe(reviewer, flags)).resolves.toEqual({}); + await expect( + routing.place({ + agent: reviewer, + flags, + context: { executionId: 'run', hostArgs: { workspace: 'agents' } }, + }), + ).rejects.toMatchObject({ + argument: 'workspace', + rejection: 'not_allowed', + }); + }); + await tenantStorage.run({ tenantId: 'tenant-a' }, async () => { + const routing = await routingFor(ownerId); + await expect( + routing.place({ + agent: reviewer, + flags, + context: { executionId: 'run', hostArgs: { workspace: 'agents' } }, + }), + ).resolves.toMatchObject({ + target: { environmentId: 'danny-vm' }, + }); + }); + }); +}); diff --git a/packages/api/src/code/targets.spec.ts b/packages/api/src/code/targets.spec.ts new file mode 100644 index 00000000000..37ade394aa2 --- /dev/null +++ b/packages/api/src/code/targets.spec.ts @@ -0,0 +1,1401 @@ +import type { SubagentExecutionContext } from '@librechat/agents'; +import type { AppConfig } from '@librechat/data-schemas'; +import type { + SubagentCodeAgent, + SubagentCodeFlags, + SubagentCodeTargets, + SubagentCodeTargetParams, +} from './targets'; +import type { CodeEnvironmentConfig } from '~/agents/execution'; +import { + SUBAGENT_MACHINE_ARG, + SUBAGENT_WORKSPACE_ARG, + buildSubagentCodeHostArgs, + createSubagentCodeRouting, + guardRoutableSubagent, + getSubagentHostArgValues, + isSubagentHostArgsSupported, + placeSubagentOnCodeTarget, + resolveSubagentCodeTargets, + selectSubagentCodeTarget, +} from './targets'; +import { + resolveCodeExecutionContext, + assertCodeExecutionApprovalBinding, + captureCodeExecutionApprovalBinding, +} from '~/agents/execution'; +import { collectCodeExecutionProfileRoutes } from '~/agents/codeFilesSession'; +import { collectAttachedCodeEnvironmentAgentIds } from '~/agents/hitl/byom'; + +/** + * Stands in for the SDK's `SubagentHostArgumentError` until LibreChat depends + * on an `@librechat/agents` release that exports it; the real class wins once + * the installed SDK provides one. + */ +jest.mock('@librechat/agents', () => { + const actual = jest.requireActual('@librechat/agents'); + class MockSubagentHostArgumentError extends Error { + constructor(name: string, reason: 'unavailable' | 'not_allowed') { + super('Subagent host argument was rejected.'); + this.name = 'SubagentHostArgumentError'; + Object.assign(this, { argument: name, rejection: reason }); + } + } + return { + ...actual, + SubagentHostArgumentError: actual.SubagentHostArgumentError ?? MockSubagentHostArgumentError, + }; +}); + +const BASE_URL = 'https://bridge.example'; +const TOKEN_ENV = 'TEST_SUBAGENT_TARGETS_TOKEN'; + +const controlPlane: CodeEnvironmentConfig = { + id: 'control-plane', + name: 'Control plane', + type: 'attached', + owner: 'deployment', + baseURL: BASE_URL, + pairing: { allowPrincipalWorkers: true, tokenEnv: TOKEN_ENV }, +}; + +const machine = (id: string, workerId: string): CodeEnvironmentConfig => ({ + id, + name: `${id} machine`, + type: 'attached', + owner: 'principal', + baseURL: BASE_URL, + workerId, + controlPlaneId: 'control-plane', +}); + +const environments: CodeEnvironmentConfig[] = [ + controlPlane, + machine('laptop', 'w-laptop'), + machine('buildbox', 'w-buildbox'), + machine('spare', 'w-spare'), +]; + +const getAppConfig = jest.fn( + async () => + ({ + endpoints: { agents: { statefulCodeSessions: { environments: [controlPlane] } } }, + }) as AppConfig, +); + +const sealed = [ + { environmentId: 'buildbox', workspaceId: 'agents' }, + { environmentId: 'laptop', workspaceId: 'librechat' }, +]; + +type WorkerState = { online: boolean; workspaces: string[] }; + +function serveWorkers(workers: Record): jest.SpyInstance { + return jest.spyOn(globalThis, 'fetch').mockImplementation(async (input) => { + const url = String(input); + const workerId = /\/bridge\/workers\/([^/]+)\/status$/.exec(url)?.[1] ?? ''; + const worker = workers[workerId]; + return new Response( + JSON.stringify({ + protocolVersion: 1, + workerId, + online: worker?.online === true, + ready: worker?.online === true, + leaseExpiresInMs: 45_000, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file', 'list_files', 'execute_command'], + workspaces: (worker?.workspaces ?? []).map((id) => ({ id })), + }, + }, + }), + ); + }); +} + +const allOnline: Record = { + 'w-laptop': { online: true, workspaces: ['librechat'] }, + 'w-buildbox': { online: true, workspaces: ['agents'] }, + 'w-spare': { online: true, workspaces: ['scratch'] }, +}; + +function params(overrides: Partial = {}): SubagentCodeTargetParams { + return { + agentId: 'agent_reviewer', + statefulSessions: true, + environment: 'conversation', + environmentId: 'laptop', + environmentIds: ['buildbox', 'spare'], + allowEnvironmentSelection: true, + persistedSelections: sealed, + environments, + userId: 'user-1', + conversationId: 'convo-1', + getAppConfig, + ...overrides, + }; +} + +function rejection(error: unknown): { argument?: string; rejection?: string } { + return error as { argument?: string; rejection?: string }; +} + +describe('subagent code targets', () => { + beforeEach(() => { + process.env[TOKEN_ENV] = `token-${Math.random()}`; + }); + + afterEach(() => { + jest.restoreAllMocks(); + delete process.env[TOKEN_ENV]; + }); + + it('detects SDK support for per-call host arguments', () => { + expect(isSubagentHostArgsSupported()).toBe(true); + }); + + it('lists only admitted, authorized machines with a ready worker', async () => { + serveWorkers(allOnline); + + const { targets, unavailableMachines } = await resolveSubagentCodeTargets(params()); + + expect(targets.map(({ environmentId, workspaceId }) => [environmentId, workspaceId])).toEqual([ + ['buildbox', 'agents'], + ['laptop', 'librechat'], + ]); + expect(unavailableMachines.size).toBe(0); + expect(targets[0].context).toMatchObject({ + environmentId: 'buildbox', + environmentType: 'attached', + bridgeWorkerId: 'w-buildbox', + codeWorkspace: { environmentId: 'buildbox', workspaceId: 'agents' }, + }); + }); + + it('excludes a machine outside the child allowlist', async () => { + serveWorkers(allOnline); + + const { targets } = await resolveSubagentCodeTargets(params({ environmentIds: ['spare'] })); + + expect(targets.map((target) => target.environmentId)).toEqual(['laptop']); + }); + + it('keeps only the default machine when per-chat machine choice is disabled', async () => { + serveWorkers(allOnline); + + const { targets } = await resolveSubagentCodeTargets( + params({ allowEnvironmentSelection: false }), + ); + + expect(targets.map((target) => target.environmentId)).toEqual(['laptop']); + }); + + it('excludes a machine missing from the principal’s list', async () => { + serveWorkers(allOnline); + + const { targets } = await resolveSubagentCodeTargets( + params({ environments: environments.filter((environment) => environment.id !== 'buildbox') }), + ); + + expect(targets.map((target) => target.environmentId)).toEqual(['laptop']); + }); + + it('never adds an allowed machine the conversation has not admitted', async () => { + serveWorkers(allOnline); + + const { targets } = await resolveSubagentCodeTargets(params({ environmentIds: ['spare'] })); + + expect(targets.some((target) => target.environmentId === 'spare')).toBe(false); + }); + + it('reports an authorized machine whose worker is offline as unavailable', async () => { + serveWorkers({ ...allOnline, 'w-buildbox': { online: false, workspaces: [] } }); + + const resolution = await resolveSubagentCodeTargets(params()); + + expect(resolution.targets.map((target) => target.environmentId)).toEqual(['laptop']); + expect([...resolution.unavailableMachines]).toEqual(['buildbox']); + expect([...resolution.unavailableWorkspaces]).toEqual(['agents']); + }); + + it('treats a missing sealed workspace as unavailable', async () => { + serveWorkers({ ...allOnline, 'w-buildbox': { online: true, workspaces: ['other'] } }); + + const resolution = await resolveSubagentCodeTargets(params()); + + expect([...resolution.unavailableMachines]).toEqual(['buildbox']); + }); + + it('pins a child to the machine the conversation assigned to it', async () => { + serveWorkers(allOnline); + + const { targets } = await resolveSubagentCodeTargets( + params({ + persistedSelections: [ + { environmentId: 'buildbox', workspaceId: 'agents', agentIds: ['agent_reviewer'] }, + { environmentId: 'laptop', workspaceId: 'librechat' }, + ], + }), + ); + + expect(targets.map((target) => target.environmentId)).toEqual(['buildbox']); + }); + + it('lists nothing without stateful sessions or an admitted decision', async () => { + const fetchSpy = serveWorkers(allOnline); + + await expect( + resolveSubagentCodeTargets(params({ statefulSessions: false })), + ).resolves.toMatchObject({ targets: [] }); + await expect( + resolveSubagentCodeTargets(params({ persistedSelections: [] })), + ).resolves.toMatchObject({ targets: [] }); + await expect( + resolveSubagentCodeTargets(params({ persistedSelections: 'tampered' })), + ).resolves.toMatchObject({ targets: [] }); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it('prefers the sealed decision over a conflicting request', async () => { + serveWorkers(allOnline); + + const { targets, unavailableMachines } = await resolveSubagentCodeTargets( + params({ requestedSelections: [{ environmentId: 'spare', workspaceId: 'scratch' }] }), + ); + + expect(targets).toEqual([]); + expect([...unavailableMachines].sort()).toEqual(['buildbox', 'laptop']); + }); + + it('resolves the same route for the same inputs', async () => { + serveWorkers(allOnline); + + const first = await resolveSubagentCodeTargets(params()); + const second = await resolveSubagentCodeTargets(params()); + + expect(second.targets.map((target) => target.context)).toEqual( + first.targets.map((target) => target.context), + ); + expect(first.targets[0].context.codeSessionKey).toBe(second.targets[0].context.codeSessionKey); + }); +}); + +describe('subagent code host arguments', () => { + const resolution = ( + targets: Array<[string, string]>, + unavailable: Array<[string, string]> = [], + ): SubagentCodeTargets => ({ + targets: targets.map(([environmentId, workspaceId]) => ({ + environmentId, + workspaceId, + context: { + baseUrl: BASE_URL, + codeSessionKey: `session-${environmentId}`, + executionProfile: 'stateful', + statefulSessions: true, + environmentId, + environmentType: 'attached', + }, + })), + unavailableMachines: new Set(unavailable.map(([environmentId]) => environmentId)), + unavailableWorkspaces: new Set(unavailable.map(([, workspaceId]) => workspaceId)), + }); + + it('declares machine and workspace enums built only from reachable targets', () => { + const hostArgs = buildSubagentCodeHostArgs( + resolution([ + ['buildbox', 'agents'], + ['laptop', 'librechat'], + ]).targets, + ); + + expect(hostArgs?.[SUBAGENT_MACHINE_ARG].enum).toEqual(['buildbox', 'laptop']); + expect(hostArgs?.[SUBAGENT_WORKSPACE_ARG].enum).toEqual(['agents', 'librechat']); + for (const spec of Object.values(hostArgs ?? {})) { + expect(spec.description.length).toBeLessThan(1024); + } + }); + + it('omits workspace when two machines share a workspace ID, and declares nothing without targets', () => { + expect( + Object.keys( + buildSubagentCodeHostArgs( + resolution([ + ['buildbox', 'agents'], + ['laptop', 'agents'], + ]).targets, + ) ?? {}, + ), + ).toEqual([SUBAGENT_MACHINE_ARG]); + expect(buildSubagentCodeHostArgs([])).toBeUndefined(); + }); + + it('returns no target for omitted arguments', () => { + expect(selectSubagentCodeTarget(undefined, resolution([['laptop', 'librechat']]))).toBe( + undefined, + ); + expect(selectSubagentCodeTarget({}, resolution([['laptop', 'librechat']]))).toBe(undefined); + }); + + it('selects by machine, by workspace, or by a matching pair', () => { + const current = resolution([ + ['buildbox', 'agents'], + ['laptop', 'librechat'], + ]); + + expect(selectSubagentCodeTarget({ machine: 'buildbox' }, current)?.workspaceId).toBe('agents'); + expect(selectSubagentCodeTarget({ workspace: 'librechat' }, current)?.environmentId).toBe( + 'laptop', + ); + expect( + selectSubagentCodeTarget({ machine: 'laptop', workspace: 'librechat' }, current) + ?.environmentId, + ).toBe('laptop'); + }); + + it.each([ + [{ machine: 'buildbox' }, SUBAGENT_MACHINE_ARG, 'unavailable'], + [{ machine: 'spare' }, SUBAGENT_MACHINE_ARG, 'not_allowed'], + [{ machine: 'laptop', workspace: 'agents' }, SUBAGENT_WORKSPACE_ARG, 'not_allowed'], + [{ workspace: 'agents' }, SUBAGENT_WORKSPACE_ARG, 'unavailable'], + [{ workspace: 'scratch' }, SUBAGENT_WORKSPACE_ARG, 'not_allowed'], + ])('rejects %j with a %s refusal', (hostArgs, argument, expected) => { + const current = resolution([['laptop', 'librechat']], [['buildbox', 'agents']]); + let thrown: unknown; + try { + selectSubagentCodeTarget(hostArgs, current); + } catch (error) { + thrown = error; + } + expect(thrown).toBeInstanceOf(Error); + expect((thrown as Error).name).toBe('SubagentHostArgumentError'); + expect(rejection(thrown)).toMatchObject({ argument, rejection: expected }); + }); + + it('reads only string host arguments from a resolver context', () => { + expect(getSubagentHostArgValues(undefined)).toBeUndefined(); + expect( + getSubagentHostArgValues({ + executionId: 'run', + hostArgs: { machine: 'laptop', other: 3 }, + } as never), + ).toEqual({ machine: 'laptop' }); + expect(getSubagentHostArgValues({ executionId: 'run', hostArgs: {} } as never)).toBe(undefined); + }); + + it('pins the routed child document to the selected machine without mutating it', () => { + const agent = { + id: 'agent_reviewer', + code_environment_id: 'laptop', + code_environment_ids: ['buildbox', 'spare'], + }; + + const routed = placeSubagentOnCodeTarget(agent, { environmentId: 'buildbox' }); + + expect(routed).toEqual({ + id: 'agent_reviewer', + code_environment_id: 'buildbox', + code_environment_ids: ['buildbox'], + }); + expect(agent).toEqual({ + id: 'agent_reviewer', + code_environment_id: 'laptop', + code_environment_ids: ['buildbox', 'spare'], + }); + }); + + it('keeps a routed child on its machine even when its parent machine would be inherited', () => { + const routed = placeSubagentOnCodeTarget( + { id: 'agent_reviewer', code_environment_id: 'laptop', code_environment_ids: ['buildbox'] }, + { environmentId: 'buildbox' }, + ); + + const context = resolveCodeExecutionContext({ + statefulSessions: true, + environment: 'conversation', + environmentId: routed.code_environment_id, + environmentIds: routed.code_environment_ids, + allowEnvironmentSelection: true, + workspaceSelections: sealed, + inheritedEnvironments: new Map([['agent_reviewer', 'laptop']]), + environments, + userId: 'user-1', + agentId: 'agent_reviewer', + conversationId: 'convo-1', + }); + const inherited = resolveCodeExecutionContext({ + statefulSessions: true, + environment: 'conversation', + environmentId: 'spare', + environmentIds: ['laptop'], + allowEnvironmentSelection: true, + workspaceSelections: sealed, + inheritedEnvironments: new Map([['agent_reviewer', 'laptop']]), + environments, + userId: 'user-1', + agentId: 'agent_reviewer', + conversationId: 'convo-1', + }); + + expect(context.environmentId).toBe('buildbox'); + expect(inherited.environmentId).toBe('laptop'); + }); +}); + +describe('subagent code routing', () => { + const request = { + allowEnvironmentSelection: true, + persistedSelections: sealed, + environments, + userId: 'user-1', + conversationId: 'convo-1', + getAppConfig, + }; + const reviewer = { + id: 'agent_reviewer', + code_environment_id: 'laptop', + code_environment_ids: ['buildbox'], + }; + const flags = { statefulCodeSessions: true, statefulCodeEnvironment: 'conversation' }; + + beforeEach(() => { + process.env[TOKEN_ENV] = `token-${Math.random()}`; + }); + + afterEach(() => { + jest.restoreAllMocks(); + delete process.env[TOKEN_ENV]; + }); + + it('describes per-child choices and their approval targets', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + + const description = await routing.describe(reviewer, flags); + const lone = await routing.describe( + { id: 'agent_writer', code_environment_id: 'laptop' }, + flags, + ); + const stateless = await routing.describe(reviewer, { statefulCodeSessions: false }); + + expect(description.subagentHostArgs?.machine.enum).toEqual(['buildbox', 'laptop']); + expect(description.codeExecutionChoices?.map((context) => context.environmentId)).toEqual([ + 'buildbox', + 'laptop', + ]); + expect(lone.subagentHostArgs?.machine.enum).toEqual(['laptop']); + expect(stateless).toEqual({}); + }); + + const call = (hostArgs?: Record, parentRunId?: string) => ({ + executionId: `run-${Math.random()}`, + ...(parentRunId == null ? {} : { parentRunId }), + ...(hostArgs == null ? {} : { hostArgs }), + }); + + it('re-validates on use and routes only to a currently reachable machine', async () => { + const routing = createSubagentCodeRouting(request); + serveWorkers(allOnline); + + const placed = await routing.place({ + agent: reviewer, + flags, + context: call({ machine: 'buildbox' }), + }); + const omitted = await createSubagentCodeRouting(request).place({ + agent: reviewer, + flags, + context: call(), + }); + + expect(placed.agent.code_environment_id).toBe('buildbox'); + expect(placed.target?.context.bridgeWorkerId).toBe('w-buildbox'); + expect(omitted).toEqual({ agent: reviewer }); + + jest.restoreAllMocks(); + process.env[TOKEN_ENV] = `token-${Math.random()}`; + serveWorkers({ ...allOnline, 'w-buildbox': { online: false, workspaces: [] } }); + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'buildbox' }) }), + ).rejects.toMatchObject({ argument: 'machine', rejection: 'unavailable' }); + }); + + it('reserves a default route before initialization so a concurrent call cannot move it', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + + const defaultCall = call(); + const placed = await routing.place({ agent: reviewer, flags, context: defaultCall }); + const concurrent = routing.place({ + agent: reviewer, + flags, + context: call({ machine: 'buildbox' }), + }); + const sameMachine = await routing.place({ + agent: reviewer, + flags, + context: call({ machine: 'laptop' }), + }); + + expect(placed).toEqual({ agent: reviewer }); + await expect(concurrent).rejects.toMatchObject({ + argument: 'machine', + rejection: 'unavailable', + }); + expect(sameMachine.target?.environmentId).toBe('laptop'); + expect(routing.routesChildren(defaultCall.executionId)).toBe(false); + }); + + it('reserves the inherited default a parent machine gives a child', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting({ + ...request, + getInheritedEnvironments: () => new Map([['agent_reviewer', 'buildbox']]), + }); + + await routing.place({ agent: reviewer, flags, context: call() }); + + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }), + ).rejects.toMatchObject({ argument: 'machine' }); + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'buildbox' }) }), + ).resolves.toMatchObject({ target: { environmentId: 'buildbox' } }); + }); + + it('settles concurrent inherited calls from differently routed parents on one machine', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const parentOn = async (machine: string) => { + const parent = { + id: `agent_parent_${machine}`, + code_environment_id: machine, + code_environment_ids: [machine], + }; + const parentCall = call({ machine }); + const placement = await routing.place({ agent: parent, flags, context: parentCall }); + routing.attach(new Map(), { + agentId: parent.id, + context: parentCall, + placement, + codeExecutionContext: placement.target?.context, + toolContext: machine, + }); + return parentCall.executionId; + }; + const fromLaptop = await parentOn('laptop'); + const fromBuildbox = await parentOn('buildbox'); + + const placements = await Promise.allSettled([ + routing.place({ agent: reviewer, flags, context: call(undefined, fromLaptop) }), + routing.place({ agent: reviewer, flags, context: call(undefined, fromBuildbox) }), + ]); + const machines = placements.map((result) => + result.status === 'fulfilled' ? result.value.target?.environmentId : 'rejected', + ); + + expect( + new Set(machines.filter((machine) => machine !== undefined && machine !== 'rejected')).size, + ).toBe(1); + }); + + it('never claims a machine for a call canceled while its targets resolve', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const controller = new AbortController(); + + const canceled = routing.place({ + agent: reviewer, + flags, + context: { ...call({ machine: 'buildbox' }), signal: controller.signal }, + }); + controller.abort(); + + await expect(canceled).rejects.toBeDefined(); + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }), + ).resolves.toMatchObject({ target: { environmentId: 'laptop' } }); + }); + + it('gives back a machine whose initialization failed unless another call still holds it', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const failed = new Error('initialization failed'); + + const first = await routing.place({ agent: reviewer, flags, context: call() }); + const second = await routing.place({ agent: reviewer, flags, context: call() }); + await expect(routing.settle(first, Promise.reject(failed))).rejects.toBe(failed); + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'buildbox' }) }), + ).rejects.toMatchObject({ argument: 'machine' }); + + await expect(routing.settle(second, Promise.reject(failed))).rejects.toBe(failed); + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'buildbox' }) }), + ).resolves.toMatchObject({ target: { environmentId: 'buildbox' } }); + }); + + it('gives back a machine when its call is canceled before the child attaches', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const controller = new AbortController(); + + await routing.place({ + agent: reviewer, + flags, + context: { ...call({ machine: 'buildbox' }), signal: controller.signal }, + }); + controller.abort(); + + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }), + ).resolves.toMatchObject({ target: { environmentId: 'laptop' } }); + }); + + it('keeps a machine once the resolution that placed it succeeds', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const controller = new AbortController(); + const routedCall = { ...call({ machine: 'buildbox' }), signal: controller.signal }; + + const placement = await routing.settleExecution(routedCall, async () => { + const placed = await routing.place({ agent: reviewer, flags, context: routedCall }); + routing.attach(new Map(), { + agentId: reviewer.id, + context: routedCall, + placement: placed, + codeExecutionContext: placed.target?.context, + toolContext: 'buildbox', + }); + return placed; + }); + controller.abort(); + await expect(routing.settle(placement, Promise.reject(new Error('late')))).rejects.toThrow(); + + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }), + ).rejects.toMatchObject({ argument: 'machine' }); + }); + + it('gives back a child and its graph members when the rest of its resolution fails', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const member = { ...reviewer, id: 'agent_member' }; + const routedCall = call({ machine: 'buildbox' }); + const memberCall = { + executionId: `${routedCall.executionId}:graph:${member.id}`, + parentRunId: routedCall.executionId, + }; + await expect( + routing.settleExecution(routedCall, async () => { + const placed = await routing.place({ agent: reviewer, flags, context: routedCall }); + routing.attach(new Map(), { + agentId: reviewer.id, + context: routedCall, + placement: placed, + codeExecutionContext: placed.target?.context, + toolContext: 'buildbox', + }); + const memberPlaced = await routing.place({ agent: member, flags, context: memberCall }); + routing.attach(new Map(), { + agentId: member.id, + context: memberCall, + placement: memberPlaced, + codeExecutionContext: memberPlaced.target?.context, + toolContext: 'buildbox', + }); + expect(memberPlaced.target?.environmentId).toBe('buildbox'); + throw new Error('a graph member failed to load'); + }), + ).rejects.toThrow('a graph member failed to load'); + + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }), + ).resolves.toMatchObject({ target: { environmentId: 'laptop' } }); + await expect( + routing.place({ agent: member, flags, context: call({ machine: 'laptop' }) }), + ).resolves.toMatchObject({ target: { environmentId: 'laptop' } }); + }); + + it('gives back the per-agent tool context a failed resolution set', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const contexts = new Map(); + const member = { ...reviewer, id: 'agent_member' }; + const failedCall = call({ machine: 'buildbox' }); + const routedMember = { + executionId: `${failedCall.executionId}:graph:${member.id}`, + parentRunId: failedCall.executionId, + }; + + await expect( + routing.settleExecution(failedCall, async () => { + const placed = await routing.place({ agent: reviewer, flags, context: failedCall }); + routing.attach(contexts, { + agentId: reviewer.id, + context: failedCall, + placement: placed, + codeExecutionContext: placed.target?.context, + toolContext: 'reviewer-on-buildbox', + }); + const memberPlaced = await routing.place({ agent: member, flags, context: routedMember }); + routing.attach(contexts, { + agentId: member.id, + context: routedMember, + placement: memberPlaced, + codeExecutionContext: memberPlaced.target?.context, + toolContext: 'member-on-buildbox', + }); + expect(contexts.get(member.id)).toBe('member-on-buildbox'); + throw new Error('canceled'); + }), + ).rejects.toThrow('canceled'); + + expect(contexts.has(member.id)).toBe(false); + expect(contexts.has(reviewer.id)).toBe(false); + }); + + it('keeps the per-agent tool context of a surviving parallel placement', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const contexts = new Map(); + const select = async (toolContext: string, fail: boolean) => { + const selected = call({ machine: 'buildbox' }); + const placement = await routing.place({ agent: reviewer, flags, context: selected }); + routing.attach(contexts, { + agentId: reviewer.id, + context: selected, + placement, + codeExecutionContext: placement.target?.context, + toolContext, + }); + return { + settle: () => + routing.settleExecution(selected, async () => { + if (fail) throw new Error('canceled'); + }), + }; + }; + + const first = await select('first', true); + const second = await select('second', false); + await second.settle(); + await expect(first.settle()).rejects.toThrow('canceled'); + expect(contexts.get(reviewer.id)).toBe('second'); + + const third = await select('third', true); + const fourth = await select('fourth', true); + await expect(third.settle()).rejects.toThrow('canceled'); + await expect(fourth.settle()).rejects.toThrow('canceled'); + expect(contexts.get(reviewer.id)).toBe('second'); + }); + + it('queues live target probes below the shared status poller limit', async () => { + const machines = Array.from({ length: 12 }, (_, index) => machine(`m${index}`, `w-m${index}`)); + let inFlight = 0; + let peak = 0; + jest.spyOn(globalThis, 'fetch').mockImplementation(async (input) => { + inFlight += 1; + peak = Math.max(peak, inFlight); + await new Promise((resolve) => setTimeout(resolve, 5)); + inFlight -= 1; + const workerId = /\/bridge\/workers\/([^/]+)\/status$/.exec(String(input))?.[1] ?? ''; + return new Response( + JSON.stringify({ + protocolVersion: 1, + workerId, + online: true, + ready: true, + leaseExpiresInMs: 45_000, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: `ws-${workerId}` }], + }, + }, + }), + ); + }); + + const { targets } = await resolveSubagentCodeTargets({ + ...request, + agentId: 'agent_reviewer', + statefulSessions: true, + environment: 'conversation', + environmentId: 'm0', + environmentIds: machines.map((entry) => entry.id), + environments: [controlPlane, ...machines], + persistedSelections: machines.map((entry) => ({ + environmentId: entry.id, + workspaceId: `ws-w-${entry.id}`, + })), + }); + + expect(targets).toHaveLength(12); + expect(peak).toBeLessThanOrEqual(8); + }); + + it('gives back a machine whose initialization landed elsewhere', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const routedCall = call({ machine: 'buildbox' }); + + const placement = await routing.place({ agent: reviewer, flags, context: routedCall }); + expect(() => + routing.attach(new Map(), { + agentId: reviewer.id, + context: routedCall, + placement, + codeExecutionContext: { environmentId: 'laptop' }, + toolContext: 'laptop', + }), + ).toThrow(); + + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }), + ).resolves.toMatchObject({ target: { environmentId: 'laptop' } }); + }); + + it('offers and accepts no attached machine while run files are shared', async () => { + const fetchSpy = serveWorkers(allOnline); + const routing = createSubagentCodeRouting({ ...request, sharedRunFiles: true }); + + await expect(routing.describe(reviewer, flags)).resolves.toEqual({}); + await expect( + routing.place({ agent: reviewer, flags, context: call({ machine: 'buildbox' }) }), + ).rejects.toMatchObject({ argument: 'machine', rejection: 'not_allowed' }); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it('skips queued target probes once the resolution is canceled', async () => { + const fetchSpy = serveWorkers(allOnline); + const controller = new AbortController(); + controller.abort(new Error('canceled')); + + await expect(resolveSubagentCodeTargets(params({ signal: controller.signal }))).rejects.toThrow( + 'canceled', + ); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it('keeps the tool context a committed placement installed over a pending one', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const contexts = new Map(); + const pendingCall = call(); + const pending = await routing.place({ agent: reviewer, flags, context: pendingCall }); + routing.attach(contexts, { + agentId: reviewer.id, + context: pendingCall, + placement: pending, + codeExecutionContext: { environmentId: 'laptop' }, + toolContext: 'pending', + }); + const shared = await routing.place({ agent: reviewer, flags, context: {} }); + routing.attach(contexts, { + agentId: reviewer.id, + context: {}, + placement: shared, + codeExecutionContext: { environmentId: 'laptop' }, + toolContext: 'shared-member', + }); + + await expect( + routing.settleExecution(pendingCall, async () => { + throw new Error('canceled'); + }), + ).rejects.toThrow('canceled'); + expect(contexts.get(reviewer.id)).toBe('shared-member'); + }); + + it('attributes a workspace-only conflict to the workspace argument', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + + await routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }); + + await expect( + routing.place({ agent: reviewer, flags, context: call({ workspace: 'agents' }) }), + ).rejects.toMatchObject({ argument: 'workspace', rejection: 'unavailable' }); + }); + + it('keeps one machine per subagent for the whole request', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const contexts = new Map(); + + const first = await routing.place({ + agent: reviewer, + flags, + context: call({ machine: 'buildbox' }), + }); + const later = await routing.place({ agent: reviewer, flags, context: call() }); + const moved = routing.place({ agent: reviewer, flags, context: call({ machine: 'laptop' }) }); + + const writer = { + id: 'agent_writer', + code_environment_id: 'laptop', + code_environment_ids: ['buildbox'], + }; + const writerCall = call(); + const writerPlacement = await routing.place({ agent: writer, flags, context: writerCall }); + routing.attach(contexts, { + agentId: writer.id, + context: writerCall, + placement: writerPlacement, + codeExecutionContext: { environmentId: 'laptop' }, + toolContext: 'laptop', + }); + const writerMoved = routing.place({ + agent: writer, + flags, + context: call({ machine: 'buildbox' }), + }); + const writerSame = await routing.place({ + agent: writer, + flags, + context: call({ machine: 'laptop' }), + }); + + expect(first.target?.environmentId).toBe('buildbox'); + expect(later.target?.environmentId).toBe('buildbox'); + await expect(moved).rejects.toMatchObject({ argument: 'machine', rejection: 'unavailable' }); + expect(writerPlacement.target).toBeUndefined(); + await expect(writerMoved).rejects.toMatchObject({ + argument: 'machine', + rejection: 'unavailable', + }); + expect(writerSame.target?.environmentId).toBe('laptop'); + }); + + it('refuses an unavailable default only when no per-call route applies', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + + await expect( + routing.place({ agent: reviewer, flags, context: call(), unavailableReason: 'missing' }), + ).rejects.toMatchObject({ name: 'CodeWorkspaceSelectionError', reason: 'missing' }); + await expect( + routing.place({ + agent: reviewer, + flags, + context: call({ machine: 'buildbox' }), + unavailableReason: 'missing', + }), + ).resolves.toMatchObject({ target: { environmentId: 'buildbox' } }); + }); + + it('passes a routed parent machine to its children, through agents without code', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const contexts = new Map(); + const routeAndAttach = async ( + agent: SubagentCodeAgent, + context: ReturnType, + agentFlags: SubagentCodeFlags = flags, + ) => { + const placement = await routing.place({ agent, flags: agentFlags, context }); + routing.attach(contexts, { + agentId: agent.id, + context, + placement, + codeExecutionContext: placement.target?.context, + toolContext: placement.agent.code_environment_id ?? 'default', + }); + return placement; + }; + + const parentCall = call({ machine: 'buildbox' }); + await routeAndAttach(reviewer, parentCall); + const writerCall = call(undefined, parentCall.executionId); + const writer = await routeAndAttach( + { id: 'agent_writer', code_environment_id: 'laptop' }, + writerCall, + { statefulCodeSessions: false }, + ); + const grandchild = await routeAndAttach( + { id: 'agent_tester', code_environment_id: 'laptop', code_environment_ids: ['buildbox'] }, + call(undefined, writerCall.executionId), + ); + const unreachable = await routing.place({ + agent: { id: 'agent_docs', code_environment_id: 'laptop' }, + flags, + context: call(undefined, parentCall.executionId), + }); + const explicit = await routing.place({ + agent: { + id: 'agent_lint', + code_environment_id: 'laptop', + code_environment_ids: ['buildbox'], + }, + flags, + context: call({ machine: 'laptop' }, parentCall.executionId), + }); + + expect(writer).toEqual({ + agent: { id: 'agent_writer', code_environment_id: 'laptop' }, + childEnvironmentId: 'buildbox', + }); + expect(grandchild.target?.environmentId).toBe('buildbox'); + expect(grandchild.agent.code_environment_ids).toEqual(['buildbox']); + expect(unreachable.target).toBeUndefined(); + expect(explicit.target?.environmentId).toBe('laptop'); + }); + + it('passes a routed parent machine on from a self-spawned run it never placed', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const parentCall = call({ machine: 'buildbox' }); + const placement = await routing.place({ agent: reviewer, flags, context: parentCall }); + routing.attach(new Map(), { + agentId: reviewer.id, + context: parentCall, + placement, + codeExecutionContext: placement.target?.context, + toolContext: 'buildbox', + }); + const tester = { + id: 'agent_tester', + code_environment_id: 'laptop', + code_environment_ids: ['buildbox'], + }; + + const fromSelfRun = await routing.place({ + agent: tester, + flags, + context: { ...call(undefined, 'self-run-never-placed'), parentAgentId: reviewer.id }, + }); + const fromUnroutedAgent = await routing.place({ + agent: { ...tester, id: 'agent_docs' }, + flags, + context: { ...call(undefined, 'other-run'), parentAgentId: 'agent_unrouted' }, + }); + + expect(fromSelfRun.target?.environmentId).toBe('buildbox'); + expect(fromUnroutedAgent.target).toBeUndefined(); + }); + + it('passes a machine on from a self-spawned run of an agent without code', async () => { + serveWorkers(allOnline); + const routing = createSubagentCodeRouting(request); + const placeAndAttach = async ( + agent: SubagentCodeAgent, + context: ReturnType, + agentFlags: SubagentCodeFlags = flags, + ) => { + const placement = await routing.place({ agent, flags: agentFlags, context }); + routing.attach(new Map(), { + agentId: agent.id, + context, + placement, + codeExecutionContext: placement.target?.context, + toolContext: 'context', + }); + }; + const writer = { id: 'agent_writer', code_environment_id: 'laptop' }; + const tester = { + id: 'agent_tester', + code_environment_id: 'laptop', + code_environment_ids: ['buildbox'], + }; + const fromWriterSelfRun = (id: string) => + routing.place({ + agent: { ...tester, id }, + flags, + context: { ...call(undefined, `self-run-${Math.random()}`), parentAgentId: writer.id }, + }); + + const parentCall = call({ machine: 'buildbox' }); + await placeAndAttach(reviewer, parentCall); + const failedWriterCall = call(undefined, parentCall.executionId); + await expect( + routing.settleExecution(failedWriterCall, async () => { + await placeAndAttach(writer, failedWriterCall, { statefulCodeSessions: false }); + throw new Error('the writer failed to resolve'); + }), + ).rejects.toThrow('the writer failed to resolve'); + expect((await fromWriterSelfRun('agent_after_failure')).target).toBeUndefined(); + + const writerCall = call(undefined, parentCall.executionId); + await routing.settleExecution(writerCall, () => + placeAndAttach(writer, writerCall, { statefulCodeSessions: false }), + ); + await expect(fromWriterSelfRun('agent_after_success')).resolves.toMatchObject({ + target: { environmentId: 'buildbox' }, + }); + + const lintCall = call({ machine: 'laptop' }); + await placeAndAttach({ ...tester, id: 'agent_lint' }, lintCall); + const secondWriterCall = call(undefined, lintCall.executionId); + await routing.settleExecution(secondWriterCall, () => + placeAndAttach(writer, secondWriterCall, { statefulCodeSessions: false }), + ); + const ambiguous = await routing.place({ + agent: { ...tester, id: 'agent_docs' }, + flags, + context: { ...call(undefined, 'self-run-ambiguous'), parentAgentId: writer.id }, + }); + expect(ambiguous.target).toBeUndefined(); + }); + + it('rejects a revoked allowlist entry and a principal that lost the machine', async () => { + serveWorkers(allOnline); + + await expect( + createSubagentCodeRouting(request).place({ + agent: { ...reviewer, code_environment_ids: [] }, + flags, + context: call({ machine: 'buildbox' }), + }), + ).rejects.toMatchObject({ argument: 'machine', rejection: 'not_allowed' }); + await expect( + createSubagentCodeRouting({ + ...request, + environments: environments.filter((environment) => environment.id !== 'buildbox'), + }).place({ agent: reviewer, flags, context: call({ machine: 'buildbox' }) }), + ).rejects.toMatchObject({ argument: 'machine', rejection: 'not_allowed' }); + }); + + it('binds every per-call choice into paused-approval targets', async () => { + serveWorkers(allOnline); + const description = await createSubagentCodeRouting(request).describe(reviewer, flags); + const descriptor = { + id: reviewer.id, + codeExecutionContext: description.codeExecutionChoices?.[1], + codeExecutionChoices: description.codeExecutionChoices, + }; + + const binding = captureCodeExecutionApprovalBinding([descriptor]); + const changedChoice = { + ...descriptor, + codeExecutionChoices: [ + { ...description.codeExecutionChoices![0], bridgeWorkerId: 'w-replaced' }, + description.codeExecutionChoices![1], + ], + }; + + const changedDefault = { + ...descriptor, + codeExecutionContext: { ...descriptor.codeExecutionContext!, bridgeWorkerId: 'w-replaced' }, + }; + + expect(binding?.targets).toHaveLength(1); + expect(() => assertCodeExecutionApprovalBinding(binding, [descriptor])).not.toThrow(); + expect(() => assertCodeExecutionApprovalBinding(binding, [changedChoice])).toThrow( + 'The attached code environment changed', + ); + expect(() => assertCodeExecutionApprovalBinding(binding, [changedDefault])).toThrow( + 'The attached code environment changed', + ); + expect( + captureCodeExecutionApprovalBinding([ + { id: reviewer.id, codeExecutionContext: descriptor.codeExecutionContext }, + ])?.targets, + ).toHaveLength(1); + }); + + it('folds many per-call choices into one approval target per agent', () => { + const choice = (environmentId: string) => ({ + baseUrl: BASE_URL, + codeSessionKey: `session-${environmentId}`, + executionProfile: 'stateful' as const, + statefulSessions: true, + environmentId, + environmentType: 'attached' as const, + }); + const agents = Array.from({ length: 5 }, (_, agentIndex) => ({ + id: `agent_${agentIndex}`, + codeExecutionChoices: Array.from({ length: 32 }, (_, index) => choice(`m${index}`)), + })); + + const binding = captureCodeExecutionApprovalBinding(agents); + + expect(binding?.targets).toHaveLength(5); + expect(() => assertCodeExecutionApprovalBinding(binding, agents)).not.toThrow(); + }); +}); + +describe('unavailable subagents with per-call choices', () => { + const hostArgs = { machine: { description: 'Machine.', enum: ['buildbox'] } }; + + it('keeps a routable subagent callable and tells the parent to pick a machine', async () => { + const resolve = jest.fn(async () => 'config'); + + const guarded = guardRoutableSubagent({ + description: 'Reviews PRs.', + codeWorkspaceUnavailable: 'missing', + subagentHostArgs: hostArgs, + resolve, + }); + + expect(guarded.description).toContain('Unavailable in this conversation'); + expect(guarded.description).toContain('Pass "machine"'); + await expect(guarded.resolve({})).resolves.toBe('config'); + }); + + it('guards a subagent without choices exactly as before', async () => { + const resolve = jest.fn(async () => 'config'); + + const guarded = guardRoutableSubagent({ + description: 'Reviews PRs.', + codeWorkspaceUnavailable: 'missing', + resolve, + }); + const available = guardRoutableSubagent({ + description: 'Reviews PRs.', + subagentHostArgs: hostArgs, + resolve, + }); + + await expect(guarded.resolve({})).rejects.toMatchObject({ reason: 'missing' }); + expect(guarded.description).not.toContain('Pass "machine"'); + expect(resolve).not.toHaveBeenCalled(); + expect(available).toEqual({ description: 'Reviews PRs.', resolve }); + }); +}); + +describe('run-wide gates for per-call choices', () => { + const choice = { + baseUrl: BASE_URL, + codeSessionKey: 'execute_code:stateful:buildbox-session', + executionProfile: 'stateful' as const, + executionRouteKey: 'stateful:buildbox', + statefulSessions: true, + environmentId: 'buildbox', + environmentType: 'attached' as const, + }; + const parent = { + id: 'agent_parent', + lazySubagentConfigs: [ + { id: 'agent_reviewer', codeEnvAvailable: false, codeExecutionChoices: [choice] }, + ], + }; + + it('installs the attached-machine policy for a child that can only be routed there', () => { + expect([...collectAttachedCodeEnvironmentAgentIds([parent])]).toEqual(['agent_reviewer']); + }); + + it('primes skill files for every per-call route', () => { + expect(collectCodeExecutionProfileRoutes([parent])).toEqual([ + { codeExecutionContext: choice, codeSessionKeys: [choice.codeSessionKey] }, + ]); + }); +}); + +describe('subagent tool contexts', () => { + const executionContext = (...entries: Array<[string, string]>): SubagentExecutionContext => ({ + rootRunId: 'root', + hookSessionId: 'hooks', + depth: entries.length, + ancestry: entries.map(([subagentRunId, subagentAgentId]) => ({ + subagentRunId, + subagentType: subagentAgentId, + subagentKind: 'agent' as const, + subagentAgentId, + parentRunId: 'root', + })), + }); + const target = (environmentId: string) => ({ + environmentId, + workspaceId: `${environmentId}-workspace`, + context: { + baseUrl: BASE_URL, + codeSessionKey: environmentId, + executionProfile: 'stateful' as const, + statefulSessions: true, + environmentId, + }, + }); + const agent = { id: 'agent_reviewer' }; + const attach = ( + routing: ReturnType>, + contexts: Map, + executionId: string, + machine?: string, + ) => + routing.attach(contexts, { + agentId: agent.id, + context: { executionId }, + placement: machine == null ? { agent } : { agent, target: target(machine) }, + codeExecutionContext: { environmentId: machine ?? 'laptop' }, + toolContext: machine ?? 'laptop', + }); + + it('keeps concurrent siblings of one agent on their own routes', () => { + const routing = createSubagentCodeRouting({}); + const contexts = new Map(); + + attach(routing, contexts, 'run-a', 'buildbox'); + attach(routing, contexts, 'run-b'); + attach(routing, contexts, 'run-c', 'spare'); + + expect(routing.getToolContext(agent.id, executionContext(['run-a', agent.id]))).toBe( + 'buildbox', + ); + expect(routing.getToolContext(agent.id, executionContext(['run-c', agent.id]))).toBe('spare'); + expect(routing.getToolContext(agent.id, executionContext(['run-b', agent.id]))).toBe(undefined); + expect(contexts.get(agent.id)).toBe('laptop'); + expect(routing.isRouted('run-a')).toBe(true); + expect(routing.isRouted('run-b')).toBe(false); + }); + + it('lets an unplaced self-spawn defer to the routed run that spawned it', () => { + const routing = createSubagentCodeRouting({}); + attach(routing, new Map(), 'run-parent', 'buildbox'); + + expect( + routing.getToolContext( + agent.id, + executionContext(['run-parent', agent.id], ['run-self', agent.id]), + ), + ).toBe('buildbox'); + }); + + it('matches the innermost execution of the asking agent only', () => { + const routing = createSubagentCodeRouting({}); + attach(routing, new Map(), 'run-parent', 'buildbox'); + + expect( + routing.getToolContext( + 'agent_helper', + executionContext(['run-parent', agent.id], ['run-child', 'agent_helper']), + ), + ).toBeUndefined(); + expect(routing.getToolContext(agent.id, undefined)).toBeUndefined(); + expect( + routing.getToolContext('agent_other', executionContext(['run-parent', agent.id])), + ).toBeUndefined(); + }); + + it('seeds the per-agent entry from a routed child only when none exists', () => { + const routing = createSubagentCodeRouting({}); + const contexts = new Map(); + + attach(routing, contexts, 'run-a', 'buildbox'); + + expect(contexts.get(agent.id)).toBe('buildbox'); + }); + + it('fails closed when initialization did not land on the routed machine', () => { + const routing = createSubagentCodeRouting({}); + + expect(() => + routing.attach(new Map(), { + agentId: agent.id, + context: { executionId: 'run-a' }, + placement: { agent, target: target('buildbox') }, + codeExecutionContext: { environmentId: 'laptop' }, + toolContext: 'laptop', + }), + ).toThrow(); + expect(routing.isRouted('run-a')).toBe(false); + }); +}); diff --git a/packages/api/src/code/targets.ts b/packages/api/src/code/targets.ts new file mode 100644 index 00000000000..ea64a66ee5c --- /dev/null +++ b/packages/api/src/code/targets.ts @@ -0,0 +1,947 @@ +import * as agentsSdk from '@librechat/agents'; +import { isCodeWorkspaceSelections, stripAgentIdSuffix } from 'librechat-data-provider'; +import type { + CodeWorkspaceSelection, + StatefulCodeEnvironment, + CodeWorkspaceSelectionErrorReason, +} from 'librechat-data-provider'; +import type { SubagentExecutionContext } from '@librechat/agents'; +import type { CodeEnvironmentConfig, CodeExecutionContext } from '~/agents/execution'; +import type { CodeCapabilityConfigLoader } from './capabilities'; +import { + resolveCodeExecutionContext, + isExecutableAttachedEnvironment, + resolveAgentCodeEnvironmentRouting, +} from '~/agents/execution'; +import { CodeWorkspaceSelectionError, describeCodeWorkspaceUnavailableSubagent } from './errors'; +import { resolveCodeExecutionWorkspaceContext } from './capabilities'; +import { guardUnavailableSubagent } from '~/agents/lazySubagents'; +import { isCodeEnvironmentSelectionEnabled } from './protocol'; +import { createConcurrencyLimiter } from '~/utils/promise'; + +/** Subagent call property naming the attached machine the child runs on. */ +export const SUBAGENT_MACHINE_ARG = 'machine'; +/** Subagent call property naming the workspace the child opens on that machine. */ +export const SUBAGENT_WORKSPACE_ARG = 'workspace'; + +const MACHINE_DESCRIPTION = + 'ID of the attached code machine the subagent runs on. Pass the machine your own workspace tools use when the subagent must see the same files; omit it to let the subagent use its default placement.'; +const WORKSPACE_DESCRIPTION = + 'ID of the workspace the subagent opens. Each machine has one workspace in this conversation, so this is an alternative to machine; when both are passed they must belong together.'; + +/** + * Host argument declaration accepted by `@librechat/agents` subagent configs + * (`SubagentConfig.hostArgs`, added after 4.0.2). Declared locally so this + * package typechecks against older SDKs, which ignore the field. + */ +export interface SubagentCodeHostArgSpec { + description: string; + enum: string[]; +} + +export type SubagentCodeHostArgSpecs = Record; + +/** Validated per-call host argument values delivered to a lazy resolver. */ +export type SubagentHostArgValues = Readonly>; + +type HostArgumentRejection = 'unavailable' | 'not_allowed'; +type HostArgumentErrorConstructor = new ( + argument: string, + rejection: HostArgumentRejection, +) => Error; + +function getHostArgumentErrorConstructor(): HostArgumentErrorConstructor | undefined { + const candidate = (agentsSdk as { SubagentHostArgumentError?: HostArgumentErrorConstructor }) + .SubagentHostArgumentError; + return typeof candidate === 'function' ? candidate : undefined; +} + +/** Whether the installed SDK accepts per-call subagent host arguments. */ +export function isSubagentHostArgsSupported(): boolean { + return getHostArgumentErrorConstructor() != null; +} + +/** + * A refusal the SDK turns into a fixed, model-visible message naming the + * argument. The value and the reason detail never reach the model. + */ +export function createSubagentHostArgumentError( + argument: string, + rejection: HostArgumentRejection, +): Error { + const HostArgumentError = getHostArgumentErrorConstructor(); + return HostArgumentError == null + ? new Error(`Subagent host argument "${argument}" was rejected.`) + : new HostArgumentError(argument, rejection); +} + +/** Reads the host argument values the SDK attached to a resolver call. */ +export function getSubagentHostArgValues( + context: object | null | undefined, +): SubagentHostArgValues | undefined { + const values = (context as { hostArgs?: unknown } | null | undefined)?.hostArgs; + if (values == null || typeof values !== 'object' || Array.isArray(values)) { + return undefined; + } + const entries = Object.entries(values).filter( + (entry): entry is [string, string] => typeof entry[1] === 'string', + ); + return entries.length === 0 ? undefined : Object.fromEntries(entries); +} + +/** One reachable machine for a subagent, with its fully resolved live route. */ +export interface SubagentCodeTarget { + environmentId: string; + workspaceId: string; + context: CodeExecutionContext; +} + +export interface SubagentCodeTargets { + targets: SubagentCodeTarget[]; + /** Authorized machines whose worker or workspace failed live checks. */ + unavailableMachines: ReadonlySet; + /** Workspaces of those machines. */ + unavailableWorkspaces: ReadonlySet; +} + +export interface SubagentCodeTargetParams { + agentId: string; + /** Whether this child runs stateful code sessions at all in this request. */ + statefulSessions: boolean; + environment?: StatefulCodeEnvironment | string | null; + /** The child's saved default machine. */ + environmentId?: string | null; + /** The child's saved machine allowlist. */ + environmentIds?: readonly string[]; + allowEnvironmentSelection?: boolean; + /** The conversation's sealed decision; authoritative over the request. */ + persistedSelections?: unknown; + requestedSelections?: unknown; + /** The requesting principal's tenant-scoped machine list. */ + environments?: readonly CodeEnvironmentConfig[]; + userId?: string | null; + conversationId?: string | null; + getAppConfig?: CodeCapabilityConfigLoader; + /** Cancels the resolution; probes still queued are skipped. */ + signal?: AbortSignal; +} + +/** + * Live target probes run through the process-wide worker status poller, which + * refuses requests past its own concurrency ceiling; queue them well below it so + * concurrent descriptors never misreport a healthy machine as unavailable. + */ +const SUBAGENT_TARGET_PROBE_CONCURRENCY = 8; +const probeTarget = createConcurrencyLimiter(SUBAGENT_TARGET_PROBE_CONCURRENCY); + +const NO_TARGETS: SubagentCodeTargets = Object.freeze({ + targets: [], + unavailableMachines: new Set(), + unavailableWorkspaces: new Set(), +}); + +function getDefaultEnvironmentId(params: SubagentCodeTargetParams): string | undefined { + return ( + params.environmentId ?? + resolveAgentCodeEnvironmentRouting({ environments: params.environments }).defaultEnvironment?.id + ); +} + +/** + * Admitted selections this child may be routed to: its own sealed choice when + * the conversation pinned one to it, otherwise every admitted machine that is + * its default or on its allowlist and that the principal can use. + */ +function getAuthorizedSelections( + params: SubagentCodeTargetParams, + selections: CodeWorkspaceSelection[], +): CodeWorkspaceSelection[] { + const stableAgentId = stripAgentIdSuffix(params.agentId); + const owned = selections.find((selection) => selection.agentIds?.includes(stableAgentId)); + const pool = owned == null ? selections : [owned]; + const defaultId = getDefaultEnvironmentId(params); + const allowlist = isCodeEnvironmentSelectionEnabled(params.allowEnvironmentSelection) + ? new Set(params.environmentIds ?? []) + : new Set(); + return pool.filter( + ({ environmentId }) => + (environmentId === defaultId || allowlist.has(environmentId)) && + isExecutableAttachedEnvironment(environmentId, params.environments), + ); +} + +type CandidateResult = + | { status: 'ready'; target: SubagentCodeTarget } + | { status: 'unavailable'; selection: CodeWorkspaceSelection } + | { status: 'unauthorized' }; + +async function resolveCandidate( + params: SubagentCodeTargetParams, + selections: CodeWorkspaceSelection[], + selection: CodeWorkspaceSelection, +): Promise { + let base: CodeExecutionContext; + try { + base = resolveCodeExecutionContext({ + statefulSessions: true, + environment: params.environment, + environmentId: selection.environmentId, + environmentIds: params.environmentIds, + allowEnvironmentSelection: params.allowEnvironmentSelection, + workspaceSelections: selections, + environments: params.environments, + userId: params.userId, + agentId: params.agentId, + conversationId: params.conversationId, + }); + } catch { + return { status: 'unauthorized' }; + } + if (base.environmentId !== selection.environmentId || base.environmentType !== 'attached') { + return { status: 'unauthorized' }; + } + try { + const context = await resolveCodeExecutionWorkspaceContext({ + context: base, + requestedSelections: params.requestedSelections, + persistedSelections: params.persistedSelections, + environments: params.environments, + getAppConfig: params.getAppConfig, + }); + if (context.codeWorkspace?.workspaceId !== selection.workspaceId) { + return { status: 'unavailable', selection }; + } + return { + status: 'ready', + target: { + environmentId: selection.environmentId, + workspaceId: selection.workspaceId, + context, + }, + }; + } catch { + return { status: 'unavailable', selection }; + } +} + +/** + * Lists the machines a parent may route this child to: admitted in the + * conversation's sealed decision, reachable by the requesting principal, on + * the child's default or allowlist, and live on a ready worker. Never adds a + * machine the conversation has not admitted. + */ +export async function resolveSubagentCodeTargets( + params: SubagentCodeTargetParams, +): Promise { + if (!params.statefulSessions) { + return NO_TARGETS; + } + const selections = params.persistedSelections ?? params.requestedSelections; + if (!isCodeWorkspaceSelections(selections) || selections.length === 0) { + return NO_TARGETS; + } + const authorized = getAuthorizedSelections(params, selections); + if (authorized.length === 0) { + return NO_TARGETS; + } + const results = await Promise.all( + authorized.map((selection) => + probeTarget(() => + params.signal?.aborted === true + ? Promise.resolve({ status: 'unauthorized' }) + : resolveCandidate(params, selections, selection), + ), + ), + ); + throwIfCanceled(params.signal); + const targets: SubagentCodeTarget[] = []; + const unavailableMachines = new Set(); + const unavailableWorkspaces = new Set(); + for (const result of results) { + if (result.status === 'ready') { + targets.push(result.target); + } else if (result.status === 'unavailable') { + unavailableMachines.add(result.selection.environmentId); + unavailableWorkspaces.add(result.selection.workspaceId); + } + } + targets.sort((left, right) => (left.environmentId < right.environmentId ? -1 : 1)); + return { targets, unavailableMachines, unavailableWorkspaces }; +} + +/** + * Declares `machine` (and `workspace` when each machine's workspace is + * distinct) as enums of exactly the reachable targets. Returns nothing when + * the child has no reachable machine, leaving its call schema unchanged. + */ +export function buildSubagentCodeHostArgs( + targets: readonly SubagentCodeTarget[], +): SubagentCodeHostArgSpecs | undefined { + if (targets.length === 0) { + return undefined; + } + const workspaces = targets.map((target) => target.workspaceId); + const distinctWorkspaces = new Set(workspaces).size === workspaces.length; + return { + [SUBAGENT_MACHINE_ARG]: { + description: MACHINE_DESCRIPTION, + enum: targets.map((target) => target.environmentId), + }, + ...(distinctWorkspaces + ? { [SUBAGENT_WORKSPACE_ARG]: { description: WORKSPACE_DESCRIPTION, enum: workspaces } } + : {}), + }; +} + +/** + * Picks the target a call asked for, or `undefined` when it asked for none. + * Throws an SDK host-argument refusal for anything not currently reachable, + * so the parent sees which argument failed instead of a silent fallback. + */ +export function selectSubagentCodeTarget( + hostArgs: SubagentHostArgValues | undefined, + resolution: SubagentCodeTargets, +): SubagentCodeTarget | undefined { + const machine = hostArgs?.[SUBAGENT_MACHINE_ARG]; + const workspace = hostArgs?.[SUBAGENT_WORKSPACE_ARG]; + if (machine == null && workspace == null) { + return undefined; + } + if (machine != null) { + const target = resolution.targets.find((candidate) => candidate.environmentId === machine); + if (target == null) { + throw createSubagentHostArgumentError( + SUBAGENT_MACHINE_ARG, + resolution.unavailableMachines.has(machine) ? 'unavailable' : 'not_allowed', + ); + } + if (workspace != null && target.workspaceId !== workspace) { + throw createSubagentHostArgumentError(SUBAGENT_WORKSPACE_ARG, 'not_allowed'); + } + return target; + } + const matches = resolution.targets.filter((candidate) => candidate.workspaceId === workspace); + if (matches.length === 1) { + return matches[0]; + } + throw createSubagentHostArgumentError( + SUBAGENT_WORKSPACE_ARG, + matches.length === 0 && workspace != null && resolution.unavailableWorkspaces.has(workspace) + ? 'unavailable' + : 'not_allowed', + ); +} + +/** + * Routes a child agent document to the selected machine without mutating it. + * The machine becomes its default and its only allowlisted choice, so neither + * parent inheritance nor another admitted selection can move it elsewhere at + * any routing site that reads this document. + */ +export function placeSubagentOnCodeTarget< + T extends { code_environment_id?: string | null; code_environment_ids?: string[] | null }, +>(agent: T, target: Pick): T { + return { + ...agent, + code_environment_id: target.environmentId, + code_environment_ids: [target.environmentId], + }; +} + +/** Fails closed when initialization did not land on the selected route. */ +export function assertSubagentCodePlacement( + context: Pick | null | undefined, + target: Pick, +): void { + if (context?.environmentId !== target.environmentId) { + throw createSubagentHostArgumentError(SUBAGENT_MACHINE_ARG, 'unavailable'); + } +} + +/** + * Lists a lazy subagent whose default machine has no usable workspace. Without + * per-call choices it stays guarded exactly as before; with them its resolver + * stays live, the description tells the parent to pick a listed machine, and + * `SubagentCodeRouting.place` still refuses a call that names none. + */ +export function guardRoutableSubagent({ + description, + codeWorkspaceUnavailable, + subagentHostArgs, + resolve, +}: { + description?: string; + codeWorkspaceUnavailable?: CodeWorkspaceSelectionErrorReason; + subagentHostArgs?: SubagentCodeHostArgSpecs; + resolve: (context: TContext) => Promise; +}): { description?: string; resolve: (context: TContext) => Promise } { + if (!codeWorkspaceUnavailable || subagentHostArgs == null) { + return guardUnavailableSubagent({ description, codeWorkspaceUnavailable, resolve }); + } + return { + description: `${describeCodeWorkspaceUnavailableSubagent( + description, + codeWorkspaceUnavailable, + )} Pass "${SUBAGENT_MACHINE_ARG}" to run it on one of its listed machines.`, + resolve, + }; +} + +/** The saved child fields that decide its code route. */ +export interface SubagentCodeAgent { + id: string; + code_environment_id?: string | null; + code_environment_ids?: string[] | null; +} + +/** Child code flags the host already resolved for this request. */ +export interface SubagentCodeFlags { + statefulCodeSessions?: boolean; + statefulCodeEnvironment?: StatefulCodeEnvironment | string | null; +} + +/** Request-level inputs shared by every child of one parent request. */ +export type SubagentCodeRequest = Omit< + SubagentCodeTargetParams, + 'agentId' | 'statefulSessions' | 'environment' | 'environmentId' | 'environmentIds' +>; + +export interface SubagentCodeDescription { + /** `SubagentConfig.hostArgs` for the child's lazy descriptor. */ + subagentHostArgs?: SubagentCodeHostArgSpecs; + /** Alternate routes covered by paused-approval bindings. */ + codeExecutionChoices?: CodeExecutionContext[]; +} + +export interface SubagentCodePlacement { + agent: T; + /** The machine this execution was routed to per call. */ + target?: SubagentCodeTarget; + /** The per-call route this execution's own subagents inherit. */ + childEnvironmentId?: string; +} + +/** The parts of an SDK resolver context that routing reads. */ +export interface SubagentCodeCallContext { + /** This child execution; its own subagents name it as their `parentRunId`. */ + executionId?: string; + parentRunId?: string; + /** The agent that dispatched this call; read when its run was never placed (a self-spawn). */ + parentAgentId?: string; + hostArgs?: SubagentHostArgValues; + /** Cancels this call; a canceled call never claims a machine. */ + signal?: AbortSignal; +} + +/** What one lazy child resolution needs to know about its call and parent. */ +export interface SubagentCodePlacementInput { + agent: T; + flags: SubagentCodeFlags; + /** The SDK resolver context; graph members initialized outside a call pass only a signal. */ + context?: SubagentCodeCallContext | null; + /** Why the child's default route is unusable this request, if it is. */ + unavailableReason?: CodeWorkspaceSelectionErrorReason; +} + +/** + * Request-scoped routing for subagents a parent may place per call. Owns the + * per-execution state, so the legacy initializer only wires it in. + */ +export interface SubagentCodeRouting { + /** Builds a child's per-call machine choices; empty on SDKs without host arguments. */ + describe( + agent: SubagentCodeAgent, + flags: SubagentCodeFlags, + signal?: AbortSignal, + ): Promise; + /** + * Routes one child execution: an explicit call choice (re-validated against + * the current request), else the machine a per-call-routed parent runs on + * when this child may reach it, else the child's default route. A child whose + * default route is unavailable is refused unless one of those applies. + */ + place( + input: SubagentCodePlacementInput, + ): Promise>; + /** + * Stores a resolved child's tool context. A routed child is kept per + * execution and only seeds the per-agent entry when none exists, so it never + * replaces the route a default-placed sibling is using. + */ + attach( + contexts: Map, + input: { + agentId: string; + context?: Pick | null; + placement: SubagentCodePlacement; + codeExecutionContext?: Pick | null; + toolContext: TContext; + }, + ): void; + /** The routed tool context for the executing child, if it was routed per call. */ + getToolContext( + agentId: string | null | undefined, + executionContext: SubagentExecutionContext | null | undefined, + ): TContext | undefined; + /** Whether an execution was routed per call (its config must not be shared). */ + isRouted(executionId: string | null | undefined): boolean; + /** Whether an execution's own subagents (including graph members) inherit a per-call route. */ + routesChildren(executionId: string | null | undefined): boolean; + /** + * Awaits a placed child's initialization. A failed initialization gives back + * the machine its placement reserved unless another execution still holds it. + */ + settle( + placement: SubagentCodePlacement, + initialization: Promise, + ): Promise; + /** + * Runs one lazy child's whole resolution, through to the inputs handed to the SDK. + * The routes and per-agent tool contexts it and its graph members hold become + * permanent only when it succeeds (then `onCommit` publishes the caller's own + * state); when it fails they are given back unless another execution holds them. + */ + settleExecution( + context: Pick | null | undefined, + resolve: () => Promise, + hooks?: { onCommit?: () => void }, + ): Promise; +} + +function throwIfCanceled(signal?: AbortSignal): void { + if (signal?.aborted === true) { + throw signal.reason ?? new Error('Subagent resolution was aborted.'); + } +} + +type RouteClaim = { + environmentId: string | null; + routed: boolean; + /** Set once an execution on this route initialized; a committed claim is never released. */ + committed: boolean; + /** Placements still initializing on this route. */ + holders: number; +}; + +export function createSubagentCodeRouting({ + getInheritedEnvironments, + sharedRunFiles, + ...request +}: SubagentCodeRequest & { + /** Request-scoped parent inheritance (#16756), read when a default route is reserved. */ + getInheritedEnvironments?: () => ReadonlyMap | undefined; + /** + * Run file sharing is active: shared children must run on managed code, so no + * attached machine is offered or accepted per call. + */ + sharedRunFiles?: boolean; +}): SubagentCodeRouting { + const routedContexts = new Map(); + const childRoutes = new Map(); + /** Machines each agent that runs no stateful code passed on, for runs it never placed. */ + const passedOnByAgent = new Map>(); + /** + * One machine per subagent per request, as with parent inheritance: run-wide + * state such as the attached-machine permission policy is keyed by agent, so + * no call may move a subagent off the machine an earlier call settled on. + */ + const routeByAgent = new Map(); + /** Placements holding a claim until they are committed or given back. */ + const holds = new WeakMap< + object, + { agentId: string; executions: readonly string[]; stopWatching: () => void } + >(); + /** Pending placements by the execution whose resolution decides them. */ + const holdsByExecution = new Map>(); + /** + * Per-agent tool context entries pending placements use: each user's own context, + * and the entry that preceded them, so the entry always belongs to a live placement. + */ + const contextOwners = new Map< + string, + { + contexts: Map; + users: Map; + previous?: { value: TContext }; + } + >(); + const settleContext = (agentId: string, placement: object, keep: boolean): void => { + const owner = contextOwners.get(agentId); + const own = owner?.users.get(placement); + if (owner == null || own === undefined) { + return; + } + owner.users.delete(placement); + if (keep) { + owner.contexts.set(agentId, own); + contextOwners.delete(agentId); + return; + } + const survivor = owner.users.values().next(); + if (!survivor.done) { + if (owner.contexts.get(agentId) === own) { + owner.contexts.set(agentId, survivor.value); + } + return; + } + contextOwners.delete(agentId); + if (owner.previous != null) { + owner.contexts.set(agentId, owner.previous.value); + } else { + owner.contexts.delete(agentId); + } + }; + /** Ends a placement's hold, keeping its claim when `keep` is true. */ + const endHold = (placement: object, keep: boolean): void => { + const held = holds.get(placement); + if (held == null) { + return; + } + holds.delete(placement); + held.stopWatching(); + settleContext(held.agentId, placement, keep); + for (const executionId of held.executions) { + const pending = holdsByExecution.get(executionId); + pending?.delete(placement); + if (pending?.size === 0) { + holdsByExecution.delete(executionId); + } + } + const claimed = routeByAgent.get(held.agentId); + if (claimed == null || claimed.committed) { + return; + } + if (keep) { + claimed.committed = true; + return; + } + claimed.holders -= 1; + if (claimed.holders <= 0) { + routeByAgent.delete(held.agentId); + } + }; + const release = (placement: object): void => endHold(placement, false); + const commit = (placement: object): void => endHold(placement, true); + const recordPassedOn = (agentId: string, environmentId: string): void => { + const passedOn = passedOnByAgent.get(agentId) ?? new Set(); + passedOn.add(environmentId); + passedOnByAgent.set(agentId, passedOn); + }; + /** Pass-throughs recorded by a call, kept only if the resolution deciding them succeeds. */ + const passOnsByExecution = new Map< + string, + Array<{ agentId: string; environmentId: string; settled: boolean }> + >(); + const settleHoldsOf = (executionId: string | undefined, keep: boolean): void => { + if (executionId == null) { + return; + } + const pending = holdsByExecution.get(executionId); + for (const placement of [...(pending ?? [])]) { + endHold(placement, keep); + } + for (const passOn of passOnsByExecution.get(executionId) ?? []) { + if (!passOn.settled) { + passOn.settled = true; + if (keep) { + recordPassedOn(passOn.agentId, passOn.environmentId); + } + } + } + passOnsByExecution.delete(executionId); + }; + const conflicts = (agentId: string, environmentId: string | null): boolean => { + const claimed = routeByAgent.get(agentId); + return claimed != null && claimed.environmentId !== environmentId; + }; + /** + * Claims (or joins) a subagent's one route for a placement. A call's placement is + * decided by its own resolution and a routed graph member's by its parent's; a + * placement outside any call is committed when it attaches. + */ + const hold = ( + placement: SubagentCodePlacement, + route: { environmentId: string | null; routed: boolean }, + context?: SubagentCodeCallContext | null, + ): SubagentCodePlacement => { + const signal = context?.signal; + const agentId = placement.agent.id; + const claimed = routeByAgent.get(agentId); + if (claimed == null) { + routeByAgent.set(agentId, { ...route, committed: false, holders: 1 }); + } else { + claimed.holders += 1; + claimed.routed = claimed.routed || route.routed; + } + const onAbort = (): void => release(placement); + signal?.addEventListener('abort', onAbort, { once: true }); + const executions = [context?.executionId, context?.parentRunId].filter( + (executionId): executionId is string => executionId != null && executionId !== '', + ); + for (const executionId of executions) { + const pending = holdsByExecution.get(executionId) ?? new Set(); + pending.add(placement); + holdsByExecution.set(executionId, pending); + } + holds.set(placement, { + agentId, + executions, + stopWatching: () => signal?.removeEventListener('abort', onAbort), + }); + return placement; + }; + /** Live targets this request may place a subagent on per call. */ + /** Every per-call target is an attached machine, so a shared-file run probes none. */ + const resolveTargets = (params: SubagentCodeTargetParams): Promise => + sharedRunFiles === true ? Promise.resolve(NO_TARGETS) : resolveSubagentCodeTargets(params); + const paramsFor = ( + agent: SubagentCodeAgent, + flags: SubagentCodeFlags, + signal?: AbortSignal, + ): SubagentCodeTargetParams => ({ + ...request, + signal, + agentId: agent.id, + statefulSessions: flags.statefulCodeSessions === true, + environment: flags.statefulCodeEnvironment, + environmentId: agent.code_environment_id, + environmentIds: agent.code_environment_ids ?? undefined, + }); + /** The machine a call that names none lands on, as `initializeAgent` will resolve it. */ + const defaultRouteOf = (agent: SubagentCodeAgent, flags: SubagentCodeFlags): string | null => { + try { + const context = resolveCodeExecutionContext({ + statefulSessions: true, + environment: flags.statefulCodeEnvironment, + environmentId: agent.code_environment_id, + environmentIds: agent.code_environment_ids ?? undefined, + allowEnvironmentSelection: request.allowEnvironmentSelection, + workspaceSelections: request.persistedSelections ?? request.requestedSelections, + inheritedEnvironments: getInheritedEnvironments?.(), + environments: request.environments, + userId: request.userId, + agentId: agent.id, + conversationId: request.conversationId, + }); + return context.environmentType === 'attached' ? (context.environmentId ?? null) : null; + } catch { + return null; + } + }; + const routeTo = ( + agent: T, + target: SubagentCodeTarget, + context?: SubagentCodeCallContext | null, + ): SubagentCodePlacement => + hold( + { + agent: placeSubagentOnCodeTarget(agent, target), + target, + childEnvironmentId: target.environmentId, + }, + { environmentId: target.environmentId, routed: true }, + context, + ); + /** + * The machine a call's parent runs on per call: its execution's recorded route, else + * the routed claim of the parent agent, since an agent keeps one machine per request + * even in runs that were never placed (an `allowSelf` spawn reuses its inputs). + */ + const parentRouteOf = (context?: SubagentCodeCallContext | null): string | undefined => { + const recorded = context?.parentRunId ? childRoutes.get(context.parentRunId) : undefined; + if (recorded != null || !context?.parentAgentId) { + return recorded; + } + const claimed = routeByAgent.get(context.parentAgentId); + if (claimed?.routed === true) { + return claimed.environmentId ?? undefined; + } + /** Only an unambiguous pass-through is followed; otherwise static inheritance applies. */ + const passedOn = passedOnByAgent.get(context.parentAgentId); + return passedOn?.size === 1 ? [...passedOn][0] : undefined; + }; + /** The machine an omitted call should follow: this subagent's earlier per-call + * route, else its routed parent's machine unless that would move it. */ + const inheritedRoute = (agentId: string, parentEnvironmentId?: string): string | undefined => { + const claimed = routeByAgent.get(agentId); + if (claimed?.routed === true && claimed.environmentId != null) { + return claimed.environmentId; + } + if (parentEnvironmentId == null || conflicts(agentId, parentEnvironmentId)) { + return undefined; + } + return parentEnvironmentId; + }; + return { + async describe(agent, flags, signal) { + if (!isSubagentHostArgsSupported() || flags.statefulCodeSessions !== true) { + return {}; + } + const { targets } = await resolveTargets(paramsFor(agent, flags, signal)); + const subagentHostArgs = buildSubagentCodeHostArgs(targets); + return subagentHostArgs == null + ? {} + : { subagentHostArgs, codeExecutionChoices: targets.map((target) => target.context) }; + }, + async place({ agent, flags, context, unavailableReason }) { + const hostArgs = getSubagentHostArgValues(context); + const parentEnvironmentId = parentRouteOf(context); + const requested = + hostArgs?.[SUBAGENT_MACHINE_ARG] != null || hostArgs?.[SUBAGENT_WORKSPACE_ARG] != null; + if (requested) { + const resolution = await resolveTargets(paramsFor(agent, flags, context?.signal)); + throwIfCanceled(context?.signal); + const target = selectSubagentCodeTarget(hostArgs, resolution); + if (target != null) { + if (conflicts(agent.id, target.environmentId)) { + throw createSubagentHostArgumentError( + hostArgs?.[SUBAGENT_MACHINE_ARG] != null + ? SUBAGENT_MACHINE_ARG + : SUBAGENT_WORKSPACE_ARG, + 'unavailable', + ); + } + return routeTo(agent, target, context); + } + } + if (flags.statefulCodeSessions !== true) { + /** An agent that runs no stateful code passes its parent's machine on. */ + return parentEnvironmentId == null + ? { agent } + : { agent, childEnvironmentId: parentEnvironmentId }; + } + if (inheritedRoute(agent.id, parentEnvironmentId) != null) { + const { targets } = await resolveTargets(paramsFor(agent, flags, context?.signal)); + throwIfCanceled(context?.signal); + /** Re-read after the await: a concurrent call may have settled this subagent meanwhile. */ + const inherited = inheritedRoute(agent.id, parentEnvironmentId); + const target = + inherited == null + ? undefined + : targets.find((candidate) => candidate.environmentId === inherited); + if (target != null) { + return routeTo(agent, target, context); + } + if (routeByAgent.get(agent.id)?.routed === true) { + throw createSubagentHostArgumentError(SUBAGENT_MACHINE_ARG, 'unavailable'); + } + } + if (unavailableReason != null) { + throw new CodeWorkspaceSelectionError(unavailableReason); + } + /** Reserved now, not after initialization, so a concurrent call cannot claim another machine. */ + const claimed = routeByAgent.get(agent.id); + return hold( + { agent }, + claimed ?? { environmentId: defaultRouteOf(agent, flags), routed: false }, + context, + ); + }, + routesChildren(executionId) { + return executionId != null && childRoutes.has(executionId); + }, + attach(contexts, { agentId, context, placement, codeExecutionContext, toolContext }) { + const executionId = context?.executionId; + if (placement.target != null) { + try { + assertSubagentCodePlacement(codeExecutionContext, placement.target); + } catch (error) { + release(placement); + throw error; + } + } + if (executionId && placement.childEnvironmentId != null) { + childRoutes.set(executionId, placement.childEnvironmentId); + } + if (placement.target == null && placement.childEnvironmentId != null) { + const executions = [context?.executionId, context?.parentRunId].filter( + (id): id is string => id != null && id !== '', + ); + if (executions.length === 0) { + recordPassedOn(agentId, placement.childEnvironmentId); + } + const passOn = { agentId, environmentId: placement.childEnvironmentId, settled: false }; + for (const id of executions) { + passOnsByExecution.set(id, [...(passOnsByExecution.get(id) ?? []), passOn]); + } + } + if ((holds.get(placement)?.executions.length ?? 0) === 0) { + commit(placement); + } + if (placement.target == null && !routeByAgent.has(agentId)) { + routeByAgent.set(agentId, { + environmentId: codeExecutionContext?.environmentId ?? null, + routed: false, + committed: true, + holders: 0, + }); + } + /** A pending placement is a user of the per-agent entry until its resolution settles. */ + const register = (replace: boolean): void => { + const owner = contextOwners.get(agentId); + if (!holds.has(placement) && (replace || !contexts.has(agentId))) { + /** A committed placement now owns the entry; nothing pending may roll it back. */ + contextOwners.delete(agentId); + } + if (holds.has(placement)) { + const record = owner ?? { + contexts, + users: new Map(), + ...(contexts.has(agentId) ? { previous: { value: contexts.get(agentId)! } } : {}), + }; + record.users.set(placement, toolContext); + contextOwners.set(agentId, record); + } + if (replace || !contexts.has(agentId)) { + contexts.set(agentId, toolContext); + } + }; + if (placement.target == null || !executionId) { + register(true); + return; + } + routedContexts.set(executionId, { agentId, toolContext }); + register(false); + }, + getToolContext(agentId, executionContext) { + if (!agentId || routedContexts.size === 0) { + return undefined; + } + const ancestry = executionContext?.ancestry ?? []; + for (let index = ancestry.length - 1; index >= 0; index--) { + const entry = ancestry[index]; + if (entry.subagentAgentId !== agentId) { + continue; + } + /** An unplaced run of the agent (an `allowSelf` spawn) defers to an enclosing one. */ + const routed = routedContexts.get(entry.subagentRunId); + if (routed?.agentId === agentId) { + return routed.toolContext; + } + } + return undefined; + }, + isRouted(executionId) { + return executionId != null && routedContexts.has(executionId); + }, + async settle(placement, initialization) { + try { + return await initialization; + } catch (error) { + release(placement); + throw error; + } + }, + async settleExecution(context, resolve, hooks) { + try { + const resolved = await resolve(); + settleHoldsOf(context?.executionId, true); + hooks?.onCommit?.(); + return resolved; + } catch (error) { + settleHoldsOf(context?.executionId, false); + throw error; + } + }, + }; +} diff --git a/packages/api/src/files/provision/callback.spec.ts b/packages/api/src/files/provision/callback.spec.ts index 18662f5573b..28846bdb928 100644 --- a/packages/api/src/files/provision/callback.spec.ts +++ b/packages/api/src/files/provision/callback.spec.ts @@ -7,6 +7,7 @@ import type { CodeFileAgent } from '../code/queued'; import type { ServerRequest } from '~/types'; import { mergeCodeFilesIntoContext } from '~/agents/codeFilesSession'; import { prepareQueuedCodeFileContext } from '../code/queued'; +import { createSubagentCodeRouting } from '~/code/targets'; import { createProvisionFilesCallback } from './callback'; jest.mock('@librechat/data-schemas', () => ({ @@ -48,10 +49,14 @@ function buildHarness({ contexts, codeImpl, vectorImpl, + resolveExecutionContext, }: { contexts: Array<[string, ProvisionToolContext]>; codeImpl?: jest.Mock; vectorImpl?: jest.Mock; + resolveExecutionContext?: Parameters< + typeof createProvisionFilesCallback + >[0]['resolveExecutionContext']; }) { const provisionToCodeEnv = codeImpl ?? @@ -101,6 +106,7 @@ function buildHarness({ updateFile, updateCodeEnvRef, addEmbeddedEntity, + resolveExecutionContext, }), }; } @@ -301,6 +307,57 @@ describe('createProvisionFilesCallback', () => { expect(agent.provisionState.codeEnvFiles).toEqual([]); }); + it("provisions a per-call routed child into its own execution's context", async () => { + const shared = { provisionState: state([makeFile({ file_id: 'shared-route' })], []) }; + const routed = { provisionState: state([makeFile({ file_id: 'routed-route' })], []) }; + const routing = createSubagentCodeRouting({}); + routing.attach(new Map(), { + agentId: 'agent-a', + context: { executionId: 'run-routed' }, + placement: { + agent: { id: 'agent-a' }, + target: { + environmentId: 'machine', + workspaceId: 'workspace', + context: { + baseUrl: 'https://bridge.example', + codeSessionKey: 'machine', + executionProfile: 'stateful', + statefulSessions: true, + environmentId: 'machine', + }, + }, + }, + codeExecutionContext: { environmentId: 'machine' }, + toolContext: routed, + }); + const executionContext = { + rootRunId: 'root', + hookSessionId: 'hooks', + depth: 1, + ancestry: [ + { + subagentRunId: 'run-routed', + subagentType: 'agent-a', + subagentKind: 'agent' as const, + subagentAgentId: 'agent-a', + parentRunId: 'root', + }, + ], + }; + const { provisionFiles, provisionToCodeEnv } = buildHarness({ + contexts: [['agent-a', shared]], + resolveExecutionContext: routing.getToolContext, + }); + + await provisionFiles([Constants.EXECUTE_CODE], 'agent-a', undefined, executionContext); + await provisionFiles([Constants.EXECUTE_CODE], 'agent-a'); + + expect( + provisionToCodeEnv.mock.calls.map(([args]) => (args as { file: TFile }).file.file_id), + ).toEqual(['routed-route', 'shared-route']); + }); + it('scopes only the agent own resource files to its identity', async () => { /* A user can attach another agent's setup file to this conversation. Uploading it under * this agent would place it in a namespace this agent's other users share, so the diff --git a/packages/api/src/files/provision/callback.ts b/packages/api/src/files/provision/callback.ts index 6e48fb01230..7aee5546053 100644 --- a/packages/api/src/files/provision/callback.ts +++ b/packages/api/src/files/provision/callback.ts @@ -1,8 +1,8 @@ import { Constants } from '@librechat/agents'; import { logger } from '@librechat/data-schemas'; import { EToolResources } from 'librechat-data-provider'; +import type { CodeEnvFile, SubagentExecutionContext } from '@librechat/agents'; import type { TFile } from 'librechat-data-provider'; -import type { CodeEnvFile } from '@librechat/agents'; import type { CodeEnvRefUpdate, ProvisionService } from './service'; import type { ProvisionToolContext } from '../code/queued'; import type { ServerRequest } from '~/types'; @@ -42,6 +42,11 @@ export interface ProvisionCallbackDeps { req: ServerRequest; agentToolContexts: Map; resolvePrimaryAgentId?: () => string | undefined; + /** A child routed per call keeps its own context, keyed by its execution. */ + resolveExecutionContext?: ( + agentId: string | undefined, + executionContext: SubagentExecutionContext | undefined, + ) => ProvisionToolContext | undefined; provisionToCodeEnv: ProvisionService['provisionToCodeEnv']; provisionToVectorDB: ProvisionService['provisionToVectorDB']; updateFile: (update: FileUpdate) => Promise; @@ -90,6 +95,7 @@ export function createProvisionFilesCallback({ req, agentToolContexts, resolvePrimaryAgentId, + resolveExecutionContext, provisionToCodeEnv, provisionToVectorDB, updateFile, @@ -99,6 +105,7 @@ export function createProvisionFilesCallback({ toolNames: string[], agentId?: string, signal?: AbortSignal, + executionContext?: SubagentExecutionContext, ) => Promise { /* Agents in a handoff or parallel graph are initialized independently over the same * request attachments, so each holds its own ProvisionState for the same file. Keyed @@ -132,6 +139,7 @@ export function createProvisionFilesCallback({ toolNames: string[], agentId?: string, signal?: AbortSignal, + executionContext?: SubagentExecutionContext, ): Promise { signal?.throwIfAborted(); /* agentId is optional on this callback and a batch for the primary agent may omit @@ -141,11 +149,15 @@ export function createProvisionFilesCallback({ * while reading state from a fallback context would upload them as user-scoped, * then reconstruct them as agent-scoped on the next turn, and the entity id used to * query those vectors would no longer match the one they were stored under. */ - const { ctx, resolvedAgentId } = resolveProvisionContext({ - agentId, - agentToolContexts, - primaryAgentId: resolvePrimaryAgentId?.(), - }); + const placedCtx = resolveExecutionContext?.(agentId, executionContext); + const { ctx, resolvedAgentId } = + placedCtx != null + ? { ctx: placedCtx, resolvedAgentId: agentId } + : resolveProvisionContext({ + agentId, + agentToolContexts, + primaryAgentId: resolvePrimaryAgentId?.(), + }); if (!ctx?.provisionState) { return []; }