feat(FN-4811): complete Step 1 — add active session registry wiring

Fusion-Task-Id: FN-4811
Fusion-Task-Lineage: f3dac123-cb46-4d31-8a7e-dfc664bfdc5f
This commit is contained in:
Fusion (runfusion.ai)
2026-05-16 16:39:04 -07:00
committed by gsxdsm
parent f4aa6d7b8b
commit 69d7b282cf
4 changed files with 199 additions and 23 deletions

View File

@@ -0,0 +1,38 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
import { activeSessionRegistry } from "../active-session-registry.js";
describe("activeSessionRegistry", () => {
beforeEach(() => {
activeSessionRegistry.clear();
});
it("registers and unregisters paths", () => {
activeSessionRegistry.registerPath("/tmp/w1", { taskId: "FN-1", kind: "executor", ownerKey: "FN-1" });
expect(activeSessionRegistry.isPathActive("/tmp/w1")).toBe(true);
activeSessionRegistry.unregisterPath("/tmp/w1");
expect(activeSessionRegistry.isPathActive("/tmp/w1")).toBe(false);
});
it("supports multiple paths for same task", () => {
activeSessionRegistry.registerPath("/tmp/w1", { taskId: "FN-1", kind: "executor", ownerKey: "FN-1" });
activeSessionRegistry.registerPath("/tmp/w2", { taskId: "FN-1", kind: "workflow-step", ownerKey: "FN-1#workflow-step" });
expect(activeSessionRegistry.pathsForTask("FN-1").sort()).toEqual(["/tmp/w1", "/tmp/w2"]);
});
it("returns null for unregistered path", () => {
expect(activeSessionRegistry.lookupByPath("/tmp/missing")).toBeNull();
});
it("overwrites duplicate registration with warning", () => {
const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {});
activeSessionRegistry.registerPath("/tmp/w1", { taskId: "FN-1", kind: "executor", ownerKey: "FN-1" });
activeSessionRegistry.registerPath("/tmp/w1", { taskId: "FN-2", kind: "workflow-step", ownerKey: "FN-2#workflow-step" });
expect(activeSessionRegistry.lookupByPath("/tmp/w1")?.taskId).toBe("FN-2");
expect(warnSpy).toHaveBeenCalledOnce();
warnSpy.mockRestore();
});
});

View File

@@ -0,0 +1,53 @@
export type ActiveSessionKind = "executor" | "step-session" | "workflow-step" | "step-session-parallel";
export interface ActiveSessionRegistration {
taskId: string;
kind: ActiveSessionKind;
ownerKey: string;
}
export interface ActiveSessionRecord extends ActiveSessionRegistration {
registeredAt: number;
}
class ActiveSessionRegistry {
private readonly records = new Map<string, ActiveSessionRecord>();
registerPath(worktreePath: string, registration: ActiveSessionRegistration): void {
if (this.records.has(worktreePath)) {
console.warn(`[active-session-registry] overwriting existing registration for ${worktreePath}`);
}
this.records.set(worktreePath, {
...registration,
registeredAt: Date.now(),
});
}
unregisterPath(worktreePath: string): void {
this.records.delete(worktreePath);
}
lookupByPath(worktreePath: string): ActiveSessionRecord | null {
return this.records.get(worktreePath) ?? null;
}
isPathActive(worktreePath: string): boolean {
return this.records.has(worktreePath);
}
pathsForTask(taskId: string): string[] {
const paths: string[] = [];
for (const [path, record] of this.records.entries()) {
if (record.taskId === taskId) {
paths.push(path);
}
}
return paths;
}
clear(): void {
this.records.clear();
}
}
export const activeSessionRegistry = new ActiveSessionRegistry();

View File

@@ -44,6 +44,7 @@ import type { SandboxBackend } from "./sandbox/types.js";
import { ModelRegistry, SessionManager, type ToolDefinition, type AgentSession } from "@mariozechner/pi-coding-agent";
import { PRIORITY_EXECUTE, type AgentSemaphore } from "./concurrency.js";
import { getRegisteredWorktreePaths, isGitRepository, isInsideWorktreesDir, isRegisteredGitWorktree, isUsableTaskWorktree, removeWorktree, type WorktreePool } from "./worktree-pool.js";
import { activeSessionRegistry } from "./active-session-registry.js";
import {
BranchConflictError,
BranchCrossContaminationError,
@@ -844,6 +845,53 @@ export class TaskExecutor {
/** Set of ephemeral spawned agent IDs with in-flight cleanup (prevents duplicate deletion attempts). */
private pendingEphemeralDeletions = new Set<string>();
private setActiveSession(taskId: string, sessionState: {
session: AgentSession;
seenSteeringIds: Set<string>;
lastResolvedModelProvider?: string;
lastResolvedModelId?: string;
lastTaskModelProvider?: string | null;
lastTaskModelId?: string | null;
lastAssignedAgentId?: string | null;
}, worktreePath: string): void {
this.activeSessions.set(taskId, sessionState);
activeSessionRegistry.registerPath(worktreePath, { taskId, kind: "executor", ownerKey: taskId });
}
private deleteActiveSession(taskId: string, worktreePath?: string): void {
this.activeSessions.delete(taskId);
const resolvedWorktreePath = worktreePath ?? this.activeWorktrees.get(taskId);
if (resolvedWorktreePath) {
activeSessionRegistry.unregisterPath(resolvedWorktreePath);
}
}
private setActiveStepExecutor(taskId: string, stepExecutor: StepSessionExecutor, worktreePath: string): void {
this.activeStepExecutors.set(taskId, stepExecutor);
activeSessionRegistry.registerPath(worktreePath, { taskId, kind: "step-session", ownerKey: `${taskId}#step-session` });
}
private deleteActiveStepExecutor(taskId: string, worktreePath?: string): void {
this.activeStepExecutors.delete(taskId);
const resolvedWorktreePath = worktreePath ?? this.activeWorktrees.get(taskId);
if (resolvedWorktreePath) {
activeSessionRegistry.unregisterPath(resolvedWorktreePath);
}
}
private setActiveWorkflowStepSession(taskId: string, session: AgentSession, worktreePath: string): void {
this.activeWorkflowStepSessions.set(taskId, session);
activeSessionRegistry.registerPath(worktreePath, { taskId, kind: "workflow-step", ownerKey: `${taskId}#workflow-step` });
}
private deleteActiveWorkflowStepSession(taskId: string, worktreePath?: string): void {
this.activeWorkflowStepSessions.delete(taskId);
const resolvedWorktreePath = worktreePath ?? this.activeWorktrees.get(taskId);
if (resolvedWorktreePath) {
activeSessionRegistry.unregisterPath(resolvedWorktreePath);
}
}
private getAutoRecoveryDispatcher(audit: RunAuditor): AutoRecoveryDispatcher {
if (this.options.autoRecoveryDispatcher) return this.options.autoRecoveryDispatcher;
const fileScopeHandler = createFileScopeAutoRecoveryHandler({
@@ -1307,7 +1355,7 @@ export class TaskExecutor {
});
}
session.dispose();
this.activeSessions.delete(task.id);
this.deleteActiveSession(task.id);
}
if (this.activeStepExecutors.has(task.id)) {
executorLog.log(`${task.id} moved from in-progress to ${to} — terminating step sessions`);
@@ -1325,7 +1373,7 @@ export class TaskExecutor {
stepExecutor.terminateAllSessions().catch((err) =>
executorLog.error(`Failed to terminate step sessions for ${task.id}:`, err),
);
this.activeStepExecutors.delete(task.id);
this.deleteActiveStepExecutor(task.id);
}
if (this.activeWorkflowStepSessions.has(task.id)) {
executorLog.log(`${task.id} moved from in-progress to ${to} — terminating workflow step session`);
@@ -1339,7 +1387,7 @@ export class TaskExecutor {
});
}
workflowSession.dispose();
this.activeWorkflowStepSessions.delete(task.id);
this.deleteActiveWorkflowStepSession(task.id);
}
// Reviewer subagents run in their own sessions outside `activeSessions`
// and `activeStepExecutors`, so the loops above don't reach them.
@@ -1408,7 +1456,7 @@ export class TaskExecutor {
);
}
workflowSession.dispose();
this.activeWorkflowStepSessions.delete(task.id);
this.deleteActiveWorkflowStepSession(task.id);
this.loopRecoveryState.delete(task.id);
this.spawnedAgents.delete(task.id);
this.stuckAborted.delete(task.id);
@@ -1628,7 +1676,7 @@ export class TaskExecutor {
});
}
workflowSession.dispose();
this.activeWorkflowStepSessions.delete(taskId);
this.deleteActiveWorkflowStepSession(taskId);
this.loopRecoveryState.delete(taskId);
this.spawnedAgents.delete(taskId);
this.stuckAborted.delete(taskId);
@@ -2187,7 +2235,7 @@ export class TaskExecutor {
if (this.activeSessions.has(task.id)) {
const { session: activeSession } = this.activeSessions.get(task.id)!;
activeSession.dispose();
this.activeSessions.delete(task.id);
this.deleteActiveSession(task.id);
}
// Untrack from stuck detector
@@ -2931,7 +2979,7 @@ export class TaskExecutor {
});
},
});
this.activeStepExecutors.set(task.id, stepExecutor);
this.setActiveStepExecutor(task.id, stepExecutor, worktreePath);
const stepWork = async () => {
const results = await stepExecutor.executeAll();
@@ -3250,7 +3298,7 @@ export class TaskExecutor {
} catch (cleanupErr) {
executorLog.warn(`StepSessionExecutor cleanup failed for ${task.id}: ${cleanupErr instanceof Error ? cleanupErr.message : String(cleanupErr)}`);
}
this.activeStepExecutors.delete(task.id);
this.deleteActiveStepExecutor(task.id);
// Stuck-requeue: clean up worktree and move to todo
if (stuckRequeue === true) {
@@ -3544,7 +3592,7 @@ export class TaskExecutor {
seenSteeringIds.add(comment.id);
}
}
this.activeSessions.set(task.id, {
this.setActiveSession(task.id, {
session,
seenSteeringIds,
lastResolvedModelProvider: executorProvider,
@@ -3552,7 +3600,7 @@ export class TaskExecutor {
lastTaskModelProvider: detail.modelProvider,
lastTaskModelId: detail.modelId,
lastAssignedAgentId: detail.assignedAgentId ?? null,
});
}, worktreePath);
let leaseRenewalTimer: ReturnType<typeof setInterval> | undefined;
if (detail.assignedAgentId && detail.checkedOutBy === detail.assignedAgentId) {
@@ -3816,7 +3864,7 @@ export class TaskExecutor {
const reclaimMessage = `${task.id}: worktree/branch reclaimed during no-fn_task_done retry — aborting retry and requeueing`;
executorLog.log(reclaimMessage);
await this.store.logEntry(task.id, reclaimMessage, undefined, this.currentRunContext);
this.activeSessions.delete(task.id);
this.deleteActiveSession(task.id);
this.tokenUsageBaselines.delete(task.id);
session.dispose();
retryAbortedDueToReclaim = true;
@@ -3852,7 +3900,7 @@ export class TaskExecutor {
// Dispose old session and create a fresh one.
// Reset lastAssistantText so the new session's text is tracked cleanly.
lastAssistantText = "";
this.activeSessions.delete(task.id);
this.deleteActiveSession(task.id);
this.tokenUsageBaselines.delete(task.id);
session.dispose();
@@ -3893,7 +3941,7 @@ export class TaskExecutor {
session = retrySession;
sessionRef.current = retrySession;
this.activeSessions.set(task.id, {
this.setActiveSession(task.id, {
session: retrySession,
seenSteeringIds,
lastResolvedModelProvider: executorProvider,
@@ -3901,7 +3949,7 @@ export class TaskExecutor {
lastTaskModelProvider: detail.modelProvider,
lastTaskModelId: detail.modelId,
lastAssignedAgentId: detail.assignedAgentId ?? null,
});
}, worktreePath);
stuckDetector?.trackTask(task.id, retrySession);
let retryPrompt: string;
@@ -3959,7 +4007,7 @@ export class TaskExecutor {
branch: null,
baseCommitSha: null,
});
this.activeSessions.delete(task.id);
this.deleteActiveSession(task.id);
this.tokenUsageBaselines.delete(task.id);
retrySession?.dispose();
retryAbortedDueToReclaim = true;
@@ -4089,7 +4137,7 @@ export class TaskExecutor {
if (leaseRenewalTimer) {
clearInterval(leaseRenewalTimer);
}
this.activeSessions.delete(task.id);
this.deleteActiveSession(task.id);
stuckDetector?.untrackTask(task.id);
await agentLogger.flush();
await this.persistTokenUsage(task.id, session).catch((err: unknown) => {
@@ -7074,7 +7122,7 @@ Backward compat fallback: if JSON is unavailable, you may still begin output wit
task.id,
`Workflow step '${workflowStep.name}' using model: ${describeModel(session)}${useOverride && attemptLabel === "primary" ? " (workflow step override)" : ""}${attemptLabel === "fallback" ? " (fallback after timeout)" : ""}`,
);
this.activeWorkflowStepSessions.set(task.id, session);
this.setActiveWorkflowStepSession(task.id, session, worktreePath);
let output = "";
session.subscribe((event) => {
@@ -7175,7 +7223,7 @@ Backward compat fallback: if JSON is unavailable, you may still begin output wit
if (timeoutHandle) clearTimeout(timeoutHandle);
const activeWorkflowStepSession = this.activeWorkflowStepSessions.get(task.id);
if (activeWorkflowStepSession === session) {
this.activeWorkflowStepSessions.delete(task.id);
this.deleteActiveWorkflowStepSession(task.id);
}
// Suppress unused-variable warning; `timedOut` documents intent.
void timedOut;

View File

@@ -49,6 +49,7 @@ import {
createTaskLogTool,
} from "./agent-tools.js";
import { removeWorktree } from "./worktree-backend.js";
import { activeSessionRegistry } from "./active-session-registry.js";
const stepExecLog = createLogger("step-session-executor");
@@ -628,6 +629,37 @@ export class StepSessionExecutor {
private aborted = false;
private maxParallel: number;
private registerActiveStepSession(stepIndex: number, handle: SessionHandle, worktreePath: string): void {
this.registerActiveStepSession(stepIndex, handle, worktreePath);
activeSessionRegistry.registerPath(worktreePath, {
taskId: this.options.taskDetail.id,
kind: "step-session",
ownerKey: `${this.options.taskDetail.id}#step-${stepIndex}`,
});
}
private unregisterActiveStepSession(stepIndex: number, worktreePath: string): void {
this.unregisterActiveStepSession(stepIndex, worktreePath);
activeSessionRegistry.unregisterPath(worktreePath);
}
private registerParallelWorktree(stepIndex: number, worktreePath: string): void {
this.registerParallelWorktree(stepIndex, worktreePath);
activeSessionRegistry.registerPath(worktreePath, {
taskId: this.options.taskDetail.id,
kind: "step-session-parallel",
ownerKey: `${this.options.taskDetail.id}#parallel-${stepIndex}`,
});
}
private unregisterParallelWorktree(stepIndex: number): void {
const path = this.parallelWorktrees.get(stepIndex);
if (path) {
activeSessionRegistry.unregisterPath(path);
}
this.parallelWorktrees.delete(stepIndex);
}
constructor(options: StepSessionExecutorOptions) {
this.options = options;
this.store = options.store ?? (NOOP_TASK_STORE as TaskStore);
@@ -734,6 +766,9 @@ export class StepSessionExecutor {
this.options.stuckTaskDetector?.untrackTask(trackingKey);
}
for (const worktreePath of activeSessionRegistry.pathsForTask(this.options.taskDetail.id)) {
activeSessionRegistry.unregisterPath(worktreePath);
}
this.activeSessions.clear();
}
@@ -778,7 +813,9 @@ export class StepSessionExecutor {
}
}
this.parallelWorktrees.clear();
for (const [stepIdx] of this.parallelWorktrees) {
this.unregisterParallelWorktree(stepIdx);
}
this.parallelBranches.clear();
stepExecLog.log(`Cleanup complete for task ${this.options.taskDetail.id}`);
@@ -1024,7 +1061,7 @@ Follow instructions precisely and avoid unrelated changes.`,
dispose: () => session?.dispose(),
abortBash: () => session?.abortBash(),
};
this.activeSessions.set(stepIndex, handle);
this.registerActiveStepSession(stepIndex, handle, worktreePath);
stuckTaskDetector?.trackTask(trackingKey, { dispose: () => session?.dispose() }, taskDetail.id);
const sessionModel = await describeAgentModel(session);
@@ -1125,7 +1162,7 @@ Follow instructions precisely and avoid unrelated changes.`,
stepExecLog.warn(`Failed to flush agent logs for step ${stepIndex}: ${flushError}`);
}
this.activeSessions.delete(stepIndex);
this.unregisterActiveStepSession(stepIndex, worktreePath);
stuckTaskDetector?.untrackTask(trackingKey);
try {
session?.dispose();
@@ -1264,7 +1301,7 @@ Follow instructions precisely and avoid unrelated changes.`,
} catch (err) {
stepExecLog.warn(`Failed to clean up worktree for step ${stepIdx}: ${err}`);
} finally {
this.parallelWorktrees.delete(stepIdx);
this.unregisterParallelWorktree(stepIdx);
this.parallelBranches.delete(stepIdx);
}
}
@@ -1305,7 +1342,7 @@ Follow instructions precisely and avoid unrelated changes.`,
throw err;
}
this.parallelWorktrees.set(stepIndex, worktreePath);
this.registerParallelWorktree(stepIndex, worktreePath);
this.parallelBranches.set(stepIndex, branchName);
return worktreePath;