diff --git a/ts/packages/defaultAgentProvider/src/mcpAgentProvider.ts b/ts/packages/defaultAgentProvider/src/mcpAgentProvider.ts index c8625e5521..bd792a2c8f 100644 --- a/ts/packages/defaultAgentProvider/src/mcpAgentProvider.ts +++ b/ts/packages/defaultAgentProvider/src/mcpAgentProvider.ts @@ -3,6 +3,7 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js"; +import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { AppAgent, AppAgentManifest } from "@typeagent/agent-sdk"; import { createActionResult } from "@typeagent/agent-sdk/helpers/action"; import { @@ -12,8 +13,8 @@ import { import { AppAgentProvider } from "agent-dispatcher"; import { ArgDefinitions, - ParameterDefinitions, ParsedCommandParams, + ParameterDefinitions, ActionContext, } from "@typeagent/agent-sdk"; import { @@ -21,6 +22,12 @@ import { getCommandInterface, } from "@typeagent/agent-sdk/helpers/command"; import { InstanceConfig, InstanceConfigProvider } from "./utils/config.js"; +import { spawn, ChildProcess } from "child_process"; +import net from "net"; +import registerDebug from "debug"; + +const debug = registerDebug("typeagent:mcp"); +const debugError = registerDebug("typeagent:mcp:error"); export type McpAppAgentInfo = { emojiChar: string; @@ -30,12 +37,16 @@ export type McpAppAgentInfo = { actionDefaultEnabled?: boolean; serverScript?: string; serverScriptArgs?: string[] | ArgDefinitions; + serverUrl?: string; + serverCommand?: string; + serverCommandArgs?: string[]; }; export type McpAppAgent = { manifest: AppAgentManifest; agent: AppAgent; - transport: StdioClientTransport | undefined; + transport: StdioClientTransport | StreamableHTTPClientTransport | undefined; + serverProcess?: ChildProcess | undefined; }; export type McpAppAgentRecord = { agentP: Promise; @@ -52,11 +63,86 @@ function convertSchema(tools: any) { return JSON.stringify(toJSONParsedActionSchema(pas)); } +// Check if anything is already listening on the port — raw TCP, no MCP handshake. +// This prevents us from launching a second server when one (even a broken one) is +// already bound to the port, which would cause compilation + bind-failure hangs. +function isPortOccupied(url: string): Promise { + return new Promise((resolve) => { + try { + const parsed = new URL(url); + const port = parseInt( + parsed.port || (parsed.protocol === "https:" ? "443" : "80"), + ); + const socket = net.createConnection(port, parsed.hostname); + socket.once("connect", () => { + socket.destroy(); + resolve(true); + }); + socket.once("error", () => resolve(false)); + socket.setTimeout(2000, () => { + socket.destroy(); + resolve(false); + }); + } catch { + resolve(false); + } + }); +} + +function launchHttpServer( + command: string, + args: string[], +): Promise { + return new Promise((resolve, reject) => { + const proc = spawn(command, args, { + stdio: ["ignore", "pipe", "pipe"], + }); + let started = false; + const timeout = setTimeout(() => { + if (!started) { + proc.kill(); + reject( + new Error(`HTTP MCP server failed to start within 180s`), + ); + } + }, 180000); + const onData = (data: Buffer) => { + const line = data.toString().trimEnd(); + debug(`[server] ${line}`); + if (!started && line.includes("Now listening on")) { + started = true; + clearTimeout(timeout); + resolve(proc); + } + }; + proc.stdout?.on("data", onData); + proc.stderr?.on("data", onData); + proc.on("error", (err) => { + clearTimeout(timeout); + reject(err); + }); + proc.on("exit", (code) => { + clearTimeout(timeout); + if (!started) { + reject( + new Error( + `HTTP MCP server exited with code ${code} before starting`, + ), + ); + } + }); + }); +} + function createMcpAppAgentTransport( appAgentName: string, info: McpAppAgentInfo, instanceConfig?: McpAppAgentConfig, -) { +): StdioClientTransport | StreamableHTTPClientTransport { + if (info.serverUrl !== undefined) { + return new StreamableHTTPClientTransport(new URL(info.serverUrl)); + } + const serverScriptPath = info.serverScript; if (serverScriptPath === undefined) { throw new Error(`Invalid app agent: ${appAgentName}`); @@ -112,28 +198,21 @@ function getMcpCommandHandlerTable( context: ActionContext, params: ParsedCommandParams, ) => { - const serverArgs: string[] = []; - if (params.args) { - for (const value of Object.values(params.args)) { - if (Array.isArray(value)) { - serverArgs.push(...value.map(String)); - } else if (value !== undefined) { - serverArgs.push(String(value)); - } - } - } const instanceConfig: InstanceConfig = structuredClone( configs.getInstanceConfig(), ); if (instanceConfig.mcpServers === undefined) { instanceConfig.mcpServers = {}; } + const serverScriptArgs = Object.keys(args).map((k) => + String((params.args as Record)[k]), + ); instanceConfig.mcpServers[appAgentName] = { - serverScriptArgs: serverArgs, + serverScriptArgs, }; configs.setInstanceConfig(instanceConfig); context.actionIO.appendDisplay( - `Server arguments set to ${serverArgs.join(" ")}. Please restart TypeAgent to reflect the change.`, + `Server arguments set to ${serverScriptArgs.join(" ")}. Please restart TypeAgent to reflect the change.`, ); }, }, @@ -170,17 +249,52 @@ function createMcpAppAgentRecord( } const createMcpAppAgent = async (): Promise => { - let transport: StdioClientTransport | undefined; + let transport: + | StdioClientTransport + | StreamableHTTPClientTransport + | undefined; + let serverProcess: ChildProcess | undefined; let agent: AppAgent; try { + if (info.serverCommand !== undefined) { + const occupied = + info.serverUrl !== undefined && + (await isPortOccupied(info.serverUrl)); + debug( + `[${appAgentName}] serverUrl=${info.serverUrl} port occupied=${occupied}`, + ); + if (!occupied) { + debug( + `[${appAgentName}] launching server: ${info.serverCommand} ${(info.serverCommandArgs ?? []).join(" ")}`, + ); + serverProcess = await launchHttpServer( + info.serverCommand, + info.serverCommandArgs ?? [], + ); + debug( + `[${appAgentName}] server process started (pid ${serverProcess.pid})`, + ); + } else { + debug( + `[${appAgentName}] server already running, skipping launch`, + ); + } + } + const transportUrl = + info.serverUrl ?? info.serverScript ?? "(stdio)"; + debug(`[${appAgentName}] connecting transport to ${transportUrl}`); transport = createMcpAppAgentTransport( appAgentName, info, instanceConfig, ); const client = new Client({ name: clientName, version }); - await client.connect(transport); + await client.connect(transport as any); + debug(`[${appAgentName}] connected, listing tools...`); const tools = (await client.listTools()).tools; + debug( + `[${appAgentName}] found ${tools.length} tool(s): ${tools.map((t) => t.name).join(", ")}`, + ); if (tools.length === 0) { throw new Error( `Invalid app agent config ${appAgentName}: No tools found`, @@ -212,10 +326,17 @@ function createMcpAppAgentRecord( }, }; } catch (error: any) { + debugError( + `[${appAgentName}] failed to connect: ${error?.message ?? error}`, + ); if (transport !== undefined) { transport.close(); transport = undefined; } + if (serverProcess !== undefined) { + serverProcess.kill(); + serverProcess = undefined; + } agent = { updateAgentContext() { // Delay throwing error until the agent is used. @@ -240,6 +361,7 @@ function createMcpAppAgentRecord( manifest, transport, agent, + serverProcess, }; }; return { @@ -255,14 +377,71 @@ export function createMcpAppAgentProvider( configs?: InstanceConfigProvider, ): AppAgentProvider { const instanceConfig = configs?.getInstanceConfig()?.mcpServers; - const mcpAppAgents = new Map(); + + // For server-command agents: background records that start the server eagerly. + // count is kept at 1 (background ref) so the server process is never killed + // by a stray unload. They are moved into mcpAppAgents on first loadAppAgent. + const backgroundRecords = new Map(); + + // Callbacks registered via onSchemaReady() + const schemaReadyCallbacks: (( + agentName: string, + manifest: AppAgentManifest, + ) => void)[] = []; + + // Manifests that are already resolved (so late-registered callbacks fire immediately) + const resolvedManifests = new Map(); + + function startBackgroundAgent(appAgentName: string) { + if ( + backgroundRecords.has(appAgentName) || + mcpAppAgents.has(appAgentName) + ) { + return; + } + const info = infos[appAgentName]; + if (info === undefined || info.serverCommand === undefined) { + return; + } + const record = createMcpAppAgentRecord( + name, + version, + appAgentName, + info, + configs, + instanceConfig?.[appAgentName], + ); + backgroundRecords.set(appAgentName, record); + + record.agentP + .then((agentData) => { + if (agentData.transport !== undefined) { + resolvedManifests.set(appAgentName, agentData.manifest); + for (const cb of schemaReadyCallbacks) { + cb(appAgentName, agentData.manifest); + } + } + }) + .catch(() => { + // errors surface when the agent is actually used + }); + } + function getMpcAppAgentRecord(appAgentName: string) { const existing = mcpAppAgents.get(appAgentName); if (existing !== undefined) { existing.count++; return existing; } + // Promote a background record (server already loading/loaded) + const background = backgroundRecords.get(appAgentName); + if (background !== undefined) { + background.count++; + backgroundRecords.delete(appAgentName); + mcpAppAgents.set(appAgentName, background); + return background; + } const info = infos[appAgentName]; if (info === undefined) { throw new Error(`Invalid app agent: ${appAgentName}`); @@ -278,19 +457,63 @@ export function createMcpAppAgentProvider( mcpAppAgents.set(appAgentName, record); return record; } + return { getAppAgentNames() { return Object.keys(infos); }, + + getLoadingAgentNames() { + return [...backgroundRecords.keys()]; + }, + + onSchemaReady(callback) { + schemaReadyCallbacks.push(callback); + // Fire immediately for any agents already resolved + for (const [agentName, manifest] of resolvedManifests) { + callback(agentName, manifest); + } + }, + async getAppAgentManifest(appAgentName: string) { + const info = infos[appAgentName]; + if (info === undefined) { + throw new Error(`Invalid app agent: ${appAgentName}`); + } + if (info.serverCommand !== undefined) { + // Non-blocking: kick off background server startup and return a + // stub manifest immediately. The real manifest (with schema) is + // delivered later via the onSchemaReady callback. + startBackgroundAgent(appAgentName); + // Include a stub schema so the agent row appears in @config + // (default view filters by schema names). The empty content + // will fail to parse, showing ❌ while loading. refreshAgentSchema + // replaces it with the real schema once the server is ready. + return { + emojiChar: info.emojiChar, + description: info.description, + defaultEnabled: info.defaultEnabled, + schema: { + description: info.description, + schemaType: entryTypeName, + schemaFile: { + format: "pas" as const, + content: "", + }, + }, + } as AppAgentManifest; + } + // Stdio agents start fast — keep blocking path. const record = getMpcAppAgentRecord(appAgentName); const manifest = (await record.agentP).manifest; await this.unloadAppAgent(appAgentName); return manifest; }, + async loadAppAgent(appAgentName: string) { return (await getMpcAppAgentRecord(appAgentName).agentP).agent; }, + async unloadAppAgent(appAgentName: string) { const record = mcpAppAgents.get(appAgentName); if (!record || record.count === 0) { @@ -301,11 +524,16 @@ export function createMcpAppAgentProvider( const agent = await record.agentP; const transport = agent.transport; if (transport !== undefined) { - return new Promise((resolve) => { - transport.onclose = resolve; - transport.close(); - }); + if (transport instanceof StreamableHTTPClientTransport) { + await transport.close(); + } else { + return new Promise((resolve) => { + transport.onclose = resolve; + transport.close(); + }); + } } + agent.serverProcess?.kill(); } }, }; diff --git a/ts/packages/dispatcher/dispatcher/src/agentProvider/agentProvider.ts b/ts/packages/dispatcher/dispatcher/src/agentProvider/agentProvider.ts index 116753c648..35dbb16772 100644 --- a/ts/packages/dispatcher/dispatcher/src/agentProvider/agentProvider.ts +++ b/ts/packages/dispatcher/dispatcher/src/agentProvider/agentProvider.ts @@ -9,6 +9,16 @@ export interface AppAgentProvider { loadAppAgent(appAgentName: string): Promise; unloadAppAgent(appAgentName: string): Promise; setTraceNamespaces?(namespaces: string): void; + // Optional: providers that start slowly can return a stub manifest from + // getAppAgentManifest and call the registered callback with the real + // manifest once the agent is ready. + onSchemaReady?: ( + callback: (agentName: string, manifest: AppAgentManifest) => void, + ) => void; + // Optional: returns the names of agents currently loading asynchronously. + // Only these agents should show ⏳ in the UI. If omitted, no agents are + // treated as loading. + getLoadingAgentNames?(): string[]; } export interface AppAgentInstaller { diff --git a/ts/packages/dispatcher/dispatcher/src/context/appAgentManager.ts b/ts/packages/dispatcher/dispatcher/src/context/appAgentManager.ts index bcf5389527..8b48148432 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/appAgentManager.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/appAgentManager.ts @@ -120,6 +120,7 @@ function loadGrammar(actionConfig: ActionConfig): Grammar | undefined { export class AppAgentManager implements ActionConfigProvider { private readonly agents = new Map(); private readonly actionConfigs = new Map(); + private readonly loadingSchemas = new Set(); private readonly flowRegistry = new Map(); private readonly transientAgents: Record = {}; private readonly actionSemanticMap?: ActionSchemaSemanticMap; @@ -149,6 +150,10 @@ export class AppAgentManager implements ActionConfigProvider { return Array.from(this.agents.keys()); } + public isSchemaLoading(schemaName: string): boolean { + return this.loadingSchemas.has(schemaName); + } + public getAppAgentDescription(appAgentName: string) { const record = this.getRecord(appAgentName); return record.manifest.description; @@ -276,6 +281,7 @@ export class AppAgentManager implements ActionConfigProvider { actionEmbeddingCache?: EmbeddingCache, agentGrammarRegistry?: AgentGrammarRegistry, useNFAGrammar?: boolean, + stateRefreshFn?: () => Promise, ) { const agentNames = provider.getAppAgentNames(); const semanticMapP: Promise[] = []; @@ -296,6 +302,94 @@ export class AppAgentManager implements ActionConfigProvider { debug("Waiting for action embeddings"); await Promise.all(semanticMapP); debug("Finish action embeddings"); + + if (provider.onSchemaReady && stateRefreshFn) { + // Mark only the agents that are actually loading asynchronously (e.g. + // serverCommand MCP agents with slow startup). Agents that failed + // synchronously should show ❌, not ⏳. + const loadingNames = new Set( + provider.getLoadingAgentNames?.() ?? [], + ); + for (const name of agentNames) { + if (!loadingNames.has(name)) { + continue; + } + const record = this.agents.get(name); + if (record) { + for (const schemaName of Object.keys( + convertToActionConfig(name, record.manifest), + )) { + this.loadingSchemas.add(schemaName); + } + } + } + + provider.onSchemaReady(async (agentName, manifest) => { + try { + const refreshSemanticMapP: Promise[] = []; + this.refreshAgentSchema( + agentName, + manifest, + refreshSemanticMapP, + actionGrammarStore, + actionEmbeddingCache, + agentGrammarRegistry, + useNFAGrammar, + ); + await Promise.all(refreshSemanticMapP); + await stateRefreshFn(); + } catch (e) { + debugError( + `Failed to refresh schema for agent '${agentName}': ${e}`, + ); + } + }); + } + } + + private refreshAgentSchema( + appAgentName: string, + manifest: AppAgentManifest, + semanticMapP: Promise[], + actionGrammarStore: GrammarStore | undefined, + actionEmbeddingCache?: EmbeddingCache, + agentGrammarRegistry?: AgentGrammarRegistry, + useNFAGrammar?: boolean, + ) { + const record = this.agents.get(appAgentName); + if (record === undefined) { + throw new Error(`Agent not found: ${appAgentName}`); + } + + // Update the manifest (emoji, schema, etc.) + record.manifest = manifest; + + const actionConfigs = convertToActionConfig(appAgentName, manifest); + for (const [schemaName, config] of Object.entries(actionConfigs)) { + debug(`Refreshing action config: ${schemaName}`); + this.actionConfigs.set(schemaName, config); + if (config.transient) { + this.transientAgents[schemaName] = false; + } + try { + const actionSchemaFile = + this.actionSchemaFileCache.getActionSchemaFile(config); + if (this.actionSemanticMap) { + semanticMapP.push( + this.actionSemanticMap.addActionSchemaFile( + config, + actionSchemaFile, + actionEmbeddingCache, + ), + ); + } + record.schemaErrors.delete(schemaName); + this.loadingSchemas.delete(schemaName); + } catch (e: any) { + record.schemaErrors.set(schemaName, e); + this.loadingSchemas.delete(schemaName); + } + } } private addAgentManifest( @@ -373,14 +467,20 @@ export class AppAgentManager implements ActionConfigProvider { // Add to NFA grammar registry if using NFA system if (useNFAGrammar && agentGrammarRegistry) { try { - // Enrich grammar with checked variables from the already-parsed action schema - enrichGrammarWithCheckedVariables( - g, - actionSchemaFile.parsedActionSchema, - ); - debug( - `Enriched grammar with checked variables for schema: ${schemaName}`, - ); + // Enrich grammar with checked variables from parsed schema + try { + enrichGrammarWithCheckedVariables( + g, + actionSchemaFile.parsedActionSchema, + ); + debug( + `Enriched grammar with checked variables for schema: ${schemaName}`, + ); + } catch (enrichError) { + debug( + `Could not enrich grammar with checked variables for ${schemaName}: ${enrichError}`, + ); + } const nfa = compileGrammarToNFA(g, schemaName); agentGrammarRegistry.registerAgent( @@ -662,16 +762,24 @@ export class AppAgentManager implements ActionConfigProvider { ); if (enableSchema !== record.schemas.has(name)) { if (enableSchema) { - const e = record.schemaErrors.get(name); - if (e !== undefined) { - failedSchemas.push([name, enableSchema, e]); - debugError( - `Schema '${name}' is not enabled because of error: ${e.message}`, + if (this.loadingSchemas.has(name)) { + // Schema is still loading (e.g. slow MCP server start). + // Skip for now; refreshAgentSchema will re-run setState once ready. + debug( + `Schema '${name}' is still loading, skipping enable`, ); } else { - record.schemas.add(name); - changedSchemas.push([name, enableSchema]); - debug(`Schema enabled ${name}`); + const e = record.schemaErrors.get(name); + if (e !== undefined) { + failedSchemas.push([name, enableSchema, e]); + debugError( + `Schema '${name}' is not enabled because of error: ${e.message}`, + ); + } else { + record.schemas.add(name); + changedSchemas.push([name, enableSchema]); + debug(`Schema enabled ${name}`); + } } } else { record.schemas.delete(name); @@ -688,21 +796,26 @@ export class AppAgentManager implements ActionConfigProvider { failedActions, ); if (enableAction !== record.actions.has(name)) { - p.push( - (async () => { - try { - await this.updateAction( - name, - record, - enableAction, - context, - ); - changedActions.push([name, enableAction]); - } catch (e: any) { - failedActions.push([name, enableAction, e]); - } - })(), - ); + if (enableAction && this.loadingSchemas.has(name)) { + // Action agent is still loading — skip to avoid blocking on server startup. + debug(`Action '${name}' is still loading, skipping enable`); + } else { + p.push( + (async () => { + try { + await this.updateAction( + name, + record, + enableAction, + context, + ); + changedActions.push([name, enableAction]); + } catch (e: any) { + failedActions.push([name, enableAction, e]); + } + })(), + ); + } } } @@ -906,30 +1019,6 @@ export class AppAgentManager implements ActionConfigProvider { debug(`Loaded dynamic schema for ${schemaName}`); } - public async reloadAgentSchema( - schemaName: string, - context: CommandHandlerContext, - ): Promise { - const appAgentName = getAppAgentName(schemaName); - const record = this.getRecord(appAgentName); - if (!record.appAgent || !record.sessionContext) { - throw new Error(`Agent '${appAgentName}' is not initialized`); - } - - await this.loadDynamicSchema( - schemaName, - record.appAgent, - record.sessionContext, - context, - ); - await this.loadDynamicGrammar( - schemaName, - record.appAgent, - record.sessionContext, - context, - ); - } - private async updateAction( schemaName: string, record: AppAgentRecord, @@ -954,7 +1043,6 @@ export class AppAgentManager implements ActionConfigProvider { schemaName, ), ); - // Load dynamic schema and grammar from agent callbacks await this.loadDynamicSchema( schemaName, @@ -1119,6 +1207,38 @@ export class AppAgentManager implements ActionConfigProvider { return this.actionSchemaFileCache.getActionSchemaFile(config); } + public async reloadAgentSchema( + appAgentName: string, + context: CommandHandlerContext, + ): Promise { + const record = this.getRecord(appAgentName); + if (record.provider === undefined) { + return; + } + + // Unload cached schema files so they get reloaded from disk + for (const schemaName of this.actionConfigs.keys()) { + if (getAppAgentName(schemaName) === appAgentName) { + this.actionSchemaFileCache.unloadActionSchemaFile(schemaName); + } + } + + // Get fresh manifest from provider and refresh schemas + const manifest = + await record.provider.getAppAgentManifest(appAgentName); + const semanticMapP: Promise[] = []; + this.refreshAgentSchema( + appAgentName, + manifest, + semanticMapP, + undefined, + ); + await Promise.all(semanticMapP); + + // Clear translator cache to force re-translation with new schema + context.translatorCache.clear(); + } + public setTraceNamespaces(namespaces: string) { const providers = new Set(); for (const { provider } of this.agents.values()) { diff --git a/ts/packages/dispatcher/dispatcher/src/context/commandHandlerContext.ts b/ts/packages/dispatcher/dispatcher/src/context/commandHandlerContext.ts index 46a39cfd73..5bd6e4034f 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/commandHandlerContext.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/commandHandlerContext.ts @@ -82,7 +82,6 @@ import { createSchemaInfoProvider } from "../translation/actionSchemaFileCache.j import { createBuiltinAppAgentProvider } from "./inlineAgentProvider.js"; import { CommandResult } from "@typeagent/dispatcher-types"; import { DispatcherName } from "./dispatcher/dispatcherUtils.js"; -import { DisplayLog } from "../displayLog.js"; import lockfile from "proper-lockfile"; import { IndexManager } from "./indexManager.js"; import { ActionContextWithClose } from "../execute/actionContext.js"; @@ -97,55 +96,15 @@ import fs from "node:fs"; import { CosmosClient, PartitionKeyBuilder } from "@azure/cosmos"; import { CosmosPartitionKeyBuilder } from "telemetry"; import { DefaultAzureCredential } from "@azure/identity"; +import { DisplayLog } from "../displayLog.js"; +import { + fromJSONParsedActionSchema, + ParsedActionSchemaJSON, +} from "@typeagent/action-schema"; const debug = registerDebug("typeagent:dispatcher:init"); const debugError = registerDebug("typeagent:dispatcher:init:error"); -function wrapClientIOWithDisplayLog( - clientIO: ClientIO, - displayLog: DisplayLog, -): ClientIO { - return { - ...clientIO, - setUserRequest(requestId, command) { - const seq = displayLog.logUserRequest(requestId, command); - clientIO.setUserRequest(requestId, command, seq); - }, - setDisplayInfo(requestId, source, actionIndex?, action?) { - const seq = displayLog.logSetDisplayInfo( - requestId, - source, - actionIndex, - action, - ); - clientIO.setDisplayInfo( - requestId, - source, - actionIndex, - action, - seq, - ); - }, - setDisplay(message) { - const seq = displayLog.logSetDisplay(message); - clientIO.setDisplay(message, seq); - }, - appendDisplay(message, mode) { - const seq = displayLog.logAppendDisplay(message, mode); - clientIO.appendDisplay(message, mode, seq); - }, - notify(notificationId, event, data, source) { - const seq = displayLog.logNotify( - notificationId, - event, - data, - source, - ); - clientIO.notify(notificationId, event, data, source, seq); - }, - }; -} - export type EmptyFunction = () => void; export type SetSettingFunction = (name: string, value: any) => void; @@ -204,10 +163,12 @@ export type CommandHandlerContext = { logger?: Logger | undefined; currentRequestId: RequestId | undefined; currentAbortSignal: AbortSignal | undefined; + activeRequests: Map; noReasoning: boolean; commandResult?: CommandResult | undefined; chatHistory: ChatHistory; constructionProvider?: ConstructionProvider | undefined; + displayLog: DisplayLog; batchMode: boolean; pendingChoiceRoutes: Map< @@ -223,16 +184,10 @@ export type CommandHandlerContext = { commandProfiler?: Profiler | undefined; promptLogger?: PromptLogger | undefined; - // Maps requestId string → AbortController for in-flight commands. - // Used by cancelCommand() to abort a running command. - activeRequests: Map; - instanceDirLock: (() => Promise) | undefined; userRequestKnowledgeExtraction: boolean; actionResultKnowledgeExtraction: boolean; - - displayLog: DisplayLog; }; export function getRequestId(context: CommandHandlerContext): RequestId { @@ -481,6 +436,7 @@ async function addAppAgentProviders( ); if (appAgentProviders) { + const stateRefreshFn = () => setAppAgentStates(context); for (const provider of appAgentProviders) { await context.agents.addProvider( provider, @@ -488,6 +444,7 @@ async function addAppAgentProviders( embeddingCache, context.agentGrammarRegistry, useNFAGrammar, + stateRefreshFn, ); } } @@ -593,11 +550,7 @@ export async function initializeCommandHandlerContext( } const sessionDirPath = session.getSessionDirPath(); debug(`Session directory: ${sessionDirPath}`); - const displayLog = await DisplayLog.load(sessionDirPath); - const clientIO = wrapClientIOWithDisplayLog( - options?.clientIO ?? nullClientIO, - displayLog, - ); + const clientIO = options?.clientIO ?? nullClientIO; const loggerSink = getLoggerSink(() => context.dblogging, clientIO); const logger = new ChildLogger(loggerSink, DispatcherName, { hostName, @@ -638,6 +591,7 @@ export async function initializeCommandHandlerContext( commandLock: createLimiter(1), // Make sure we process one command at a time. currentRequestId: undefined, currentAbortSignal: undefined, + activeRequests: new Map(), noReasoning: false, pendingToggleTransientAgents: [], agentCache: await getAgentCache( @@ -654,12 +608,12 @@ export async function initializeCommandHandlerContext( chatHistory: createChatHistory( session.getConfig().execution.history, ), + displayLog: await DisplayLog.load(persistDir), logger, metricsManager: metrics ? new RequestMetricsManager() : undefined, promptLogger: createPromptLogger(getCosmosFactories()), batchMode: false, pendingChoiceRoutes: new Map(), - activeRequests: new Map(), instanceDirLock, constructionProvider, collectCommandResult: options?.collectCommandResult ?? false, @@ -674,8 +628,6 @@ export async function initializeCommandHandlerContext( actionResultKnowledgeExtraction: options?.conversationMemorySettings ?.actionResultKnowledgeExtraction ?? true, - - displayLog, }; await initializeMemory(context, sessionDirPath); @@ -811,20 +763,63 @@ async function setupGrammarGeneration(context: CommandHandlerContext) { // Enable auto-save await grammarStore.setAutoSave(config.cache.autoSave); + // Import getPackageFilePath for resolving schema paths + const { getPackageFilePath } = await import( + "../utils/getPackageFilePath.js" + ); + // Configure agent cache with grammar generation support context.agentCache.configureGrammarGeneration( context.agentGrammarRegistry, grammarStore, true, (schemaName: string) => { - const actionSchemaFile = - context.agents.tryGetActionSchemaFile(schemaName); - if (!actionSchemaFile) { + // Get compiled schema file path (.pas.json) from action config for grammar generation + const actionConfig = context.agents.tryGetActionConfig(schemaName); + if (!actionConfig) { throw new Error( - `Action schema file not found for schema: ${schemaName}`, + `Action config not found for schema: ${schemaName}`, ); } - return actionSchemaFile.parsedActionSchema; + + let schemaPath: string | undefined; + + // Use schemaFilePath directly if it's already a .pas.json file + if ( + actionConfig.schemaFilePath && + actionConfig.schemaFilePath.endsWith(".pas.json") + ) { + schemaPath = getPackageFilePath(actionConfig.schemaFilePath); + } else if ( + actionConfig.schemaFilePath && + actionConfig.schemaFilePath.endsWith(".ts") + ) { + // Fallback: try to derive .pas.json path from .ts schemaFilePath + // Try common pattern: ./src/schema.ts -> ../dist/schema.pas.json + const derivedPath = actionConfig.schemaFilePath + .replace(/^\.\/src\//, "../dist/") + .replace(/\.ts$/, ".pas.json"); + debug( + `Attempting fallback .pas.json path for ${schemaName}: ${derivedPath}`, + ); + try { + schemaPath = getPackageFilePath(derivedPath); + } catch { + // Fallback path doesn't exist, continue to error + } + } + + if (!schemaPath) { + throw new Error( + `Compiled schema file path (.pas.json) not found for schema: ${schemaName}. ` + + `Please ensure the schema is compiled to a .pas.json file.`, + ); + } + + const content = fs.readFileSync(schemaPath, "utf-8"); + return fromJSONParsedActionSchema( + JSON.parse(content) as ParsedActionSchemaJSON, + ); }, ); @@ -908,7 +903,6 @@ export async function closeCommandHandlerContext( ) { // Save the session because the token count is in it. context.session.save(); - await context.displayLog.save(); await context.agents.close(); if (context.instanceDirLock) { await context.instanceDirLock(); @@ -919,10 +913,7 @@ export async function setSessionOnCommandHandlerContext( context: CommandHandlerContext, session: Session, ) { - // Persist the old session's display log before switching - await context.displayLog.save(); context.session = session; - context.displayLog = await DisplayLog.load(session.getSessionDirPath()); await context.agents.close(); await initializeMemory(context, session.getSessionDirPath()); diff --git a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts index 8bfe229ba3..5344d4d54b 100644 --- a/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts +++ b/ts/packages/dispatcher/dispatcher/src/context/system/handlers/configCommandHandlers.ts @@ -255,16 +255,23 @@ function showAgentStatus( if (showSchema || showAction) { for (const name of agents.getSchemaNames()) { + const loading = agents.isSchemaLoading(name); if (showSchema) { const state = agents.isSchemaEnabled(name); const active = agents.isSchemaActive(name); setStatus(status, "schemas", name, state, active, changes); + if (loading && status[name] !== undefined) { + status[name]["schemas"] = "⏳"; + } } if (showAction) { const state = agents.isActionEnabled(name); const active = agents.isActionActive(name); setStatus(status, "actions", name, state, active, changes); + if (loading && status[name] !== undefined) { + status[name]["actions"] = "⏳"; + } } } } diff --git a/ts/packages/shell/playwright.config.ts b/ts/packages/shell/playwright.config.ts index 4102afa60a..aa6cf4910c 100644 --- a/ts/packages/shell/playwright.config.ts +++ b/ts/packages/shell/playwright.config.ts @@ -16,6 +16,8 @@ import { defineConfig, devices } from "@playwright/test"; */ export default defineConfig({ testDir: "./test", + /* Exclude Jest unit tests from Playwright test discovery */ + testIgnore: /\/partialCompletion\//, /* Run tests sequentially otherwise the client will complain about locked session file */ fullyParallel: false, /* Fail the build on CI if you accidentally left test.only in the source code. */ diff --git a/ts/packages/shell/test/testHelper.ts b/ts/packages/shell/test/testHelper.ts index 3dac9985b9..ca701d1412 100644 --- a/ts/packages/shell/test/testHelper.ts +++ b/ts/packages/shell/test/testHelper.ts @@ -217,6 +217,12 @@ export async function sendUserRequest(prompt: string, page: Page) { await locator.waitFor({ timeout: 30000, state: "visible" }); await locator.focus({ timeout: 30000 }); await locator.fill(prompt, { timeout: 30000 }); + + // robgruen - dismiss completion suggestion since it doesn't auto-dismiss on input and would cause the Enter key press to not submit the request but instead accept the suggestion + // TODO: fix completion to not need this workaround + await locator.press("ArrowLeft", { timeout: 30000 }); + + // send the request await locator.press("Enter", { timeout: 30000 }); }