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:
committed by
gsxdsm
parent
f4aa6d7b8b
commit
69d7b282cf
@@ -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();
|
||||
});
|
||||
});
|
||||
53
packages/engine/src/active-session-registry.ts
Normal file
53
packages/engine/src/active-session-registry.ts
Normal 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();
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user