feat: Implement ProcessTrackingService (#740)

Author: charlesvien

apps/twig/src/main/di/container.ts

@@ -14,6 +14,7 @@ import { FoldersService } from "../services/folders/service.js";
 import { FsService } from "../services/fs/service.js";
 import { GitService } from "../services/git/service.js";
 import { OAuthService } from "../services/oauth/service.js";
+import { ProcessTrackingService } from "../services/process-tracking/service.js";
 import { ShellService } from "../services/shell/service.js";
 import { TaskLinkService } from "../services/task-link/service.js";
 import { UIService } from "../services/ui/service.js";
@@ -39,6 +40,7 @@ container.bind(MAIN_TOKENS.FoldersService).to(FoldersService);
 container.bind(MAIN_TOKENS.FsService).to(FsService);
 container.bind(MAIN_TOKENS.GitService).to(GitService);
 container.bind(MAIN_TOKENS.OAuthService).to(OAuthService);
+container.bind(MAIN_TOKENS.ProcessTrackingService).to(ProcessTrackingService);
 container.bind(MAIN_TOKENS.ShellService).to(ShellService);
 container.bind(MAIN_TOKENS.UIService).to(UIService);
 container.bind(MAIN_TOKENS.UpdatesService).to(UpdatesService);

apps/twig/src/main/di/tokens.ts

@@ -20,6 +20,7 @@ export const MAIN_TOKENS = Object.freeze({
   GitService: Symbol.for("Main.GitService"),
   DeepLinkService: Symbol.for("Main.DeepLinkService"),
   OAuthService: Symbol.for("Main.OAuthService"),
+  ProcessTrackingService: Symbol.for("Main.ProcessTrackingService"),
   ShellService: Symbol.for("Main.ShellService"),
   UIService: Symbol.for("Main.UIService"),
   UpdatesService: Symbol.for("Main.UpdatesService"),

apps/twig/src/main/lib/process-utils.ts

@@ -0,0 +1,46 @@
+import { execSync } from "node:child_process";
+import { platform } from "node:os";
+import { logger } from "./logger.js";
+
+const log = logger.scope("process-utils");
+
+/**
+ * Kill a process and all its children by killing the process group.
+ * On Unix, we use process.kill(-pid) to kill the entire process group.
+ * On Windows, we use taskkill with /T flag to kill the process tree.
+ */
+export function killProcessTree(pid: number): void {
+  try {
+    if (platform() === "win32") {
+      // Windows: use taskkill with /T to kill process tree
+      execSync(`taskkill /PID ${pid} /T /F`, { stdio: "ignore" });
+    } else {
+      // Unix: kill the process group by using negative PID
+      // This sends SIGTERM to all processes in the group
+      try {
+        process.kill(-pid, "SIGTERM");
+      } catch {
+        // If SIGTERM fails (process may have already exited), try SIGKILL
+        try {
+          process.kill(-pid, "SIGKILL");
+        } catch (err) {
+          log.warn(`Failed to kill process group for PID ${pid}`, err);
+        }
+      }
+    }
+  } catch (err) {
+    log.warn(`Failed to kill process tree for PID ${pid}`, err);
+  }
+}
+
+/**
+ * Check if a process is alive using signal 0.
+ */
+export function isProcessAlive(pid: number): boolean {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch {
+    return false;
+  }
+}

apps/twig/src/main/services/agent/service.ts

@@ -19,11 +19,13 @@ import {
 import { getLlmGatewayUrl } from "@posthog/agent/posthog-api";
 import type { OnLogCallback } from "@posthog/agent/types";
 import { app } from "electron";
-import { injectable, preDestroy } from "inversify";
+import { inject, injectable, preDestroy } from "inversify";
 import type { ExecutionMode } from "@/shared/types.js";
 import type { AcpMessage } from "../../../shared/types/session-events.js";
+import { MAIN_TOKENS } from "../../di/tokens.js";
 import { logger } from "../../lib/logger.js";
 import { TypedEventEmitter } from "../../lib/typed-event-emitter.js";
+import type { ProcessTrackingService } from "../process-tracking/service.js";
 import {
   AgentServiceEvent,
   type AgentServiceEvents,
@@ -185,6 +187,7 @@ interface ManagedSession {
   config: SessionConfig;
   interruptReason?: InterruptReason;
   needsRecreation: boolean;
+  recreationPromise?: Promise<ManagedSession>;
   promptPending: boolean;
   pendingContext?: string;
   availableModels?: Array<{
@@ -214,6 +217,15 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
   private sessions = new Map<string, ManagedSession>();
   private currentToken: string | null = null;
   private pendingPermissions = new Map<string, PendingPermission>();
+  private processTracking: ProcessTrackingService;
+
+  constructor(
+    @inject(MAIN_TOKENS.ProcessTrackingService)
+    processTracking: ProcessTrackingService,
+  ) {
+    super();
+    this.processTracking = processTracking;
+  }
 
   public updateToken(newToken: string): void {
     this.currentToken = newToken;
@@ -394,6 +406,9 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
         return existing;
       }
 
+      // Kill any lingering processes from previous runs of this task
+      this.processTracking.killByTaskId(taskId);
+
       // Clean up any prior session for this taskRunId before creating a new one
       await this.cleanupSession(taskRunId);
     }
@@ -413,7 +428,26 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
     });
 
     try {
-      const acpConnection = await agent.run(taskId, taskRunId);
+      const acpConnection = await agent.run(taskId, taskRunId, {
+        processCallbacks: {
+          onProcessSpawned: (info) => {
+            this.processTracking.register(
+              info.pid,
+              "agent",
+              `agent:${taskRunId}`,
+              {
+                taskRunId,
+                taskId,
+                command: info.command,
+              },
+              taskId,
+            );
+          },
+          onProcessExited: (pid) => {
+            this.processTracking.unregister(pid, "agent-exited");
+          },
+        },
+      });
       const { clientStreams } = acpConnection;
 
       const connection = this.createClientConnection(
@@ -572,10 +606,18 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
 
     // Recreate session if marked (token was refreshed while session was active)
     if (session.needsRecreation) {
-      log.info("Recreating session before prompt (token refreshed)", {
-        sessionId,
-      });
-      session = await this.recreateSession(sessionId);
+      if (!session.recreationPromise) {
+        log.info("Recreating session before prompt (token refreshed)", {
+          sessionId,
+        });
+        session.recreationPromise = this.recreateSession(sessionId).finally(
+          () => {
+            const s = this.sessions.get(sessionId);
+            if (s) s.recreationPromise = undefined;
+          },
+        );
+      }
+      session = await session.recreationPromise;
     }
 
     // Prepend pending context if present
@@ -636,6 +678,14 @@ export class AgentService extends TypedEventEmitter<AgentServiceEvents> {
     }
   }
 
+  async cancelSessionsByTaskId(taskId: string): Promise<void> {
+    for (const [taskRunId, session] of this.sessions) {
+      if (session.taskId === taskId) {
+        await this.cleanupSession(taskRunId);
+      }
+    }
+  }
+
   async cancelPrompt(
     sessionId: string,
     reason?: InterruptReason,

apps/twig/src/main/services/app-lifecycle/service.ts

@@ -2,9 +2,11 @@ import { app } from "electron";
 import { injectable } from "inversify";
 import { ANALYTICS_EVENTS } from "../../../types/analytics.js";
 import { container } from "../../di/container.js";
+import { MAIN_TOKENS } from "../../di/tokens.js";
 import { withTimeout } from "../../lib/async.js";
 import { logger } from "../../lib/logger.js";
 import { shutdownPostHog, trackAppEvent } from "../posthog-analytics.js";
+import type { ProcessTrackingService } from "../process-tracking/service.js";
 
 const log = logger.scope("app-lifecycle");
 
@@ -53,7 +55,38 @@ export class AppLifecycleService {
   }
 
   private async doShutdown(): Promise<void> {
-    log.info("Shutdown started: unbinding container");
+    log.info("Shutdown started");
+
+    try {
+      const processTracking = container.get<ProcessTrackingService>(
+        MAIN_TOKENS.ProcessTrackingService,
+      );
+      const snapshot = await processTracking.getSnapshot(true);
+      log.info("Process snapshot at shutdown", {
+        tracked: {
+          shell: snapshot.tracked.shell.length,
+          agent: snapshot.tracked.agent.length,
+          child: snapshot.tracked.child.length,
+        },
+        discovered: snapshot.discovered?.length ?? 0,
+        untrackedDiscovered:
+          snapshot.discovered?.filter((p) => !p.tracked).length ?? 0,
+      });
+
+      if (
+        snapshot.tracked.shell.length +
+          snapshot.tracked.agent.length +
+          snapshot.tracked.child.length >
+        0
+      ) {
+        log.info("Killing all tracked processes before container unbind");
+        processTracking.killAll();
+      }
+    } catch (error) {
+      log.warn("Failed to get process snapshot at shutdown", error);
+    }
+
+    log.info("Unbinding container");
     try {
       await container.unbindAll();
       log.info("Container unbound successfully");