diff --git a/api/server/controllers/agents/__tests__/callbacks.spec.js b/api/server/controllers/agents/__tests__/callbacks.spec.js index 318adaf8bc8..b94e7f3e7a6 100644 --- a/api/server/controllers/agents/__tests__/callbacks.spec.js +++ b/api/server/controllers/agents/__tests__/callbacks.spec.js @@ -25,12 +25,15 @@ jest.mock('@librechat/api', () => ({ isCodeArtifactToolOutput: jest.requireActual('@librechat/api').isCodeArtifactToolOutput, isCodeSessionToolName: jest.requireActual('@librechat/api').isCodeSessionToolName, collectToolCallIds: jest.requireActual('@librechat/api').collectToolCallIds, + captureSubagentIdentity: jest.requireActual('@librechat/api').captureSubagentIdentity, + createToolTimingAdapter: jest.requireActual('@librechat/api').createToolTimingAdapter, })); jest.mock('@librechat/data-schemas', () => ({ logger: { debug: jest.fn(), error: jest.fn(), + warn: jest.fn(), }, })); @@ -361,6 +364,139 @@ describe('resumable event generation fencing', () => { expect(resumedPublish.mock.calls[0][0].activityEventId).not.toBe(firstUpdate.activityEventId); }); + it('publishes tool preparation and handoff into event-child activity', async () => { + const { GraphEvents, createContentAggregator } = jest.requireActual('@librechat/agents'); + const { getDefaultHandlers } = require('../callbacks'); + const publish = jest.fn().mockResolvedValue(undefined); + const { contentParts, stepMap, aggregateContent } = createContentAggregator(); + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + aggregateContent, + contentParts, + stepMap, + toolEndCallback: jest.fn(), + collectedUsage: [], + streamId: 'event-thread', + eventChildActivity: { + runId: 'event-thread', + parentRunId: 'parent-conversation', + subagentRunId: 'child-1', + subagentType: 'researcher', + subagentAgentId: 'agent-1', + parentAgentId: 'director', + publish, + }, + }); + const step = { + id: 'step-child', + index: 0, + type: 'tool_calls', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-child', name: 'query', args: '{}' }], + }, + }; + await handlers[GraphEvents.ON_RUN_STEP].handle(GraphEvents.ON_RUN_STEP, step); + await handlers[GraphEvents.ON_RUN_STEP_DELTA].handle(GraphEvents.ON_RUN_STEP_DELTA, { + id: 'step-child', + observed_at: 100, + delta: { type: 'tool_calls', tool_calls: [{ id: 'call-child', index: 0, args: '{' }] }, + }); + await handlers[StepEvents.ON_TOOL_CALLS_DISPATCHED].handle( + StepEvents.ON_TOOL_CALLS_DISPATCHED, + { + dispatched_at: 500, + toolCalls: [{ id: 'call-child', name: 'query', stepId: 'step-child' }], + }, + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(publish.mock.calls.map(([value]) => value.phase)).toEqual([ + 'run_step', + 'tool_preparation', + 'run_step_delta', + 'tool_calls_dispatched', + ]); + expect(publish.mock.calls[1][0].data).toEqual({ + id: 'step-child', + index: 0, + toolCallId: 'call-child', + observed_at: 100, + }); + expect(publish.mock.calls[3][0].data.toolCalls[0]).not.toHaveProperty('args'); + }); + + it('folds child dispatch and result into the parent-owned subagent tool part', async () => { + const { GraphEvents } = jest.requireActual('@librechat/agents'); + const { getDefaultHandlers } = require('../callbacks'); + const aggregators = new Map(); + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + aggregateContent: jest.fn(), + toolEndCallback: jest.fn(), + collectedUsage: [], + subagentAggregatorsByToolCallId: aggregators, + }); + const base = { + parentToolCallId: 'parent-call', + parentRunId: 'parent-run', + subagentRunId: 'child-run', + subagentType: 'researcher', + subagentAgentId: 'child-agent', + runId: 'parent-run', + }; + for (const event of [ + { + phase: 'run_step', + data: { + id: 'child-step', + index: 0, + type: 'tool_calls', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'child-call', name: 'query', args: '{}' }], + }, + }, + }, + { + phase: 'tool_preparation', + data: { id: 'child-step', toolCallId: 'child-call', observed_at: 100 }, + }, + { + phase: 'tool_calls_dispatched', + data: { + dispatched_at: 500, + toolCalls: [{ id: 'child-call', stepId: 'child-step', name: 'query' }], + }, + }, + { + phase: 'run_step_completed', + data: { + result: { + id: 'child-step', + index: 0, + type: 'tool_call', + completed_at: 540, + tool_call: { id: 'child-call', name: 'query', args: '{}', output: 'ok', progress: 1 }, + }, + }, + }, + ]) { + await handlers[GraphEvents.ON_SUBAGENT_UPDATE].handle(GraphEvents.ON_SUBAGENT_UPDATE, { + ...base, + ...event, + }); + } + expect(jest.requireMock('@librechat/data-schemas').logger.warn).not.toHaveBeenCalled(); + expect(aggregators.get('parent-call')?.contentParts[0]?.tool_call).toMatchObject({ + id: 'child-call', + toolPreparationStartedAt: 100, + toolPreparationDurationMs: 400, + toolDispatchedAt: 500, + toolExecutionDurationMs: 40, + output: 'ok', + }); + }); + it('forwards the originating job epoch with deferred attachments', () => { const { GenerationJobManager } = require('@librechat/api'); const { createAttachmentEmitter } = require('../callbacks'); @@ -1308,6 +1444,145 @@ describe('createToolEndCallback', () => { }); }); +describe('tool dispatch timing', () => { + it('forwards the SDK handoff and stores preparation and result intervals independently', async () => { + const { GraphEvents, createContentAggregator } = jest.requireActual('@librechat/agents'); + const { GenerationJobManager } = require('@librechat/api'); + const { getDefaultHandlers } = require('../callbacks'); + const { contentParts, stepMap, aggregateContent } = createContentAggregator(); + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + contentParts, + stepMap, + aggregateContent, + toolEndCallback: jest.fn(), + collectedUsage: [], + streamId: 'run', + }); + const step = { + id: 'step-1', + index: 0, + type: 'tool_calls', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-1', name: 'query', args: '{}' }], + }, + }; + await handlers[GraphEvents.ON_RUN_STEP].handle(GraphEvents.ON_RUN_STEP, step); + await handlers[GraphEvents.ON_RUN_STEP_DELTA].handle(GraphEvents.ON_RUN_STEP_DELTA, { + id: 'step-1', + observed_at: 1_000, + delta: { type: 'tool_calls', tool_calls: [{ id: 'call-1', index: 0, args: '{' }] }, + }); + const dispatched = { + dispatched_at: 248_000, + toolCalls: [{ id: 'call-1', name: 'query', stepId: 'step-1' }], + }; + await handlers[StepEvents.ON_TOOL_CALLS_DISPATCHED].handle( + StepEvents.ON_TOOL_CALLS_DISPATCHED, + dispatched, + ); + await handlers[GraphEvents.ON_RUN_STEP_COMPLETED].handle(GraphEvents.ON_RUN_STEP_COMPLETED, { + result: { + id: 'step-1', + completed_at: 248_340, + index: 0, + tool_call: { id: 'call-1', name: 'query', args: '{}', output: 'ok' }, + }, + }); + await handlers[GraphEvents.ON_RUN_STEP_CLOSED].handle(GraphEvents.ON_RUN_STEP_CLOSED, { + id: 'step-1', + index: 0, + type: 'tool_calls', + status: 'completed', + created_at: 1_000, + closed_at: 248_340, + }); + expect(GenerationJobManager.emitChunk).toHaveBeenCalledWith( + 'run', + { + event: StepEvents.ON_TOOL_PREPARATION, + data: { + id: 'step-1', + index: 0, + toolCallId: 'call-1', + observed_at: 1_000, + }, + }, + expect.anything(), + ); + expect(GenerationJobManager.emitChunk).toHaveBeenCalledWith( + 'run', + { event: StepEvents.ON_TOOL_CALLS_DISPATCHED, data: dispatched }, + expect.anything(), + ); + expect(contentParts[0].tool_call).toMatchObject({ + runStepDurationMs: 247_340, + runStepClosedAt: 248_340, + toolPreparationDurationMs: 247_000, + toolExecutionDurationMs: 340, + }); + }); + + it('keeps preparation across a new handler created after HITL approval', async () => { + const { GraphEvents, createContentAggregator } = jest.requireActual('@librechat/agents'); + const { getDefaultHandlers } = require('../callbacks'); + const { contentParts, stepMap, aggregateContent } = createContentAggregator(); + const handlers = getDefaultHandlers({ + res: { write: jest.fn() }, + contentParts, + stepMap, + aggregateContent, + toolEndCallback: jest.fn(), + collectedUsage: [], + toolTimingReplayEvents: [ + { + event: StepEvents.ON_TOOL_PREPARATION, + data: { + id: 'step-1', + index: 0, + toolCallId: 'call-1', + observed_at: 1_000, + }, + }, + ], + }); + await handlers[GraphEvents.ON_RUN_STEP].handle(GraphEvents.ON_RUN_STEP, { + id: 'step-1', + index: 0, + type: 'tool_calls', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-1', name: 'query', args: '{}' }], + }, + }); + await handlers[StepEvents.ON_TOOL_CALLS_DISPATCHED].handle( + StepEvents.ON_TOOL_CALLS_DISPATCHED, + { dispatched_at: 51_000, toolCalls: [{ id: 'call-1', name: 'query', stepId: 'step-1' }] }, + ); + await handlers[GraphEvents.ON_RUN_STEP_COMPLETED].handle(GraphEvents.ON_RUN_STEP_COMPLETED, { + result: { + id: 'step-1', + index: 0, + completed_at: 51_200, + tool_call: { id: 'call-1', name: 'query', args: '{}', output: 'ok' }, + }, + }); + await handlers[GraphEvents.ON_RUN_STEP_CLOSED].handle(GraphEvents.ON_RUN_STEP_CLOSED, { + id: 'step-1', + index: 0, + type: 'tool_calls', + status: 'completed', + created_at: 1_000, + closed_at: 51_200, + }); + expect(contentParts[0].tool_call).toMatchObject({ + toolPreparationDurationMs: 50_000, + toolExecutionDurationMs: 200, + }); + }); +}); + describe('tool input validation marker', () => { it('marks the streamed result and persisted content part out of band', async () => { const { GraphEvents, createContentAggregator } = jest.requireActual('@librechat/agents'); diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 9858b4d69ee..d762976f31b 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -34,6 +34,7 @@ const { getToolInputValidationDetails, captureSubagentIdentity, collectToolCallIds, + createToolTimingAdapter, } = require('@librechat/api'); const { processFileCitations } = require('~/server/services/Files/Citations'); const { processCodeOutput, runPreviewFinalize } = require('~/server/services/Files/Code/process'); @@ -342,10 +343,11 @@ function subagentPhaseToGraphEvent(event) { * @param {{ aggregateContent: Function, contentParts?: Array, stepMap?: Map }} aggregator * @param {SubagentUpdateEvent} event */ -function feedSubagentAggregator(aggregator, event) { +function feedSubagentAggregator(aggregator, event, applyChildTiming) { const graphEvent = subagentPhaseToGraphEvent(event); + if (graphEvent) aggregator.aggregateContent({ event: graphEvent, data: event.data }); + applyChildTiming(aggregator, event); if (!graphEvent) return; - aggregator.aggregateContent({ event: graphEvent, data: event.data }); /** The SDK aggregator intentionally projects run-step tool calls onto its * public content shape, so host-only routing metadata is not copied. Restore @@ -416,6 +418,7 @@ function getDefaultHandlers({ usageEmitSink = null, eventChildActivity = null, resolveMcpServerName = null, + toolTimingReplayEvents = [], }) { if (!res || !aggregateContent) { throw new Error( @@ -425,6 +428,8 @@ function getDefaultHandlers({ const eventActivityPhases = { [GraphEvents.ON_RUN_STEP]: 'run_step', [GraphEvents.ON_RUN_STEP_DELTA]: 'run_step_delta', + [StepEvents.ON_TOOL_PREPARATION]: 'tool_preparation', + [StepEvents.ON_TOOL_CALLS_DISPATCHED]: 'tool_calls_dispatched', [GraphEvents.ON_RUN_STEP_COMPLETED]: 'run_step_completed', [GraphEvents.ON_RUN_STEP_CLOSED]: 'run_step_closed', [GraphEvents.ON_MESSAGE_DELTA]: 'message_delta', @@ -514,7 +519,12 @@ function getDefaultHandlers({ } return emitForJob({ event: UsageEvents.ON_TOKEN_USAGE, data: payload }); }; + const toolTiming = createToolTimingAdapter({ + replayEvents: toolTimingReplayEvents, + emit: emitForJob, + }); const handlers = { + [StepEvents.ON_TOOL_CALLS_DISPATCHED]: toolTiming.dispatch, [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler( collectedUsage, collectedThoughtSignatures, @@ -595,6 +605,7 @@ function getDefaultHandlers({ const index = stepMap?.get(stepId)?.index; const part = typeof index === 'number' ? contentParts[index] : undefined; if (part?.type === ContentTypes.TOOL_CALL && part.tool_call) { + toolTiming.close(part.tool_call, stepId); part.tool_call.runStepStatus = data.status; Object.assign(part.tool_call, getRunStepCloseMetadata(data)); /** @@ -622,6 +633,7 @@ function getDefaultHandlers({ */ handle: async (event, data, metadata) => { aggregateContent({ event, data }); + await toolTiming.delta(data); if (data?.delta.type === StepTypes.TOOL_CALLS) { await emitForJob({ event, data }); } else if (checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node)) { @@ -657,6 +669,7 @@ function getDefaultHandlers({ agentId: metadata?.agent_id, }); } + toolTiming.completed(data); aggregateContent({ event, data }); const stepId = data?.result?.id; const runStep = stepMap?.get(stepId); @@ -769,7 +782,7 @@ function getDefaultHandlers({ } try { captureSubagentIdentity(aggregator, data); - feedSubagentAggregator(aggregator, data); + feedSubagentAggregator(aggregator, data, toolTiming.child); } catch (err) { logger.warn( `[ON_SUBAGENT_UPDATE] Failed to aggregate phase "${data?.phase}" for tool_call ${key}: ${err?.message ?? err}`, diff --git a/api/server/controllers/agents/resume.js b/api/server/controllers/agents/resume.js index 4caeceb2597..fe4e200798a 100644 --- a/api/server/controllers/agents/resume.js +++ b/api/server/controllers/agents/resume.js @@ -1854,6 +1854,7 @@ const ResumeAgentController = async (req, res, next, initializeClient, addTitle) checkpointNamespace, foregroundRunId: mcpRequestBody.messageId, requestBody: mcpRequestBody, + toolTimingReplayEvents: resumeState?.replayEvents, }); client = result.client; diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index d6137f281b9..9c3ab1d2ade 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -67,6 +67,7 @@ const { EModelEndpoint, resolveAllowedStatefulCodeEnvironments, removeCodeExecutionCaller, + pickUserToggleOptions, hasActivePiiFields, hasActivePiiPatterns, openapiToFunction, @@ -1018,6 +1019,8 @@ const getAgentHandler = async (req, res, expandProperties = false) => { provider: agent.provider, model: agent.model, model_parameters: getSafeModelParameters(agent.model_parameters), + tools: agent.tools, + tool_options: pickUserToggleOptions(agent.tool_options), isPublic: agent.isPublic, version: agent.version, // Safe metadata diff --git a/api/server/controllers/agents/v1.spec.js b/api/server/controllers/agents/v1.spec.js index bf704940ba7..8e6cda9d4de 100644 --- a/api/server/controllers/agents/v1.spec.js +++ b/api/server/controllers/agents/v1.spec.js @@ -1407,6 +1407,30 @@ describe('Agent Controllers - Mass Assignment Protection', () => { const response = mockRes.json.mock.calls[0][0]; expect(response.conversation_starters).toEqual(starters); }); + + test('should include the tools and only their user switches in the basic VIEW response', async () => { + const agent = await Agent.create({ + id: `agent_${uuidv4()}`, + name: 'Switchable Agent', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['web_search', 'search_mcp_docs'], + tool_options: { + web_search: { user_toggle: 'off', defer_loading: true }, + search_mcp_docs: { allowed_callers: ['direct'] }, + }, + }); + + mockReq.params = { id: agent.id }; + + await getAgentHandler(mockReq, mockRes); + + expect(mockRes.status).toHaveBeenCalledWith(200); + const response = mockRes.json.mock.calls[0][0]; + expect(response.tools).toEqual(['web_search', 'search_mcp_docs']); + expect(response.tool_options).toEqual({ web_search: { user_toggle: 'off' } }); + }); }); describe('getAgentVersionsHandler', () => { diff --git a/api/server/services/Endpoints/agents/build.js b/api/server/services/Endpoints/agents/build.js index 7c5ce99434c..ce8f1f8c533 100644 --- a/api/server/services/Endpoints/agents/build.js +++ b/api/server/services/Endpoints/agents/build.js @@ -20,6 +20,7 @@ const buildOptions = (req, endpoint, parsedBody, endpointType) => { agent_id: isAgentsEndpoint(endpoint) ? agent_id : Constants.EPHEMERAL_AGENT_ID, endpoint, model_parameters, + userToolSwitches: req.body?.ephemeralAgent ?? null, }).catch((error) => { logger.error(`[/agents/:${agent_id}] Error retrieving agent during build options step`, error); return undefined; diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 94dc2e5b339..68de908eb13 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -204,6 +204,7 @@ const initializeClientWithProvider = async ({ checkpointNamespace, foregroundRunId, requestBody, + toolTimingReplayEvents, upstreamTokenProvider, upstreamTokenProviderResolver, }) => { @@ -1824,6 +1825,7 @@ const initializeClientWithProvider = async ({ usageEmitSink, eventChildActivity, resolveMcpServerName, + toolTimingReplayEvents, }); const client = new AgentClient({ diff --git a/api/server/services/Files/Code/process.js b/api/server/services/Files/Code/process.js index a5c2bbdfe97..6f0409ba0d5 100644 --- a/api/server/services/Files/Code/process.js +++ b/api/server/services/Files/Code/process.js @@ -1068,6 +1068,7 @@ async function readSandboxFile({ * @param {string} params.file_path * @param {string} params.workspace_id * @param {string} [params.workspace_instance_id] + * @param {boolean} [params.linked_worktrees] * @param {number} params.start_line * @param {number} params.max_lines * @param {string} params.codeApiBaseUrl @@ -1080,6 +1081,7 @@ async function readWorkspaceFile({ file_path, workspace_id, workspace_instance_id, + linked_worktrees, start_line, max_lines, codeApiBaseUrl, @@ -1093,6 +1095,7 @@ async function readWorkspaceFile({ }) { return executeWorkspaceTool({ baseURL: codeApiBaseUrl, + linkedWorktrees: linked_worktrees, maxQueueWaitMs, codeApiMaxRetryWaitMs: req?.config?.endpoints?.agents?.codeApiMaxRetryWaitMs, maxRequestTimeoutMs, @@ -1122,6 +1125,7 @@ async function readWorkspaceFile({ * @param {string} params.query * @param {string} params.workspace_id * @param {string} [params.workspace_instance_id] + * @param {boolean} [params.linked_worktrees] * @param {string} [params.path] * @param {number} params.max_results * @param {string} params.codeApiBaseUrl @@ -1134,6 +1138,7 @@ async function searchWorkspace({ query, workspace_id, workspace_instance_id, + linked_worktrees, path, max_results, codeApiBaseUrl, @@ -1147,6 +1152,7 @@ async function searchWorkspace({ }) { return executeWorkspaceTool({ baseURL: codeApiBaseUrl, + linkedWorktrees: linked_worktrees, maxQueueWaitMs, codeApiMaxRetryWaitMs: req?.config?.endpoints?.agents?.codeApiMaxRetryWaitMs, maxRequestTimeoutMs, @@ -1175,6 +1181,7 @@ async function searchWorkspace({ * @param {Object} params * @param {string} params.workspace_id * @param {string} [params.workspace_instance_id] + * @param {boolean} [params.linked_worktrees] * @param {string} [params.path] * @param {string} [params.after_path] * @param {number} params.max_results @@ -1187,6 +1194,7 @@ async function searchWorkspace({ async function listWorkspaceFiles({ workspace_id, workspace_instance_id, + linked_worktrees, path, after_path, max_results, @@ -1201,6 +1209,7 @@ async function listWorkspaceFiles({ }) { return executeWorkspaceTool({ baseURL: codeApiBaseUrl, + linkedWorktrees: linked_worktrees, maxQueueWaitMs, codeApiMaxRetryWaitMs: req?.config?.endpoints?.agents?.codeApiMaxRetryWaitMs, maxRequestTimeoutMs, @@ -1230,6 +1239,7 @@ async function writeWorkspaceFile({ overwrite, workspace_id, workspace_instance_id, + linked_worktrees, codeApiBaseUrl, executionProfile, bridgeWorkerId, @@ -1241,6 +1251,7 @@ async function writeWorkspaceFile({ }) { return executeWorkspaceTool({ baseURL: codeApiBaseUrl, + linkedWorktrees: linked_worktrees, maxQueueWaitMs, codeApiMaxRetryWaitMs: req?.config?.endpoints?.agents?.codeApiMaxRetryWaitMs, maxRequestTimeoutMs, @@ -1270,6 +1281,7 @@ async function editWorkspaceFile({ expected_base_sha256, workspace_id, workspace_instance_id, + linked_worktrees, codeApiBaseUrl, executionProfile, bridgeWorkerId, @@ -1281,6 +1293,7 @@ async function editWorkspaceFile({ }) { return executeWorkspaceTool({ baseURL: codeApiBaseUrl, + linkedWorktrees: linked_worktrees, maxQueueWaitMs, codeApiMaxRetryWaitMs: req?.config?.endpoints?.agents?.codeApiMaxRetryWaitMs, maxRequestTimeoutMs, @@ -1309,6 +1322,7 @@ async function previewWorkspaceEdit({ edits, workspace_id, workspace_instance_id, + linked_worktrees, codeApiBaseUrl, executionProfile, bridgeWorkerId, @@ -1320,6 +1334,7 @@ async function previewWorkspaceEdit({ }) { return executeWorkspaceTool({ baseURL: codeApiBaseUrl, + linkedWorktrees: linked_worktrees, maxQueueWaitMs, codeApiMaxRetryWaitMs: req?.config?.endpoints?.agents?.codeApiMaxRetryWaitMs, maxRequestTimeoutMs, diff --git a/api/server/services/Files/Code/process.spec.js b/api/server/services/Files/Code/process.spec.js index 65175294c02..54564260f1b 100644 --- a/api/server/services/Files/Code/process.spec.js +++ b/api/server/services/Files/Code/process.spec.js @@ -2199,6 +2199,7 @@ describe('Code Process', () => { await expect( listWorkspaceFiles({ workspace_id: 'primary', + linked_worktrees: true, path: 'src', after_path: 'src/app.ts', max_results: 20, @@ -2226,6 +2227,7 @@ describe('Code Process', () => { expect(getCodeApiAuthHeaders).toHaveBeenNthCalledWith(2, mockReq, 'worker-user-1'); expect(mockExecuteWorkspaceTool).toHaveBeenCalledWith({ baseURL: 'https://attached-code.example.com/v1', + linkedWorktrees: true, authHeaders: expect.any(Function), maxQueueWaitMs: 0, codeApiMaxRetryWaitMs: undefined, diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index 5ce4e4d603b..51b97957cef 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -2308,6 +2308,7 @@ async function loadToolsForExecution({ baseUrl: codeExecutionContext.baseUrl, workspaceId: codeExecutionContext.codeWorkspace.workspaceId, workspaceInstanceId: codeExecutionContext.codeWorkspace.workspaceInstanceId, + linkedWorktrees: codeExecutionContext.codeWorkspace.linkedWorktrees, environment: codeExecutionContext.codeWorkspace.environment, gitIdentity: agent?.git_identity, maxTimeoutMs: resolveAttachedWorkspaceCommandTimeoutMax( diff --git a/client/src/Providers/BadgeRowContext.tsx b/client/src/Providers/BadgeRowContext.tsx index 59388825a28..7b81b73ec64 100644 --- a/client/src/Providers/BadgeRowContext.tsx +++ b/client/src/Providers/BadgeRowContext.tsx @@ -1,7 +1,7 @@ import React, { createContext, useContext, useEffect, useMemo, useRef } from 'react'; import { useSetRecoilState } from 'recoil'; import { Tools, Constants, LocalStorageKeys, AgentCapabilities } from 'librechat-data-provider'; -import type { TAgentsEndpoint, TEphemeralAgent } from 'librechat-data-provider'; +import type { AgentToolSwitches, TAgentsEndpoint, TEphemeralAgent } from 'librechat-data-provider'; import { useMCPServerManager, useSearchApiKeyForm, @@ -16,6 +16,8 @@ interface BadgeRowContextType { conversationId?: string | null; storageContextKey?: string; agentsConfig?: TAgentsEndpoint | null; + /** Present only in a saved agent's chat: the tools its creator made switchable. */ + agentToolSwitches?: AgentToolSwitches; skills: ReturnType; memory: ReturnType; webSearch: ReturnType; @@ -38,6 +40,7 @@ interface BadgeRowProviderProps { conversationId?: string | null; specName?: string | null; observeToolAuthorization?: boolean; + agentToolSwitches?: AgentToolSwitches; } export default function BadgeRowProvider({ @@ -46,6 +49,7 @@ export default function BadgeRowProvider({ conversationId, specName, observeToolAuthorization = false, + agentToolSwitches, }: BadgeRowProviderProps) { const lastContextKeyRef = useRef(''); const hasInitializedRef = useRef(false); @@ -272,14 +276,35 @@ export default function BadgeRowProvider({ isAuthenticated: true, }); + const agentServers = useMemo( + () => (agentToolSwitches ? Object.keys(agentToolSwitches.mcp) : undefined), + [agentToolSwitches], + ); const mcpServerManager = useMCPServerManager({ conversationId, storageContextKey, specName, + agentServers, ownsChatSelection: true, observeToolAuthorization, }); + /** In a saved agent's chat the MCP menu offers exactly the servers its creator made + * switchable, including ones `chatMenu: false` hides from plain chats. */ + const chatMcpServerManager = useMemo(() => { + if (!agentToolSwitches) { + return mcpServerManager; + } + const isSwitchable = ({ serverName }: { serverName: string }) => + serverName in agentToolSwitches.mcp; + const agentMCPServers = mcpServerManager.availableMCPServers.filter(isSwitchable); + return { + ...mcpServerManager, + availableMCPServers: agentMCPServers, + selectableServers: agentMCPServers, + }; + }, [mcpServerManager, agentToolSwitches]); + const value: BadgeRowContextType = { skills, memory, @@ -287,11 +312,12 @@ export default function BadgeRowProvider({ artifacts, fileSearch, agentsConfig, + agentToolSwitches, conversationId, storageContextKey, codeInterpreter, searchApiKeyForm, - mcpServerManager, + mcpServerManager: chatMcpServerManager, }; return {children}; diff --git a/client/src/Providers/__tests__/BadgeRowContext.agentSwitches.spec.tsx b/client/src/Providers/__tests__/BadgeRowContext.agentSwitches.spec.tsx new file mode 100644 index 00000000000..386c8457d58 --- /dev/null +++ b/client/src/Providers/__tests__/BadgeRowContext.agentSwitches.spec.tsx @@ -0,0 +1,71 @@ +import React from 'react'; +import { RecoilRoot } from 'recoil'; +import { render } from '@testing-library/react'; +import type { AgentToolSwitches } from 'librechat-data-provider'; +import BadgeRowProvider, { useBadgeRowContext } from '../BadgeRowContext'; + +const mockUseMCPServerManager = jest.fn(); + +const server = (serverName: string, chatMenu?: boolean) => ({ + serverName, + config: { type: 'sse', url: 'http://mcp', chatMenu }, +}); + +jest.mock('~/hooks', () => ({ + useMCPServerManager: (args: unknown) => mockUseMCPServerManager(args), + useSearchApiKeyForm: () => ({}), + useGetAgentsConfig: () => ({ agentsConfig: undefined }), + useToolToggle: () => ({}), +})); + +jest.mock('~/data-provider', () => ({ + useGetStartupConfig: () => ({ data: undefined }), +})); + +let context: ReturnType; +const Consumer = () => { + context = useBadgeRowContext(); + return null; +}; + +const renderProvider = (agentToolSwitches?: AgentToolSwitches) => + render( + + + + + , + ); + +describe('BadgeRowProvider agent switches', () => { + beforeEach(() => { + mockUseMCPServerManager.mockReset(); + const hidden = server('hidden', false); + const visible = server('visible'); + const other = server('other'); + mockUseMCPServerManager.mockReturnValue({ + availableMCPServers: [hidden, visible, other], + selectableServers: [visible, other], + }); + }); + + it("offers the agent's switchable servers, including one hidden from the chat menu", () => { + renderProvider({ builtins: {}, mcp: { hidden: true, visible: false } }); + + expect(mockUseMCPServerManager).toHaveBeenCalledWith( + expect.objectContaining({ agentServers: ['hidden', 'visible'] }), + ); + const names = (list: Array<{ serverName: string }>) => list.map((s) => s.serverName); + expect(names(context!.mcpServerManager.selectableServers)).toEqual(['hidden', 'visible']); + expect(names(context!.mcpServerManager.availableMCPServers)).toEqual(['hidden', 'visible']); + }); + + it('leaves the manager untouched outside a saved agent chat', () => { + renderProvider(); + + expect(mockUseMCPServerManager).toHaveBeenCalledWith( + expect.objectContaining({ agentServers: undefined }), + ); + expect(context!.mcpServerManager.selectableServers).toHaveLength(2); + }); +}); diff --git a/client/src/components/Chat/Input/BadgeRow.tsx b/client/src/components/Chat/Input/BadgeRow.tsx index 2a0e0100123..1b715dc684d 100644 --- a/client/src/components/Chat/Input/BadgeRow.tsx +++ b/client/src/components/Chat/Input/BadgeRow.tsx @@ -10,12 +10,14 @@ import React, { } from 'react'; import { Badge } from '@librechat/client'; import { useRecoilValue, useRecoilCallback } from 'recoil'; +import { Tools, getAgentToolSwitches } from 'librechat-data-provider'; import type { LucideIcon } from 'lucide-react'; import type { BadgeItem } from '~/common'; +import { useChatBadges, useApplyAgentToolSwitches } from '~/hooks'; +import { useGetAgentByIdQuery } from '~/data-provider'; import CodeInterpreter from './CodeInterpreter'; import { BadgeRowProvider } from '~/Providers'; import ToolsDropdown from './ToolsDropdown'; -import { useChatBadges } from '~/hooks'; import ToolDialogs from './ToolDialogs'; import FileSearch from './FileSearch'; import Artifacts from './Artifacts'; @@ -30,6 +32,7 @@ interface BadgeRowProps { onChange: (badges: Pick[]) => void; onToggle?: (badgeId: string, currentActive: boolean) => void; conversationId?: string | null; + agentId?: string | null; specName?: string | null; isSubmitting?: boolean; isInChat: boolean; @@ -145,6 +148,7 @@ const dragReducer = (state: DragState, action: DragAction): DragState => { function BadgeRow({ showEphemeralBadges, conversationId, + agentId, specName, isSubmitting, onChange, @@ -166,6 +170,13 @@ function BadgeRow({ const containerRectRef = useRef(null); const allBadges = useChatBadges(); + const { data: agent } = useGetAgentByIdQuery(agentId); + useApplyAgentToolSwitches({ agent, conversationId }); + const agentToolSwitches = useMemo( + () => (agentId && agent ? getAgentToolSwitches(agent) : undefined), + [agentId, agent], + ); + const hasSwitchableServer = Object.keys(agentToolSwitches?.mcp ?? {}).length > 0; const isEditing = useRecoilValue(store.isEditingBadges); const badges = useMemo( @@ -329,9 +340,10 @@ function BadgeRow({ specName={specName} isSubmitting={isSubmitting} observeToolAuthorization={showEphemeralBadges === true} + agentToolSwitches={agentToolSwitches} >
- {showEphemeralBadges === true && } + {(showEphemeralBadges === true || agentToolSwitches != null) && } {tempBadges.map((badge, index) => ( {dragState.draggedBadge && dragState.insertIndex === index && ghostBadge && ( @@ -382,6 +394,14 @@ function BadgeRow({ )} + {showEphemeralBadges !== true && agentToolSwitches != null && ( + <> + {agentToolSwitches.builtins[Tools.web_search] != null && } + {agentToolSwitches.builtins[Tools.execute_code] != null && } + {agentToolSwitches.builtins[Tools.file_search] != null && } + {hasSwitchableServer && } + + )} {ghostBadge && (
{ const canUseMemory = useHasMemoryAccess(); const showMemory = canUseMemory && memoryEnabled && user?.personalization?.memories !== false; + const agentToolSwitches = context?.agentToolSwitches; + const isSwitchable = (tool: Tools) => + agentToolSwitches == null || agentToolSwitches.builtins[tool] != null; + const [isPopoverActive, setIsPopoverActive] = useState(false); const isDisabled = disabled ?? false; const { @@ -162,7 +167,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => { const dropdownItems: MenuItemProps[] = []; - if (fileSearchEnabled && canUseFileSearch) { + if (fileSearchEnabled && canUseFileSearch && isSwitchable(Tools.file_search)) { dropdownItems.push({ onClick: handleFileSearchToggle, hideOnClick: false, @@ -194,7 +199,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => { }); } - if (canUseWebSearch && webSearchEnabled) { + if (canUseWebSearch && webSearchEnabled && isSwitchable(Tools.web_search)) { dropdownItems.push({ onClick: handleWebSearchToggle, hideOnClick: false, @@ -248,7 +253,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => { }); } - if (canUseSkills && skillsEnabled) { + if (canUseSkills && skillsEnabled && agentToolSwitches == null) { dropdownItems.push({ onClick: handleSkillsToggle, hideOnClick: false, @@ -280,7 +285,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => { }); } - if (showMemory) { + if (showMemory && agentToolSwitches == null) { dropdownItems.push({ onClick: handleMemoryToggle, hideOnClick: false, @@ -312,7 +317,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => { }); } - if (canRunCode && codeEnabled) { + if (canRunCode && codeEnabled && isSwitchable(Tools.execute_code)) { dropdownItems.push({ onClick: handleCodeInterpreterToggle, hideOnClick: false, @@ -346,7 +351,7 @@ const ToolsDropdown = ({ disabled }: ToolsDropdownProps) => { }); } - if (artifactsEnabled && setIsArtifactsPinned != null) { + if (artifactsEnabled && setIsArtifactsPinned != null && agentToolSwitches == null) { dropdownItems.push({ hideOnClick: false, render: (props) => ( diff --git a/client/src/components/Chat/Input/__tests__/BadgeRow.agentSwitches.test.tsx b/client/src/components/Chat/Input/__tests__/BadgeRow.agentSwitches.test.tsx new file mode 100644 index 00000000000..1945a6e286b --- /dev/null +++ b/client/src/components/Chat/Input/__tests__/BadgeRow.agentSwitches.test.tsx @@ -0,0 +1,176 @@ +import React from 'react'; +import { RecoilRoot } from 'recoil'; +import userEvent from '@testing-library/user-event'; +import { render, screen } from '@testing-library/react'; +import { Tools, mcpServerToggleKey } from 'librechat-data-provider'; +import type { Agent } from 'librechat-data-provider'; +import BadgeRow from '../BadgeRow'; + +let mockAgent: Partial | undefined; +let mockContext: Record = {}; +const mockApplySwitches = jest.fn(); +const mockProvider = jest.fn(); + +jest.mock('~/data-provider', () => ({ + useGetStartupConfig: () => ({ data: undefined }), + useGetAgentByIdQuery: () => ({ data: mockAgent }), +})); + +jest.mock('~/hooks', () => ({ + useChatBadges: () => [], + useApplyAgentToolSwitches: (args: unknown) => mockApplySwitches(args), + useLocalize: () => (key: string) => key, + useHasAccess: () => true, + useAuthContext: () => ({ user: undefined }), + useHasMemoryAccess: () => false, + useAgentCapabilities: () => ({ + codeEnabled: true, + memoryEnabled: true, + webSearchEnabled: true, + artifactsEnabled: false, + fileSearchEnabled: true, + skillsEnabled: true, + }), +})); + +jest.mock('~/Providers', () => ({ + useBadgeRowContext: () => mockContext, + BadgeRowProvider: (props: { children: React.ReactNode }) => { + mockProvider(props); + return <>{props.children}; + }, +})); + +jest.mock('~/components/Chat/Input/MCPSubMenu', () => ({ + __esModule: true, + default: () =>
, +})); +jest.mock('~/components/Chat/Input/ArtifactsSubMenu', () => ({ + __esModule: true, + default: () => null, +})); + +const stub = (testId: string) => ({ + __esModule: true, + default: () =>
, +}); +jest.mock('../CodeInterpreter', () => stub('code-badge')); +jest.mock('../FileSearch', () => stub('file-badge')); +jest.mock('../WebSearch', () => stub('web-badge')); +jest.mock('../MCPSelect', () => stub('mcp-badge')); +jest.mock('../Artifacts', () => stub('artifacts-badge')); +jest.mock('../Memory', () => stub('memory-badge')); +jest.mock('../Skills', () => stub('skills-badge')); +jest.mock('../ToolDialogs', () => stub('tool-dialogs')); + +const savedAgent: Partial = { + id: 'agent_1', + tools: [Tools.web_search, Tools.execute_code, 'search_mcp_docs'], + tool_options: { + [Tools.web_search]: { user_toggle: 'on' }, + [mcpServerToggleKey('docs')]: { user_toggle: 'off' }, + }, +}; + +const renderRow = (props: Partial>) => + render( + + + , + ); + +describe('BadgeRow agent switches', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockAgent = savedAgent; + mockContext = {}; + }); + + it('shows only the switchable tools of a saved agent and seeds their defaults', () => { + renderRow({ agentId: 'agent_1', showEphemeralBadges: false }); + + expect(screen.getByTestId('web-badge')).toBeInTheDocument(); + expect(screen.getByTestId('mcp-badge')).toBeInTheDocument(); + expect(screen.queryByTestId('code-badge')).not.toBeInTheDocument(); + expect(screen.queryByTestId('file-badge')).not.toBeInTheDocument(); + expect(screen.queryByTestId('skills-badge')).not.toBeInTheDocument(); + expect(screen.queryByTestId('memory-badge')).not.toBeInTheDocument(); + expect(mockApplySwitches).toHaveBeenCalledWith({ + agent: savedAgent, + conversationId: 'convo_1', + }); + expect(mockProvider).toHaveBeenCalledWith( + expect.objectContaining({ + agentToolSwitches: { builtins: { [Tools.web_search]: true }, mcp: { docs: false } }, + }), + ); + }); + + it('renders no switch badges for a saved agent without switches', () => { + mockAgent = { id: 'agent_1', tools: [Tools.web_search] }; + renderRow({ agentId: 'agent_1', showEphemeralBadges: false }); + + expect(screen.queryByTestId('web-badge')).not.toBeInTheDocument(); + expect(screen.queryByTestId('mcp-badge')).not.toBeInTheDocument(); + }); + + it('keeps the full badge set for ephemeral chats', () => { + renderRow({ showEphemeralBadges: true }); + + for (const testId of ['web-badge', 'code-badge', 'file-badge', 'skills-badge', 'mcp-badge']) { + expect(screen.getByTestId(testId)).toBeInTheDocument(); + } + expect(mockProvider).toHaveBeenCalledWith( + expect.objectContaining({ agentToolSwitches: undefined }), + ); + }); +}); + +describe('ToolsDropdown agent switches', () => { + const renderDropdown = async () => { + const ToolsDropdown = jest.requireActual('../ToolsDropdown').default; + render( +
+ +
, + ); + await userEvent.click(screen.getByRole('button', { name: 'Tools Options' })); + }; + + it('lists only the switchable tools for a saved agent', async () => { + mockContext = { + agentToolSwitches: { builtins: { [Tools.web_search]: true }, mcp: {} }, + webSearch: {}, + codeInterpreter: {}, + fileSearch: {}, + skills: {}, + memory: {}, + mcpServerManager: { availableMCPServers: [] }, + }; + await renderDropdown(); + + expect(await screen.findByText('com_ui_web_search')).toBeInTheDocument(); + expect(screen.queryByText('com_ui_run_code')).not.toBeInTheDocument(); + expect(screen.queryByText('com_assistants_file_search')).not.toBeInTheDocument(); + expect(screen.queryByText('com_ui_skills')).not.toBeInTheDocument(); + expect(screen.queryByTestId('mcp-submenu')).not.toBeInTheDocument(); + }); + + it('lists every enabled tool without agent switches', async () => { + mockContext = { + webSearch: {}, + codeInterpreter: {}, + fileSearch: {}, + skills: {}, + memory: {}, + mcpServerManager: { availableMCPServers: [{ serverName: 'docs' }] }, + }; + await renderDropdown(); + + expect(await screen.findByText('com_ui_web_search')).toBeInTheDocument(); + expect(screen.getByText('com_ui_run_code')).toBeInTheDocument(); + expect(screen.getByText('com_assistants_file_search')).toBeInTheDocument(); + expect(screen.getByText('com_ui_skills')).toBeInTheDocument(); + expect(screen.getByTestId('mcp-submenu')).toBeInTheDocument(); + }); +}); diff --git a/client/src/components/Chat/Messages/Content/Part.tsx b/client/src/components/Chat/Messages/Content/Part.tsx index 480f81b3406..775cc1efdb8 100644 --- a/client/src/components/Chat/Messages/Content/Part.tsx +++ b/client/src/components/Chat/Messages/Content/Part.tsx @@ -466,6 +466,10 @@ const Part = memo(function Part({ toolCall.backgroundTask?.cancelled === true ? 'cancelled' : toolCall.runStepStatus } runStepDurationMs={toolCall.runStepDurationMs} + toolPreparationStartedAt={toolCall.toolPreparationStartedAt} + toolDispatchedAt={toolCall.toolDispatchedAt} + toolPreparationDurationMs={toolCall.toolPreparationDurationMs} + toolExecutionDurationMs={toolCall.toolExecutionDurationMs} /> ); })(); diff --git a/client/src/components/Chat/Messages/Content/ProgressText.tsx b/client/src/components/Chat/Messages/Content/ProgressText.tsx index fba3dcc3ccc..d4dc702a2eb 100644 --- a/client/src/components/Chat/Messages/Content/ProgressText.tsx +++ b/client/src/components/Chat/Messages/Content/ProgressText.tsx @@ -8,6 +8,7 @@ import type { ToolCallPhase } from '~/utils/toolCallPhase'; import { cn, getRunStepDurationLabels } from '~/utils'; import CancelledIcon from './CancelledIcon'; import { useFailedReveal } from './reveal'; +import { ElapsedTimer } from '../Elapsed'; import { ROW_GLYPH_SLOT } from './rows'; import { useLocalize } from '~/hooks'; @@ -69,6 +70,9 @@ export default function ProgressText({ icon: iconProp, subtitle, durationMs, + toolPreparationDurationMs, + toolExecutionDurationMs, + phaseStartAt, hasInput = true, popover = false, isExpanded = false, @@ -89,8 +93,11 @@ export default function ProgressText({ authText?: string; icon?: React.ReactNode; subtitle?: string; - /** Wall-clock duration of the run step, from `PartMetadata.runStepDurationMs`. */ + /** Total run-step lifetime, not necessarily tool execution. */ durationMs?: number; + toolPreparationDurationMs?: number; + toolExecutionDurationMs?: number; + phaseStartAt?: number; hasInput?: boolean; popover?: boolean; isExpanded?: boolean; @@ -99,6 +106,15 @@ export default function ProgressText({ /** For locale-aware decimal formatting of the sub-10s duration value. */ const { i18n } = useTranslation(); const isRunning = phase === 'running'; + /** A server-authored phase stamp is an identity, not a browser clock origin. + * On reconnect we can only time from local receipt, never infer cross-host skew. */ + const phaseTimer = useRef<{ stamp?: number; receivedAt: number } | null>(null); + if (!isRunning || phaseStartAt == null) { + phaseTimer.current = null; + } else if (phaseTimer.current?.stamp !== phaseStartAt) { + phaseTimer.current = { stamp: phaseStartAt, receivedAt: Date.now() }; + } + const localPhaseStart = phaseTimer.current?.receivedAt; const rootRef = useRef(null); /** A header above asked for its failures. This control is the disclosure * every card renders, so answering here reaches a failed bash, code, @@ -139,10 +155,41 @@ export default function ProgressText({ * failed card "how long it took" is not the fact the reader needs — that * slot already carries the cancelled icon or the failure suffix. */ - const duration = - phase === 'completed' && isReportableRunStepDuration(durationMs) - ? getRunStepDurationLabels(durationMs, i18n.language) - : undefined; + const measured = toolPreparationDurationMs != null || toolExecutionDurationMs != null; + /** The formatter rounds to tenths below ten seconds. Shorter than 50 ms reads as 0.0s. */ + const showToolCallTime = + toolExecutionDurationMs != null && + Number.isFinite(toolExecutionDurationMs) && + toolExecutionDurationMs >= 50; + const durationParts = + phase !== 'completed' + ? [] + : [ + ...(isReportableRunStepDuration(toolPreparationDurationMs) + ? [ + { + label: localize('com_ui_tool_preparation_time'), + duration: getRunStepDurationLabels(toolPreparationDurationMs, i18n.language), + }, + ] + : []), + ...(showToolCallTime + ? [ + { + label: localize('com_ui_tool_call_time'), + duration: getRunStepDurationLabels(toolExecutionDurationMs, i18n.language), + }, + ] + : []), + ...(!measured && isReportableRunStepDuration(durationMs) + ? [ + { + label: localize('com_ui_tool_total_time'), + duration: getRunStepDurationLabels(durationMs, i18n.language), + }, + ] + : []), + ]; return ( @@ -194,22 +241,19 @@ export default function ProgressText({ {errorSuffix && ( · {errorSuffix} )} - {duration && ( - <> - {/* The compact form is the readable one on screen but a poor - thing to hear ("one point four s"), so it is hidden from - assistive technology and paired with a spoken equivalent. - Both live inside the button, so its accessible name carries - the duration — this is not an `aria-live` region and does not - re-announce. */} -
); } diff --git a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx index e52b62a2864..098bf5eaa5f 100644 --- a/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx +++ b/client/src/components/SidePanel/Agents/Tools/ItemDialog/sections/McpSection.tsx @@ -6,6 +6,7 @@ import { AgentCapabilities, Constants, splitMCPToolKey, + mcpServerToggleKey, normalizeServerName, buildServerNameAliases, stripServerNamePrefix, @@ -28,6 +29,7 @@ import MCPConfigDialog from '~/components/MCP/MCPConfigDialog'; import McpOAuthDialog from '~/components/MCP/McpOAuthDialog'; import { useAgentPanelContext } from '~/Providers'; import { getIconForItem } from '../../items/icons'; +import UserToggleSelect from '../UserToggleSelect'; import OptionToggle from '../../../OptionToggle'; import MCPToolItem from '../../../MCPToolItem'; import { Collapse } from '~/components/ui'; @@ -610,6 +612,11 @@ export default function McpSection({ item }: Props) {
+ + {configDialogProps && } ) : null} - {(canFavorite || canConfigure || showInfoOnly) && ( + {(canFavorite || canConfigure) && (
- {(canConfigure || showInfoOnly) && ( + {canConfigure && ( )} {canFavorite && ( diff --git a/client/src/components/SidePanel/Agents/Tools/hooks.ts b/client/src/components/SidePanel/Agents/Tools/hooks.ts index ad78248132e..73354fa3cbc 100644 --- a/client/src/components/SidePanel/Agents/Tools/hooks.ts +++ b/client/src/components/SidePanel/Agents/Tools/hooks.ts @@ -49,21 +49,6 @@ export function useBuiltinAuthMap(): Map { }, [data, isLoading]); } -/** - * Whether `web_search` uses USER_PROVIDED auth (a user-managed key). When false - * the deployment uses SYSTEM_DEFINED keys, so there is nothing for the user to - * configure. Shares the `useBuiltinAuthMap` React Query key, so it adds no - * request. Threaded into `buildCatalog` so the card/row affordance is decided - * synchronously (cog vs info) without a per-row hook. - */ -export function useWebSearchUserProvided(): boolean { - const { data } = useVerifyAgentToolAuth({ toolId: Tools.web_search }, { retry: 1 }); - return useMemo( - () => data?.authTypes?.some(([, authType]) => authType === AuthType.USER_PROVIDED) ?? false, - [data], - ); -} - /** * Resolves whether the Memory capability should be offered in the builder. * Mirrors the legacy `AgentConfig` gate: the admin must enable the `memory` @@ -176,7 +161,6 @@ export function useAgentItems({ permission: Permissions.USE, }); const showMemory = useShowMemory(); - const webSearchUserProvided = useWebSearchUserProvided(); const builtinAuthMap = useBuiltinAuthMap(); const toolsField = useWatch({ control, name: 'tools' }); @@ -211,7 +195,6 @@ export function useAgentItems({ fileSearch: hasFileSearchAccess, }, showMemory, - webSearchUserProvided, builtinAuthMap, }), [ @@ -226,7 +209,6 @@ export function useAgentItems({ hasFileSearchAccess, skillsPermission, showMemory, - webSearchUserProvided, builtinAuthMap, ], ); diff --git a/client/src/components/SidePanel/Agents/Tools/items/__tests__/catalog.spec.ts b/client/src/components/SidePanel/Agents/Tools/items/__tests__/catalog.spec.ts index 94c0945bc2d..67062642ac3 100644 --- a/client/src/components/SidePanel/Agents/Tools/items/__tests__/catalog.spec.ts +++ b/client/src/components/SidePanel/Agents/Tools/items/__tests__/catalog.spec.ts @@ -1,6 +1,7 @@ import { AgentCapabilities } from 'librechat-data-provider'; import type { BuildCatalogInputs } from '../catalog'; import { makePlugin, makeSkill, makeAction } from 'test/itemFactories'; +import { deriveSelectedItems } from '../selectors'; import { buildCatalog } from '../catalog'; const emptyInputs: BuildCatalogInputs = { @@ -47,19 +48,34 @@ describe('buildCatalog', () => { expect(buildCatalog({ ...emptyInputs, showMemory: true }).find(memoryId)).toBeDefined(); }); - test('flags web_search userProvidedAuth from the webSearchUserProvided input', () => { - const findWebSearch = (inputs: BuildCatalogInputs) => - buildCatalog(inputs).find( - (i) => i.kind === 'builtin' && i.id === AgentCapabilities.web_search, - ); - const base = { - ...emptyInputs, - agentsConfig: { capabilities: [AgentCapabilities.web_search] }, + test('offers no web search item, and so no switch setting, when capabilities or the role block it', () => { + const form = { + execute_code: false, + web_search: true, + file_search: false, + memory: false, + artifacts: undefined, + tools: [], + skills: [], + context_files: [], + knowledge_files: [], + code_files: [], }; - const userProvided = findWebSearch({ ...base, webSearchUserProvided: true }); - const systemDefined = findWebSearch({ ...base, webSearchUserProvided: false }); - expect(userProvided?.kind === 'builtin' && userProvided.userProvidedAuth).toBe(true); - expect(systemDefined?.kind === 'builtin' && systemDefined.userProvidedAuth).toBe(false); + const webSearchEnabled = { capabilities: [AgentCapabilities.web_search] }; + const selectedIds = (inputs: BuildCatalogInputs) => + deriveSelectedItems(form, buildCatalog(inputs), []).map((item) => item.id); + + expect(selectedIds(emptyInputs)).toEqual([]); + expect( + selectedIds({ + ...emptyInputs, + agentsConfig: webSearchEnabled, + permissions: { ...emptyInputs.permissions, webSearch: false }, + }), + ).toEqual([]); + expect(selectedIds({ ...emptyInputs, agentsConfig: webSearchEnabled })).toEqual([ + AgentCapabilities.web_search, + ]); }); test('surfaces ask_user_question as a BUILTIN (not a plugin) when the server lists it', () => { diff --git a/client/src/components/SidePanel/Agents/Tools/items/__tests__/configurable.spec.ts b/client/src/components/SidePanel/Agents/Tools/items/__tests__/configurable.spec.ts index 52682526877..fa7550ab1af 100644 --- a/client/src/components/SidePanel/Agents/Tools/items/__tests__/configurable.spec.ts +++ b/client/src/components/SidePanel/Agents/Tools/items/__tests__/configurable.spec.ts @@ -14,9 +14,9 @@ describe('hasConfigurableSettings', () => { expect(hasConfigurableSettings(builtin('memory'))).toBe(true); }); - test('web_search is configurable only when auth is user-provided', () => { - expect(hasConfigurableSettings(builtin('web_search'))).toBe(false); - expect(hasConfigurableSettings(builtin('web_search', { userProvidedAuth: false }))).toBe(false); + test('web_search is always configurable because it carries the user toggle', () => { + expect(hasConfigurableSettings(builtin('web_search'))).toBe(true); + expect(hasConfigurableSettings(builtin('web_search', { userProvidedAuth: false }))).toBe(true); expect(hasConfigurableSettings(builtin('web_search', { userProvidedAuth: true }))).toBe(true); }); diff --git a/client/src/components/SidePanel/Agents/Tools/items/catalog.ts b/client/src/components/SidePanel/Agents/Tools/items/catalog.ts index 6080c9c31e8..5c7093a0530 100644 --- a/client/src/components/SidePanel/Agents/Tools/items/catalog.ts +++ b/client/src/components/SidePanel/Agents/Tools/items/catalog.ts @@ -56,12 +56,6 @@ export interface BuildCatalogInputs { * `agentsConfig.capabilities` here. */ showMemory?: boolean; - /** - * Whether `web_search` uses USER_PROVIDED auth (a user-managed key). Drives the - * cog-vs-info affordance on the web_search card/row: configurable only when a - * user key exists; SYSTEM_DEFINED deployments have nothing to configure. - */ - webSearchUserProvided?: boolean; } interface BuiltinDef { @@ -137,8 +131,6 @@ export function buildCatalog(inputs: BuildCatalogInputs): AgentItem[] { name: def.nameKey, description: def.descriptionKey, status: inputs.builtinAuthMap?.get(def.id) === true ? 'needs_setup' : undefined, - userProvidedAuth: - def.id === AgentCapabilities.web_search ? inputs.webSearchUserProvided === true : undefined, }); } diff --git a/client/src/components/SidePanel/Agents/Tools/items/configurable.ts b/client/src/components/SidePanel/Agents/Tools/items/configurable.ts index cc63ac7f80d..9d535101f7c 100644 --- a/client/src/components/SidePanel/Agents/Tools/items/configurable.ts +++ b/client/src/components/SidePanel/Agents/Tools/items/configurable.ts @@ -16,7 +16,7 @@ export function hasConfigurableSettings(item: AgentItem): boolean { item.id === 'file_search' || item.id === 'context' || item.id === 'memory' || - (item.id === 'web_search' && item.userProvidedAuth === true) + item.id === 'web_search' ); case 'tool': return pluginNeedsAuth(item.plugin); diff --git a/client/src/components/SidePanel/Agents/Tools/items/types.ts b/client/src/components/SidePanel/Agents/Tools/items/types.ts index d3c563bb0d6..785a0a73a32 100644 --- a/client/src/components/SidePanel/Agents/Tools/items/types.ts +++ b/client/src/components/SidePanel/Agents/Tools/items/types.ts @@ -43,12 +43,6 @@ interface ItemBase { export interface BuiltinItem extends ItemBase { kind: 'builtin'; id: BuiltinId; - /** - * True when `web_search` auth is USER_PROVIDED (a user-managed key exists to - * configure). Undefined/false means SYSTEM_DEFINED — nothing to configure, so - * the card/row shows an info icon instead of a settings cog. - */ - userProvidedAuth?: boolean; } export interface ToolItem extends ItemBase { diff --git a/client/src/hooks/Agents/__tests__/useApplyAgentToolSwitches.test.tsx b/client/src/hooks/Agents/__tests__/useApplyAgentToolSwitches.test.tsx new file mode 100644 index 00000000000..e9642c6749e --- /dev/null +++ b/client/src/hooks/Agents/__tests__/useApplyAgentToolSwitches.test.tsx @@ -0,0 +1,259 @@ +import React from 'react'; +import { RecoilRoot } from 'recoil'; +import { Provider, createStore } from 'jotai'; +import { useRecoilValue, useSetRecoilState } from 'recoil'; +import { act, render, waitFor } from '@testing-library/react'; +import { Constants, LocalStorageKeys, mcpServerToggleKey } from 'librechat-data-provider'; +import type { TEphemeralAgent } from 'librechat-data-provider'; +import type { Agent } from 'librechat-data-provider'; +import type { MCPServerDefinition } from '~/hooks/MCP/useMCPServerManager'; +import { useApplyAgentToolSwitches } from '../useApplyAgentToolSwitches'; +import { cleanupTimestampedStorage } from '~/utils/timestamps'; +import { useMCPSelect } from '~/hooks/MCP/useMCPSelect'; +import { ephemeralAgentByConvoId } from '~/store'; + +jest.mock('~/data-provider', () => ({ + ...jest.requireActual('~/data-provider'), + useGetStartupConfig: jest.fn(() => ({ data: undefined })), +})); + +const serverName = 'mock-elicitation'; +const servers = [ + { serverName, config: { type: 'sse', url: 'http://localhost' }, effectivePermissions: 1 }, +] as unknown as MCPServerDefinition[]; + +type SavedAgent = Pick; + +const withToggle = (toggle?: 'on' | 'off'): SavedAgent => ({ + id: 'agent_1', + tools: [`echo_mcp_${serverName}`], + tool_options: toggle ? { [mcpServerToggleKey(serverName)]: { user_toggle: toggle } } : {}, +}); + +let selected: string[] = []; +let select: (value: string[]) => void = () => undefined; + +function Picker({ conversationId }: { conversationId: string | null }) { + const { mcpValues, setMCPValues } = useMCPSelect({ + conversationId, + servers, + ownsChatSelection: true, + }); + selected = mcpValues; + select = setMCPValues; + return null; +} + +function Chat({ + agent, + conversationId = null, +}: { + agent: SavedAgent; + conversationId?: string | null; +}) { + useApplyAgentToolSwitches({ agent, conversationId }); + return ; +} + +describe('useApplyAgentToolSwitches', () => { + beforeEach(() => { + localStorage.clear(); + selected = []; + }); + + it('selects a default-on MCP server in a new chat', async () => { + render( + + + + + , + ); + await waitFor(() => expect(selected).toEqual([serverName])); + }); + + it.each([ + ['not switchable', undefined], + ['switchable, starts off', 'off' as const], + ])( + 'selects a default-on MCP server once the saved agent changes from %s', + async (_label, previousToggle) => { + const store = createStore(); + const tree = (agent: SavedAgent) => ( + + + + + + ); + const { rerender } = render(tree(withToggle(previousToggle))); + await waitFor(() => expect(selected).toEqual([])); + + rerender(tree(withToggle('on'))); + await waitFor(() => expect(selected).toEqual([serverName])); + }, + ); + + describe('switching agents inside one new chat', () => { + const agentX: SavedAgent = { + id: 'agent_x', + tools: ['web_search', 'echo_mcp_x-server'], + tool_options: { + web_search: { user_toggle: 'on' }, + [mcpServerToggleKey('x-server')]: { user_toggle: 'on' }, + }, + }; + let ephemeralAgent: TEphemeralAgent | null = null; + let store = createStore(); + beforeEach(() => { + ephemeralAgent = null; + store = createStore(); + }); + function Probe({ agent }: { agent: SavedAgent }) { + useApplyAgentToolSwitches({ agent, conversationId: null }); + ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return null; + } + const tree = (agent: SavedAgent) => ( + + + + + + ); + + it("shows the next agent's defaults and clears the previous agent's switches", async () => { + const { rerender } = render(tree(agentX)); + await waitFor(() => expect(ephemeralAgent).toEqual({ web_search: true, mcp: ['x-server'] })); + + rerender( + tree({ + id: 'agent_y', + tools: ['execute_code'], + tool_options: { execute_code: { user_toggle: 'off' } }, + }), + ); + await waitFor(() => expect(ephemeralAgent).toEqual({ execute_code: false, mcp: [] })); + }); + + it("clears the previous agent's switches when the next agent has none", async () => { + const { rerender } = render(tree(agentX)); + await waitFor(() => expect(ephemeralAgent).toEqual({ web_search: true, mcp: ['x-server'] })); + + rerender(tree({ id: 'agent_z', tools: ['web_search'], tool_options: {} })); + await waitFor(() => expect(ephemeralAgent).toEqual({ mcp: [] })); + }); + }); + + it('keeps a default-on server the user turned off once the new chat gets its real id', async () => { + const store = createStore(); + let applyTemplate: (agent: TEphemeralAgent) => void = () => undefined; + function TemplateWriter() { + applyTemplate = useSetRecoilState(ephemeralAgentByConvoId('real1')); + return null; + } + const tree = (conversationId: string | null) => ( + + + + + + + ); + const { rerender } = render(tree(null)); + await waitFor(() => expect(selected).toEqual([serverName])); + + act(() => select([])); + await waitFor(() => expect(selected).toEqual([])); + + /** What the SSE handlers do on the first response: copy the submitted state to the real id. */ + act(() => { + applyTemplate({ mcp: [] }); + rerender(tree('real1')); + }); + await waitFor(() => + expect(localStorage.getItem(`${LocalStorageKeys.LAST_MCP_}real1`)).not.toBeNull(), + ); + expect(selected).toEqual([]); + expect( + JSON.parse(localStorage.getItem(`${LocalStorageKeys.LAST_MCP_}real1`) ?? 'null'), + ).toEqual([]); + }); + + it('keeps a default-on server the user turned off after the app reloads', async () => { + let applyTemplate: (agent: TEphemeralAgent) => void = () => undefined; + function TemplateWriter() { + applyTemplate = useSetRecoilState(ephemeralAgentByConvoId('real1')); + return null; + } + const store = createStore(); + const tree = (conversationId: string | null) => ( + + + + + + + ); + const { rerender, unmount } = render(tree(null)); + await waitFor(() => expect(selected).toEqual([serverName])); + + act(() => select([])); + await waitFor(() => expect(selected).toEqual([])); + + act(() => { + applyTemplate({ mcp: [] }); + rerender(tree('real1')); + }); + await waitFor(() => + expect(localStorage.getItem(`${LocalStorageKeys.LAST_MCP_}real1`)).not.toBeNull(), + ); + unmount(); + + /** App startup prunes conversation keys without a timestamp. */ + cleanupTimestampedStorage(); + render( + + + + + , + ); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(selected).toEqual([]); + }); + + it('keeps a default-on server turned off in an existing chat after a reload near the cleanup age', async () => { + const hour = 60 * 60 * 1000; + const storageKey = `${LocalStorageKeys.LAST_MCP_}c1`; + const tree = () => ( + + + + + + ); + const { unmount } = render(tree()); + await waitFor(() => expect(selected).toEqual([serverName])); + + /** The last non-empty selection was stamped 47 hours ago. */ + localStorage.setItem(`${storageKey}_TIMESTAMP`, String(Date.now() - 47 * hour)); + act(() => select([])); + await waitFor(() => expect(selected).toEqual([])); + unmount(); + + /** Reload three hours later: startup cleanup runs before the chat mounts. */ + const now = Date.now() + 3 * hour; + const dateNow = jest.spyOn(Date, 'now').mockReturnValue(now); + try { + cleanupTimestampedStorage(); + } finally { + dateNow.mockRestore(); + } + expect(JSON.parse(localStorage.getItem(storageKey) ?? 'null')).toEqual([]); + + render(tree()); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect(selected).toEqual([]); + }); +}); diff --git a/client/src/hooks/Agents/index.ts b/client/src/hooks/Agents/index.ts index ede43629732..f5b5efb5999 100644 --- a/client/src/hooks/Agents/index.ts +++ b/client/src/hooks/Agents/index.ts @@ -20,3 +20,4 @@ export type { } from './useCodeWorkspace'; export { default as useMCPToolOptions } from './useMCPToolOptions'; export * from './useApplyModelSpecAgents'; +export * from './useApplyAgentToolSwitches'; diff --git a/client/src/hooks/Agents/useApplyAgentToolSwitches.ts b/client/src/hooks/Agents/useApplyAgentToolSwitches.ts new file mode 100644 index 00000000000..50348b52718 --- /dev/null +++ b/client/src/hooks/Agents/useApplyAgentToolSwitches.ts @@ -0,0 +1,131 @@ +import { useEffect, useRef } from 'react'; +import { useAtom } from 'jotai'; +import { useRecoilCallback, useSetRecoilState } from 'recoil'; +import { + Constants, + LocalStorageKeys, + getAgentToolSwitches, + switchableBuiltinTools, +} from 'librechat-data-provider'; +import type { Agent, TEphemeralAgent, SwitchableBuiltinTool } from 'librechat-data-provider'; +import { ephemeralAgentByConvoId, mcpValuesAtomFamily } from '~/store'; +import { applyAgentToolSwitchDefaults } from '~/utils'; +import { setTimestamp } from '~/utils/timestamps'; + +export function useApplyAgentToolSwitches({ + agent, + conversationId, +}: { + agent?: Pick | null; + conversationId?: string | null; +}) { + const convoId = conversationId ?? Constants.NEW_CONVO; + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(convoId)); + const [mcpValues, setMCPValues] = useAtom(mcpValuesAtomFamily(convoId)); + const getEphemeralAgent = useRecoilCallback( + ({ snapshot }) => + () => + snapshot.getLoadable(ephemeralAgentByConvoId(convoId)).contents as TEphemeralAgent | null, + [convoId], + ); + const seededRef = useRef<{ + convoId: string; + seedKey: string; + builtins: SwitchableBuiltinTool[]; + servers: string[]; + } | null>(null); + const agentRef = useRef(agent); + agentRef.current = agent; + const agentId = agent?.id; + /** The saved agent's cache entry is replaced in place when its creator edits it, + * so reseed on a change in the switch config itself, not just the agent id. */ + const agentSwitches = agent ? getAgentToolSwitches(agent) : null; + const switchesKey = agentSwitches ? JSON.stringify(agentSwitches) : ''; + const hasMCPSwitches = agentSwitches != null && Object.keys(agentSwitches.mcp).length > 0; + + useEffect(() => { + const current = agentRef.current; + if (!current) { + return; + } + const switches = getAgentToolSwitches(current); + const switchableBuiltins = switchableBuiltinTools.filter((tool) => tool in switches.builtins); + const switchableServers = Object.keys(switches.mcp); + const previousSeed = seededRef.current; + const seedKey = `${agentId}:${switchesKey}`; + seededRef.current = { + convoId, + seedKey, + builtins: switchableBuiltins, + servers: switchableServers, + }; + /** Another agent picked in this same chat leaves its switches behind; drop them. */ + const stale = + previousSeed?.convoId === convoId && previousSeed.seedKey !== seedKey ? previousSeed : null; + const withoutStale = (previous: TEphemeralAgent | null): TEphemeralAgent | null => { + if (!stale || !previous) { + return previous; + } + const next: TEphemeralAgent = { ...previous }; + for (const tool of stale.builtins) { + delete next[tool]; + } + if (Array.isArray(next.mcp)) { + next.mcp = next.mcp.filter((name) => !stale.servers.includes(name)); + } + return next; + }; + if (switchableBuiltins.length === 0 && switchableServers.length === 0) { + if (stale) { + setEphemeralAgent(withoutStale); + } + return; + } + /** A new chat that just received its real id already carries the user's + * choices, copied over from the submission; persist them instead of reseeding. */ + const carried = + previousSeed?.convoId === Constants.NEW_CONVO && + previousSeed.seedKey === seedKey && + convoId !== Constants.NEW_CONVO + ? getEphemeralAgent() + : null; + if (carried) { + if (Array.isArray(carried.mcp)) { + setMCPValues(carried.mcp); + } + return; + } + const seeded = applyAgentToolSwitchDefaults({ + agent: current, + convoId, + isNewConvo: convoId === Constants.NEW_CONVO, + }); + setEphemeralAgent((current) => { + const previous = withoutStale(current); + return { + ...previous, + ...seeded, + ...(seeded.mcp && { + mcp: [ + ...(previous?.mcp ?? []).filter((name) => !switchableServers.includes(name)), + ...seeded.mcp, + ], + }), + }; + }); + }, [agentId, switchesKey, convoId, setEphemeralAgent, setMCPValues, getEphemeralAgent]); + + /** An explicit empty list overrides the creator's default-on servers, so it needs a + * fresh timestamp to survive startup cleanup; plain chats only stamp non-empty lists. */ + useEffect(() => { + const storageKey = `${LocalStorageKeys.LAST_MCP_}${convoId}`; + if ( + !hasMCPSwitches || + convoId === Constants.NEW_CONVO || + localStorage.getItem(storageKey) === null + ) { + return; + } + setTimestamp(storageKey); + }, [hasMCPSwitches, convoId, mcpValues]); +} diff --git a/client/src/hooks/Agents/useUserToolToggle.ts b/client/src/hooks/Agents/useUserToolToggle.ts new file mode 100644 index 00000000000..edddf4907b9 --- /dev/null +++ b/client/src/hooks/Agents/useUserToolToggle.ts @@ -0,0 +1,35 @@ +import { useCallback } from 'react'; +import { useFormContext, useWatch } from 'react-hook-form'; +import type { AgentToolOptions, UserToolToggle } from 'librechat-data-provider'; +import type { AgentForm } from '~/common'; + +interface UseUserToolToggleReturn { + value: UserToolToggle | undefined; + setValue: (next: UserToolToggle | undefined) => void; +} + +/** Reads/writes one tool's `user_toggle`; `undefined` (locked) removes the key and any emptied entry. */ +export default function useUserToolToggle(toolOptionKey: string): UseUserToolToggleReturn { + const { getValues, setValue: setFormValue, control } = useFormContext(); + const toolOptions = useWatch({ control, name: 'tool_options' }); + + const setValue = useCallback( + (next: UserToolToggle | undefined) => { + const options: AgentToolOptions = { ...getValues('tool_options') }; + if (next != null) { + options[toolOptionKey] = { ...options[toolOptionKey], user_toggle: next }; + } else if (options[toolOptionKey] != null) { + const { user_toggle: _omit, ...rest } = options[toolOptionKey]; + if (Object.keys(rest).length === 0) { + delete options[toolOptionKey]; + } else { + options[toolOptionKey] = rest; + } + } + setFormValue('tool_options', options, { shouldDirty: true }); + }, + [getValues, setFormValue, toolOptionKey], + ); + + return { value: toolOptions?.[toolOptionKey]?.user_toggle, setValue }; +} diff --git a/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx b/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx index 93ca85f519d..bfde05cdb83 100644 --- a/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx +++ b/client/src/hooks/MCP/__tests__/useMCPSelect.test.tsx @@ -965,6 +965,33 @@ describe('useMCPSelect', () => { }); }); + it("keeps a saved agent's switchable server, even though the menu hides it", async () => { + const { Wrapper } = createWrapper(); + + const TestComponent = () => { + const mcpHook = useMCPSelect({ + ownsChatSelection: true, + servers: createMCPServers(['visible']), + allServers: createMCPServers(['visible', 'agent-server', 'stale-hidden']), + agentServers: ['agent-server'], + }); + const setEphemeralAgent = useSetRecoilState(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + const ephemeralAgent = useRecoilValue(ephemeralAgentByConvoId(Constants.NEW_CONVO)); + return { mcpHook, ephemeralAgent, setEphemeralAgent }; + }; + + const { result } = renderHook(() => TestComponent(), { wrapper: Wrapper }); + + act(() => { + result.current.setEphemeralAgent({ mcp: ['visible', 'agent-server', 'stale-hidden'] }); + }); + + await waitFor(() => { + expect(result.current.ephemeralAgent?.mcp).toEqual(['visible', 'agent-server']); + expect(result.current.mcpHook.mcpValues).toEqual(['visible', 'agent-server']); + }); + }); + it("lets a catalog-only instance leave the picker instance's selection alone", async () => { mockStartupConfig = { modelSpecs: { list: [{ name: 'pins-hidden', mcpServers: ['spec-server'] }] }, diff --git a/client/src/hooks/MCP/useMCPSelect.ts b/client/src/hooks/MCP/useMCPSelect.ts index f4e08efaee3..4fa77495c1c 100644 --- a/client/src/hooks/MCP/useMCPSelect.ts +++ b/client/src/hooks/MCP/useMCPSelect.ts @@ -18,6 +18,7 @@ export function useMCPSelect({ servers, allServers, specName, + agentServers, ownsChatSelection = false, }: { conversationId?: string | null; @@ -28,6 +29,8 @@ export function useMCPSelect({ allServers?: MCPServerDefinition[]; /** Active model spec, whose pinned servers are exempt from pruning. */ specName?: string | null; + /** The saved agent's switchable servers, exempt from pruning like a spec's. */ + agentServers?: string[]; /** * Whether this instance drives the chat picker and may therefore rewrite the * shared selection. Off by default: every instance keyed to a conversation @@ -57,21 +60,23 @@ export function useMCPSelect({ const { data: startupConfig } = useGetStartupConfig(); /** * Selections that survive pruning: what the dropdown offers, plus whatever the - * active model spec pins. `chatMenu` hides a server from the picker; it does - * not override an admin's spec, so a spec-assigned server stays selected even - * when the picker would never have offered it. + * active model spec pins or the saved agent makes switchable. `chatMenu` hides a + * server from the picker; it does not override an admin's spec or an agent + * creator's switch, so such a server stays selected even when the picker would + * never have offered it. */ const retainedServers = useMemo(() => { - const specServers = getModelSpec({ specName, startupConfig })?.mcpServers; - if (!specServers?.length) { + const specServers = getModelSpec({ specName, startupConfig })?.mcpServers ?? []; + const exempt = [...specServers, ...(agentServers ?? [])]; + if (!exempt.length) { return configuredServers; } const retained = new Set(configuredServers); - for (const serverName of specServers) { + for (const serverName of exempt) { retained.add(serverName); } return retained; - }, [configuredServers, specName, startupConfig]); + }, [configuredServers, specName, startupConfig, agentServers]); /** * For new conversations, key the MCP atom by environment (spec or defaults) diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index 160ea9ec1e6..47c9d97d429 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -65,12 +65,15 @@ export function useMCPServerManager({ conversationId, storageContextKey, specName, + agentServers, ownsChatSelection = false, observeToolAuthorization = false, }: { conversationId?: string | null; storageContextKey?: string; specName?: string | null; + /** The saved agent's switchable servers, kept selected even when the menu hides them. */ + agentServers?: string[]; /** * Opt in to managing the chat MCP selection. Most callers mount this hook for * the catalog, the server actions, or the status icons and never read the @@ -150,6 +153,7 @@ export function useMCPServerManager({ servers: selectableServers, allServers: availableMCPServers, specName, + agentServers, ownsChatSelection, }); const mcpValuesRef = useRef(mcpValues); diff --git a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts index ca8d4a7b5b9..fd5fa24dd80 100644 --- a/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts +++ b/client/src/hooks/SSE/__tests__/useStepHandler.spec.ts @@ -1817,6 +1817,338 @@ describe('useStepHandler', () => { }); describe('on_run_step_delta event', () => { + it('separates streamed preparation, dispatch, and this call’s result', () => { + mockGetMessages.mockReturnValue([createResponseMessage()]); + const { result } = renderHook(() => useStepHandler(createHookParams())); + const submission = createSubmission(); + const latest = () => { + const messages = mockSetMessages.mock.lastCall?.[0] as TMessage[] | undefined; + const part = messages?.find((message) => message.messageId === 'response-msg-1') + ?.content?.[0]; + return part?.type === ContentTypes.TOOL_CALL ? part.tool_call : undefined; + }; + act(() => { + result.current.stepHandler( + { event: StepEvents.ON_RUN_STEP, data: createToolCallRunStep() }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_DELTA, + data: { + id: 'step-tool-1', + observed_at: 1_000, + delta: { + type: StepTypes.TOOL_CALLS, + tool_calls: [{ id: 'tool-call-1', index: 0, args: '{' }], + }, + }, + }, + submission, + ); + }); + expect(latest()?.toolPreparationStartedAt).toBe(1_000); + expect(latest()?.toolDispatchedAt).toBeUndefined(); + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_TOOL_CALLS_DISPATCHED, + data: { + dispatched_at: 248_000, + toolCalls: [{ id: 'tool-call-1', name: 'test_tool', stepId: 'step-tool-1' }], + }, + }, + submission, + ); + }); + expect(latest()?.toolDispatchedAt).toBe(248_000); + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_COMPLETED, + data: { + result: { + id: 'step-tool-1', + index: 0, + completed_at: 248_340, + tool_call: { id: 'tool-call-1', name: 'test_tool', args: '{}', output: 'done' }, + }, + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_CLOSED, + data: { + id: 'step-tool-1', + index: 0, + type: StepTypes.TOOL_CALLS, + status: 'completed', + created_at: 1_000, + closed_at: 248_340, + }, + }, + submission, + ); + }); + expect(latest()).toMatchObject({ + runStepStatus: 'completed', + runStepDurationMs: 247_340, + runStepClosedAt: 248_340, + toolPreparationDurationMs: 247_000, + toolExecutionDurationMs: 340, + }); + }); + + it('retains an ID-less first-fragment timestamp when the call receives its ID later', () => { + mockGetMessages.mockReturnValue([createResponseMessage()]); + const { result } = renderHook(() => useStepHandler(createHookParams())); + const submission = createSubmission(); + act(() => { + result.current.stepHandler( + { event: StepEvents.ON_RUN_STEP, data: createToolCallRunStep() }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_DELTA, + data: { + id: 'step-tool-1', + observed_at: 100, + delta: { type: StepTypes.TOOL_CALLS, tool_calls: [{ index: 0, args: '{' }] }, + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_DELTA, + data: { + id: 'step-tool-1', + observed_at: 200, + delta: { + type: StepTypes.TOOL_CALLS, + tool_calls: [{ id: 'tool-call-1', index: 0, args: '"x":1}' }], + }, + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_TOOL_CALLS_DISPATCHED, + data: { + dispatched_at: 500, + toolCalls: [{ id: 'tool-call-1', name: 'test_tool', stepId: 'step-tool-1' }], + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_COMPLETED, + data: { + result: { + id: 'step-tool-1', + index: 0, + completed_at: 530, + tool_call: { id: 'tool-call-1', name: 'test_tool', args: '{}', output: 'done' }, + }, + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_CLOSED, + data: { + id: 'step-tool-1', + index: 0, + type: StepTypes.TOOL_CALLS, + status: 'completed', + created_at: 100, + closed_at: 530, + }, + }, + submission, + ); + }); + const messages = mockSetMessages.mock.lastCall?.[0] as TMessage[] | undefined; + const part = messages?.find((message) => message.messageId === 'response-msg-1') + ?.content?.[0]; + expect(part?.type === ContentTypes.TOOL_CALL ? part.tool_call : undefined).toMatchObject({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 30, + }); + }); + + it('restores the Preparing and Calling phases from compact replay events', () => { + mockGetMessages.mockReturnValue([createResponseMessage()]); + const { result } = renderHook(() => useStepHandler(createHookParams())); + const submission = createSubmission(); + const currentPart = () => { + const response = (mockSetMessages.mock.lastCall?.[0] as TMessage[] | undefined)?.find( + (message) => message.messageId === 'response-msg-1', + ); + const part = response?.content?.[0]; + return part?.type === ContentTypes.TOOL_CALL ? part.tool_call : undefined; + }; + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_TOOL_PREPARATION, + data: { id: 'step-tool-1', toolCallId: 'tool-call-1', index: 0, observed_at: 100 }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP, + data: createToolCallRunStep(), + }, + submission, + ); + }); + expect(currentPart()?.toolPreparationStartedAt).toBe(100); + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_TOOL_CALLS_DISPATCHED, + data: { + dispatched_at: 500, + toolCalls: [{ id: 'tool-call-1', name: 'test_tool', stepId: 'step-tool-1' }], + }, + }, + submission, + ); + }); + expect(currentPart()?.toolDispatchedAt).toBe(500); + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_COMPLETED, + data: { + result: { + id: 'step-tool-1', + index: 0, + completed_at: 540, + tool_call: { id: 'tool-call-1', name: 'test_tool', args: '{}', output: 'done' }, + }, + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_CLOSED, + data: { + id: 'step-tool-1', + type: StepTypes.TOOL_CALLS, + index: 0, + status: 'completed', + created_at: 90, + closed_at: 540, + }, + }, + submission, + ); + }); + expect(currentPart()).toMatchObject({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 40, + }); + }); + + it('isolates duplicate provider IDs across simultaneous run steps', () => { + mockGetMessages.mockReturnValue([createResponseMessage()]); + const { result } = renderHook(() => useStepHandler(createHookParams())); + const submission = createSubmission(); + for (const [id, index, observedAt, dispatchedAt, completedAt] of [ + ['step-tool-1', 0, 100, 300, 310], + ['step-tool-2', 1, 200, 500, 550], + ] as const) { + act(() => { + result.current.stepHandler( + { event: StepEvents.ON_RUN_STEP, data: createToolCallRunStep({ id, index }) }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_DELTA, + data: { + id, + observed_at: observedAt, + delta: { + type: StepTypes.TOOL_CALLS, + tool_calls: [{ id: 'tool-call-1', index: 0, args: '{' }], + }, + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_TOOL_CALLS_DISPATCHED, + data: { + dispatched_at: dispatchedAt, + toolCalls: [{ id: 'tool-call-1', name: 'test_tool', stepId: id }], + }, + }, + submission, + ); + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_COMPLETED, + data: { + result: { + id, + index, + completed_at: completedAt, + tool_call: { id: 'tool-call-1', name: 'test_tool', args: '{}', output: 'done' }, + }, + }, + }, + submission, + ); + }); + } + for (const [id, index, closedAt] of [ + ['step-tool-1', 0, 310], + ['step-tool-2', 1, 550], + ] as const) { + act(() => { + result.current.stepHandler( + { + event: StepEvents.ON_RUN_STEP_CLOSED, + data: { + id, + index, + type: StepTypes.TOOL_CALLS, + status: 'completed', + created_at: 100, + closed_at: closedAt, + }, + }, + submission, + ); + }); + } + const parts = (mockSetMessages.mock.lastCall?.[0] as TMessage[] | undefined)?.find( + (message) => message.messageId === 'response-msg-1', + )?.content; + const first = parts?.[0]; + const second = parts?.[1]; + expect(first?.type === ContentTypes.TOOL_CALL ? first.tool_call : undefined).toMatchObject({ + toolPreparationDurationMs: 200, + toolExecutionDurationMs: 10, + }); + expect(second?.type === ContentTypes.TOOL_CALL ? second.tool_call : undefined).toMatchObject({ + toolPreparationDurationMs: 300, + toolExecutionDurationMs: 50, + }); + }); + it('should update tool call with delta args', () => { const responseMessage = createResponseMessage(); mockGetMessages.mockReturnValue([responseMessage]); diff --git a/client/src/hooks/SSE/useStepHandler.ts b/client/src/hooks/SSE/useStepHandler.ts index f6a08bf63f1..4fd07ffc7cb 100644 --- a/client/src/hooks/SSE/useStepHandler.ts +++ b/client/src/hooks/SSE/useStepHandler.ts @@ -10,6 +10,7 @@ import { getNonEmptyValue, getRunStepDurationMs, getRunStepCloseMetadata, + getToolTimingDurations, } from 'librechat-data-provider'; import type { Agents, @@ -61,12 +62,16 @@ type TUseStepHandler = { onSubagentIndexChange?: (conversationId: string) => void; }; +const toolTimingKey = (stepId: string, callId: string) => `${stepId}\u0000${callId}`; + type TStepEvent = | { event: StepEvents.ON_RUN_STEP; data: Agents.RunStep } | { event: StepEvents.ON_AGENT_UPDATE; data: Agents.AgentUpdate } | { event: StepEvents.ON_MESSAGE_DELTA; data: Agents.MessageDeltaEvent } | { event: StepEvents.ON_REASONING_DELTA; data: Agents.ReasoningDeltaEvent } | { event: StepEvents.ON_RUN_STEP_DELTA; data: Agents.RunStepDeltaEvent } + | { event: StepEvents.ON_TOOL_CALLS_DISPATCHED; data: Agents.ToolCallsDispatchedEvent } + | { event: StepEvents.ON_TOOL_PREPARATION; data: Agents.ToolPreparationMarker } | { event: StepEvents.ON_RUN_STEP_COMPLETED; data: { result: Agents.ToolEndEvent } } | { event: StepEvents.ON_RUN_STEP_CLOSED; data: Agents.RunStepClosedEvent } | { event: StepEvents.ON_SUMMARIZE_START; data: Agents.SummarizeStartEvent } @@ -165,6 +170,10 @@ export default function useStepHandler({ }: TUseStepHandler) { const subagentStore = useStore(); const toolCallIdMap = useRef(new Map()); + const firstFragmentByCall = useRef(new Map()); + const firstFragmentByStep = useRef(new Map()); + const dispatchedByCall = useRef(new Map()); + const completedByCall = useRef(new Map()); const messageMap = useRef(new Map()); const stepMap = useRef(new Map()); /** Buffer for deltas that arrive before their corresponding run step */ @@ -622,7 +631,9 @@ export default function useStepHandler({ }; } else if (contentType === ContentTypes.TOOL_CALL && 'tool_call' in contentPart) { const existingContent = updatedContent[index] as Agents.ToolCallContent | undefined; - const existingToolCall = existingContent?.tool_call; + const existingToolCall = existingContent?.tool_call as + | (Agents.ToolCall & PartMetadata) + | undefined; const toolCallArgs = (contentPart.tool_call as Agents.ToolCall).args; /** When args are a valid object, they are likely already invoked */ let args = @@ -647,6 +658,11 @@ export default function useStepHandler({ type: ToolCallTypes.TOOL_CALL, auth: contentPart.tool_call.auth, expires_at: contentPart.tool_call.expires_at, + toolPreparationStartedAt: + contentPart.tool_call.toolPreparationStartedAt ?? + existingToolCall?.toolPreparationStartedAt, + toolDispatchedAt: + contentPart.tool_call.toolDispatchedAt ?? existingToolCall?.toolDispatchedAt, }; if (finalUpdate) { @@ -956,6 +972,7 @@ export default function useStepHandler({ // Store tool call IDs if present if (runStep.stepDetails.type === StepTypes.TOOL_CALLS) { + const streamedCallCount = runStep.stepDetails.tool_calls?.length ?? 0; let updatedResponse = { ...response }; ((runStep.stepDetails.tool_calls ?? []) as Agents.ToolCall[]).forEach((toolCall) => { const toolCallId = toolCall.id ?? ''; @@ -970,6 +987,14 @@ export default function useStepHandler({ args: toolCall.args, id: toolCallId, stepId: runStep.id, + toolPreparationStartedAt: + firstFragmentByCall.current.get(toolTimingKey(runStep.id, toolCallId)) ?? + (streamedCallCount <= 1 + ? firstFragmentByStep.current.get(runStep.id) + : undefined), + toolDispatchedAt: dispatchedByCall.current.get( + toolTimingKey(runStep.id, toolCallId), + ), }, }; @@ -1200,8 +1225,92 @@ export default function useStepHandler({ scheduleCoalescedMessagesFlush(responseMessageId); } } + } else if (stepEvent.event === StepEvents.ON_TOOL_PREPARATION) { + const { id, index, toolCallId, observed_at: at } = stepEvent.data; + if (!id || typeof at !== 'number' || !Number.isFinite(at) || at < 0) return; + const runStep = stepMap.current.get(id); + const declaredCalls = + runStep?.stepDetails.type === StepTypes.TOOL_CALLS + ? runStep.stepDetails.tool_calls + : undefined; + const resolvedId = + toolCallId ?? + (index === 0 && declaredCalls?.length === 1 ? declaredCalls[0]?.id : undefined); + if (resolvedId) { + const key = toolTimingKey(id, resolvedId); + firstFragmentByCall.current.set( + key, + Math.min(firstFragmentByCall.current.get(key) ?? at, at), + ); + } else if (index === 0) { + firstFragmentByStep.current.set( + id, + Math.min(firstFragmentByStep.current.get(id) ?? at, at), + ); + } + if (!runStep?.runId || (runStep.status && runStep.status !== 'in_progress')) return; + const responseId = + runStep.runId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID + ? (submission.initialResponse?.messageId ?? '') + : runStep.runId; + const response = messageMap.current.get(responseId); + const contentIndex = runStep.index + editPrefixOffset; + const part = response?.content?.[contentIndex]; + if ( + !response || + part?.type !== ContentTypes.TOOL_CALL || + part.tool_call.runStepStatus != null || + (resolvedId && part.tool_call.id !== resolvedId) + ) + return; + const content = [...(response.content ?? [])]; + content[contentIndex] = { + ...part, + tool_call: { + ...part.tool_call, + toolPreparationStartedAt: Math.min(part.tool_call.toolPreparationStartedAt ?? at, at), + }, + }; + const updated = { ...response, content }; + messageMap.current.set(responseId, updated); + setMessages( + mergeResponseMessage(messages, updated, responseId, { ensureUserMessage: true }), + ); } else if (stepEvent.event === StepEvents.ON_RUN_STEP_DELTA) { const runStepDelta = stepEvent.data; + const at = runStepDelta.observed_at; + if (typeof at === 'number' && Number.isFinite(at) && at >= 0) { + for (const chunk of runStepDelta.delta.tool_calls ?? []) { + if (chunk.id) { + const first = + chunk.index === 0 ? firstFragmentByStep.current.get(runStepDelta.id) : undefined; + const key = toolTimingKey(runStepDelta.id, chunk.id); + firstFragmentByCall.current.set( + key, + Math.min(firstFragmentByCall.current.get(key) ?? at, first ?? at, at), + ); + if (first != null) firstFragmentByStep.current.delete(runStepDelta.id); + } else if (chunk.index === 0 && runStepDelta.delta.tool_calls?.length === 1) { + const declared = stepMap.current.get(runStepDelta.id)?.stepDetails; + const firstCallId = + declared?.type === StepTypes.TOOL_CALLS && declared.tool_calls?.length === 1 + ? declared.tool_calls[0]?.id + : undefined; + if (firstCallId) { + const key = toolTimingKey(runStepDelta.id, firstCallId); + firstFragmentByCall.current.set( + key, + Math.min(firstFragmentByCall.current.get(key) ?? at, at), + ); + } else { + firstFragmentByStep.current.set( + runStepDelta.id, + Math.min(firstFragmentByStep.current.get(runStepDelta.id) ?? at, at), + ); + } + } + } + } const runStep = stepMap.current.get(runStepDelta.id); let responseMessageId = runStep?.runId ?? ''; if (responseMessageId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID) { @@ -1225,7 +1334,7 @@ export default function useStepHandler({ let updatedResponse = { ...response }; runStepDelta.delta.tool_calls.forEach((toolCallDelta) => { - const toolCallId = toolCallIdMap.current.get(runStepDelta.id) ?? ''; + const toolCallId = toolCallDelta.id || toolCallIdMap.current.get(runStepDelta.id) || ''; const contentPart: Agents.MessageContentComplex = { type: ContentTypes.TOOL_CALL, @@ -1234,6 +1343,14 @@ export default function useStepHandler({ args: toolCallDelta.args ?? '', id: toolCallId, stepId: runStepDelta.id, + toolPreparationStartedAt: + firstFragmentByCall.current.get(toolTimingKey(runStepDelta.id, toolCallId)) ?? + (toolCallDelta.index === 0 + ? firstFragmentByStep.current.get(runStepDelta.id) + : undefined), + toolDispatchedAt: dispatchedByCall.current.get( + toolTimingKey(runStepDelta.id, toolCallId), + ), }, }; @@ -1260,11 +1377,55 @@ export default function useStepHandler({ }), ); } + } else if (stepEvent.event === StepEvents.ON_TOOL_CALLS_DISPATCHED) { + const { dispatched_at: at, toolCalls } = stepEvent.data; + if (typeof at !== 'number' || !Number.isFinite(at) || at < 0) return; + for (const call of toolCalls ?? []) { + if (!call.id) continue; + if (!call.stepId) continue; + const key = toolTimingKey(call.stepId, call.id); + const dispatchedAt = Math.min(dispatchedByCall.current.get(key) ?? at, at); + dispatchedByCall.current.set(key, dispatchedAt); + const runStep = stepMap.current.get(call.stepId ?? ''); + if (!runStep?.runId) continue; + const responseId = + runStep.runId === Constants.USE_PRELIM_RESPONSE_MESSAGE_ID + ? (submission.initialResponse?.messageId ?? '') + : runStep.runId; + const response = messageMap.current.get(responseId); + const index = runStep.index + editPrefixOffset; + const part = response?.content?.[index]; + if ( + !response || + part?.type !== ContentTypes.TOOL_CALL || + part.tool_call.id !== call.id || + part.tool_call.runStepStatus != null + ) + continue; + const content = [...(response.content ?? [])]; + content[index] = { + ...part, + tool_call: { ...part.tool_call, toolDispatchedAt: dispatchedAt }, + }; + const updated = { ...response, content }; + messageMap.current.set(responseId, updated); + setMessages( + mergeResponseMessage(messages, updated, responseId, { ensureUserMessage: true }), + ); + } } else if (stepEvent.event === StepEvents.ON_RUN_STEP_COMPLETED) { const { result } = stepEvent.data; const { id: stepId } = result; - clearSandboxStarting(result.tool_call?.id); + const completedCallId = result.tool_call?.id; + if ( + completedCallId && + typeof result.completed_at === 'number' && + Number.isFinite(result.completed_at) + ) { + completedByCall.current.set(toolTimingKey(stepId, completedCallId), result.completed_at); + } + clearSandboxStarting(completedCallId); const runStep = stepMap.current.get(stepId); let responseMessageId = runStep?.runId ?? ''; @@ -1351,6 +1512,25 @@ export default function useStepHandler({ * server already stamped in place, rather than overwriting it with * `undefined`. */ const durationMs = getRunStepDurationMs(closed); + const callId = existingToolCall.id ?? ''; + const key = toolTimingKey(closed.id, callId); + const singleCallStep = + runStep.stepDetails.type === StepTypes.TOOL_CALLS && + (runStep.stepDetails.tool_calls?.length ?? 0) <= 1; + const observedAt = Math.min( + firstFragmentByCall.current.get(key) ?? Infinity, + existingToolCall.toolPreparationStartedAt ?? Infinity, + singleCallStep ? (firstFragmentByStep.current.get(closed.id) ?? Infinity) : Infinity, + ); + const timing = getToolTimingDurations({ + observedAt: Number.isFinite(observedAt) ? observedAt : undefined, + dispatchedAt: existingToolCall.toolDispatchedAt ?? dispatchedByCall.current.get(key), + completedAt: completedByCall.current.get(key), + }); + firstFragmentByCall.current.delete(key); + firstFragmentByStep.current.delete(closed.id); + dispatchedByCall.current.delete(key); + completedByCall.current.delete(key); const updatedContent = [...(response.content ?? [])]; updatedContent[currentIndex] = { ...existing, @@ -1359,6 +1539,7 @@ export default function useStepHandler({ runStepStatus: closed.status, ...getRunStepCloseMetadata(closed), ...(durationMs != null && { runStepDurationMs: durationMs }), + ...timing, }, }; @@ -1561,6 +1742,10 @@ export default function useStepHandler({ } cancelPendingDeltaFlush(); toolCallIdMap.current.clear(); + firstFragmentByCall.current.clear(); + firstFragmentByStep.current.clear(); + dispatchedByCall.current.clear(); + completedByCall.current.clear(); messageMap.current.clear(); stepMap.current.clear(); pendingDeltaBuffer.current.clear(); diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index ced97bbc6db..92b51c671d9 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -1468,6 +1468,11 @@ "com_ui_duration_announced_seconds_one": "took {{count}} second", "com_ui_duration_minutes": "{{0}}m {{1}}s", "com_ui_duration_seconds": "{{0}}s", + "com_ui_tool_preparing": "Preparing {{0}}", + "com_ui_tool_calling": "Calling {{0}}", + "com_ui_tool_preparation_time": "Preparation", + "com_ui_tool_call_time": "Tool call", + "com_ui_tool_total_time": "Total elapsed", "com_ui_during_run_actions": "More send options", "com_ui_edit": "Edit", "com_ui_edit_editing_image": "Editing image", @@ -2713,6 +2718,11 @@ "com_ui_tool_name_image_gen": "Image Generation", "com_ui_tool_name_set_memory": "Save Memory", "com_ui_tool_name_web_search": "Web Search", + "com_ui_tool_toggle_info": "Locked keeps the tool always on. On or Off lets users switch it in chat, starting in that state.", + "com_ui_tool_toggle_label": "User control", + "com_ui_tool_toggle_locked": "Locked", + "com_ui_tool_toggle_off": "Starts off", + "com_ui_tool_toggle_on": "Starts on", "com_ui_tool_response_placeholder": "Type a substitute result to return to the agent", "com_ui_tools": "Tools", "com_ui_tools_configure": "Configure", diff --git a/client/src/utils/__tests__/applyAgentToolSwitchDefaults.test.ts b/client/src/utils/__tests__/applyAgentToolSwitchDefaults.test.ts new file mode 100644 index 00000000000..624e1c082ad --- /dev/null +++ b/client/src/utils/__tests__/applyAgentToolSwitchDefaults.test.ts @@ -0,0 +1,74 @@ +import { Constants, LocalStorageKeys, Tools, mcpServerToggleKey } from 'librechat-data-provider'; +import type { Agent } from 'librechat-data-provider'; +import { applyAgentToolSwitchDefaults } from '../endpoints'; +import { setTimestamp } from '../timestamps'; + +const createAgent = (overrides: Partial = {}): Agent => + ({ + id: 'agent_1', + tools: [Tools.web_search, Tools.execute_code, 'search_mcp_alpha', 'lookup_mcp_beta'], + tool_options: { + [Tools.web_search]: { user_toggle: 'on' }, + [Tools.execute_code]: { user_toggle: 'off' }, + [mcpServerToggleKey('alpha')]: { user_toggle: 'on' }, + [mcpServerToggleKey('beta')]: { user_toggle: 'off' }, + }, + ...overrides, + }) as Agent; + +function writeToggle(storagePrefix: string, convoId: string, value: unknown): void { + const key = `${storagePrefix}${convoId}`; + localStorage.setItem(key, JSON.stringify(value)); + setTimestamp(key); +} + +describe('applyAgentToolSwitchDefaults', () => { + beforeEach(() => { + localStorage.clear(); + }); + + it('seeds a new chat from the creator defaults', () => { + const result = applyAgentToolSwitchDefaults({ + agent: createAgent(), + convoId: `${Constants.NEW_CONVO}`, + isNewConvo: true, + }); + expect(result).toEqual({ web_search: true, execute_code: false, mcp: ['alpha'] }); + }); + + it('ignores stored overrides for a new chat', () => { + writeToggle(LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_, 'convo-1', false); + const result = applyAgentToolSwitchDefaults({ + agent: createAgent(), + convoId: 'convo-1', + isNewConvo: true, + }); + expect(result.web_search).toBe(true); + }); + + it('layers stored overrides on top of defaults for an existing chat', () => { + writeToggle(LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_, 'convo-1', false); + writeToggle(LocalStorageKeys.LAST_CODE_TOGGLE_, 'convo-1', true); + localStorage.setItem(`${LocalStorageKeys.LAST_MCP_}convo-1`, JSON.stringify(['beta'])); + const result = applyAgentToolSwitchDefaults({ + agent: createAgent(), + convoId: 'convo-1', + isNewConvo: false, + }); + expect(result).toEqual({ web_search: false, execute_code: true, mcp: ['beta'] }); + }); + + it('never stores keys outside the switchable set', () => { + writeToggle(LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_, 'convo-1', true); + localStorage.setItem( + `${LocalStorageKeys.LAST_MCP_}convo-1`, + JSON.stringify(['alpha', 'gamma']), + ); + const result = applyAgentToolSwitchDefaults({ + agent: createAgent({ tools: [Tools.web_search], tool_options: undefined }), + convoId: 'convo-1', + isNewConvo: false, + }); + expect(result).toEqual({}); + }); +}); diff --git a/client/src/utils/__tests__/subagentContent.test.ts b/client/src/utils/__tests__/subagentContent.test.ts index dd8782dcfcb..e861303b0f9 100644 --- a/client/src/utils/__tests__/subagentContent.test.ts +++ b/client/src/utils/__tests__/subagentContent.test.ts @@ -160,6 +160,105 @@ describe('aggregateSubagentContent', () => { expect((parts[0] as { tool_call: { progress: number } }).tool_call.progress).toBe(0.1); }); + it('folds a child handoff into its own running tool card and measured call time', () => { + const events = [ + makeEvent({ + phase: 'run_step', + data: { + id: 'step-child', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-child', name: 'query', args: '{}' }], + }, + }, + }), + makeEvent({ + phase: 'tool_preparation', + data: { + id: 'step-child', + toolCallId: 'call-child', + observed_at: 100, + }, + }), + makeEvent({ + phase: 'tool_calls_dispatched', + data: { + dispatched_at: 500, + toolCalls: [ + { id: 'call-child', stepId: 'step-child', name: 'query' }, + { id: 'other-child', stepId: 'other-step', name: 'query' }, + ], + }, + }), + ]; + const running = aggregateSubagentContent(events); + expect(running[0]).toMatchObject({ + tool_call: { + id: 'call-child', + stepId: 'step-child', + toolPreparationStartedAt: 100, + toolDispatchedAt: 500, + progress: 0.1, + }, + }); + const completed = aggregateSubagentContent([ + ...events, + makeEvent({ + phase: 'run_step_completed', + data: { + result: { + id: 'step-child', + completed_at: 540, + type: 'tool_call', + tool_call: { id: 'call-child', name: 'query', args: '{}', output: 'ok', progress: 1 }, + }, + }, + }), + ]); + expect(completed[0]).toMatchObject({ + tool_call: { + id: 'call-child', + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 40, + output: 'ok', + progress: 1, + }, + }); + }); + + it('ignores a late handoff for a settled child tool', () => { + const completed = aggregateSubagentContent([ + makeEvent({ + phase: 'run_step', + data: { + id: 'step-child', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-child', name: 'query', args: '{}' }], + }, + }, + }), + makeEvent({ + phase: 'run_step_completed', + data: { + result: { + id: 'step-child', + completed_at: 540, + tool_call: { id: 'call-child', name: 'query', output: 'ok', progress: 1 }, + }, + }, + }), + makeEvent({ + phase: 'tool_calls_dispatched', + data: { + dispatched_at: 500, + toolCalls: [{ id: 'call-child', stepId: 'step-child' }], + }, + }), + ]); + expect(completed[0]).not.toHaveProperty('tool_call.toolDispatchedAt'); + }); + it('finalizes a TOOL_CALL part on run_step_completed with output and progress=1', () => { const parts = aggregateSubagentContent([ makeEvent({ diff --git a/client/src/utils/endpoints.ts b/client/src/utils/endpoints.ts index b39fbdf23f5..d7b245a84d4 100644 --- a/client/src/utils/endpoints.ts +++ b/client/src/utils/endpoints.ts @@ -1,10 +1,12 @@ import { + Tools, Constants, EModelEndpoint, defaultEndpoints, modularEndpoints, LocalStorageKeys, getEndpointField, + getAgentToolSwitches, isAgentsEndpoint, isEphemeralAgentId, isAssistantsEndpoint, @@ -394,6 +396,60 @@ export function applyModelSpecEphemeralAgent({ updateEphemeralAgent(key, agent); } +export function applyAgentToolSwitchDefaults({ + agent, + convoId, + isNewConvo, +}: { + agent: Pick; + convoId: string; + isNewConvo: boolean; +}): t.TEphemeralAgent { + const switches = getAgentToolSwitches(agent); + const builtins = { ...switches.builtins }; + const serverNames = Object.keys(switches.mcp); + let enabledServers = serverNames.filter((name) => switches.mcp[name]); + + if (!isNewConvo) { + const storageMap: Array<[keyof typeof switches.builtins, string]> = [ + [Tools.execute_code, LocalStorageKeys.LAST_CODE_TOGGLE_], + [Tools.web_search, LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_], + [Tools.file_search, LocalStorageKeys.LAST_FILE_SEARCH_TOGGLE_], + ]; + for (const [toolKey, storagePrefix] of storageMap) { + if (!(toolKey in builtins)) { + continue; + } + const raw = getTimestampedValue(`${storagePrefix}${convoId}`); + if (raw === null) { + continue; + } + try { + const stored = JSON.parse(raw); + if (typeof stored === 'boolean') { + builtins[toolKey] = stored; + } + } catch { + // ignore parse errors + } + } + + const mcpRaw = localStorage.getItem(`${LocalStorageKeys.LAST_MCP_}${convoId}`); + if (mcpRaw !== null && serverNames.length > 0) { + try { + const stored = JSON.parse(mcpRaw); + if (Array.isArray(stored)) { + enabledServers = serverNames.filter((name) => stored.includes(name)); + } + } catch { + // ignore parse errors + } + } + } + + return serverNames.length > 0 ? { ...builtins, mcp: enabledServers } : { ...builtins }; +} + /** * Resolves the default model spec for a new chat. Priority: hard admin default → * the most recent conversation's own selection → soft default → legacy first spec. diff --git a/client/src/utils/subagentContent.ts b/client/src/utils/subagentContent.ts index 393c67278e5..ccab985ddd0 100644 --- a/client/src/utils/subagentContent.ts +++ b/client/src/utils/subagentContent.ts @@ -1,4 +1,4 @@ -import { ContentTypes, ToolCallTypes } from 'librechat-data-provider'; +import { ContentTypes, ToolCallTypes, getToolTimingDurations } from 'librechat-data-provider'; import type { SubagentUpdateEvent } from 'librechat-data-provider'; /** @@ -35,6 +35,8 @@ type RunStepData = { type RunStepCompletedData = { result?: { + id?: string; + completed_at?: number; type?: string; tool_call?: { id?: string; @@ -51,6 +53,17 @@ type RunStepClosedData = { id?: string; }; +type ToolDispatchData = { + dispatched_at?: number; + toolCalls?: Array<{ id?: string; stepId?: string }>; +}; + +type ToolPreparationData = { + id?: string; + toolCallId?: string; + observed_at?: number; +}; + type MessageDeltaData = { id?: string; delta?: { @@ -81,6 +94,11 @@ type ToolCallPart = { progress: number; inputValidationError?: true; type?: string; + stepId?: string; + toolPreparationStartedAt?: number; + toolDispatchedAt?: number; + toolPreparationDurationMs?: number; + toolExecutionDurationMs?: number; }; }; @@ -269,6 +287,7 @@ export function foldSubagentEvent( type: ContentTypes.TOOL_CALL, tool_call: { id: tc.id, + ...(typeof data?.id === 'string' ? { stepId: data.id } : {}), name: tc.name ?? '', args: stringifyArgs(tc.args), progress: 0.1, @@ -291,6 +310,57 @@ export function foldSubagentEvent( }; } + if (event.phase === 'tool_preparation') { + const data = event.data as ToolPreparationData | undefined; + const id = data?.toolCallId; + const at = data?.observed_at; + if (!id || typeof at !== 'number' || !Number.isFinite(at) || at < 0) return { parts, state }; + const idx = state.toolCallIndexById[id]; + const part = idx == null ? undefined : parts[idx]; + if ( + part?.type !== ContentTypes.TOOL_CALL || + part.tool_call.stepId !== data.id || + part.tool_call.progress >= 1 + ) + return { parts, state }; + const next = parts.slice(); + next[idx] = { + ...part, + tool_call: { + ...part.tool_call, + toolPreparationStartedAt: Math.min(part.tool_call.toolPreparationStartedAt ?? at, at), + }, + }; + return { parts: next, state }; + } + + if (event.phase === 'tool_calls_dispatched') { + const data = event.data as ToolDispatchData | undefined; + const at = data?.dispatched_at; + if (typeof at !== 'number' || !Number.isFinite(at) || at < 0) return { parts, state }; + let next = parts; + for (const call of data?.toolCalls ?? []) { + if (!call.id || !call.stepId) continue; + const idx = state.toolCallIndexById[call.id]; + const part = idx == null ? undefined : next[idx]; + if ( + part?.type !== ContentTypes.TOOL_CALL || + part.tool_call.stepId !== call.stepId || + part.tool_call.progress >= 1 + ) + continue; + if (next === parts) next = parts.slice(); + next[idx] = { + ...part, + tool_call: { + ...part.tool_call, + toolDispatchedAt: Math.min(part.tool_call.toolDispatchedAt ?? at, at), + }, + }; + } + return { parts: next, state }; + } + if (event.phase === 'run_step_completed') { const data = event.data as RunStepCompletedData | undefined; const tc = data?.result?.tool_call; @@ -298,10 +368,19 @@ export function foldSubagentEvent( const existingIdx = state.toolCallIndexById[tc.id]; if (existingIdx != null) { const existing = parts[existingIdx] as ToolCallPart; + const timings = + data?.result?.id === existing.tool_call.stepId + ? getToolTimingDurations({ + observedAt: existing.tool_call.toolPreparationStartedAt, + dispatchedAt: existing.tool_call.toolDispatchedAt, + completedAt: data?.result?.completed_at, + }) + : {}; const merged: ToolCallPart = { type: ContentTypes.TOOL_CALL, tool_call: { ...existing.tool_call, + ...timings, ...(tc.name ? { name: tc.name } : {}), ...(tc.args != null ? { args: stringifyArgs(tc.args) } : {}), ...(tc.output != null ? { output: tc.output } : {}), diff --git a/package-lock.json b/package-lock.json index 85504bcc0e4..47d470a9d83 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10622,9 +10622,9 @@ } }, "node_modules/@librechat/agents": { - "version": "3.9.7", - "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.9.7.tgz", - "integrity": "sha512-r3JSh7ACjKVn/ZzDlXtRU/l16L19pxjlPzcTxa+h+Myol3qLDdx/ldmSiNpSFMUQcGDcCxE07A+ZJWhChcaBZg==", + "version": "3.9.8", + "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.9.8.tgz", + "integrity": "sha512-mVunoG+xQ7dXTDCLF2CmeDnrWOs0BCpZzNRqZZwoXEfLToXGkIjZiGFPiffhsQ6FkOn6DvMipzBFLKA9k5G2YA==", "license": "MIT", "dependencies": { "@anthropic-ai/sdk": "^0.115.0", @@ -42428,7 +42428,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.9.7", + "@librechat/agents": "^3.9.8", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/openapi/agents.openapi.json b/packages/api/openapi/agents.openapi.json index df429025b71..a2e4d216fab 100644 --- a/packages/api/openapi/agents.openapi.json +++ b/packages/api/openapi/agents.openapi.json @@ -2263,6 +2263,13 @@ }, "describe_intent": { "type": "boolean" + }, + "user_toggle": { + "type": "string", + "enum": [ + "on", + "off" + ] } } } @@ -2702,6 +2709,13 @@ }, "describe_intent": { "type": "boolean" + }, + "user_toggle": { + "type": "string", + "enum": [ + "on", + "off" + ] } } } @@ -3165,6 +3179,13 @@ }, "describe_intent": { "type": "boolean" + }, + "user_toggle": { + "type": "string", + "enum": [ + "on", + "off" + ] } } } diff --git a/packages/api/package.json b/packages/api/package.json index bf9bb076194..d9259129df6 100644 --- a/packages/api/package.json +++ b/packages/api/package.json @@ -120,7 +120,7 @@ "@azure/storage-blob": "^12.30.0", "@google/genai": "^2.8.0", "@keyv/redis": "5.1.6", - "@librechat/agents": "^3.9.7", + "@librechat/agents": "^3.9.8", "@librechat/data-schemas": "*", "@modelcontextprotocol/sdk": "^1.30.0", "@opentelemetry/api": "^1.9.0", diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index 43bac3fd1db..c20c0af9e53 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -3020,6 +3020,7 @@ describe('initializeAgent — execute_code capability expansion', () => { workspaceId: 'project-a', operations: ['read_file', 'list_files', 'execute_command'], environment: { fingerprint: 'a'.repeat(64), repo: 'owner/project', actions: ['check'] }, + linkedWorktrees: true, }, }; if (protectedEdit) codeExecutionContext.codeWorkspace!.operations.push('edit_file'); @@ -3063,6 +3064,10 @@ describe('initializeAgent — execute_code capability expansion', () => { (bashTool?.parameters as { properties?: { timeoutMs?: { maximum?: number } } })?.properties ?.timeoutMs?.maximum, ).toBe(120_000); + expect( + (bashTool?.parameters as { properties?: { cwd?: { description?: string } } })?.properties + ?.cwd?.description, + ).toContain('.worktrees/'); }, ); diff --git a/packages/api/src/agents/__tests__/load.spec.ts b/packages/api/src/agents/__tests__/load.spec.ts index 18d199cedad..7fa8f411d7a 100644 --- a/packages/api/src/agents/__tests__/load.spec.ts +++ b/packages/api/src/agents/__tests__/load.spec.ts @@ -11,6 +11,8 @@ import type { } from 'librechat-data-provider'; import type { AppConfig } from '@librechat/data-schemas'; import type { LoadAgentParams, LoadAgentDeps } from '../load'; +import { applyResumeContext, pickResumeContext } from '../hitl/policy'; +import { extractMCPServers } from '../context'; import { loadAddedAgent } from '../added'; import { loadAgent } from '../load'; @@ -1101,6 +1103,160 @@ describe('loadAgent', () => { } }); + describe('user tool switches', () => { + const createSwitchableAgent = async () => { + const agentId = `agent_${uuidv4()}`; + await createAgent({ + id: agentId, + name: 'Switchable Agent', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['web_search', 'execute_code', 'search_mcp_docs', 'read_mcp_files'], + tool_options: { + web_search: { user_toggle: 'on' }, + execute_code: { user_toggle: 'off' }, + sys__server__sys_mcp_docs: { user_toggle: 'on' }, + }, + }); + return agentId; + }; + + const load = (agentId: string, userToolSwitches?: LoadAgentParams['userToolSwitches']) => + loadAgent( + { + req: { user: { id: 'user123' } }, + agent_id: agentId, + endpoint: 'agents', + userToolSwitches, + }, + deps, + ); + + test('keeps creator defaults when the chat sends no switches', async () => { + const agentId = await createSwitchableAgent(); + const result = await load(agentId, {}); + expect(result!.tools).toEqual(['web_search', 'search_mcp_docs', 'read_mcp_files']); + }); + + test('drops a built-in and an MCP server the chat switched off', async () => { + const agentId = await createSwitchableAgent(); + const result = await load(agentId, { web_search: false, mcp: [] }); + expect(result!.tools).toEqual(['read_mcp_files']); + }); + + test('keeps a default-off built-in the chat switched on, without adding unattached tools', async () => { + const agentId = await createSwitchableAgent(); + const result = await load(agentId, { execute_code: true, file_search: true, mcp: ['docs'] }); + expect(result!.tools).toEqual([ + 'web_search', + 'execute_code', + 'search_mcp_docs', + 'read_mcp_files', + ]); + }); + + test('does not filter when no switch state is supplied at all', async () => { + const agentId = await createSwitchableAgent(); + const result = await load(agentId); + expect(result!.tools).toEqual([ + 'web_search', + 'execute_code', + 'search_mcp_docs', + 'read_mcp_files', + ]); + }); + + test('applies creator defaults when the switch state is null', async () => { + const agentId = await createSwitchableAgent(); + const result = await load(agentId, null); + expect(result!.tools).toEqual(['web_search', 'search_mcp_docs', 'read_mcp_files']); + }); + + test.each([ + [['C'], ['B']], + [['A'], ['A', 'B']], + ])( + 'leaves the request server list intact while the filtered tools name the servers (%j)', + async (requestedServers, expectedServers) => { + const agentId = `agent_${uuidv4()}`; + await createAgent({ + id: agentId, + name: 'Mixed Servers', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['search_mcp_A', 'read_mcp_B'], + tool_options: { sys__server__sys_mcp_A: { user_toggle: 'on' } }, + }); + const ephemeralAgent = { mcp: requestedServers }; + const result = await load(agentId, ephemeralAgent); + expect(ephemeralAgent).toEqual({ mcp: requestedServers }); + const toolDefinitions = result!.tools!.map((name) => ({ name })); + expect(extractMCPServers({ ...result, toolDefinitions } as never)).toEqual(expectedServers); + }, + ); + + test.each([ + ['on' as const, [], ['web_search']], + ['off' as const, ['A'], ['search_mcp_A', 'web_search']], + ])( + 'rebuilds the same tools after a pause and resume (server starts %s, chat sends %j)', + async (userToggle, requestedServers, expectedTools) => { + const agentId = `agent_${uuidv4()}`; + await createAgent({ + id: agentId, + name: 'Paused Agent', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['search_mcp_A', 'web_search'], + tool_options: { sys__server__sys_mcp_A: { user_toggle: userToggle } }, + }); + const pausedBody: Record = { + agent_id: agentId, + ephemeralAgent: { mcp: requestedServers }, + }; + const paused = await load(agentId, pausedBody.ephemeralAgent as TEphemeralAgent); + const resumedBody: Record = { agent_id: agentId }; + applyResumeContext(resumedBody, pickResumeContext(pausedBody)); + const resumed = await load(agentId, resumedBody.ephemeralAgent as TEphemeralAgent); + expect(paused!.tools).toEqual(expectedTools); + expect(resumed!.tools).toEqual(expectedTools); + }, + ); + + test('leaves the request server list alone when no server is switchable', async () => { + const agentId = `agent_${uuidv4()}`; + await createAgent({ + id: agentId, + name: 'Locked Servers', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['web_search', 'read_mcp_B'], + tool_options: { web_search: { user_toggle: 'on' } }, + }); + const ephemeralAgent = { mcp: ['B'] }; + await load(agentId, ephemeralAgent); + expect(ephemeralAgent.mcp).toEqual(['B']); + }); + + test('leaves tools untouched for an agent without user toggles', async () => { + const agentId = `agent_${uuidv4()}`; + await createAgent({ + id: agentId, + name: 'Plain Agent', + provider: 'openai', + model: 'gpt-4', + author: new mongoose.Types.ObjectId(), + tools: ['web_search', 'search_mcp_docs'], + }); + const result = await load(agentId, { web_search: false, mcp: [] }); + expect(result!.tools).toEqual(['web_search', 'search_mcp_docs']); + }); + }); + describe('Edge Cases', () => { test('should handle loadAgent with malformed req object', async () => { const result = await loadAgent( diff --git a/packages/api/src/agents/activity.spec.ts b/packages/api/src/agents/activity.spec.ts index 7113971ca06..6b2938bb065 100644 --- a/packages/api/src/agents/activity.spec.ts +++ b/packages/api/src/agents/activity.spec.ts @@ -6,6 +6,49 @@ import { } from './activity'; describe('durable subagent activity projection', () => { + it('drops negative or non-finite persisted child timings', () => { + const result = projectPersistedMessageActivity([ + { + type: 'tool', + toolCallId: 'call-child', + name: 'query', + runStepStatus: 'completed', + toolPreparationStartedAt: -1, + toolDispatchedAt: Infinity, + toolPreparationDurationMs: -400, + toolExecutionDurationMs: NaN, + }, + ]); + expect(result.activity[0]).toEqual({ + type: 'tool', + toolCallId: 'call-child', + name: 'query', + status: 'completed', + }); + }); + + it('preserves split tool timing through a bounded persisted child view', () => { + const result = projectPersistedMessageActivity([ + { + type: 'tool', + toolCallId: 'call-child', + name: 'query', + runStepStatus: 'completed', + toolPreparationStartedAt: 100, + toolDispatchedAt: 500, + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 40, + }, + ]); + expect(result.activity[0]).toMatchObject({ + type: 'tool', + toolCallId: 'call-child', + toolPreparationStartedAt: 100, + toolDispatchedAt: 500, + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 40, + }); + }); it('projects ordinary persisted chat content into the shared activity vocabulary', () => { const projection = projectPersistedMessageActivity([ { type: 'reasoning' }, diff --git a/packages/api/src/agents/activity.ts b/packages/api/src/agents/activity.ts index 7d0aa864d7b..0f2df76684b 100644 --- a/packages/api/src/agents/activity.ts +++ b/packages/api/src/agents/activity.ts @@ -139,6 +139,11 @@ const visibleStatus = (value: unknown): 'running' | 'completed' | 'failed' | 'ca const finiteNumber = (value: unknown): number | undefined => typeof value === 'number' && Number.isFinite(value) ? value : undefined; +const nonNegativeTiming = (value: unknown): number | undefined => { + const parsed = finiteNumber(value); + return parsed != null && parsed >= 0 ? parsed : undefined; +}; + const stringArray = (value: unknown): string[] | undefined => { if (!Array.isArray(value)) return undefined; const result = value.filter((candidate): candidate is string => typeof candidate === 'string'); @@ -227,6 +232,10 @@ export function projectPersistedMessageActivity( finiteNumber(candidate.progress) != null && finiteNumber(candidate.progress)! >= 1; const output = typeof candidate.output === 'string' ? candidate.output : undefined; const runStepStatus = visibleStatus(candidate.runStepStatus); + const toolPreparationStartedAt = nonNegativeTiming(candidate.toolPreparationStartedAt); + const toolDispatchedAt = nonNegativeTiming(candidate.toolDispatchedAt); + const toolPreparationDurationMs = nonNegativeTiming(candidate.toolPreparationDurationMs); + const toolExecutionDurationMs = nonNegativeTiming(candidate.toolExecutionDurationMs); let status: MutableToolActivity['status'] = runStepStatus; if (candidate.runStepStatus == null) { status = completed || output != null ? 'completed' : 'running'; @@ -241,6 +250,10 @@ export function projectPersistedMessageActivity( : {}), ...(output == null || output === '' ? {} : { output }), status, + ...(toolPreparationStartedAt == null ? {} : { toolPreparationStartedAt }), + ...(toolDispatchedAt == null ? {} : { toolDispatchedAt }), + ...(toolPreparationDurationMs == null ? {} : { toolPreparationDurationMs }), + ...(toolExecutionDurationMs == null ? {} : { toolExecutionDurationMs }), ...(candidate.inputValidationError === true ? { inputValidationError: true } : {}), ...(candidate.inputTruncated === true ? { inputTruncated: true } : {}), ...(candidate.outputTruncated === true ? { outputTruncated: true } : {}), diff --git a/packages/api/src/agents/context.spec.ts b/packages/api/src/agents/context.spec.ts index 4094f73638d..c97116400d2 100644 --- a/packages/api/src/agents/context.spec.ts +++ b/packages/api/src/agents/context.spec.ts @@ -368,6 +368,36 @@ describe('Agent Context Utilities', () => { expect(agent.additional_instructions).toBe('Context'); }); + it('follows the agent tools, not the request list, when the agent has switchable servers', async () => { + const agent: AgentWithTools = { + id: 'test-agent', + instructions: 'Base', + tool_options: { + [`${Constants.mcp_server}${Constants.mcp_delimiter}A`]: { user_toggle: 'on' }, + }, + tools: [ + new DynamicStructuredTool({ + name: `tool${Constants.mcp_delimiter}B`, + description: 'Locked server tool', + schema: testSchema, + func: async () => 'result', + }), + ], + }; + + mockMCPManager.formatInstructionsForContext.mockResolvedValue('B MCP'); + + await applyContextToAgent({ + agent, + sharedRunContext: '', + mcpManager: mockMCPManager, + ephemeralAgent: { mcp: ['C'] }, + logger: mockLogger, + }); + + expect(mockMCPManager.formatInstructionsForContext).toHaveBeenCalledWith(['B'], undefined); + }); + it('should prefer agent tools over empty ephemeral MCP array', async () => { const agent: AgentWithTools = { id: 'test-agent', diff --git a/packages/api/src/agents/context.ts b/packages/api/src/agents/context.ts index f27875aeb1d..3082d7bfc77 100644 --- a/packages/api/src/agents/context.ts +++ b/packages/api/src/agents/context.ts @@ -1,5 +1,5 @@ -import { Constants } from 'librechat-data-provider'; import { DynamicStructuredTool } from '@librechat/agents/langchain/tools'; +import { Constants, getMCPSwitchServerNames } from 'librechat-data-provider'; import type { Agent, TEphemeralAgent } from 'librechat-data-provider'; import type { LCTool } from '@librechat/agents'; import type { Logger } from 'winston'; @@ -160,7 +160,13 @@ export async function applyContextToAgent({ const additionalInstructions = agent.additional_instructions || ''; try { - const mcpServers = ephemeralAgent?.mcp?.length ? ephemeralAgent.mcp : extractMCPServers(agent); + /** An agent with switchable servers already had its tools filtered by the chat's + * list, which omits its locked servers, so its tools name the servers. */ + const requestServers = ephemeralAgent?.mcp; + const mcpServers = + requestServers?.length && getMCPSwitchServerNames(agent.tool_options).length === 0 + ? requestServers + : extractMCPServers(agent); const mcpInstructions = await getMCPInstructionsForServers( mcpServers, mcpManager, diff --git a/packages/api/src/agents/execution.ts b/packages/api/src/agents/execution.ts index bf1768815bc..431beecfc80 100644 --- a/packages/api/src/agents/execution.ts +++ b/packages/api/src/agents/execution.ts @@ -44,6 +44,8 @@ export interface CodeExecutionContext { codeWorkspace?: CodeWorkspaceSelection & { operations: CodeWorkspaceOperation[]; workspaceInstanceId?: string; + /** The worker schedules each `.worktrees/` of this root as its own lane. */ + linkedWorktrees?: boolean; /** Live Code API execution ceiling. Omitted by older deployments. */ maxCommandTimeoutMs?: number; instructions?: CodeWorkspaceDescriptor['instructions']; diff --git a/packages/api/src/agents/handlers.ts b/packages/api/src/agents/handlers.ts index a239557eb96..0afc2299db9 100644 --- a/packages/api/src/agents/handlers.ts +++ b/packages/api/src/agents/handlers.ts @@ -563,6 +563,7 @@ export interface ToolExecuteOptions { file_path: string; workspace_id: string; workspace_instance_id?: string; + linked_worktrees?: boolean; start_line: number; max_lines: number; codeApiBaseUrl: string; @@ -579,6 +580,7 @@ export interface ToolExecuteOptions { query: string; workspace_id: string; workspace_instance_id?: string; + linked_worktrees?: boolean; path?: string; max_results: number; codeApiBaseUrl: string; @@ -594,6 +596,7 @@ export interface ToolExecuteOptions { listWorkspaceFiles?: (params: { workspace_id: string; workspace_instance_id?: string; + linked_worktrees?: boolean; path?: string; after_path?: string; max_results: number; @@ -613,6 +616,7 @@ export interface ToolExecuteOptions { overwrite: boolean; workspace_id: string; workspace_instance_id?: string; + linked_worktrees?: boolean; codeApiBaseUrl: string; executionProfile: CodeExecutionContext['executionProfile']; bridgeWorkerId?: string; @@ -628,6 +632,7 @@ export interface ToolExecuteOptions { edits: Array<{ oldText: string; newText: string }>; workspace_id: string; workspace_instance_id?: string; + linked_worktrees?: boolean; codeApiBaseUrl: string; executionProfile: CodeExecutionContext['executionProfile']; bridgeWorkerId?: string; @@ -644,6 +649,7 @@ export interface ToolExecuteOptions { expected_base_sha256?: string; workspace_id: string; workspace_instance_id?: string; + linked_worktrees?: boolean; codeApiBaseUrl: string; executionProfile: CodeExecutionContext['executionProfile']; bridgeWorkerId?: string; @@ -2474,6 +2480,7 @@ async function handleWorkspaceFileRead( ...(codeExecutionContext.codeWorkspace?.workspaceInstanceId ? { workspace_instance_id: codeExecutionContext.codeWorkspace.workspaceInstanceId } : {}), + ...(codeExecutionContext.codeWorkspace?.linkedWorktrees ? { linked_worktrees: true } : {}), start_line: startLine, max_lines: maxLines, codeApiBaseUrl: codeExecutionContext.baseUrl, @@ -2580,6 +2587,7 @@ async function handleWorkspaceSearchCall( ...(codeExecutionContext.codeWorkspace?.workspaceInstanceId ? { workspace_instance_id: codeExecutionContext.codeWorkspace.workspaceInstanceId } : {}), + ...(codeExecutionContext.codeWorkspace?.linkedWorktrees ? { linked_worktrees: true } : {}), ...(typeof args.path === 'string' && args.path.length > 0 ? { path: args.path } : {}), max_results: Number(maxResults), codeApiBaseUrl: codeExecutionContext.baseUrl, @@ -2668,6 +2676,7 @@ async function handleWorkspaceListCall( ...(codeExecutionContext.codeWorkspace?.workspaceInstanceId ? { workspace_instance_id: codeExecutionContext.codeWorkspace.workspaceInstanceId } : {}), + ...(codeExecutionContext.codeWorkspace?.linkedWorktrees ? { linked_worktrees: true } : {}), ...(typeof args.path === 'string' && args.path.length > 0 ? { path: args.path } : {}), ...(typeof args.after_path === 'string' && args.after_path.length > 0 ? { after_path: args.after_path } @@ -3939,6 +3948,7 @@ function attachedWorkspaceMutationParams( ): { workspace_id: string; workspace_instance_id?: string; + linked_worktrees?: boolean; codeApiBaseUrl: string; executionProfile: CodeExecutionContext['executionProfile']; bridgeWorkerId?: string; @@ -3954,6 +3964,7 @@ function attachedWorkspaceMutationParams( ...(codeExecutionContext.codeWorkspace?.workspaceInstanceId ? { workspace_instance_id: codeExecutionContext.codeWorkspace.workspaceInstanceId } : {}), + ...(codeExecutionContext.codeWorkspace?.linkedWorktrees ? { linked_worktrees: true } : {}), codeApiBaseUrl: codeExecutionContext.baseUrl, ...limits, ...(limits.maxRequestTimeoutMs == null diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index 5869bdf5824..e4976ebf54b 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -88,6 +88,7 @@ export * from './workspace'; export * from './reasoningLabels'; export * from './refusal'; export * from './toolValidation'; +export * from './toolTiming'; export * from './remote'; export * from './queuedTurns'; export * from './queuedTurnHttp'; diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 9a09ac949c2..6f0af0c7c74 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -2080,6 +2080,7 @@ export async function initializeAgent( workspaceOperations: attachedWorkspaceOperations, workspaceCommandTimeoutMaxMs: attachedWorkspaceCommandTimeoutMaxMs, workspaceEnvironment: trustedCodeExecutionContext.codeWorkspace?.environment, + workspaceLinkedWorktrees: trustedCodeExecutionContext.codeWorkspace?.linkedWorktrees, }); toolDefinitions = codeExecResult.toolDefinitions; recordCapabilityToolNames(AgentCapabilities.execute_code, codeExecResult.toolNames); @@ -2325,6 +2326,7 @@ export async function initializeAgent( userId: user?.id, workspaceCommandTimeoutMaxMs: attachedWorkspaceCommandTimeoutMaxMs, workspaceEnvironment: trustedCodeExecutionContext.codeWorkspace?.environment, + workspaceLinkedWorktrees: trustedCodeExecutionContext.codeWorkspace?.linkedWorktrees, skillStates: params.skillStates, defaultActiveOnShare: params.defaultActiveOnShare, maxCatalogSkills: getMaxCatalogSkills(runtime), diff --git a/packages/api/src/agents/load.ts b/packages/api/src/agents/load.ts index 7ae16333b5a..e0c26084462 100644 --- a/packages/api/src/agents/load.ts +++ b/packages/api/src/agents/load.ts @@ -6,6 +6,7 @@ import { isEphemeralAgentId, getEphemeralSender, encodeEphemeralAgentId, + applyAgentToolSwitches, } from 'librechat-data-provider'; import type { AgentModelParameters, @@ -62,6 +63,8 @@ export interface LoadAgentParams { agent_id: string; endpoint: string; model_parameters?: AgentModelParameters & { model?: string }; + /** The chat's tool switch state; only applied to the primary saved agent. */ + userToolSwitches?: TEphemeralAgent | null; } /** @@ -232,7 +235,7 @@ export async function loadAgent( params: LoadAgentParams, deps: LoadAgentDeps, ): Promise { - const { req, spec, agent_id, endpoint, model_parameters } = params; + const { req, spec, agent_id, endpoint, model_parameters, userToolSwitches } = params; if (!agent_id) { return null; } @@ -246,5 +249,8 @@ export async function loadAgent( } agent.version ??= agent.versions?.length ?? 0; + if (userToolSwitches !== undefined) { + agent.tools = applyAgentToolSwitches(agent, userToolSwitches).tools; + } return agent; } diff --git a/packages/api/src/agents/skills.ts b/packages/api/src/agents/skills.ts index d080628cc31..f74fee50ea3 100644 --- a/packages/api/src/agents/skills.ts +++ b/packages/api/src/agents/skills.ts @@ -441,6 +441,8 @@ export interface InjectSkillCatalogParams { /** Deployment ceiling advertised on attached Bash tool definitions. */ workspaceCommandTimeoutMaxMs?: number; workspaceEnvironment?: CodeWorkspaceDescriptor['environment']; + /** The worker runs `.worktrees/` in its own lane; advertise `cwd` routing to the model. */ + workspaceLinkedWorktrees?: boolean; /** Current user ID — used to determine skill ownership for active-state resolution. */ userId?: string; /** Per-user skill overrides: `{ [skillId]: boolean }`. Missing entries use the default. */ @@ -677,6 +679,7 @@ export async function injectSkillCatalog( workspaceOperations, workspaceCommandTimeoutMaxMs, workspaceEnvironment, + workspaceLinkedWorktrees, userId, skillStates, defaultActiveOnShare = false, @@ -863,6 +866,7 @@ export async function injectSkillCatalog( workspaceOperations, workspaceCommandTimeoutMaxMs, workspaceEnvironment, + workspaceLinkedWorktrees, }); workingDefs = codeExecResult.toolDefinitions; diff --git a/packages/api/src/agents/toolTiming.spec.ts b/packages/api/src/agents/toolTiming.spec.ts new file mode 100644 index 00000000000..f84a28e5b30 --- /dev/null +++ b/packages/api/src/agents/toolTiming.spec.ts @@ -0,0 +1,292 @@ +import { StepEvents } from 'librechat-data-provider'; +import { createToolTimingAdapter, createToolTimingTracker } from './toolTiming'; + +describe('createToolTimingTracker', () => { + it('projects child dispatch and completion only onto its own persisted tool part', () => { + const part = { type: 'tool_call', tool_call: { id: 'child-call' } }; + const aggregator = { contentParts: [part], stepMap: new Map([['child-step', { index: 0 }]]) }; + const adapter = createToolTimingAdapter({ emit: async () => undefined }); + adapter.child(aggregator, { + phase: 'tool_preparation', + data: { + id: 'child-step', + toolCallId: 'child-call', + observed_at: 100, + }, + }); + adapter.child(aggregator, { + phase: 'tool_calls_dispatched', + data: { + dispatched_at: 1_000, + toolCalls: [ + { id: 'different-call', stepId: 'child-step' }, + { id: 'child-call', stepId: 'other-step' }, + { id: 'child-call', stepId: 'child-step' }, + ], + }, + }); + expect(part.tool_call).toMatchObject({ toolDispatchedAt: 1_000 }); + const completedPart = { type: 'tool_call', tool_call: { id: 'child-call', output: 'ok' } }; + aggregator.contentParts[0] = completedPart; + adapter.child(aggregator, { + phase: 'run_step_completed', + data: { + result: { id: 'child-step', completed_at: 1_340, tool_call: { id: 'child-call' } }, + }, + }); + expect(completedPart.tool_call).toMatchObject({ + toolPreparationStartedAt: 100, + toolPreparationDurationMs: 900, + toolDispatchedAt: 1_000, + toolExecutionDurationMs: 340, + }); + }); + + it('publishes each first fragment once without forwarding tool arguments', async () => { + const publish = jest.fn(async () => undefined); + const adapter = createToolTimingAdapter({ emit: publish }); + const fragment = { + id: 'step-1', + observed_at: 100, + delta: { + type: 'tool_calls' as const, + tool_calls: [{ id: 'call-1', index: 0, args: 'secret' }], + }, + }; + await adapter.delta(fragment); + await adapter.delta({ ...fragment, observed_at: 200 }); + expect(publish).toHaveBeenCalledTimes(1); + expect(publish).toHaveBeenCalledWith({ + event: StepEvents.ON_TOOL_PREPARATION, + data: { id: 'step-1', index: 0, toolCallId: 'call-1', observed_at: 100 }, + }); + }); + it('records each call separately even when results complete out of order', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'step_a', + observed_at: 100, + delta: { + type: 'tool_calls', + tool_calls: [{ id: 'a', index: 0, args: '{' }], + }, + }); + timing.observe({ + id: 'step_b', + observed_at: 200, + delta: { + type: 'tool_calls', + tool_calls: [{ id: 'b', index: 1, args: '{' }], + }, + }); + timing.dispatched({ + dispatched_at: 500, + toolCalls: [ + { id: 'a', stepId: 'step_a' }, + { id: 'b', stepId: 'step_b' }, + ], + }); + timing.completed('step_b', 'b', 530); + timing.completed('step_a', 'a', 700); + expect(timing.take('b', 'step_b')).toEqual({ + toolPreparationDurationMs: 300, + toolExecutionDurationMs: 30, + }); + expect(timing.take('a', 'step_a')).toEqual({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 200, + }); + expect(timing.take('a', 'step_a')).toEqual({}); + }); + + it('keeps an idless initial fragment for the sole first call without assigning sibling fragments', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'sole_step', + observed_at: 100, + delta: { + type: 'tool_calls', + tool_calls: [{ index: 0, args: '{' }], + }, + }); + timing.observe({ + id: 'multi_step', + observed_at: 90, + delta: { + type: 'tool_calls', + tool_calls: [ + { index: 0, args: '{' }, + { index: 1, args: '{' }, + ], + }, + }); + timing.dispatched({ + dispatched_at: 400, + toolCalls: [ + { id: 'sole', stepId: 'sole_step' }, + { id: 'sibling', stepId: 'multi_step' }, + ], + }); + timing.completed('sole_step', 'sole', 405); + timing.completed('multi_step', 'sibling', 420); + expect(timing.take('sole', 'sole_step')).toEqual({ + toolPreparationDurationMs: 300, + toolExecutionDurationMs: 5, + }); + expect(timing.take('sibling', 'multi_step')).toEqual({ toolExecutionDurationMs: 20 }); + }); + + it('keeps the earliest ID-less fragment when the first call gains an ID later', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'step', + observed_at: 100, + delta: { type: 'tool_calls', tool_calls: [{ index: 0, args: '{' }] }, + }); + timing.observe({ + id: 'step', + observed_at: 200, + delta: { type: 'tool_calls', tool_calls: [{ id: 'call', index: 0, args: '"x":1}' }] }, + }); + timing.dispatched({ dispatched_at: 500, toolCalls: [{ id: 'call', stepId: 'step' }] }); + timing.completed('step', 'call', 530); + expect(timing.take('call', 'step')).toEqual({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 30, + }); + }); + + it('keeps the earliest timestamp when a named fragment is delivered before the ID-less one', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'step', + observed_at: 200, + delta: { type: 'tool_calls', tool_calls: [{ id: 'call', index: 0, args: '"x":1}' }] }, + }); + timing.observe({ + id: 'step', + observed_at: 100, + delta: { type: 'tool_calls', tool_calls: [{ index: 0, args: '{' }] }, + }); + timing.dispatched({ dispatched_at: 500, toolCalls: [{ id: 'call', stepId: 'step' }] }); + timing.completed('step', 'call', 530); + expect(timing.take('call', 'step')).toEqual({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 30, + }); + }); + + it('uses the ID-less start for an indexless named fragment only when it owns the step', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'step', + observed_at: 100, + delta: { type: 'tool_calls', tool_calls: [{ index: 0, args: '{' }] }, + }); + timing.observe({ + id: 'step', + observed_at: 200, + delta: { type: 'tool_calls', tool_calls: [{ id: 'call', args: '"x":1}' }] }, + }); + timing.dispatched({ dispatched_at: 500, toolCalls: [{ id: 'call', stepId: 'step' }] }); + timing.completed('step', 'call', 530); + expect(timing.take('call', 'step')).toEqual({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 30, + }); + }); + + it('does not give an ID-less first call’s preparation to an unobserved sibling', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'step', + observed_at: 100, + delta: { type: 'tool_calls', tool_calls: [{ index: 0, args: '{' }] }, + }); + timing.dispatched({ + dispatched_at: 500, + toolCalls: [ + { id: 'first', stepId: 'step' }, + { id: 'second', stepId: 'step' }, + ], + }); + timing.completed('step', 'second', 540); + expect(timing.take('second', 'step')).toEqual({ toolExecutionDurationMs: 40 }); + }); + + it('keeps a first-call ID-less fragment after a faster sibling settles', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'step', + observed_at: 100, + delta: { + type: 'tool_calls', + tool_calls: [{ index: 0, args: '{' }], + }, + }); + timing.observe({ + id: 'step', + observed_at: 200, + delta: { + type: 'tool_calls', + tool_calls: [{ id: 'first', index: 0, args: '"x":1}' }], + }, + }); + timing.dispatched({ + dispatched_at: 500, + toolCalls: [ + { id: 'first', stepId: 'step' }, + { id: 'second', stepId: 'step' }, + ], + }); + timing.completed('step', 'second', 515); + expect(timing.take('second', 'step')).toEqual({ toolExecutionDurationMs: 15 }); + timing.completed('step', 'first', 560); + expect(timing.take('first', 'step')).toEqual({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 60, + }); + }); + + it('isolates simultaneous duplicate provider call IDs by step', () => { + const timing = createToolTimingTracker(); + for (const [stepId, observedAt, dispatchedAt, completedAt] of [ + ['step_1', 100, 300, 310], + ['step_2', 150, 500, 550], + ] as const) { + timing.observe({ + id: stepId, + observed_at: observedAt, + delta: { + type: 'tool_calls', + tool_calls: [{ id: 'shared', index: 0, args: '{' }], + }, + }); + timing.dispatched({ dispatched_at: dispatchedAt, toolCalls: [{ id: 'shared', stepId }] }); + timing.completed(stepId, 'shared', completedAt); + } + expect(timing.take('shared', 'step_1')).toEqual({ + toolPreparationDurationMs: 200, + toolExecutionDurationMs: 10, + }); + expect(timing.take('shared', 'step_2')).toEqual({ + toolPreparationDurationMs: 350, + toolExecutionDurationMs: 50, + }); + }); + + it('does not manufacture execution duration on abort or an unobserved dispatch', () => { + const timing = createToolTimingTracker(); + timing.observe({ + id: 'step', + observed_at: 100, + delta: { + type: 'tool_calls', + tool_calls: [{ id: 'denied', index: 0, args: '{}' }], + }, + }); + expect(timing.take('denied', 'step')).toEqual({}); + timing.dispatched({ dispatched_at: 500, toolCalls: [{ id: 'cancelled', stepId: 'step' }] }); + expect(timing.take('cancelled', 'step')).toEqual({}); + }); +}); diff --git a/packages/api/src/agents/toolTiming.ts b/packages/api/src/agents/toolTiming.ts new file mode 100644 index 00000000000..c35c597e11d --- /dev/null +++ b/packages/api/src/agents/toolTiming.ts @@ -0,0 +1,284 @@ +import { StepEvents, getToolTimingDurations } from 'librechat-data-provider'; +import type { Agents } from 'librechat-data-provider'; + +/** SDK handoff, not proof that the host or database started executing. */ +type DispatchedCall = { id: string; stepId?: string }; +type Dispatch = { dispatched_at: number; toolCalls: DispatchedCall[] }; +type Fragment = Pick; +export type ToolPreparationMarker = Agents.ToolPreparationMarker; +type StepTiming = { + firstByCall: Map; + firstUnidentified?: number; + firstCallId?: string; + dispatchedByCall: Map; + completedByCall: Map; + pendingIds: Set; +}; + +export type ToolTimingTracker = { + observe(fragment: Fragment): ToolPreparationMarker[]; + prepare(marker: ToolPreparationMarker): void; + dispatched(event: Dispatch): void; + completed(stepId: string, callId: string, at?: number): void; + take(callId: string, stepId: string): ReturnType; +}; + +const validTime = (value: number | undefined): value is number => + typeof value === 'number' && Number.isFinite(value) && value >= 0; + +/** One instance per response. Only bounded IDs and timestamps are retained. */ +type TimingEvent = + | { event: StepEvents.ON_TOOL_PREPARATION; data: ToolPreparationMarker } + | { event: StepEvents.ON_TOOL_CALLS_DISPATCHED; data: Agents.ToolCallsDispatchedEvent }; + +type ToolCompletion = { + result?: { id?: string; completed_at?: number; tool_call?: { id?: string } }; +}; + +type ChildTimingStamps = { observedAt?: number; dispatchedAt?: number }; + +type ChildTimingAggregator = { + stepMap?: Map; + contentParts?: Array<{ + type?: string; + tool_call?: { + id?: string; + toolPreparationStartedAt?: number; + toolDispatchedAt?: number; + toolPreparationDurationMs?: number; + toolExecutionDurationMs?: number; + }; + }>; +}; + +export type ToolTimingAdapter = { + dispatch: { handle(event: string, data: Agents.ToolCallsDispatchedEvent): Promise }; + delta(fragment: Fragment): Promise; + completed(event: ToolCompletion): void; + close(toolCall: { id?: string }, stepId: string): void; + child(aggregator: ChildTimingAggregator, update: { phase?: string; data?: unknown }): void; +}; + +/** Owns timing event transitions; legacy Express only wires these handlers. */ +export function createToolTimingAdapter({ + replayEvents, + emit, +}: { + replayEvents?: readonly { event: string; data?: unknown }[]; + emit(event: TimingEvent): Promise | void; +}): ToolTimingAdapter { + const tracker = createToolTimingTracker(replayEvents); + const childStamps = new WeakMap>(); + return { + dispatch: { + async handle(_event, data) { + tracker.dispatched(data); + await emit({ event: StepEvents.ON_TOOL_CALLS_DISPATCHED, data }); + }, + }, + async delta(fragment) { + for (const marker of tracker.observe(fragment)) { + await emit({ event: StepEvents.ON_TOOL_PREPARATION, data: marker }); + } + }, + completed(event) { + const result = event?.result; + if (result?.id && result.tool_call?.id) { + tracker.completed(result.id, result.tool_call.id, result.completed_at); + } + }, + close(toolCall, stepId) { + Object.assign(toolCall, tracker.take(toolCall.id ?? '', stepId)); + }, + child(aggregator, update) { + let stamps = childStamps.get(aggregator); + if (!stamps) { + stamps = new Map(); + childStamps.set(aggregator, stamps); + } + applySubagentToolTiming(aggregator, update, stamps); + }, + }; +} + +function applySubagentToolTiming( + aggregator: ChildTimingAggregator, + update: { phase?: string; data?: unknown }, + stamps: Map, +): void { + if ( + update.phase !== 'tool_preparation' && + update.phase !== 'tool_calls_dispatched' && + update.phase !== 'run_step_completed' + ) + return; + if (update.data == null || typeof update.data !== 'object') return; + const data = update.data as { + id?: string; + toolCallId?: string; + observed_at?: number; + dispatched_at?: number; + toolCalls?: Array<{ id?: string; stepId?: string }>; + result?: { id?: string; completed_at?: number; tool_call?: { id?: string } }; + }; + const apply = ( + stepId: string | undefined, + callId: string | undefined, + fn: (tool: { + id?: string; + toolPreparationStartedAt?: number; + toolDispatchedAt?: number; + toolPreparationDurationMs?: number; + toolExecutionDurationMs?: number; + }) => void, + ) => { + if (!stepId || !callId) return; + const index = aggregator.stepMap?.get(stepId)?.index; + const part = index == null ? undefined : aggregator.contentParts?.[index]; + if (part?.type !== 'tool_call' || part.tool_call?.id !== callId) return; + fn(part.tool_call); + }; + if (update.phase === 'tool_preparation' && validTime(data.observed_at)) { + if (!data.id || !data.toolCallId) return; + const key = `${data.id}\u0000${data.toolCallId}`; + const previous = stamps.get(key) ?? {}; + const observedAt = Math.min(previous.observedAt ?? data.observed_at, data.observed_at); + stamps.set(key, { ...previous, observedAt }); + apply(data.id, data.toolCallId, (tool) => { + tool.toolPreparationStartedAt = observedAt; + }); + } + if (update.phase === 'tool_calls_dispatched' && validTime(data.dispatched_at)) { + for (const call of data.toolCalls ?? []) { + if (!call.stepId || !call.id) continue; + const key = `${call.stepId}\u0000${call.id}`; + const previous = stamps.get(key) ?? {}; + const dispatchedAt = Math.min( + previous.dispatchedAt ?? data.dispatched_at!, + data.dispatched_at!, + ); + stamps.set(key, { ...previous, dispatchedAt }); + apply(call.stepId, call.id, (tool) => { + tool.toolDispatchedAt = dispatchedAt; + if (previous.observedAt != null) tool.toolPreparationStartedAt = previous.observedAt; + }); + } + } + if (update.phase === 'run_step_completed' && validTime(data.result?.completed_at)) { + const stepId = data.result?.id; + const callId = data.result?.tool_call?.id; + if (!stepId || !callId) return; + const key = `${stepId}\u0000${callId}`; + const recorded = stamps.get(key); + apply(stepId, callId, (tool) => { + if (recorded?.dispatchedAt != null) tool.toolDispatchedAt = recorded.dispatchedAt; + if (recorded?.observedAt != null) tool.toolPreparationStartedAt = recorded.observedAt; + Object.assign( + tool, + getToolTimingDurations({ + observedAt: tool.toolPreparationStartedAt, + dispatchedAt: tool.toolDispatchedAt, + completedAt: data.result?.completed_at, + }), + ); + }); + stamps.delete(key); + } +} + +export function createToolTimingTracker( + replayEvents: readonly { event: string; data?: unknown }[] = [], +): ToolTimingTracker { + const steps = new Map(); + const getStep = (id: string): StepTiming => { + let step = steps.get(id); + if (!step) { + step = { + firstByCall: new Map(), + dispatchedByCall: new Map(), + completedByCall: new Map(), + pendingIds: new Set(), + }; + steps.set(id, step); + } + return step; + }; + const prepare = ({ id, index, toolCallId, observed_at: at }: ToolPreparationMarker): void => { + if (!id || !validTime(at)) return; + const step = getStep(id); + if (toolCallId) { + step.firstByCall.set(toolCallId, Math.min(step.firstByCall.get(toolCallId) ?? at, at)); + if (index === 0) step.firstCallId = toolCallId; + } else if (index === 0) { + step.firstUnidentified = Math.min(step.firstUnidentified ?? at, at); + if (step.firstCallId) { + const first = step.firstCallId; + step.firstByCall.set(first, Math.min(step.firstByCall.get(first) ?? at, at)); + } + } + }; + + const tracker: ToolTimingTracker = { + observe({ id, delta, observed_at: at }: Fragment): ToolPreparationMarker[] { + if (!id || !validTime(at) || delta.type !== 'tool_calls') return []; + const markers: ToolPreparationMarker[] = []; + const step = getStep(id); + for (const chunk of delta.tool_calls ?? []) { + const index = typeof chunk.index === 'number' ? chunk.index : undefined; + const callId = typeof chunk.id === 'string' && chunk.id ? chunk.id : undefined; + if (!callId && (index !== 0 || delta.tool_calls?.length !== 1)) continue; + const prior = callId ? step.firstByCall.get(callId) : step.firstUnidentified; + const marker = { id, index, ...(callId && { toolCallId: callId }), observed_at: at }; + prepare(marker); + if (prior == null) markers.push(marker); + } + return markers; + }, + prepare, + dispatched({ dispatched_at: at, toolCalls }: Dispatch): void { + if (!validTime(at) || !Array.isArray(toolCalls)) return; + for (const call of toolCalls) { + if (!call.id || !call.stepId) continue; + const step = getStep(call.stepId); + step.pendingIds.add(call.id); + step.dispatchedByCall.set(call.id, Math.min(step.dispatchedByCall.get(call.id) ?? at, at)); + } + }, + completed(stepId: string, callId: string, at?: number): void { + if (stepId && callId && validTime(at)) getStep(stepId).completedByCall.set(callId, at); + }, + take(callId: string, stepId: string): ReturnType { + const step = steps.get(stepId); + if (!step) return {}; + const sole = step.pendingIds.size === 1 && step.pendingIds.has(callId); + const unidentified = + step.firstCallId === callId || (step.firstCallId == null && sole) + ? step.firstUnidentified + : undefined; + const start = Math.min(step.firstByCall.get(callId) ?? Infinity, unidentified ?? Infinity); + const duration = getToolTimingDurations({ + observedAt: Number.isFinite(start) ? start : undefined, + dispatchedAt: step.dispatchedByCall.get(callId), + completedAt: step.completedByCall.get(callId), + }); + step.firstByCall.delete(callId); + step.dispatchedByCall.delete(callId); + step.completedByCall.delete(callId); + step.pendingIds.delete(callId); + if (step.pendingIds.size === 0) steps.delete(stepId); + return duration; + }, + }; + for (const event of replayEvents) { + if (event.data == null || typeof event.data !== 'object') continue; + if (event.event === StepEvents.ON_TOOL_PREPARATION) { + const marker = event.data as Partial; + if (typeof marker.id === 'string' && typeof marker.observed_at === 'number') { + tracker.prepare(marker as ToolPreparationMarker); + } + } else if (event.event === StepEvents.ON_TOOL_CALLS_DISPATCHED) { + tracker.dispatched(event.data as Dispatch); + } + } + return tracker; +} diff --git a/packages/api/src/agents/tools.spec.ts b/packages/api/src/agents/tools.spec.ts index 0455c7cb0fb..1852ea5a1a0 100644 --- a/packages/api/src/agents/tools.spec.ts +++ b/packages/api/src/agents/tools.spec.ts @@ -704,6 +704,24 @@ describe('registerCodeExecutionTools', () => { }); expect(bash?.description).toContain('owner/app'); }); + + it('tells the model to route worktree commands through cwd only when lanes are available', () => { + const cwdDescription = (workspaceLinkedWorktrees: boolean): string | undefined => { + const bash = registerCodeExecutionTools({ + toolRegistry: undefined, + toolDefinitions: [], + includeBash: true, + workspaceTools: true, + workspaceOperations: new Set(['execute_command']), + workspaceLinkedWorktrees, + }).toolDefinitions.find((def) => def.name === 'bash_tool'); + return (bash?.parameters as { properties?: { cwd?: { description?: string } } })?.properties + ?.cwd?.description; + }; + + expect(cwdDescription(true)).toContain('.worktrees/'); + expect(cwdDescription(false)).not.toContain('.worktrees'); + }); const makeRegistry = (): LCToolRegistry => new Map() as unknown as LCToolRegistry; describe('fresh run (no pre-existing defs or registry entries)', () => { diff --git a/packages/api/src/agents/tools.ts b/packages/api/src/agents/tools.ts index 8dd9742dbdd..3c8d85c24e0 100644 --- a/packages/api/src/agents/tools.ts +++ b/packages/api/src/agents/tools.ts @@ -439,6 +439,8 @@ export interface RegisterCodeExecutionToolsParams { /** Deployment ceiling advertised on attached Bash tool definitions. */ workspaceCommandTimeoutMaxMs?: number; workspaceEnvironment?: CodeWorkspaceDescriptor['environment']; + /** The worker runs `.worktrees/` in its own lane; advertise `cwd` routing to the model. */ + workspaceLinkedWorktrees?: boolean; /** * When `true`, the registered `bash_tool` description includes the * LLM-facing `{{toolturn}}` reference syntax guide so the @@ -1040,6 +1042,7 @@ function createBashToolDef( workspaceTools = false, workspaceCommandTimeoutMaxMs?: number, workspaceEnvironment?: CodeWorkspaceDescriptor['environment'], + workspaceLinkedWorktrees = false, ): LCTool { /* Passed as a variable (not an inline literal) so the extra * `statefulSessions` key stays assignable against pinned SDK versions @@ -1052,7 +1055,11 @@ function createBashToolDef( ? buildAttachedWorkspaceBashDescription(enableToolOutputReferences, workspaceEnvironment) : buildBashExecutionToolDescription(descriptionOpts), parameters: (workspaceTools - ? buildAttachedWorkspaceBashSchema(workspaceCommandTimeoutMaxMs, workspaceEnvironment) + ? buildAttachedWorkspaceBashSchema( + workspaceCommandTimeoutMaxMs, + workspaceEnvironment, + workspaceLinkedWorktrees, + ) : BashExecutionToolDefinition.schema) as unknown as LCTool['parameters'], }) as LCTool; } @@ -1066,6 +1073,8 @@ function buildBashToolDef(opts: { workspaceTools?: boolean; workspaceCommandTimeoutMaxMs?: number; workspaceEnvironment?: CodeWorkspaceDescriptor['environment']; + /** The worker runs `.worktrees/` in its own lane; advertise `cwd` routing to the model. */ + workspaceLinkedWorktrees?: boolean; }): LCTool { /* Stateful defs are built on demand: the stateless pair covers the * default path, and per-run construction is negligible next to init. */ @@ -1076,6 +1085,7 @@ function buildBashToolDef(opts: { opts.workspaceTools === true, opts.workspaceCommandTimeoutMaxMs, opts.workspaceEnvironment, + opts.workspaceLinkedWorktrees === true, ); } return opts.enableToolOutputReferences @@ -1109,6 +1119,7 @@ export function registerCodeExecutionTools( workspaceOperations, workspaceCommandTimeoutMaxMs, workspaceEnvironment, + workspaceLinkedWorktrees, enableToolOutputReferences = false, statefulSessions = false, } = params; @@ -1129,6 +1140,7 @@ export function registerCodeExecutionTools( workspaceTools, workspaceCommandTimeoutMaxMs, workspaceEnvironment, + workspaceLinkedWorktrees, }), ); } diff --git a/packages/api/src/agents/validation.spec.ts b/packages/api/src/agents/validation.spec.ts index 8b3a75637e1..c2b5bcedb59 100644 --- a/packages/api/src/agents/validation.spec.ts +++ b/packages/api/src/agents/validation.spec.ts @@ -452,3 +452,24 @@ describe('validateAgentModel', () => { expect(logViolation).toHaveBeenCalledTimes(1); }); }); + +describe('tool_options user_toggle', () => { + const base = { provider: 'openAI', model: 'gpt-4o-mini', tools: ['web_search'] }; + + it.each(['on', 'off'])('keeps %s on create and update', (user_toggle) => { + const tool_options = { web_search: { user_toggle } }; + expect(agentCreateSchema.parse({ ...base, tool_options }).tool_options).toEqual(tool_options); + expect(agentUpdateSchema.parse({ tool_options }).tool_options).toEqual(tool_options); + }); + + it('rejects any other value', () => { + const tool_options = { web_search: { user_toggle: 'auto' } }; + expect(agentCreateSchema.safeParse({ ...base, tool_options }).success).toBe(false); + expect(agentUpdateSchema.safeParse({ tool_options }).success).toBe(false); + }); + + it('treats a missing toggle as locked', () => { + const parsed = agentCreateSchema.parse({ ...base, tool_options: { web_search: {} } }); + expect(parsed.tool_options?.web_search.user_toggle).toBeUndefined(); + }); +}); diff --git a/packages/api/src/agents/validation.ts b/packages/api/src/agents/validation.ts index 49ada6e2121..b7a03c7db31 100644 --- a/packages/api/src/agents/validation.ts +++ b/packages/api/src/agents/validation.ts @@ -149,13 +149,14 @@ export const graphEdgeSchema: z.ZodObject< .transform((v) => (v === '' ? undefined : v)), }); -/** Per-tool options schema (defer_loading, allowed_callers, run_in_background, describe_intent) */ +/** Per-tool options schema (defer_loading, allowed_callers, run_in_background, describe_intent, user_toggle) */ export const toolOptionsSchema: z.ZodObject< { defer_loading: z.ZodOptional; allowed_callers: z.ZodOptional, 'many'>>; run_in_background: z.ZodOptional; describe_intent: z.ZodOptional; + user_toggle: z.ZodOptional>; }, 'strip' > = z.object({ @@ -163,6 +164,7 @@ export const toolOptionsSchema: z.ZodObject< allowed_callers: z.array(z.enum(['direct', 'code_execution'])).optional(), run_in_background: z.boolean().optional(), describe_intent: z.boolean().optional(), + user_toggle: z.enum(['on', 'off']).optional(), }); /** Agent tool options - map of tool_id to tool options */ @@ -175,6 +177,7 @@ export const agentToolOptionsSchema: z.ZodOptional< allowed_callers: z.ZodOptional, 'many'>>; run_in_background: z.ZodOptional; describe_intent: z.ZodOptional; + user_toggle: z.ZodOptional>; }, 'strip', z.ZodTypeAny, @@ -183,12 +186,14 @@ export const agentToolOptionsSchema: z.ZodOptional< allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; }, { defer_loading?: boolean | undefined; allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; } > > @@ -525,6 +530,7 @@ export const agentBaseSchema: z.ZodObject< >; run_in_background: z.ZodOptional; describe_intent: z.ZodOptional; + user_toggle: z.ZodOptional>; }, 'strip', z.ZodTypeAny, @@ -533,12 +539,14 @@ export const agentBaseSchema: z.ZodObject< allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; }, { defer_loading?: boolean | undefined; allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; } > > @@ -704,6 +712,7 @@ export const agentCreateSchema: z.ZodObject< >; run_in_background: z.ZodOptional; describe_intent: z.ZodOptional; + user_toggle: z.ZodOptional>; }, 'strip', z.ZodTypeAny, @@ -712,12 +721,14 @@ export const agentCreateSchema: z.ZodObject< allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; }, { defer_loading?: boolean | undefined; allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; } > > @@ -841,6 +852,7 @@ export const agentUpdateSchema: z.ZodObject< >; run_in_background: z.ZodOptional; describe_intent: z.ZodOptional; + user_toggle: z.ZodOptional>; }, 'strip', z.ZodTypeAny, @@ -849,12 +861,14 @@ export const agentUpdateSchema: z.ZodObject< allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; }, { defer_loading?: boolean | undefined; allowed_callers?: ('direct' | 'code_execution')[] | undefined; run_in_background?: boolean | undefined; describe_intent?: boolean | undefined; + user_toggle?: 'on' | 'off' | undefined; } > > diff --git a/packages/api/src/code/bridge.spec.ts b/packages/api/src/code/bridge.spec.ts index bd9dae4bf0f..9fe3e8990fe 100644 --- a/packages/api/src/code/bridge.spec.ts +++ b/packages/api/src/code/bridge.spec.ts @@ -97,6 +97,7 @@ describe('getCodeBridgeWorkerStatus', () => { id: 'project-a', name: 'Project A', workspaceInstances: ['git_worktree'], + workspaceScopes: ['git_linked_worktree'], }, { id: 'docs', operations: ['read_file'] }, ], @@ -128,6 +129,7 @@ describe('getCodeBridgeWorkerStatus', () => { id: 'project-a', name: 'Project A', workspaceInstances: ['git_worktree'], + workspaceScopes: ['git_linked_worktree'], }, { id: 'docs', operations: ['read_file'] }, ], @@ -249,6 +251,53 @@ describe('getCodeBridgeWorkerStatus', () => { }, }, }, + { + online: true, + ready: true, + leaseExpiresInMs: 5_000, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'project-a', workspaceScopes: ['subdirectory'] }], + }, + }, + }, + { + online: true, + ready: true, + leaseExpiresInMs: 5_000, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [{ id: 'project-a', workspaceScopes: 'git_linked_worktree' }], + }, + }, + }, + { + online: true, + ready: true, + leaseExpiresInMs: 5_000, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: 1, + operations: ['read_file'], + workspaces: [ + { id: 'project-a', workspaceScopes: ['git_linked_worktree', 'git_linked_worktree'] }, + ], + }, + }, + }, ])('rejects an invalid upstream status response: %p', async (invalid) => { const fetchImpl = jest .fn() diff --git a/packages/api/src/code/bridge.ts b/packages/api/src/code/bridge.ts index 0820901eca4..17f31aa1f5e 100644 --- a/packages/api/src/code/bridge.ts +++ b/packages/api/src/code/bridge.ts @@ -282,6 +282,7 @@ function validWorkspaceCapabilities(value: unknown): value is { key !== 'name' && key !== 'operations' && key !== 'workspaceInstances' && + key !== 'workspaceScopes' && key !== 'environment' && key !== 'instructions', ) || @@ -297,6 +298,10 @@ function validWorkspaceCapabilities(value: unknown): value is { (!Array.isArray(workspace.workspaceInstances) || workspace.workspaceInstances.length !== 1 || workspace.workspaceInstances[0] !== 'git_worktree')) || + (workspace.workspaceScopes !== undefined && + (!Array.isArray(workspace.workspaceScopes) || + workspace.workspaceScopes.length !== 1 || + workspace.workspaceScopes[0] !== 'git_linked_worktree')) || (workspace.name !== undefined && (typeof workspace.name !== 'string' || workspace.name.trim().length === 0 || diff --git a/packages/api/src/code/capabilities.spec.ts b/packages/api/src/code/capabilities.spec.ts index 746d68b8d1f..4365e0d24e5 100644 --- a/packages/api/src/code/capabilities.spec.ts +++ b/packages/api/src/code/capabilities.spec.ts @@ -461,6 +461,49 @@ describe('resolveCodeExecutionWorkspaceContext', () => { expect(legacy.codeWorkspace).not.toHaveProperty('workspaceInstanceId'); }); + it('routes linked worktrees into lanes only when configured and no conversation instance owns the checkout', async () => { + const workspaceInstanceId = 'c'.repeat(64); + jest.spyOn(globalThis, 'fetch').mockImplementation(async () => + workspaceStatus([ + { + id: 'lanes', + workspaceInstances: ['git_worktree'], + workspaceScopes: ['git_linked_worktree'], + }, + { id: 'legacy' }, + ]), + ); + const resolve = (workspaceId: string, instanceId?: string, linkedWorktrees = true) => + resolveCodeExecutionWorkspaceContext({ + context: { + ...context, + codeEnvironmentConfigSchema: { workspaces: { linkedWorktrees } }, + ...(instanceId ? { conversationWorkspaceInstanceId: instanceId } : {}), + }, + requestedSelections: [{ environmentId: 'personal', workspaceId }], + environments, + getAppConfig, + }); + + const lanes = await resolve('lanes'); + const disabled = await resolve('lanes', undefined, false); + const unconfigured = await resolveCodeExecutionWorkspaceContext({ + context, + requestedSelections: [{ environmentId: 'personal', workspaceId: 'lanes' }], + environments, + getAppConfig, + }); + const instance = await resolve('lanes', workspaceInstanceId); + const legacy = await resolve('legacy'); + + expect(lanes.codeWorkspace?.linkedWorktrees).toBe(true); + expect(disabled.codeWorkspace).not.toHaveProperty('linkedWorktrees'); + expect(unconfigured.codeWorkspace).not.toHaveProperty('linkedWorktrees'); + expect(instance.codeWorkspace?.workspaceInstanceId).toBe(workspaceInstanceId); + expect(instance.codeWorkspace).not.toHaveProperty('linkedWorktrees'); + expect(legacy.codeWorkspace).not.toHaveProperty('linkedWorktrees'); + }); + it('admits native workspace tools without enabling programmatic runtime execution', async () => { jest .spyOn(globalThis, 'fetch') diff --git a/packages/api/src/code/capabilities.ts b/packages/api/src/code/capabilities.ts index 85b8f7ed72a..32f1d37ad22 100644 --- a/packages/api/src/code/capabilities.ts +++ b/packages/api/src/code/capabilities.ts @@ -182,6 +182,14 @@ export async function resolveCodeExecutionWorkspaceContext({ workspace.workspaceInstances?.includes('git_worktree') ? { workspaceInstanceId: context.conversationWorkspaceInstanceId } : {}), + ...(context.codeEnvironmentConfigSchema?.workspaces?.linkedWorktrees === true && + workspace.workspaceScopes?.includes('git_linked_worktree') && + !( + context.conversationWorkspaceInstanceId && + workspace.workspaceInstances?.includes('git_worktree') + ) + ? { linkedWorktrees: true } + : {}), ...(status.maxCommandTimeoutMs == null ? {} : { maxCommandTimeoutMs: status.maxCommandTimeoutMs }), diff --git a/packages/api/src/code/command.spec.ts b/packages/api/src/code/command.spec.ts index b027583ef6a..7a625819a62 100644 --- a/packages/api/src/code/command.spec.ts +++ b/packages/api/src/code/command.spec.ts @@ -316,6 +316,23 @@ describe('createAttachedWorkspaceBashTool', () => { }); }); + test('runs a command whose working directory is a linked worktree in that lane', async () => { + const fetchImpl: CodeBridgeFetch = jest.fn(async () => commandResponse()); + const bashTool = createAttachedWorkspaceBashTool({ + baseUrl: 'https://code.example.com/v1', + authHeaders: () => ({}), + workspaceId: 'project-a', + linkedWorktrees: true, + fetchImpl, + }); + + expect(JSON.stringify(bashTool.schema)).toContain('.worktrees/'); + await bashTool.func({ command: 'npm test', cwd: '.worktrees/fix-a/api' }, undefined, {}); + + const request = JSON.parse(String((fetchImpl as jest.Mock).mock.calls[0][1]?.body)); + expect(request).toMatchObject({ command: 'npm test', cwd: 'api', worktree: 'fix-a' }); + }); + test('forwards a bounded per-call execution timeout', async () => { const fetchImpl: CodeBridgeFetch = jest.fn(async () => commandResponse()); const bashTool = createAttachedWorkspaceBashTool({ diff --git a/packages/api/src/code/command.ts b/packages/api/src/code/command.ts index 7981204f2ad..d9754ee6704 100644 --- a/packages/api/src/code/command.ts +++ b/packages/api/src/code/command.ts @@ -173,9 +173,15 @@ export function resolveAttachedWorkspaceRequestTimeoutMs( return configured; } +const linkedWorktreeWorkingDirectorySchema: BoundedWorkingDirectorySchema = { + ...attachedWorkingDirectorySchema, + description: `${attachedWorkingDirectorySchema.description} To work in a linked worktree, pass its directory here (for example ".worktrees/fix-auth") instead of running cd inside the command: commands in different .worktrees/ directories then run in parallel, while a cd from the workspace root waits for all of them.`, +}; + export function buildAttachedWorkspaceBashSchema( maxTimeoutMs: number = WORKSPACE_COMMAND_DEFAULT_TIMEOUT_MS, environment?: CodeWorkspaceDescriptor['environment'], + linkedWorktrees = false, ): NonNullable { const effectiveMaxTimeoutMs = normalizeAttachedWorkspaceCommandTimeoutMax(maxTimeoutMs); return { @@ -183,7 +189,10 @@ export function buildAttachedWorkspaceBashSchema( properties: { ...bashSchema.properties, command: attachedCommandSchema, - cwd: attachedWorkingDirectorySchema, + cwd: + linkedWorktrees === true + ? linkedWorktreeWorkingDirectorySchema + : attachedWorkingDirectorySchema, timeoutMs: buildAttachedTimeoutSchema(effectiveMaxTimeoutMs), ...(environment?.actions.length ? { @@ -322,12 +331,15 @@ export function createAttachedWorkspaceBashTool({ codeApiMaxRetryWaitMs, maxRequestTimeoutMs, minCommandAdmissionMs, + linkedWorktrees = false, fetchImpl, }: { baseUrl: string; authHeaders: () => Promise> | Record; workspaceId: string; workspaceInstanceId?: string; + /** The worker runs each `.worktrees/` in its own lane; a matching `cwd` is routed there. */ + linkedWorktrees?: boolean; environment?: CodeWorkspaceDescriptor['environment']; gitIdentity?: AgentGitIdentity | null; /** Effective admin/upstream ceiling already intersected with the protocol hard cap. */ @@ -347,7 +359,7 @@ export function createAttachedWorkspaceBashTool({ minCommandAdmissionMs, ); const schema = structuredClone( - buildAttachedWorkspaceBashSchema(effectiveMaxTimeoutMs, environment), + buildAttachedWorkspaceBashSchema(effectiveMaxTimeoutMs, environment, linkedWorktrees), ); const actions = environment?.actions ?? []; return tool( @@ -405,6 +417,7 @@ export function createAttachedWorkspaceBashTool({ baseURL: baseUrl, /** Passed as a supplier: a queued call outlives its minted token. */ authHeaders, + linkedWorktrees, request: { protocolVersion: 1, operation: 'execute_command', diff --git a/packages/api/src/code/workspace.lanes.spec.ts b/packages/api/src/code/workspace.lanes.spec.ts new file mode 100644 index 00000000000..46ae78b4293 --- /dev/null +++ b/packages/api/src/code/workspace.lanes.spec.ts @@ -0,0 +1,183 @@ +import type { WorkspaceToolRequest, WorkspaceToolResult } from './workspace'; +import { fromLinkedWorktreeResult, toLinkedWorktreeRequest } from './workspace'; + +const base = { protocolVersion: 1, workspaceId: 'librechat' } as const; + +describe('toLinkedWorktreeRequest', () => { + it.each([ + ['read_file', { operation: 'read_file', path: '.worktrees/fix-a/src/index.ts' }], + ['write_file', { operation: 'write_file', path: '.worktrees/fix-a/src/index.ts', content: '' }], + ['edit_file', { operation: 'edit_file', path: '.worktrees/fix-a/src/index.ts', edits: [] }], + [ + 'preview_edit', + { operation: 'preview_edit', path: '.worktrees/fix-a/src/index.ts', edits: [] }, + ], + ])('routes a %s path into its worktree lane', (_operation, fields) => { + const request = { ...base, ...fields } as WorkspaceToolRequest; + expect(toLinkedWorktreeRequest(request)).toEqual({ + worktree: 'fix-a', + request: { ...request, path: 'src/index.ts', worktree: 'fix-a' }, + }); + }); + + it('scopes searches and listings to the worktree, omitting the path at its root', () => { + expect( + toLinkedWorktreeRequest({ + ...base, + operation: 'search_text', + query: 'TODO', + path: '.worktrees/fix-a', + }), + ).toEqual({ + worktree: 'fix-a', + request: { ...base, operation: 'search_text', query: 'TODO', worktree: 'fix-a' }, + }); + expect( + toLinkedWorktreeRequest({ + ...base, + operation: 'list_files', + path: '.worktrees/fix-a/src', + afterPath: '.worktrees/fix-a/src/a.ts', + }), + ).toEqual({ + worktree: 'fix-a', + request: { + ...base, + operation: 'list_files', + path: 'src', + afterPath: 'src/a.ts', + worktree: 'fix-a', + }, + }); + }); + + it('runs a command whose cwd is inside a worktree in that lane', () => { + expect( + toLinkedWorktreeRequest({ + ...base, + operation: 'execute_command', + command: 'npm test', + cwd: '.worktrees/fix-a', + }), + ).toEqual({ + worktree: 'fix-a', + request: { ...base, operation: 'execute_command', command: 'npm test', worktree: 'fix-a' }, + }); + expect( + toLinkedWorktreeRequest({ + ...base, + operation: 'execute_command', + command: 'npx jest', + cwd: '.worktrees/fix-a/packages/api', + })?.request, + ).toMatchObject({ cwd: 'packages/api', worktree: 'fix-a' }); + }); + + it.each<[string, WorkspaceToolRequest]>([ + ['a root path', { ...base, operation: 'read_file', path: 'src/index.ts' }], + ['the worktree directory itself', { ...base, operation: 'read_file', path: '.worktrees/a' }], + ['the worktrees container', { ...base, operation: 'list_files', path: '.worktrees' }], + ['an invalid name', { ...base, operation: 'read_file', path: '.worktrees/-bad/x.ts' }], + ['a lock name', { ...base, operation: 'read_file', path: '.worktrees/a.lock/x.ts' }], + [ + 'an instruction read', + { + ...base, + operation: 'read_file', + path: '.worktrees/a/AGENTS.md', + instructionSha256: 'a'.repeat(64), + }, + ], + [ + 'a conversation instance', + { + ...base, + operation: 'read_file', + path: '.worktrees/a/x.ts', + workspaceInstanceId: 'b'.repeat(64), + }, + ], + [ + 'a listing cursor in another worktree', + { + ...base, + operation: 'list_files', + path: '.worktrees/a', + afterPath: '.worktrees/b/x.ts', + }, + ], + [ + 'a listing cursor at the worktree root', + { ...base, operation: 'list_files', path: '.worktrees/a', afterPath: '.worktrees/a' }, + ], + [ + 'a command without cwd', + { ...base, operation: 'execute_command', command: 'cd .worktrees/a' }, + ], + [ + 'an environment action', + { + ...base, + operation: 'execute_command', + command: 'setup', + cwd: '.worktrees/a', + environmentAction: { name: 'setup', fingerprint: 'c'.repeat(64) }, + }, + ], + ])('leaves %s root-scoped', (_case, request) => { + expect(toLinkedWorktreeRequest(request)).toBeUndefined(); + }); +}); + +describe('fromLinkedWorktreeResult', () => { + it('restores the worktree prefix on every reported path', () => { + const listing: WorkspaceToolResult = { + ...base, + operation: 'list_files', + paths: ['src/a.ts', 'src/b.ts'], + truncated: true, + nextAfterPath: 'src/b.ts', + }; + expect(fromLinkedWorktreeResult(listing, 'fix-a')).toEqual({ + ...listing, + paths: ['.worktrees/fix-a/src/a.ts', '.worktrees/fix-a/src/b.ts'], + nextAfterPath: '.worktrees/fix-a/src/b.ts', + }); + + const search: WorkspaceToolResult = { + ...base, + operation: 'search_text', + matches: [{ path: 'src/a.ts', line: 3, column: 1, text: 'TODO' }], + truncated: false, + }; + expect(fromLinkedWorktreeResult(search, 'fix-a')).toEqual({ + ...search, + matches: [{ path: '.worktrees/fix-a/src/a.ts', line: 3, column: 1, text: 'TODO' }], + }); + + const written: WorkspaceToolResult = { + ...base, + operation: 'write_file', + path: 'src/a.ts', + created: true, + bytesWritten: 1, + }; + expect(fromLinkedWorktreeResult(written, 'fix-a')).toEqual({ + ...written, + path: '.worktrees/fix-a/src/a.ts', + }); + }); + + it('returns command results unchanged', () => { + const result: WorkspaceToolResult = { + ...base, + operation: 'execute_command', + stdout: '', + stderr: '', + exitCode: 0, + timedOut: false, + truncated: false, + }; + expect(fromLinkedWorktreeResult(result, 'fix-a')).toBe(result); + }); +}); diff --git a/packages/api/src/code/workspace.spec.ts b/packages/api/src/code/workspace.spec.ts index fb48d989333..35fe63292c4 100644 --- a/packages/api/src/code/workspace.spec.ts +++ b/packages/api/src/code/workspace.spec.ts @@ -2176,3 +2176,72 @@ describe('executeWorkspaceTool', () => { expect(fetchImpl).toHaveBeenCalledTimes(2); }); }); + +describe('linked worktree lanes', () => { + const listing = (paths: string[]) => + new Response( + JSON.stringify({ + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'librechat', + paths, + truncated: false, + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ); + const request: WorkspaceToolRequest = { + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'librechat', + path: '.worktrees/fix-a/src', + }; + + test('sends a worktree request and reports paths relative to the checkout', async () => { + const fetchImpl = jest.fn().mockResolvedValue(listing(['src/a.ts'])); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example/v1', + authHeaders: {}, + fetchImpl, + request, + linkedWorktrees: true, + }), + ).resolves.toMatchObject({ paths: ['.worktrees/fix-a/src/a.ts'] }); + expect(JSON.parse(fetchImpl.mock.calls[0][1].body)).toEqual({ + protocolVersion: 1, + operation: 'list_files', + workspaceId: 'librechat', + path: 'src', + worktree: 'fix-a', + }); + }); + + test('keeps the checkout-scoped request when the worker has no lanes', async () => { + const fetchImpl = jest.fn().mockResolvedValue(listing(['.worktrees/fix-a/src/a.ts'])); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example/v1', + authHeaders: {}, + fetchImpl, + request, + }), + ).resolves.toMatchObject({ paths: ['.worktrees/fix-a/src/a.ts'] }); + expect(JSON.parse(fetchImpl.mock.calls[0][1].body)).toEqual(request); + }); + + test('rejects lane results that escape the requested worktree scope', async () => { + const fetchImpl = jest.fn().mockResolvedValue(listing(['lib/a.ts'])); + + await expect( + executeWorkspaceTool({ + baseURL: 'https://code.example/v1', + authHeaders: {}, + fetchImpl, + request, + linkedWorktrees: true, + }), + ).rejects.toMatchObject({ reason: 'invalid' }); + }); +}); diff --git a/packages/api/src/code/workspace.ts b/packages/api/src/code/workspace.ts index c953dffbfa9..80ab78a4dc1 100644 --- a/packages/api/src/code/workspace.ts +++ b/packages/api/src/code/workspace.ts @@ -858,6 +858,142 @@ function getWorkspaceAuthHeaders( }); } +/** Linked worktrees live at `.worktrees/` beneath a registered checkout. */ +const LINKED_WORKTREE_DIRECTORY = '.worktrees/'; +/** Mirrors Code API's single-segment worktree name rule. */ +const LINKED_WORKTREE_NAME_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$/; + +export type LinkedWorktreeRequest = WorkspaceToolRequest & { worktree: string }; + +interface LinkedWorktreePath { + worktree: string; + /** Path relative to the worktree root; empty for the worktree root itself. */ + rest: string; +} + +function splitLinkedWorktreePath(path: string | undefined): LinkedWorktreePath | undefined { + if (path == null || !path.startsWith(LINKED_WORKTREE_DIRECTORY)) return undefined; + const remainder = path.slice(LINKED_WORKTREE_DIRECTORY.length); + const slash = remainder.indexOf('/'); + const worktree = slash === -1 ? remainder : remainder.slice(0, slash); + if (!LINKED_WORKTREE_NAME_PATTERN.test(worktree) || worktree.endsWith('.lock')) { + return undefined; + } + return { worktree, rest: slash === -1 ? '' : remainder.slice(slash + 1) }; +} + +function prefixed(worktree: string, path: string): string { + return `${LINKED_WORKTREE_DIRECTORY}${worktree}/${path}`; +} + +/** + * Route a request that targets `.worktrees//…` into that worktree's own + * scheduling lane, so work in sibling worktrees runs concurrently. Only + * structural targets are routed: a file path, a search or listing scope, or a + * command `cwd`. Anything else, including a command that `cd`s into a worktree + * from the root, stays root-scoped because its reach cannot be bounded. + */ +export function toLinkedWorktreeRequest( + request: WorkspaceToolRequest, +): { request: LinkedWorktreeRequest; worktree: string } | undefined { + if (request.workspaceInstanceId != null) return undefined; + switch (request.operation) { + case 'read_file': + case 'write_file': + case 'edit_file': + case 'preview_edit': { + if (request.operation === 'read_file' && request.instructionSha256 != null) return undefined; + const target = splitLinkedWorktreePath(request.path); + if (target == null || target.rest === '') return undefined; + return { + worktree: target.worktree, + request: { ...request, path: target.rest, worktree: target.worktree }, + }; + } + case 'search_text': { + const target = splitLinkedWorktreePath(request.path); + if (target == null) return undefined; + const { path: _path, ...rest } = request; + return { + worktree: target.worktree, + request: { + ...rest, + ...(target.rest === '' ? {} : { path: target.rest }), + worktree: target.worktree, + }, + }; + } + case 'list_files': { + const target = splitLinkedWorktreePath(request.path); + if (target == null) return undefined; + const after = + request.afterPath == null ? undefined : splitLinkedWorktreePath(request.afterPath); + if ( + request.afterPath != null && + (after == null || after.worktree !== target.worktree || after.rest === '') + ) { + return undefined; + } + const { path: _path, afterPath: _afterPath, ...rest } = request; + return { + worktree: target.worktree, + request: { + ...rest, + ...(target.rest === '' ? {} : { path: target.rest }), + ...(after == null ? {} : { afterPath: after.rest }), + worktree: target.worktree, + }, + }; + } + case 'execute_command': { + if (request.environmentAction != null) return undefined; + const target = splitLinkedWorktreePath(request.cwd); + if (target == null) return undefined; + const { cwd: _cwd, ...rest } = request; + return { + worktree: target.worktree, + request: { + ...rest, + ...(target.rest === '' ? {} : { cwd: target.rest }), + worktree: target.worktree, + }, + }; + } + } +} + +/** Restore the `.worktrees//` prefix on every path a lane result reports. */ +export function fromLinkedWorktreeResult( + result: WorkspaceToolResult, + worktree: string, +): WorkspaceToolResult { + switch (result.operation) { + case 'read_file': + case 'write_file': + case 'edit_file': + case 'preview_edit': + return { ...result, path: prefixed(worktree, result.path) }; + case 'search_text': + return { + ...result, + matches: result.matches.map((match) => ({ + ...match, + path: prefixed(worktree, match.path), + })), + }; + case 'list_files': + return { + ...result, + paths: result.paths.map((path) => prefixed(worktree, path)), + ...(result.nextAfterPath == null + ? {} + : { nextAfterPath: prefixed(worktree, result.nextAfterPath) }), + }; + case 'execute_command': + return result; + } +} + export async function executeWorkspaceTool({ baseURL, authHeaders, @@ -868,6 +1004,7 @@ export async function executeWorkspaceTool({ codeApiMaxRetryWaitMs = CODE_API_RATE_LIMIT_WAIT_DEFAULT_MS, maxRequestTimeoutMs, deadlineAtMs, + linkedWorktrees = false, }: { baseURL: string; authHeaders: WorkspaceToolAuthHeaders; @@ -881,6 +1018,8 @@ export async function executeWorkspaceTool({ maxRequestTimeoutMs?: number; /** Optional earlier caller deadline; a signal alone has no remaining-time value. */ deadlineAtMs?: number; + /** The worker runs each `.worktrees/` in its own lane; route matching requests there. */ + linkedWorktrees?: boolean; }): Promise { if ( !isValidRequest(request) || @@ -898,16 +1037,18 @@ export async function executeWorkspaceTool({ ) { throw new WorkspaceToolHttpError('invalid'); } - const executionBudgetMs = getWorkspaceExecutionBudgetMs(request); + const lane = linkedWorktrees === true ? toLinkedWorktreeRequest(request) : undefined; + const wireRequest: WorkspaceToolRequest = lane?.request ?? request; + const executionBudgetMs = getWorkspaceExecutionBudgetMs(wireRequest); const completionReserveMs = executionBudgetMs + WORKSPACE_COMMAND_TRANSPORT_GRACE_MS; - const perAttemptTimeoutMs = maxRequestTimeoutMs ?? getWorkspaceToolTimeoutMs(request); + const perAttemptTimeoutMs = maxRequestTimeoutMs ?? getWorkspaceToolTimeoutMs(wireRequest); const callerDeadlineAt = Math.min( deadlineAtMs ?? Infinity, maxRequestTimeoutMs == null ? Infinity : Date.now() + maxRequestTimeoutMs, ); const queueDeadlineAt = Date.now() + maxQueueWaitMs; const callerRetryDeadlineAt = callerDeadlineAt - completionReserveMs; - const body = JSON.stringify(request); + const body = JSON.stringify(wireRequest); let lastAdmissionRejection: WorkspaceToolHttpError | undefined; let lastRetryDeadlineAt = Infinity; let rateLimitWaitedMs = 0; @@ -1011,10 +1152,10 @@ export async function executeWorkspaceTool({ continue; } const result = await readBoundedJson(response, requestSignal); - if (!isValidResult(request, result)) { + if (!isValidResult(wireRequest, result)) { throw new WorkspaceToolHttpError('invalid'); } - return result; + return lane ? fromLinkedWorktreeResult(result, lane.worktree) : result; } catch (error) { if (error instanceof WorkspaceToolHttpError) throw error; if ( diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index a54ee838beb..749db957bed 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -4,6 +4,7 @@ import { Constants, ContentTypes, StepTypes, + StepEvents, UsageEvents, ApprovalEvents, SteerEvents, @@ -443,10 +444,44 @@ function getReplayStepId(event: t.ServerSentEvent): unknown { const result = 'result' in event.data ? event.data.result : undefined; return result != null && typeof result === 'object' && 'id' in result ? result.id : undefined; } + if (event.event === StepEvents.ON_TOOL_PREPARATION) { + const marker = event.data as { id?: unknown; toolCallId?: unknown; index?: unknown }; + return typeof marker.id === 'string' + ? `${marker.id}:${typeof marker.toolCallId === 'string' ? marker.toolCallId : `#${marker.index}`}` + : undefined; + } + if (event.event === StepEvents.ON_TOOL_CALLS_DISPATCHED) { + const calls = (event.data as { toolCalls?: unknown }).toolCalls; + return Array.isArray(calls) + ? JSON.stringify(calls.map((call) => [call?.stepId, call?.id])) + : undefined; + } return undefined; } +function isToolTimingReplayEvent(event: t.ServerSentEvent): boolean { + if (!('event' in event) || event.data == null || typeof event.data !== 'object') return false; + if (event.event === StepEvents.ON_TOOL_PREPARATION) { + const marker = event.data as { id?: unknown; observed_at?: unknown }; + return ( + typeof marker.id === 'string' && + marker.id !== '' && + typeof marker.observed_at === 'number' && + Number.isFinite(marker.observed_at) + ); + } + if (event.event === StepEvents.ON_TOOL_CALLS_DISPATCHED) { + const dispatch = event.data as { dispatched_at?: unknown; toolCalls?: unknown }; + return ( + typeof dispatch.dispatched_at === 'number' && + Number.isFinite(dispatch.dispatched_at) && + Array.isArray(dispatch.toolCalls) + ); + } + return false; +} + function isOAuthReplayEvent(event: t.ServerSentEvent): boolean { if (!('event' in event) || !event.data || typeof event.data !== 'object') { return false; @@ -7468,6 +7503,9 @@ class GenerationJobManagerClass { if (event.event === UsageEvents.ON_TOKEN_USAGE) { return this.trackTokenUsage(streamId, event, expectedCreatedAt); } + if (isToolTimingReplayEvent(event)) { + return this.trackReplayEvent(streamId, event, expectedCreatedAt); + } if ( (event.event === 'on_run_step' || event.event === 'on_run_step_delta' || @@ -7758,7 +7796,7 @@ class GenerationJobManagerClass { event: t.ServerSentEvent, expectedCreatedAt: number, ): Promise { - if (!isOAuthReplayEvent(event)) { + if (!isOAuthReplayEvent(event) && !isToolTimingReplayEvent(event)) { return; } diff --git a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts index 7e67fec1228..5456045ebbe 100644 --- a/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts +++ b/packages/api/src/stream/__tests__/GenerationJobManager.resumeReplay.spec.ts @@ -1,9 +1,11 @@ +import { StepEvents } from 'librechat-data-provider'; import type { StandardGraph } from '@librechat/agents'; import type { Agents } from 'librechat-data-provider'; import type { ServerSentEvent } from '~/types'; import { InMemoryEventTransport } from '~/stream/implementations/InMemoryEventTransport'; import { InMemoryJobStore } from '~/stream/implementations/InMemoryJobStore'; import { GenerationJobManagerClass } from '~/stream/GenerationJobManager'; +import { createToolTimingTracker } from '~/agents/toolTiming'; jest.spyOn(console, 'log').mockImplementation(); @@ -83,6 +85,46 @@ describe('GenerationJobManager resume replay events', () => { }); }); + test('replays bounded first-fragment and dispatch evidence without tool arguments', async () => { + manager = createInMemoryManager(); + const streamId = `tool-timing-resume-${Date.now()}`; + await manager.createJob(streamId, 'user-1', streamId); + const first = { + event: StepEvents.ON_TOOL_PREPARATION, + data: { + id: 'step-1', + index: 0, + observed_at: 100, + }, + }; + const named = { + event: StepEvents.ON_TOOL_PREPARATION, + data: { + id: 'step-1', + index: 0, + toolCallId: 'call-1', + observed_at: 200, + }, + }; + const dispatched = { + event: StepEvents.ON_TOOL_CALLS_DISPATCHED, + data: { + dispatched_at: 500, + toolCalls: [{ id: 'call-1', name: 'query', stepId: 'step-1' }], + }, + }; + for (const event of [first, named, dispatched, named]) await manager.emitChunk(streamId, event); + const state = await manager.getResumeState(streamId); + expect(state?.replayEvents).toEqual([first, named, dispatched]); + expect(JSON.stringify(state?.replayEvents)).not.toContain('args'); + const rebuilt = createToolTimingTracker(state?.replayEvents ?? []); + rebuilt.completed('step-1', 'call-1', 560); + expect(rebuilt.take('call-1', 'step-1')).toEqual({ + toolPreparationDurationMs: 400, + toolExecutionDurationMs: 60, + }); + }); + test('includes OAuth run step and delta replay events in resume state', async () => { manager = createInMemoryManager(); const streamId = `oauth-delta-resume-${Date.now()}`; diff --git a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts index df15322b643..30068cea15e 100644 --- a/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts +++ b/packages/api/src/stream/__tests__/RedisJobStore.stream_integration.spec.ts @@ -1971,6 +1971,95 @@ describe('RedisJobStore Integration Tests', () => { await store.destroy(); }); + test('reconstructs split tool timings and close metadata across instances', async () => { + if (!ioredisClient) return; + const { RedisJobStore } = await import('../implementations/RedisJobStore'); + const producer = new RedisJobStore(ioredisClient); + const consumer = new RedisJobStore(ioredisClient); + await producer.initialize(); + await consumer.initialize(); + const streamId = `tool-timing-recon-${Date.now()}`; + const job = await producer.createJob(streamId, 'user-1', streamId); + const events = [ + { + event: 'on_run_step', + data: { + id: 'step-1', + runId: 'response-1', + index: 0, + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-1', name: 'lookup', args: '{}' }], + }, + }, + }, + { + event: 'on_run_step_delta', + data: { + id: 'step-1', + observed_at: 1_000, + delta: { type: 'tool_calls', tool_calls: [{ id: 'call-1', index: 0, args: '{' }] }, + }, + }, + { + event: 'on_tool_preparation', + data: { + id: 'step-1', + toolCallId: 'call-1', + index: 0, + observed_at: 1_000, + }, + }, + { + event: 'on_tool_calls_dispatched', + data: { + dispatched_at: 248_000, + toolCalls: [{ id: 'call-1', name: 'lookup', stepId: 'step-1' }], + }, + }, + { + event: 'on_run_step_completed', + data: { + result: { + id: 'step-1', + index: 0, + type: 'tool_call', + completed_at: 248_340, + tool_call: { id: 'call-1', name: 'lookup', args: '{}', output: 'done', progress: 1 }, + }, + }, + }, + { + event: 'on_run_step_closed', + data: { + id: 'step-1', + type: 'tool_calls', + index: 0, + status: 'completed', + created_at: 1_000, + closed_at: 248_340, + }, + }, + ]; + for (const event of events) { + await producer.appendChunk(streamId, event); + } + const result = await consumer.getContentParts(streamId, job.createdAt, { durableOnly: true }); + expect(result?.content[0]).toMatchObject({ + type: 'tool_call', + tool_call: { + id: 'call-1', + runStepStatus: 'completed', + runStepClosedAt: 248_340, + runStepDurationMs: 247_340, + toolPreparationDurationMs: 247_000, + toolExecutionDurationMs: 340, + }, + }); + await producer.destroy(); + await consumer.destroy(); + }); + test('should share run steps between instances', async () => { if (!ioredisClient) { return; diff --git a/packages/api/src/stream/__tests__/RedisJobStore.toolTiming.spec.ts b/packages/api/src/stream/__tests__/RedisJobStore.toolTiming.spec.ts new file mode 100644 index 00000000000..d867c190fca --- /dev/null +++ b/packages/api/src/stream/__tests__/RedisJobStore.toolTiming.spec.ts @@ -0,0 +1,90 @@ +import { RedisJobStore } from '../implementations/RedisJobStore'; + +/** Exercise the actual reconstruction method against a durable chunk snapshot, not a live host cache. */ +describe('RedisJobStore tool timing reconstruction', () => { + it('restores preparation, SDK handoff, result time, and the upstream close stamp', async () => { + const chunks = [ + { + event: 'on_run_step', + data: { + id: 'step-1', + index: 0, + runId: 'response-1', + stepDetails: { + type: 'tool_calls', + tool_calls: [{ id: 'call-1', name: 'lookup', args: '{}' }], + }, + }, + }, + { + event: 'on_run_step_delta', + data: { + id: 'step-1', + observed_at: 1_000, + delta: { type: 'tool_calls', tool_calls: [{ id: 'call-1', index: 0, args: '{' }] }, + }, + }, + { + event: 'on_tool_preparation', + data: { + id: 'step-1', + toolCallId: 'call-1', + index: 0, + observed_at: 1_000, + }, + }, + { + event: 'on_tool_calls_dispatched', + data: { + dispatched_at: 248_000, + toolCalls: [{ id: 'call-1', name: 'lookup', stepId: 'step-1' }], + }, + }, + { + event: 'on_run_step_completed', + data: { + result: { + id: 'step-1', + index: 0, + type: 'tool_call', + completed_at: 248_340, + tool_call: { id: 'call-1', name: 'lookup', args: '{}', output: 'done', progress: 1 }, + }, + }, + }, + { + event: 'on_run_step_closed', + data: { + id: 'step-1', + type: 'tool_calls', + index: 0, + status: 'completed', + created_at: 1_000, + closed_at: 248_340, + }, + }, + ]; + const store = Object.create(RedisJobStore.prototype) as RedisJobStore; + const snapshot = jest.fn(async () => ({ chunks, durableEventCount: chunks.length })); + Object.defineProperty(store, 'getChunkSnapshot', { value: snapshot }); + const result = await store.getContentParts('tool-run', undefined, { durableOnly: true }); + expect(snapshot).toHaveBeenCalledWith('tool-run', undefined, true); + expect(result).toMatchObject({ + reconstructedEventCount: chunks.length, + durableEventCount: chunks.length, + content: [ + { + type: 'tool_call', + tool_call: { + id: 'call-1', + runStepStatus: 'completed', + runStepClosedAt: 248_340, + runStepDurationMs: 247_340, + toolPreparationDurationMs: 247_000, + toolExecutionDurationMs: 340, + }, + }, + ], + }); + }); +}); diff --git a/packages/api/src/stream/implementations/RedisJobStore.ts b/packages/api/src/stream/implementations/RedisJobStore.ts index c616170e577..4d1f74db770 100644 --- a/packages/api/src/stream/implementations/RedisJobStore.ts +++ b/packages/api/src/stream/implementations/RedisJobStore.ts @@ -2,6 +2,7 @@ import { logger } from '@librechat/data-schemas'; import { createContentAggregator } from '@librechat/agents'; import { ContentTypes, + StepEvents, getRunStepDurationMs, getRunStepCloseMetadata, } from 'librechat-data-provider'; @@ -53,6 +54,7 @@ import { import { instrumentIORedisClient, RedisUseCases } from '~/cache/redisTelemetry'; import { RecoveredSteerPayloadMismatchError } from '~/stream/SteerRecovery'; import { createCheckpointNamespace } from '~/stream/checkpoints'; +import { createToolTimingTracker } from '~/agents/toolTiming'; import { evalScript } from '~/cache/redisScript'; const CLIENT_REQUEST_ID_PATTERN = /^[A-Za-z0-9:_-]{1,128}$/; @@ -4120,6 +4122,7 @@ export class RedisJobStore implements IJobStoreV2 { // Use the same content aggregator as live streaming const { contentParts, aggregateContent } = createContentAggregator(); + const toolTiming = createToolTimingTracker(); // Step ID -> content index, rebuilt from the replayed `on_run_step` // payloads. Those carry the index the offset wrappers shifted, whereas a @@ -4233,6 +4236,24 @@ export class RedisJobStore implements IJobStoreV2 { continue; } + if (event.event === StepEvents.ON_TOOL_PREPARATION) { + toolTiming.prepare(event.data as Agents.ToolPreparationMarker); + continue; + } + if (event.event === StepEvents.ON_TOOL_CALLS_DISPATCHED) { + toolTiming.dispatched(event.data as Agents.ToolCallsDispatchedEvent); + continue; + } + if (event.event === StepEvents.ON_RUN_STEP_DELTA) { + toolTiming.observe(event.data as Agents.RunStepDeltaEvent); + } + if (event.event === StepEvents.ON_RUN_STEP_COMPLETED) { + const completion = (event.data as { result?: Agents.ToolEndEvent }).result; + if (completion?.tool_call?.id) { + toolTiming.completed(completion.id, completion.tool_call.id, completion.completed_at); + } + } + // Step closures are host-authored like steers and labels: the SDK // aggregator has no notion of the event, so the terminal status is // stamped onto the part the replayed steps already rebuilt. Resolved by @@ -4250,6 +4271,7 @@ export class RedisJobStore implements IJobStoreV2 { const part = index != null ? contentParts[index] : undefined; if (closed.status && part?.type === ContentTypes.TOOL_CALL && part.tool_call) { part.tool_call.runStepStatus = closed.status; + Object.assign(part.tool_call, toolTiming.take(part.tool_call.id ?? '', closed.id ?? '')); Object.assign(part.tool_call, getRunStepCloseMetadata(closed)); const durationMs = getRunStepDurationMs(closed); if (durationMs != null) { diff --git a/packages/data-provider/src/agentToolOptions.spec.ts b/packages/data-provider/src/agentToolOptions.spec.ts index 70ac60238f9..c6ae177c62b 100644 --- a/packages/data-provider/src/agentToolOptions.spec.ts +++ b/packages/data-provider/src/agentToolOptions.spec.ts @@ -1,5 +1,12 @@ import type { AgentToolOptions } from './types/tools'; -import { normalizeActionToolName, removeCodeExecutionCaller } from './agentToolOptions'; +import { + mcpServerToggleKey, + getAgentToolSwitches, + applyAgentToolSwitches, + pickUserToggleOptions, + normalizeActionToolName, + removeCodeExecutionCaller, +} from './agentToolOptions'; describe('normalizeActionToolName', () => { it('normalizes only the encoded action domain', () => { @@ -51,3 +58,179 @@ describe('removeCodeExecutionCaller', () => { expect(input.search.allowed_callers).toEqual(['code_execution']); }); }); + +describe('mcpServerToggleKey', () => { + it('builds the server-level placeholder key', () => { + expect(mcpServerToggleKey('alpha')).toBe('sys__server__sys_mcp_alpha'); + }); +}); + +describe('getAgentToolSwitches', () => { + it('reports only attached tools that carry a user_toggle', () => { + expect( + getAgentToolSwitches({ + tools: ['web_search', 'execute_code', 'search_mcp_alpha', 'sys__all__sys_mcp_beta'], + tool_options: { + web_search: { user_toggle: 'on' }, + execute_code: { user_toggle: 'off' }, + file_search: { user_toggle: 'on' }, + [mcpServerToggleKey('alpha')]: { user_toggle: 'off' }, + [mcpServerToggleKey('beta')]: { defer_loading: true }, + [mcpServerToggleKey('absent')]: { user_toggle: 'on' }, + }, + }), + ).toEqual({ + builtins: { web_search: true, execute_code: false }, + mcp: { alpha: false }, + }); + }); + + it('is empty without tool_options', () => { + expect(getAgentToolSwitches({ tools: ['web_search'] })).toEqual({ builtins: {}, mcp: {} }); + }); + + it('keeps a server whose name contains the MCP delimiter', () => { + expect( + getAgentToolSwitches({ + tools: ['search_mcp_Google_mcp_Workspace'], + tool_options: { [mcpServerToggleKey('Google_mcp_Workspace')]: { user_toggle: 'on' } }, + }), + ).toEqual({ builtins: {}, mcp: { Google_mcp_Workspace: true } }); + }); + + it('keys a server by its configured name when its tool keys use the normalized name', () => { + expect( + getAgentToolSwitches({ + tools: ['search_mcp_My_Docs'], + tool_options: { [mcpServerToggleKey('My Docs')]: { user_toggle: 'off' } }, + }), + ).toEqual({ builtins: {}, mcp: { 'My Docs': false } }); + }); +}); + +describe('applyAgentToolSwitches', () => { + const agent = { + tools: [ + 'web_search', + 'execute_code', + 'file_search', + 'a_mcp_alpha', + 'b_mcp_alpha', + 'c_mcp_beta', + 'plain', + ], + tool_options: { + web_search: { user_toggle: 'on' as const }, + execute_code: { user_toggle: 'off' as const }, + [mcpServerToggleKey('alpha')]: { user_toggle: 'on' as const }, + }, + }; + + it('uses the defaults without a request', () => { + expect(applyAgentToolSwitches(agent, null)).toEqual({ + tools: ['web_search', 'file_search', 'a_mcp_alpha', 'b_mcp_alpha', 'c_mcp_beta', 'plain'], + mcp: ['alpha', 'beta'], + }); + }); + + it('honours requested booleans and server lists, dropping every tool of an off server', () => { + expect( + applyAgentToolSwitches(agent, { web_search: false, execute_code: true, mcp: [] }), + ).toEqual({ + tools: ['execute_code', 'file_search', 'c_mcp_beta', 'plain'], + mcp: ['beta'], + }); + }); + + it('ignores non-boolean values and a non-array mcp', () => { + const requested = { web_search: 'no', mcp: 'alpha' } as unknown as Parameters< + typeof applyAgentToolSwitches + >[1]; + expect(applyAgentToolSwitches(agent, requested).tools).toEqual([ + 'web_search', + 'file_search', + 'a_mcp_alpha', + 'b_mcp_alpha', + 'c_mcp_beta', + 'plain', + ]); + }); + + it('never adds a locked tool or an unattached one', () => { + const result = applyAgentToolSwitches( + { tools: ['file_search'], tool_options: {} }, + { web_search: true, execute_code: true, mcp: ['alpha'] }, + ); + expect(result).toEqual({ tools: ['file_search'], mcp: [] }); + }); + + it('tells apart a server named bar from one named foo_mcp_bar', () => { + const tool_options = { + [mcpServerToggleKey('bar')]: { user_toggle: 'on' as const }, + [mcpServerToggleKey('foo_mcp_bar')]: { user_toggle: 'on' as const }, + }; + const tools = ['a_mcp_bar', 'b_mcp_foo_mcp_bar']; + expect(applyAgentToolSwitches({ tools, tool_options }, { mcp: ['bar'] }).tools).toEqual([ + 'a_mcp_bar', + ]); + expect(applyAgentToolSwitches({ tools, tool_options }, { mcp: ['foo_mcp_bar'] }).tools).toEqual( + ['b_mcp_foo_mcp_bar'], + ); + }); + + it('matches the chat list by configured name for a normalized server', () => { + const agent = { + tools: ['search_mcp_My_Docs'], + tool_options: { [mcpServerToggleKey('My Docs')]: { user_toggle: 'off' as const } }, + }; + expect(applyAgentToolSwitches(agent, { mcp: [] }).tools).toEqual([]); + expect(applyAgentToolSwitches(agent, { mcp: ['My Docs'] }).tools).toEqual([ + 'search_mcp_My_Docs', + ]); + }); + + it('drops the server placeholder and wildcard tokens with the server', () => { + const result = applyAgentToolSwitches( + { + tools: ['sys__server__sys_mcp_alpha', 'sys__all__sys_mcp_alpha'], + tool_options: { [mcpServerToggleKey('alpha')]: { user_toggle: 'on' } }, + }, + { mcp: [] }, + ); + expect(result.tools).toEqual([]); + }); + + it('drops raw-name placeholder and wildcard tokens of a normalized server', () => { + const agent = { + tools: ['sys__server__sys_mcp_My Docs', 'sys__all__sys_mcp_My Docs', 'search_mcp_My_Docs'], + tool_options: { [mcpServerToggleKey('My Docs')]: { user_toggle: 'on' as const } }, + }; + expect(applyAgentToolSwitches(agent, { mcp: [] }).tools).toEqual([]); + expect( + applyAgentToolSwitches({ ...agent, tools: ['sys__all__sys_mcp_My Docs'] }, { mcp: [] }).tools, + ).toEqual([]); + expect(applyAgentToolSwitches(agent, { mcp: ['My Docs'] })).toEqual({ + tools: agent.tools, + mcp: ['My_Docs'], + }); + }); +}); + +describe('pickUserToggleOptions', () => { + it('keeps only user_toggle entries and nothing else from the options', () => { + expect( + pickUserToggleOptions({ + web_search: { user_toggle: 'off', defer_loading: true }, + search_mcp_docs: { allowed_callers: ['direct'] }, + [mcpServerToggleKey('docs')]: { user_toggle: 'on' }, + }), + ).toEqual({ + web_search: { user_toggle: 'off' }, + [mcpServerToggleKey('docs')]: { user_toggle: 'on' }, + }); + }); + + it('returns undefined without options', () => { + expect(pickUserToggleOptions(undefined)).toBeUndefined(); + }); +}); diff --git a/packages/data-provider/src/agentToolOptions.ts b/packages/data-provider/src/agentToolOptions.ts index 8de99fd47bb..759ff9a3c36 100644 --- a/packages/data-provider/src/agentToolOptions.ts +++ b/packages/data-provider/src/agentToolOptions.ts @@ -1,10 +1,14 @@ +import type { TEphemeralAgent } from './types'; +import type { Agent } from './types/agents'; import { + Tools, actionDelimiter, actionDomainSeparator, isActionTool, type AgentToolOptions, type AllowedCaller, } from './types/tools'; +import { Constants, splitMCPToolKey, normalizeMCPToolKey, normalizeServerName } from './config'; const actionDomainSeparatorRegex = new RegExp(actionDomainSeparator, 'g'); @@ -56,3 +60,127 @@ export function removeCodeExecutionCaller( return normalized; } + +export const switchableBuiltinTools = [ + Tools.web_search, + Tools.execute_code, + Tools.file_search, +] as const; + +export type SwitchableBuiltinTool = (typeof switchableBuiltinTools)[number]; + +/** `tool_options` key for a whole MCP server (`sys__server__sys_mcp_`). */ +export function mcpServerToggleKey(serverName: string): string { + return `${Constants.mcp_server}${Constants.mcp_delimiter}${serverName}`; +} + +export interface AgentToolSwitches { + /** Switchable built-ins attached to the agent; value = starts on. */ + builtins: Partial>; + /** Switchable attached MCP servers by name; value = starts on. */ + mcp: Record; +} + +type SwitchableAgent = Pick; + +const mcpServerTogglePrefix = mcpServerToggleKey(''); + +/** Configured names of the servers carrying a `user_toggle`, attached or not. */ +export function getMCPSwitchServerNames(toolOptions: AgentToolOptions | undefined): string[] { + return Object.entries(toolOptions ?? {}) + .filter( + ([key, options]) => key.startsWith(mcpServerTogglePrefix) && options.user_toggle != null, + ) + .map(([key]) => key.slice(mcpServerTogglePrefix.length)); +} + +/** Only the `user_toggle` of each entry, so a view-only reader learns nothing else. */ +export function pickUserToggleOptions( + toolOptions: AgentToolOptions | undefined, +): AgentToolOptions | undefined { + if (toolOptions == null) { + return undefined; + } + const picked: AgentToolOptions = {}; + for (const [key, options] of Object.entries(toolOptions)) { + if (options.user_toggle != null) { + picked[key] = { user_toggle: options.user_toggle }; + } + } + return picked; +} + +/** Tool keys carry the normalized server name, which may itself contain the MCP + * delimiter; the switch servers' key names disambiguate the split. Placeholder and + * wildcard tokens may still carry the raw configured name, so that is normalized first. */ +function toolServerName( + tool: string, + serverNames: string[], + keyServerNames: string[], +): string | undefined { + return splitMCPToolKey(normalizeMCPToolKey(tool, serverNames), keyServerNames)[1]; +} + +/** The creator's switchable set and defaults. Locked or unattached tools never appear. */ +export function getAgentToolSwitches(agent: SwitchableAgent): AgentToolSwitches { + const tools = agent.tools ?? []; + const options = agent.tool_options ?? {}; + const switches: AgentToolSwitches = { builtins: {}, mcp: {} }; + for (const tool of switchableBuiltinTools) { + const toggle = options[tool]?.user_toggle; + if (toggle != null && tools.includes(tool)) { + switches.builtins[tool] = toggle === 'on'; + } + } + const serverNames = getMCPSwitchServerNames(options); + const keyServerNames = serverNames.map(normalizeServerName); + const attached = new Set(tools.map((tool) => toolServerName(tool, serverNames, keyServerNames))); + serverNames.forEach((serverName, index) => { + if (attached.has(keyServerNames[index])) { + switches.mcp[serverName] = options[mcpServerToggleKey(serverName)]?.user_toggle === 'on'; + } + }); + return switches; +} + +/** The agent's tools after applying the chat's untrusted switch state. Only ever removes. */ +export function applyAgentToolSwitches( + agent: SwitchableAgent, + requested: TEphemeralAgent | null | undefined, +): { tools: string[]; mcp: string[] } { + const tools = agent.tools ?? []; + const { builtins, mcp } = getAgentToolSwitches(agent); + const dropped = new Set(); + for (const tool of switchableBuiltinTools) { + const requestedValue = requested?.[tool]; + const defaultValue = builtins[tool]; + const isOn = typeof requestedValue === 'boolean' ? requestedValue : defaultValue; + if (defaultValue != null && !isOn) { + dropped.add(tool); + } + } + const requestedServers = Array.isArray(requested?.mcp) ? requested.mcp : undefined; + const serverNames = getMCPSwitchServerNames(agent.tool_options); + const keyServerNames = serverNames.map(normalizeServerName); + const offServers = new Set( + Object.entries(mcp) + .filter( + ([name, isDefaultOn]) => + !(requestedServers ? requestedServers.includes(name) : isDefaultOn), + ) + .map(([name]) => normalizeServerName(name)), + ); + const kept: string[] = []; + const keptServers = new Set(); + for (const tool of tools) { + const serverName = toolServerName(tool, serverNames, keyServerNames); + if (dropped.has(tool) || (serverName != null && offServers.has(serverName))) { + continue; + } + kept.push(tool); + if (serverName != null) { + keptServers.add(serverName); + } + } + return { tools: kept, mcp: Array.from(keptServers) }; +} diff --git a/packages/data-provider/src/code/workspace.ts b/packages/data-provider/src/code/workspace.ts index 33e0592fe4b..5afb6567ab0 100644 --- a/packages/data-provider/src/code/workspace.ts +++ b/packages/data-provider/src/code/workspace.ts @@ -25,6 +25,8 @@ export const CODE_WORKSPACE_OPERATIONS = [ 'execute_command', ] as const; export const CODE_WORKSPACE_INSTANCE_TYPES = ['git_worktree'] as const; +/** Scheduling scopes a worker can admit beneath one registered root. */ +export const CODE_WORKSPACE_SCOPES = ['git_linked_worktree'] as const; export const CODE_WORKSPACE_SELECTION_ERROR_REASONS = [ 'required', 'invalid', @@ -37,6 +39,7 @@ export const CODE_ENVIRONMENT_MODES = ['attached', 'without_attached'] as const; export type CodeWorkspaceOperation = (typeof CODE_WORKSPACE_OPERATIONS)[number]; export type CodeWorkspaceInstanceType = (typeof CODE_WORKSPACE_INSTANCE_TYPES)[number]; +export type CodeWorkspaceScope = (typeof CODE_WORKSPACE_SCOPES)[number]; export type CodeWorkspaceSelectionErrorReason = (typeof CODE_WORKSPACE_SELECTION_ERROR_REASONS)[number]; export type CodeEnvironmentMode = (typeof CODE_ENVIRONMENT_MODES)[number]; @@ -50,6 +53,8 @@ export interface CodeWorkspaceDescriptor { operations?: CodeWorkspaceOperation[]; /** Optional worker-managed isolation modes available beneath this root. */ workspaceInstances?: CodeWorkspaceInstanceType[]; + /** `git_linked_worktree`: each `.worktrees/` runs in its own scheduling lane. */ + workspaceScopes?: CodeWorkspaceScope[]; environment?: { fingerprint: string; repo?: string; diff --git a/packages/data-provider/src/config.spec.ts b/packages/data-provider/src/config.spec.ts index 6ef043dbe16..c4e874f424b 100644 --- a/packages/data-provider/src/config.spec.ts +++ b/packages/data-provider/src/config.spec.ts @@ -736,6 +736,35 @@ describe('attached code environment user config schema', () => { }); }); + it.each([ + [{ linkedWorktrees: true }, true], + [{ linkedWorktrees: 'yes' }, false], + [{ linkedWorktrees: true, subdirectories: true }, false], + ])('validates the linked worktree lane toggle %p', (workspaces, valid) => { + const result = configSchema.safeParse({ + version: '1.0', + endpoints: { + agents: { + statefulCodeSessions: { + allowedEnvironments: ['user'], + environments: [ + { + id: 'personal-vm', + name: 'Personal VM', + type: 'attached', + baseURL: 'https://code.example.com/v1', + default: true, + configSchema: { workspaces }, + }, + ], + }, + }, + }, + }); + + expect(result.success).toBe(valid); + }); + it('rejects an attached command timeout above the protocol hard cap', () => { const result = configSchema.safeParse({ version: '1.0', diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 815ec69d87e..b7ecf82b33f 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1290,6 +1290,15 @@ export const codeEnvironmentUserConfigSchema = z }); }) .optional(), + workspaces: z + .object({ + /** Run requests aimed at `.worktrees/` in that worktree's own lane when the + * worker advertises linked-worktree lanes. Omission keeps every request scoped to + * its checkout. */ + linkedWorktrees: z.boolean().optional(), + }) + .strict() + .optional(), }) .strict(); diff --git a/packages/data-provider/src/index.ts b/packages/data-provider/src/index.ts index 92cb7a6cc86..e8e8e3f3d3a 100644 --- a/packages/data-provider/src/index.ts +++ b/packages/data-provider/src/index.ts @@ -13,6 +13,7 @@ export * from './messages'; export * from './errors'; /* run steps */ export * from './runSteps'; +export * from './toolTiming'; /* artifacts */ export * from './artifacts'; /* schema helpers */ diff --git a/packages/data-provider/src/toolTiming.spec.ts b/packages/data-provider/src/toolTiming.spec.ts new file mode 100644 index 00000000000..23cdb261815 --- /dev/null +++ b/packages/data-provider/src/toolTiming.spec.ts @@ -0,0 +1,31 @@ +import { getToolTimingDurations } from './toolTiming'; + +describe('getToolTimingDurations', () => { + it('separates minutes of argument preparation from a fast tool result', () => { + expect( + getToolTimingDurations({ observedAt: 1_000, dispatchedAt: 248_000, completedAt: 248_340 }), + ).toEqual({ toolPreparationDurationMs: 247_000, toolExecutionDurationMs: 340 }); + }); + + it('omits unknown phases independently rather than inventing a start', () => { + expect(getToolTimingDurations({ dispatchedAt: 10_000, completedAt: 11_000 })).toEqual({ + toolExecutionDurationMs: 1_000, + }); + expect(getToolTimingDurations({ observedAt: 1_000, completedAt: 11_000 })).toEqual({}); + expect(getToolTimingDurations({ observedAt: 1_000, dispatchedAt: 10_000 })).toEqual({ + toolPreparationDurationMs: 9_000, + }); + }); + + it('does not mislabel a reordered or non-finite clock pair', () => { + expect( + getToolTimingDurations({ observedAt: 12_000, dispatchedAt: 10_000, completedAt: 11_000 }), + ).toEqual({ toolExecutionDurationMs: 1_000 }); + expect( + getToolTimingDurations({ observedAt: 1_000, dispatchedAt: 10_000, completedAt: 9_000 }), + ).toEqual({ toolPreparationDurationMs: 9_000 }); + expect( + getToolTimingDurations({ observedAt: NaN, dispatchedAt: Infinity, completedAt: 12_000 }), + ).toEqual({}); + }); +}); diff --git a/packages/data-provider/src/toolTiming.ts b/packages/data-provider/src/toolTiming.ts new file mode 100644 index 00000000000..3a42cd68fac --- /dev/null +++ b/packages/data-provider/src/toolTiming.ts @@ -0,0 +1,29 @@ +/** Wall-clock phases of one tool call, never batch or database execution time. */ +export type ToolTimingStamps = { + observedAt?: number; + dispatchedAt?: number; + completedAt?: number; +}; + +const validTime = (value: number | undefined): value is number => + typeof value === 'number' && Number.isFinite(value) && value >= 0; + +/** Omit an interval rather than presenting clock skew or a missing stamp as measured time. */ +export function getToolTimingDurations({ + observedAt, + dispatchedAt, + completedAt, +}: ToolTimingStamps): { + toolPreparationDurationMs?: number; + toolExecutionDurationMs?: number; +} { + if (!validTime(dispatchedAt)) return {}; + return { + ...(validTime(observedAt) && observedAt <= dispatchedAt + ? { toolPreparationDurationMs: dispatchedAt - observedAt } + : {}), + ...(validTime(completedAt) && dispatchedAt <= completedAt + ? { toolExecutionDurationMs: completedAt - dispatchedAt } + : {}), + }; +} diff --git a/packages/data-provider/src/types/agents.ts b/packages/data-provider/src/types/agents.ts index 98f9591cb6e..aa88059fb0b 100644 --- a/packages/data-provider/src/types/agents.ts +++ b/packages/data-provider/src/types/agents.ts @@ -137,6 +137,7 @@ export namespace Agents { export type ToolEndEvent = { /** The Step Id of the Tool Call */ id: string; + completed_at?: number; /** The Completed Tool Call */ tool_call?: ToolCall; /** The content index of the tool call */ @@ -258,6 +259,19 @@ export namespace Agents { * because the caller aborted — which is the only signal that distinguishes * a stopped step from one still in flight. */ + export type ToolPreparationMarker = { + id: string; + index?: number; + toolCallId?: string; + observed_at: number; + }; + + export type ToolCallsDispatchedEvent = { + dispatched_at: number; + runId?: string; + toolCalls: Array<{ id: string; name: string; stepId?: string }>; + }; + export type RunStepClosedEvent = { id: string; index: number; @@ -361,6 +375,8 @@ export namespace Agents { * The delta containing the fields that have changed on the run step. */ delta: ToolCallDelta; + /** SDK receipt timestamp before awaited step dispatch, when available. */ + observed_at?: number; } export type StepDetails = MessageCreationDetails | ToolCallsDetails; export type MessageCreationDetails = { diff --git a/packages/data-provider/src/types/content.ts b/packages/data-provider/src/types/content.ts index d751a15afd2..e0d3cce1bd5 100644 --- a/packages/data-provider/src/types/content.ts +++ b/packages/data-provider/src/types/content.ts @@ -162,6 +162,14 @@ export type PartMetadata = { runStepDurationMs?: number; /** Host-reported close time in epoch milliseconds, when valid. Absent on older content. */ runStepClosedAt?: number; + /** First observed argument fragment, while this tool call is still being prepared. */ + toolPreparationStartedAt?: number; + /** SDK handoff to direct invocation or host dispatch, not the MCP round trip. */ + toolDispatchedAt?: number; + /** Time from the first observed argument fragment to SDK handoff, if both are known. */ + toolPreparationDurationMs?: number; + /** Time from SDK handoff to this call's result, not database execution time. */ + toolExecutionDurationMs?: number; /** * Stamped by the background harvester when a detached task's final output * replaces the dispatch handle in `tool_call.output`. The handle JSON and diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index 45a5b2ac829..8cebf418bc2 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -38,6 +38,10 @@ export enum StepEvents { ON_MESSAGE_DELTA = 'on_message_delta', ON_REASONING_DELTA = 'on_reasoning_delta', ON_RUN_STEP_DELTA = 'on_run_step_delta', + /** SDK handoff from generated arguments to host or direct tool execution. */ + ON_TOOL_CALLS_DISPATCHED = 'on_tool_calls_dispatched', + /** Compact first-fragment marker, retained for resumable generations. */ + ON_TOOL_PREPARATION = 'on_tool_preparation', ON_RUN_STEP_COMPLETED = 'on_run_step_completed', /** Terminal signal for a run step: closed with a status and timestamps. */ ON_RUN_STEP_CLOSED = 'on_run_step_closed', @@ -539,6 +543,8 @@ export type SubagentUpdatePhase = | 'start' | 'run_step' | 'run_step_delta' + | 'tool_preparation' + | 'tool_calls_dispatched' | 'run_step_completed' | 'run_step_closed' | 'message_delta' diff --git a/packages/data-provider/src/types/subagents.ts b/packages/data-provider/src/types/subagents.ts index 1199a6b2ddb..86fa5326404 100644 --- a/packages/data-provider/src/types/subagents.ts +++ b/packages/data-provider/src/types/subagents.ts @@ -79,6 +79,10 @@ export type SubagentActivityItem = input?: string; output?: string; status: 'running' | 'completed' | 'failed' | 'cancelled'; + toolPreparationStartedAt?: number; + toolDispatchedAt?: number; + toolPreparationDurationMs?: number; + toolExecutionDurationMs?: number; inputValidationError?: true; inputTruncated?: boolean; outputTruncated?: boolean; diff --git a/packages/data-provider/src/types/tools.ts b/packages/data-provider/src/types/tools.ts index 26db147d12e..dbbda766144 100644 --- a/packages/data-provider/src/types/tools.ts +++ b/packages/data-provider/src/types/tools.ts @@ -77,6 +77,9 @@ export interface FileSearchResource { */ export type AllowedCaller = 'direct' | 'code_execution'; +/** Switchable by the chat user; the value is the state a new chat starts in. */ +export type UserToolToggle = 'on' | 'off'; + /** * Per-tool configuration options stored at the agent level. * Keyed by tool_id (e.g., "search_mcp_github"). @@ -109,6 +112,8 @@ export type ToolOptions = { * @default false */ describe_intent?: boolean; + /** Chat-user switch for this tool on a saved agent. Absent = locked. */ + user_toggle?: UserToolToggle; }; /**