feat(KB-161): terminate active agent sessions on engine pause

- Add enginePaused agent termination to executor (kills sessions, moves tasks back to todo)
- Add enginePaused agent termination to triage processor (kills sessions, clears specifying status)
- Update enginePaused JSDoc to reflect new termination behavior vs old graceful drain
- Add unit tests for executor/triage pause termination and integration tests for restart flow
- Include changeset for engine pause behavior change
This commit is contained in:
Dustin Byrne
2026-03-28 09:15:50 -04:00
parent c4f3c39c64
commit ceb379d3c9
8 changed files with 531 additions and 22 deletions

View File

@@ -344,6 +344,48 @@ describe("runDashboard — immediate resume on unpause", () => {
});
});
describe("runDashboard — engine pause/unpause cycle", () => {
let mockStore: ReturnType<typeof makeMockStore>;
beforeEach(async () => {
capturedExecutorOpts = undefined;
vi.clearAllMocks();
mockStore = makeMockStore();
const { TaskStore } = await import("@kb/core");
(TaskStore as ReturnType<typeof vi.fn>).mockImplementation(() => mockStore);
const engine = await import("@kb/engine");
(engine.aiMergeTask as ReturnType<typeof vi.fn>).mockImplementation(() =>
Promise.resolve({ merged: true }),
);
});
it("calls executor.resumeOrphaned() when enginePaused transitions true → false", async () => {
const { TaskExecutor } = await import("@kb/engine");
const resumeOrphaned = vi.fn().mockResolvedValue(undefined);
(TaskExecutor as unknown as ReturnType<typeof vi.fn>).mockImplementation(
(_store: unknown, _cwd: unknown, opts: unknown) => {
capturedExecutorOpts = opts as Record<string, unknown>;
return { resumeOrphaned };
},
);
await runDashboard(0, { open: false });
// Clear the startup call to resumeOrphaned
resumeOrphaned.mockClear();
// Trigger engine unpause event
mockStore.emit("settings:updated", {
settings: { enginePaused: false, maxConcurrent: 1, autoMerge: false },
previous: { enginePaused: true },
});
await new Promise((r) => setTimeout(r, 50));
expect(resumeOrphaned).toHaveBeenCalled();
});
});
describe("runDashboard — port fallback on EADDRINUSE", () => {
let consoleSpy: ReturnType<typeof vi.spyOn>;

View File

@@ -101,12 +101,17 @@ export interface Settings {
* global emergency stop for the entire AI engine.
* Individual per-task pause flags are unaffected. */
globalPause?: boolean;
/** Soft pause: when true, the scheduler and triage processor stop
/** Engine pause: when true, the scheduler and triage processor stop
* dispatching **new** work (scheduling, triage specification, and
* auto-merge), but in-flight agent sessions are allowed to finish
* gracefully. Use this to drain the queue without killing active work.
* Has no effect when {@link globalPause} is also true (hard stop
* takes precedence). */
* auto-merge), and all active agent sessions (executor and triage) are
* terminated — matching the {@link globalPause} behavior for agent
* lifecycle. Terminated tasks are moved back to `todo` (executor) or
* have their `specifying` status cleared (triage) so they can resume
* cleanly when the engine is unpaused. Remains independent from
* `globalPause` for scheduling control: `globalPause` is the emergency
* stop, while `enginePaused` is the normal on/off toggle for the AI
* engine. Has no additional effect when {@link globalPause} is also
* true (hard stop already covers agent termination). */
enginePaused?: boolean;
/** Maximum number of concurrent AI agents across all activity types
* (triage specification, task execution, and merge operations). */

View File

@@ -1316,6 +1316,142 @@ describe("TaskExecutor global pause behavior", () => {
});
});
describe("TaskExecutor enginePaused agent termination", () => {
beforeEach(() => {
vi.clearAllMocks();
mockedExistsSync.mockReturnValue(true);
});
it("disposes all active sessions when enginePaused transitions false→true", async () => {
const store = createMockStore();
const disposeFn1 = vi.fn();
const disposeFn2 = vi.fn();
let callCount = 0;
mockedCreateHaiAgent.mockImplementation(async () => {
callCount++;
const dispose = callCount === 1 ? disposeFn1 : disposeFn2;
return {
session: {
prompt: vi.fn().mockImplementation(async () => {
// Fire engine pause when the second task starts
if (callCount === 2) {
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
}
throw new Error("Session terminated");
}),
dispose,
},
} as any;
});
const executor = new TaskExecutor(store, "/tmp/test");
await Promise.all([
executor.execute({
id: "KB-001", title: "T1", description: "T", column: "in-progress",
dependencies: [], steps: [], currentStep: 0, log: [],
createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(),
}),
executor.execute({
id: "KB-002", title: "T2", description: "T", column: "in-progress",
dependencies: [], steps: [], currentStep: 0, log: [],
createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(),
}),
]);
// Both tasks should be moved to todo (not marked as failed)
expect(store.moveTask).toHaveBeenCalledWith("KB-001", "todo");
expect(store.moveTask).toHaveBeenCalledWith("KB-002", "todo");
expect(store.updateTask).not.toHaveBeenCalledWith("KB-001", { status: "failed" });
expect(store.updateTask).not.toHaveBeenCalledWith("KB-002", { status: "failed" });
});
it("moves terminated tasks to todo (not failed)", async () => {
const store = createMockStore();
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
throw new Error("Session terminated");
}),
dispose: vi.fn(),
},
} as any));
const executor = new TaskExecutor(store, "/tmp/test");
await executor.execute({
id: "KB-001", title: "Test", description: "T", column: "in-progress",
dependencies: [], steps: [], currentStep: 0, log: [],
createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(),
});
expect(store.moveTask).toHaveBeenCalledWith("KB-001", "todo");
expect(store.updateTask).not.toHaveBeenCalledWith("KB-001", { status: "failed" });
});
it("takes no action when enginePaused stays false (false→false)", async () => {
const store = createMockStore();
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
store._trigger("settings:updated", {
settings: { enginePaused: false },
previous: { enginePaused: false },
});
}),
dispose: vi.fn(),
},
} as any));
const executor = new TaskExecutor(store, "/tmp/test");
await executor.execute({
id: "KB-001", title: "Test", description: "T", column: "in-progress",
dependencies: [], steps: [], currentStep: 0, log: [],
createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(),
});
// Should move to in-review (normal completion), not todo
expect(store.moveTask).toHaveBeenCalledWith("KB-001", "in-review");
expect(store.moveTask).not.toHaveBeenCalledWith("KB-001", "todo");
});
it("takes no action when enginePaused stays true (true→true)", async () => {
const store = createMockStore();
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: true },
});
}),
dispose: vi.fn(),
},
} as any));
const executor = new TaskExecutor(store, "/tmp/test");
await executor.execute({
id: "KB-001", title: "Test", description: "T", column: "in-progress",
dependencies: [], steps: [], currentStep: 0, log: [],
createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(),
});
// Should move to in-review (normal completion), not todo
expect(store.moveTask).toHaveBeenCalledWith("KB-001", "in-review");
expect(store.moveTask).not.toHaveBeenCalledWith("KB-001", "todo");
});
});
// ── Code review verdict enforcement tests ────────────────────────────
const mockedReviewStep = vi.mocked(mockedReviewStepFn);

View File

@@ -172,7 +172,7 @@ export class TaskExecutor {
* Listens for `task:moved` to auto-execute tasks moved to `in-progress`,
* `task:updated` to terminate agent sessions when individual tasks are paused,
* and `settings:updated` to terminate **all** active agent sessions when
* `globalPause` transitions from `false` to `true` (emergency stop).
* `globalPause` or `enginePaused` transitions from `false` to `true`.
* Paused tasks are moved back to `todo` rather than marked as `failed`.
*/
constructor(
@@ -208,6 +208,18 @@ export class TaskExecutor {
}
}
});
// When enginePaused transitions from false → true, terminate all active agent sessions.
// Same pattern as globalPause: agents are killed, tasks moved to todo (not failed).
store.on("settings:updated", ({ settings, previous }) => {
if (settings.enginePaused && !previous.enginePaused) {
for (const [taskId, session] of this.activeSessions) {
executorLog.log(`Engine pause — terminating agent session for ${taskId}`);
this.pausedAborted.add(taskId);
session.dispose();
}
}
});
}
/**

View File

@@ -80,6 +80,9 @@ function createMockStore(overrides: Record<string, any> = {}) {
return makeTask("KB-NEW", "triage");
}),
deleteTask: vi.fn().mockResolvedValue(undefined),
_trigger(event: string, ...args: any[]) {
for (const fn of listeners.get(event) || []) fn(...args);
},
_listeners: listeners,
...overrides,
} as any;
@@ -910,3 +913,185 @@ describe("Edge case: worktree deleted between scan and acquire", () => {
expect(pool.size).toBe(0);
});
});
// ── Engine pause/unpause cycle integration tests ──────────────────────────
describe("Engine pause/unpause cycle", () => {
it("executor: agents terminated on pause, tasks moved to todo, resume picks them up", async () => {
const store = createMockStore();
const task = makeTask("KB-EP1", "in-progress");
store.getTask.mockResolvedValue(makeTaskDetail("KB-EP1", "in-progress"));
// Agent waits for engine pause, then throws (simulating session termination)
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
// Trigger engine pause
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
throw new Error("Session terminated");
}),
dispose: vi.fn(),
},
} as any));
const executor = new TaskExecutor(store, "/tmp/test");
await executor.execute(task);
// Task should be moved to todo (not failed)
expect(store.moveTask).toHaveBeenCalledWith("KB-EP1", "todo");
expect(store.updateTask).not.toHaveBeenCalledWith("KB-EP1", { status: "failed" });
// Now simulate unpause: resumeOrphaned picks up in-progress tasks
store.moveTask.mockClear();
mockedCreateHaiAgent.mockClear();
// After unpause, the task would be in todo; scheduler would move it to in-progress
// Here we test that resumeOrphaned can pick up orphaned in-progress tasks
const resumeTask = makeTask("KB-EP1", "in-progress");
store.listTasks.mockResolvedValue([resumeTask]);
mockedCreateHaiAgent.mockResolvedValue({
session: {
prompt: vi.fn().mockResolvedValue(undefined),
dispose: vi.fn(),
},
} as any);
await executor.resumeOrphaned();
await new Promise((r) => setTimeout(r, 50));
// Agent should be created again for the resumed task
expect(mockedCreateHaiAgent).toHaveBeenCalled();
});
it("triage: agents terminated on pause, status cleared, resume picks up tasks", async () => {
const store = createMockStore();
// Agent waits for engine pause, then throws
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
throw new Error("Session terminated");
}),
dispose: vi.fn(),
},
} as any));
const triage = new TriageProcessor(store, "/tmp/test");
await triage.specifyTask(makeTask("KB-EP2", "triage"));
// Status should be cleared (not reported as error)
expect(store.updateTask).toHaveBeenCalledWith("KB-EP2", { status: null });
// Now simulate unpause: triage poll picks up unspecified tasks
store.updateTask.mockClear();
mockedCreateHaiAgent.mockClear();
const triageTask = makeTask("KB-EP2", "triage");
store.listTasks.mockResolvedValue([triageTask]);
store.getSettings.mockResolvedValue({ ...DEFAULT_SETTINGS, enginePaused: false });
mockedCreateHaiAgent.mockResolvedValue({
session: {
prompt: vi.fn().mockResolvedValue(undefined),
dispose: vi.fn(),
},
} as any);
// Trigger a poll — triage should process the task again
(triage as any).running = true;
await (triage as any).poll();
await new Promise((r) => setTimeout(r, 50));
// Task should be picked up for specification
expect(store.updateTask).toHaveBeenCalledWith("KB-EP2", { status: "specifying" });
});
it("scheduler resumes on unpause: schedule() runs when enginePaused goes true→false", async () => {
const store = createMockStore();
const todoTask = makeTask("KB-EP3", "todo");
store.listTasks.mockResolvedValue([todoTask]);
store.getSettings.mockResolvedValue({ ...DEFAULT_SETTINGS, enginePaused: false });
store.parseFileScopeFromPrompt.mockResolvedValue([]);
const onSchedule = vi.fn();
const scheduler = new Scheduler(store, {
maxConcurrent: 2,
maxWorktrees: 4,
onSchedule,
});
scheduler.start();
await new Promise((r) => setTimeout(r, 50));
// Scheduler should have moved todo task to in-progress
expect(store.moveTask).toHaveBeenCalledWith("KB-EP3", "in-progress");
// Now simulate engine pause then unpause
store.moveTask.mockClear();
onSchedule.mockClear();
// During pause, scheduler halts
store.getSettings.mockResolvedValue({ ...DEFAULT_SETTINGS, enginePaused: true });
// Add a new todo task
const newTask = makeTask("KB-EP4", "todo");
store.listTasks.mockResolvedValue([newTask]);
// Unpause — trigger settings:updated to wake the scheduler
store.getSettings.mockResolvedValue({ ...DEFAULT_SETTINGS, enginePaused: false });
store._trigger("settings:updated", {
settings: { ...DEFAULT_SETTINGS, enginePaused: false },
previous: { ...DEFAULT_SETTINGS, enginePaused: true },
});
await new Promise((r) => setTimeout(r, 100));
scheduler.stop();
// The new task should have been scheduled after unpause
expect(store.moveTask).toHaveBeenCalledWith("KB-EP4", "in-progress");
});
it("concurrency slots freed on pause: semaphore released when agents are terminated", async () => {
const sem = new AgentSemaphore(1); // Only 1 concurrent slot
const store = createMockStore();
store.getTask.mockResolvedValue(makeTaskDetail("KB-EP5", "in-progress"));
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
// Trigger engine pause while the agent holds the semaphore slot
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
throw new Error("Session terminated");
}),
dispose: vi.fn(),
},
} as any));
const executor = new TaskExecutor(store, "/tmp/test", { semaphore: sem });
// Execute — this acquires the semaphore slot, runs the agent, then releases on termination
await executor.execute(makeTask("KB-EP5", "in-progress"));
// After termination, the semaphore slot should be freed.
// Verify by running a new task through the semaphore — it should not block.
let secondSlotAcquired = false;
const runPromise = sem.run(async () => {
secondSlotAcquired = true;
}, 10);
// If the slot was properly released, this resolves immediately
await runPromise;
expect(secondSlotAcquired).toBe(true);
});
});

View File

@@ -1813,3 +1813,119 @@ describe("TriageProcessor global pause agent kill", () => {
);
});
});
describe("TriageProcessor enginePaused agent termination", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("terminates active triage sessions when enginePaused transitions false→true", async () => {
const store = createMockStore();
const disposeFn = vi.fn();
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
// Trigger engine pause while the session is active
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
throw new Error("Session terminated");
}),
dispose: disposeFn,
},
} as any));
const triage = new TriageProcessor(store, "/tmp/test");
await triage.specifyTask({
id: "KB-001",
title: "Test",
description: "Test",
column: "triage",
dependencies: [],
steps: [],
currentStep: 0,
log: [],
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
});
// dispose is called by the engine pause listener and again in finally
expect(disposeFn).toHaveBeenCalled();
});
it("clears specifying status on terminated tasks", async () => {
const store = createMockStore();
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
throw new Error("Session terminated");
}),
dispose: vi.fn(),
},
} as any));
const triage = new TriageProcessor(store, "/tmp/test");
await triage.specifyTask({
id: "KB-001",
title: "Test",
description: "Test",
column: "triage",
dependencies: [],
steps: [],
currentStep: 0,
log: [],
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
});
// Status should be cleared (not reported as error)
expect(store.updateTask).toHaveBeenCalledWith("KB-001", { status: null });
});
it("does not report errors for engine-pause-aborted tasks", async () => {
const store = createMockStore();
const onError = vi.fn();
mockedCreateHaiAgent.mockImplementation(async () => ({
session: {
prompt: vi.fn().mockImplementation(async () => {
store._trigger("settings:updated", {
settings: { enginePaused: true },
previous: { enginePaused: false },
});
throw new Error("Session terminated");
}),
dispose: vi.fn(),
},
} as any));
const triage = new TriageProcessor(store, "/tmp/test", { onSpecifyError: onError });
await triage.specifyTask({
id: "KB-001",
title: "Test",
description: "Test",
column: "triage",
dependencies: [],
steps: [],
currentStep: 0,
log: [],
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
});
// onSpecifyError should NOT be called for engine-pause aborted tasks
expect(onError).not.toHaveBeenCalled();
// Status should be cleared
expect(store.updateTask).toHaveBeenCalledWith("KB-001", { status: null });
});
});

View File

@@ -185,19 +185,19 @@ export class TriageProcessor {
private processing = new Set<string>();
private wasGlobalPaused = false;
private wasEnginePaused = false;
/** Active agent sessions per task, used to terminate on global pause. */
/** Active agent sessions per task, used to terminate on pause. */
private activeSessions = new Map<string, { dispose: () => void }>();
/** Tasks that were aborted due to global pause (to avoid reporting as errors). */
private globalPauseAborted = new Set<string>();
/** Tasks aborted due to globalPause or enginePaused (to avoid reporting as errors). */
private pauseAborted = new Set<string>();
/**
* @param store — Task store instance (also used to listen for `settings:updated` events)
* @param rootDir — Project root directory
* @param options — Processor configuration
*
* Listens for `settings:updated` events: when `globalPause` transitions from
* `false` to `true`, all active triage specification sessions are immediately
* terminated so the engine acts as a true emergency stop.
* Listens for `settings:updated` events: when `globalPause` or `enginePaused`
* transitions from `false` to `true`, all active triage specification sessions
* are immediately terminated so the engine stops all AI activity.
*/
constructor(
private store: TaskStore,
@@ -209,7 +209,19 @@ export class TriageProcessor {
if (settings.globalPause && !previous.globalPause) {
for (const [taskId, session] of this.activeSessions) {
triageLog.log(`Global pause — terminating triage session for ${taskId}`);
this.globalPauseAborted.add(taskId);
this.pauseAborted.add(taskId);
session.dispose();
}
}
});
// When enginePaused transitions from false → true, terminate all active triage sessions.
// Same pattern as globalPause: agents are killed, status cleared (not reported as error).
store.on("settings:updated", ({ settings, previous }) => {
if (settings.enginePaused && !previous.enginePaused) {
for (const [taskId, session] of this.activeSessions) {
triageLog.log(`Engine pause — terminating triage session for ${taskId}`);
this.pauseAborted.add(taskId);
session.dispose();
}
}
@@ -231,14 +243,10 @@ export class TriageProcessor {
});
/**
* Immediate soft-unpause resume: when `enginePaused` transitions from
* Immediate engine-unpause resume: when `enginePaused` transitions from
* `true` to `false`, trigger a triage poll right away instead of
* waiting for the next poll interval. Same pattern as the globalPause
* unpause handler above.
*
* Note: the agent-kill listener above only fires on `globalPause`
* transitions (hard stop). `enginePaused` (soft pause) lets in-flight
* agents finish gracefully.
*/
store.on("settings:updated", ({ settings, previous }) => {
if (previous.enginePaused && !settings.enginePaused && this.running) {
@@ -453,10 +461,10 @@ export class TriageProcessor {
// and specifyTask(). The file is gone, so just log and skip — no point retrying.
if (err.code === "ENOENT") {
triageLog.log(`${task.id} no longer exists — skipping`);
} else if (this.globalPauseAborted.has(task.id)) {
// Global pause — clear specifying status without reporting an error
this.globalPauseAborted.delete(task.id);
triageLog.log(`${task.id} aborted by global pause — clearing status`);
} else if (this.pauseAborted.has(task.id)) {
// Pause (global or engine) — clear specifying status without reporting an error
this.pauseAborted.delete(task.id);
triageLog.log(`${task.id} aborted by pause — clearing status`);
await this.store.updateTask(task.id, { status: null }).catch(() => {});
} else {
// Check if the error is a usage-limit error and trigger global pause