Ref count agent process (#426)

Kill the agent processes once all the referencing dispatchers are closed.
microsoft · Nov 22, 2024 · 02673d9 · 02673d9
1 parent 1e82a4e
commit 02673d9
Show file tree

Hide file tree

Showing 9 changed files with 75 additions and 25 deletions.
diff --git a/ts/packages/api/src/typeAgentServer.ts b/ts/packages/api/src/typeAgentServer.ts
@@ -54,5 +54,6 @@ export class TypeAgentServer {
     stop() {
         this.webServer?.stop();
         this.webSocketServer?.stop();
+        this.dispatcher?.close();
     }
 }
diff --git a/ts/packages/cli/src/commands/run/explain.ts b/ts/packages/cli/src/commands/run/explain.ts
@@ -166,8 +166,5 @@ export default class ExplainCommand extends Command {
             printProcessRequestActionResult(result);
         }
         await closeCommandHandlerContext(context);
-
-        // Some background network (like monogo) might keep the process live, exit explicitly.
-        process.exit(0);
     }
 }
diff --git a/ts/packages/cli/src/commands/run/translate.ts b/ts/packages/cli/src/commands/run/translate.ts
@@ -50,7 +50,5 @@ export default class TranslateCommand extends Command {
             `@dispatcher translate ${args.request}`,
         );
         await dispatcher.close();
-        // Some background network (like monogo) might keep the process live, exit explicitly.
-        process.exit(0);
     }
 }
diff --git a/ts/packages/dispatcher/src/agent/agentConfig.ts b/ts/packages/dispatcher/src/agent/agentConfig.ts
@@ -13,7 +13,7 @@ import {
 import { createRequire } from "module";
 import path from "node:path";
 
-import { createAgentProcessShim } from "./agentProcessShim.js";
+import { AgentProcess, createAgentProcess } from "./agentProcessShim.js";
 import { AppAgentProvider } from "./agentProvider.js";
 import { CommandHandlerContext } from "../handlers/common/commandHandlerContext.js";
 import { loadInlineAgent } from "./inlineAgentHandlers.js";
@@ -128,10 +128,12 @@ function enableExecutionMode() {
     return process.env.TYPEAGENT_EXECMODE !== "0";
 }
 
-async function loadModuleAgent(info: ModuleAppAgentInfo): Promise<AppAgent> {
+async function loadModuleAgent(
+    info: ModuleAppAgentInfo,
+): Promise<AgentProcess> {
     const execMode = info.execMode ?? ExecutionMode.SeparateProcess;
     if (enableExecutionMode() && execMode === ExecutionMode.SeparateProcess) {
-        return createAgentProcessShim(`${info.name}/agent/handlers`);
+        return createAgentProcess(`${info.name}/agent/handlers`);
     }
 
     const module = await import(`${info.name}/agent/handlers`);
@@ -140,12 +142,16 @@ async function loadModuleAgent(info: ModuleAppAgentInfo): Promise<AppAgent> {
             `Failed to load module agent ${info.name}: missing 'instantiate' function.`,
         );
     }
-    return module.instantiate();
+    return {
+        appAgent: module.instantiate(),
+        process: undefined,
+        count: 1,
+    };
 }
 
 async function loadExternalModuleAgent(
     info: ModuleAppAgentInfo,
-): Promise<AppAgent> {
+): Promise<AgentProcess> {
     const pkgpath = path.join(
         getUserProfileDir(),
         "externalagents/package.json",
@@ -155,7 +161,7 @@ async function loadExternalModuleAgent(
 
     const execMode = info.execMode ?? ExecutionMode.SeparateProcess;
     if (enableExecutionMode() && execMode === ExecutionMode.SeparateProcess) {
-        return createAgentProcessShim(`file://${handlerPath}`);
+        return createAgentProcess(`file://${handlerPath}`);
     }
 
     const module = await import(`${handlerPath}`);
@@ -164,24 +170,31 @@ async function loadExternalModuleAgent(
             `Failed to load module agent ${info.name}: missing 'instantiate' function.`,
         );
     }
-    return module.instantiate();
+    return {
+        appAgent: module.instantiate(),
+        process: undefined,
+        count: 1,
+    };
 }
 
 // Load on demand, doesn't unload for now
-const moduleAgents = new Map<string, AppAgent>();
+const moduleAgents = new Map<string, AgentProcess>();
 async function getModuleAgent(appAgentName: string) {
     const existing = moduleAgents.get(appAgentName);
-    if (existing) return existing;
+    if (existing) {
+        existing.count++;
+        return existing.appAgent;
+    }
     const config = getDispatcherConfig().agents[appAgentName];
     if (config === undefined || config.type !== "module") {
         throw new Error(`Unable to load app agent name: ${appAgentName}`);
     }
     const agent = await loadModuleAgent(config);
     moduleAgents.set(appAgentName, agent);
-    return agent;
+    return agent.appAgent;
 }
 
-const externalAgents = new Map<string, AppAgent>();
+const externalAgents = new Map<string, AgentProcess>();
 export function getExternalAppAgentProvider(
     context: CommandHandlerContext,
 ): AppAgentProvider {
@@ -203,19 +216,31 @@ export function getExternalAppAgentProvider(
                 ? await getExternalModuleAgent(appAgentName)
                 : loadInlineAgent(appAgentName, context);
         },
+        unloadAppAgent(appAgentName: string) {
+            const agent = externalAgents.get(appAgentName);
+            if (agent) {
+                if (--agent.count === 0) {
+                    agent.process?.kill();
+                    externalAgents.delete(appAgentName);
+                }
+            }
+        },
     };
 }
 
-async function getExternalModuleAgent(appAgentName: string) {
+async function getExternalModuleAgent(appAgentName: string): Promise<AppAgent> {
     const existing = moduleAgents.get(appAgentName);
-    if (existing) return existing;
+    if (existing) {
+        existing.count++;
+        return existing.appAgent;
+    }
     const config = getExternalAgentsConfig().agents[appAgentName];
     if (config === undefined || config.type !== "module") {
         throw new Error(`Unable to load app agent name: ${appAgentName}`);
     }
     const agent = await loadExternalModuleAgent(config);
     moduleAgents.set(appAgentName, agent);
-    return agent;
+    return agent.appAgent;
 }
 
 export function getBuiltinAppAgentProvider(
@@ -239,5 +264,14 @@ export function getBuiltinAppAgentProvider(
                 ? await getModuleAgent(appAgentName)
                 : loadInlineAgent(appAgentName, context);
         },
+        unloadAppAgent(appAgentName: string) {
+            const agent = moduleAgents.get(appAgentName);
+            if (agent) {
+                if (--agent.count === 0) {
+                    agent.process?.kill();
+                    externalAgents.delete(appAgentName);
+                }
+            }
+        },
     };
 }
diff --git a/ts/packages/dispatcher/src/agent/agentProcessShim.ts b/ts/packages/dispatcher/src/agent/agentProcessShim.ts
@@ -68,9 +68,16 @@ function createContextMap<T>() {
         close,
     };
 }
-export async function createAgentProcessShim(
+
+export type AgentProcess = {
+    appAgent: AppAgent;
+    process: child_process.ChildProcess | undefined;
+    count: number;
+};
+
+export async function createAgentProcess(
     modulePath: string,
-): Promise<AppAgent> {
+): Promise<AgentProcess> {
     const process = child_process.fork(
         fileURLToPath(new URL(`./agentProcess.js`, import.meta.url)),
         [modulePath],
@@ -407,5 +414,9 @@ export async function createAgentProcessShim(
         contextMap.close(context);
         return result;
     };
-    return result;
+    return {
+        process,
+        appAgent: result,
+        count: 1,
+    };
 }
diff --git a/ts/packages/dispatcher/src/agent/agentProvider.ts b/ts/packages/dispatcher/src/agent/agentProvider.ts
@@ -7,4 +7,5 @@ export interface AppAgentProvider {
     getAppAgentNames(): string[];
     getAppAgentManifest(appAgentName: string): Promise<AppAgentManifest>;
     loadAppAgent(appAgentName: string): Promise<AppAgent>;
+    unloadAppAgent(appAgentName: string): void;
 }
diff --git a/ts/packages/dispatcher/src/handlers/common/appAgentManager.ts b/ts/packages/dispatcher/src/handlers/common/appAgentManager.ts
@@ -426,6 +426,10 @@ export class AppAgentManager implements ActionConfigProvider {
             record.actions.clear();
             record.commands = false;
             await this.closeSessionContext(record);
+            if (record.appAgent !== undefined) {
+                record.provider.unloadAppAgent(record.name);
+            }
+            record.appAgent = undefined;
         }
     }
 

diff --git a/ts/packages/dispatcher/src/session/session.ts b/ts/packages/dispatcher/src/session/session.ts
@@ -72,10 +72,9 @@ async function loadSessions(): Promise<Sessions> {
                 },
             });
         } catch (e) {
-            console.error(
-                `ERROR: Unable to lock profile ${userProfileDir}. Only one client per profile can be active at a time.`,
+            throw new Error(
+                `Unable to lock profile ${userProfileDir}. Only one client per profile can be active at a time.`,
             );
-            process.exit(-1);
         }
     }
 

diff --git a/ts/packages/shell/src/main/agent.ts b/ts/packages/shell/src/main/agent.ts
@@ -243,4 +243,9 @@ export const shellAgentProvider: AppAgentProvider = {
         }
         return agent;
     },
+    unloadAppAgent: (appAgentName: string) => {
+        if (appAgentName !== "shell") {
+            throw new Error(`Unknown app agent: ${appAgentName}`);
+        }
+    },
 };