feat(monitoring): add onBeforeCleanup handler to save session state

Add new workflow handler to persist session and monitoring IDs before cleanup when workflow is interrupted. This enables resuming workflows even if interrupted before first step completion.

Author: moazbuilds

src/agents/monitoring/cleanup.ts

@@ -18,13 +18,19 @@ export class MonitoringCleanup {
   private static workflowHandlers: {
     onStop?: () => void;
     onExit?: () => void;
+    onBeforeCleanup?: () => Promise<void>;
   } = {};
 
   /**
    * Register callbacks invoked during the two-stage Ctrl+C flow.
+   * Merges with existing handlers (new handlers override existing ones for the same key).
    */
-  static registerWorkflowHandlers(handlers: { onStop?: () => void; onExit?: () => void }): void {
-    this.workflowHandlers = handlers;
+  static registerWorkflowHandlers(handlers: {
+    onStop?: () => void;
+    onExit?: () => void;
+    onBeforeCleanup?: () => Promise<void>;
+  }): void {
+    this.workflowHandlers = { ...this.workflowHandlers, ...handlers };
   }
 
   static clearWorkflowHandlers(): void {
@@ -139,6 +145,15 @@ export class MonitoringCleanup {
     // Emit workflow:skip to abort the currently running step
     (process as NodeJS.EventEmitter).emit('workflow:skip');
 
+    // Save session state for active agents before cleanup (for resume on restart)
+    if (this.workflowHandlers.onBeforeCleanup) {
+      try {
+        await this.workflowHandlers.onBeforeCleanup();
+      } catch (error) {
+        logger.debug('onBeforeCleanup failed:', error);
+      }
+    }
+
     // Stop active agents
     await this.stopActiveAgents();
 

src/workflows/run.ts

@@ -19,7 +19,7 @@ import {
 } from '../shared/workflows/index.js';
 import { StepIndexManager } from './indexing/index.js';
 import { registry } from '../infra/engines/index.js';
-import { MonitoringCleanup } from '../agents/monitoring/index.js';
+import { MonitoringCleanup, AgentMonitorService } from '../agents/monitoring/index.js';
 import { WorkflowEventBus, WorkflowEventEmitter } from './events/index.js';
 import { validateSpecification } from '../runtime/services/index.js';
 import { WorkflowRunner } from './runner/index.js';
@@ -54,8 +54,33 @@ export async function runWorkflow(options: RunWorkflowOptions = {}): Promise<voi
   // Set up cleanup handlers
   MonitoringCleanup.setup();
 
-  // Load template
+  // Load template (needed before we can set up the before-cleanup handler)
   const cmRoot = path.join(cwd, '.codemachine');
+  const indexManager = new StepIndexManager(cmRoot);
+
+  // Register callback to save session state before cleanup on Ctrl+C
+  // This ensures session/monitoring IDs are persisted even if the first turn hasn't completed
+  MonitoringCleanup.registerWorkflowHandlers({
+    onBeforeCleanup: async () => {
+      const monitor = AgentMonitorService.getInstance();
+      const activeAgents = monitor.getActiveAgents();
+
+      // Find root agents (no parentId) - these are the main step agents
+      const rootAgents = activeAgents.filter((agent) => !agent.parentId);
+
+      for (const agent of rootAgents) {
+        // Only save if agent has a sessionId (needed for resume)
+        if (agent.sessionId) {
+          const stepIndex = indexManager.currentStepIndex;
+          debug('[Workflow] Saving session state on Ctrl+C: step=%d, sessionId=%s, monitoringId=%d',
+            stepIndex, agent.sessionId, agent.id);
+          await indexManager.stepSessionInitialized(stepIndex, agent.sessionId, agent.id);
+        }
+      }
+    },
+  });
+
+  // Load template
   const templatePath = options.templatePath || (await getTemplatePathFromTracking(cmRoot));
   const { template } = await loadTemplateWithPath(cwd, templatePath);
 
@@ -108,9 +133,6 @@ export async function runWorkflow(options: RunWorkflowOptions = {}): Promise<voi
     globalThis.__workflowEventBus = eventBus;
   }
 
-  // Create step index manager
-  const indexManager = new StepIndexManager(cmRoot);
-
   // Get resume info
   const resumeInfo = await indexManager.getResumeInfo();
   const startIndex = resumeInfo.startIndex;