fix(KB-148): fix triage concurrency and scheduler double-counting

- Add re-entrance guard to triage poll() to prevent overlapping poll cycles
- Move status update to 'specifying' inside semaphore callback so queued tasks don't appear active
- Fix scheduler double-counting specifying tasks when semaphore is present
- Add tests for poll re-entrance guard, semaphore-aware status transitions, and scheduler slot counting
- Add changeset for patch release
This commit is contained in:
Dustin Byrne
2026-03-28 01:24:02 -04:00
parent 90764b9657
commit 72a8953e4e
5 changed files with 246 additions and 7 deletions

View File

@@ -1,5 +1,6 @@
import { describe, it, expect, vi, beforeEach } from "vitest";
import { Scheduler } from "./scheduler.js";
import { AgentSemaphore } from "./concurrency.js";
function makeTask(overrides: Record<string, unknown> = {}) {
return {
@@ -962,3 +963,97 @@ describe("Scheduler in-review worktrees do not count against maxWorktrees", () =
expect(store.moveTask).not.toHaveBeenCalled();
});
});
describe("Scheduler semaphore-aware slot counting", () => {
beforeEach(() => {
vi.clearAllMocks();
});
async function runSchedule(scheduler: Scheduler): Promise<void> {
(scheduler as any).running = true;
await scheduler.schedule();
}
it("with semaphore, specifying tasks do not double-count against agentSlots", async () => {
// Scenario: maxConcurrent=2, 1 specifying task holding a semaphore slot,
// 0 in-progress. Semaphore has 1 available slot. Should allow 1 new task.
const sem = new AgentSemaphore(2);
// Simulate a specifying task holding a slot
await sem.acquire();
const tasks = [
makeTask({ id: "KB-001", column: "triage", status: "specifying" }),
makeTask({ id: "KB-002", column: "todo" }),
];
const store = createMockStore(tasks);
store.getSettings.mockResolvedValue({
maxConcurrent: 2,
maxWorktrees: 4,
pollIntervalMs: 15000,
groupOverlappingFiles: false,
autoMerge: false,
});
const scheduler = new Scheduler(store, { maxConcurrent: 2, semaphore: sem });
await runSchedule(scheduler);
// With semaphore: agentSlots = inProgress(0), available = min(2-0, 4-0, 1) = 1
// KB-002 should be scheduled
expect(store.moveTask).toHaveBeenCalledWith("KB-002", "in-progress");
sem.release();
});
it("with semaphore, fully occupied semaphore blocks scheduling even with no in-progress tasks", async () => {
const sem = new AgentSemaphore(2);
// Both slots held (e.g., by two specifying agents)
await sem.acquire();
await sem.acquire();
const tasks = [
makeTask({ id: "KB-001", column: "triage", status: "specifying" }),
makeTask({ id: "KB-002", column: "triage", status: "specifying" }),
makeTask({ id: "KB-003", column: "todo" }),
];
const store = createMockStore(tasks);
store.getSettings.mockResolvedValue({
maxConcurrent: 2,
maxWorktrees: 4,
pollIntervalMs: 15000,
groupOverlappingFiles: false,
autoMerge: false,
});
const scheduler = new Scheduler(store, { maxConcurrent: 2, semaphore: sem });
await runSchedule(scheduler);
// semaphoreAvailable = 0 so nothing can start
expect(store.moveTask).not.toHaveBeenCalled();
sem.release();
sem.release();
});
it("without semaphore, specifying tasks still reduce available slots (backward compat)", async () => {
const tasks = [
makeTask({ id: "KB-001", column: "triage", status: "specifying" }),
makeTask({ id: "KB-002", column: "triage", status: "specifying" }),
makeTask({ id: "KB-003", column: "todo" }),
];
const store = createMockStore(tasks);
store.getSettings.mockResolvedValue({
maxConcurrent: 2,
maxWorktrees: 4,
pollIntervalMs: 15000,
groupOverlappingFiles: false,
autoMerge: false,
});
// No semaphore provided — fallback path
const scheduler = new Scheduler(store, { maxConcurrent: 2 });
await runSchedule(scheduler);
// agentSlots = 0 + 2(specifying) = 2, available = min(2-2, 4-0, Inf) = 0
expect(store.moveTask).not.toHaveBeenCalled();
});
});

View File

@@ -218,7 +218,15 @@ export class Scheduler {
const specifying = tasks.filter(
(t) => t.column === "triage" && t.status === "specifying" && !t.paused,
);
const agentSlots = inProgress.length + specifying.length;
// When a semaphore is provided, it is the single source of truth for
// global concurrency — its availableCount already accounts for ALL
// slot holders (executors, specifiers, mergers). Counting specifying
// tasks in agentSlots as well would double-count them. Without a
// semaphore (fallback mode), count specifying tasks directly.
const agentSlots = this.options.semaphore
? inProgress.length
: inProgress.length + specifying.length;
// When a semaphore is provided, factor in its available slots so we
// don't schedule more tasks than the global limit allows. Triage and

View File

@@ -179,6 +179,128 @@ describe("TriageProcessor with semaphore", () => {
expect(maxConcurrent).toBe(1);
expect(sem.activeCount).toBe(0);
});
it("does not set status 'specifying' until semaphore slot is acquired", async () => {
const sem = new AgentSemaphore(1);
const store = createMockStore();
// Acquire the only slot so specifyTask must wait
await sem.acquire();
let agentStarted = false;
mockedCreateHaiAgent.mockImplementation(async () => {
agentStarted = true;
return {
session: {
prompt: vi.fn().mockResolvedValue(undefined),
dispose: vi.fn(),
},
} as any;
});
const triage = new TriageProcessor(store, "/tmp/test", { semaphore: sem });
const task = {
id: "KB-001",
title: "Test",
description: "Test",
column: "triage" as const,
dependencies: [],
steps: [],
currentStep: 0,
log: [],
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
};
// Start specifyTask — it will queue on the semaphore
const specPromise = triage.specifyTask(task);
await new Promise((r) => setTimeout(r, 20));
// While queued, status should NOT have been set to "specifying"
const specifyingCalls = store.updateTask.mock.calls.filter(
(c: any[]) => c[1]?.status === "specifying",
);
expect(specifyingCalls).toHaveLength(0);
expect(agentStarted).toBe(false);
// Release the slot — now specifyTask should proceed
sem.release();
await specPromise;
// Now status should have been set to "specifying"
expect(store.updateTask).toHaveBeenCalledWith("KB-001", { status: "specifying" });
expect(agentStarted).toBe(true);
});
});
describe("TriageProcessor poll re-entrance guard", () => {
beforeEach(() => {
vi.clearAllMocks();
});
it("prevents overlapping poll() calls — second call is a no-op", async () => {
const store = createMockStore([
{
id: "KB-001",
title: "Test",
description: "Test",
column: "triage",
dependencies: [],
steps: [],
currentStep: 0,
log: [],
createdAt: new Date().toISOString(),
updatedAt: new Date().toISOString(),
},
]);
// Make specifyTask slow so the first poll is still running when the second fires
let resolveAgent: (() => void) | undefined;
mockedCreateHaiAgent.mockImplementation(async () => {
await new Promise<void>((r) => {
resolveAgent = r;
});
return {
session: {
prompt: vi.fn().mockResolvedValue(undefined),
dispose: vi.fn(),
},
} as any;
});
const triage = new TriageProcessor(store, "/tmp/test");
(triage as any).running = true;
// Start first poll (will block in createKbAgent)
const poll1 = (triage as any).poll();
// Allow microtasks to run so poll1 gets into specifyTask
await new Promise((r) => setTimeout(r, 10));
// Start second poll — should return immediately due to guard
const poll2 = (triage as any).poll();
await poll2;
// listTasks should only have been called once (first poll)
expect(store.listTasks).toHaveBeenCalledTimes(1);
// Resolve the agent to let the first poll finish
resolveAgent?.();
await poll1;
});
it("allows a new poll() after the previous one completes", async () => {
const store = createMockStore([]);
const triage = new TriageProcessor(store, "/tmp/test");
(triage as any).running = true;
await (triage as any).poll();
await (triage as any).poll();
// Both polls should have called listTasks (sequentially, guard released)
expect(store.listTasks).toHaveBeenCalledTimes(2);
});
});
describe("TriageProcessor dynamic poll interval", () => {
@@ -660,13 +782,15 @@ describe("TriageProcessor deleted task handling", () => {
await triage.specifyTask(dummyTask);
expect(onError).not.toHaveBeenCalled();
// updateTask called once for "specifying", but NOT for status reset (ENOENT path skips it)
expect(store.updateTask).toHaveBeenCalledTimes(1);
// getTask throws ENOENT before updateTask(status: "specifying") is reached
// (status update moved inside agentWork, after semaphore acquisition)
expect(store.updateTask).toHaveBeenCalledTimes(0);
});
it("cleans up processing Set on ENOENT so task is not stuck", async () => {
const store = createMockStore();
store.updateTask.mockRejectedValueOnce(createEnoentError());
// getTask throws ENOENT (task deleted between poll and specify)
store.getTask.mockRejectedValueOnce(createEnoentError());
const triage = new TriageProcessor(store, "/tmp/test", {});
@@ -868,7 +992,7 @@ describe("TriageProcessor dependency parsing", () => {
// Verify updateTask was called with dependencies, size, and reviewLevel
const updateCalls = store.updateTask.mock.calls;
// First call is { status: "specifying" }, second is the post-parse call
// First call is { status: "specifying" } (inside agentWork), second is the post-parse call
expect(updateCalls.length).toBeGreaterThanOrEqual(2);
const postParseCAll = updateCalls[1];
expect(postParseCAll[0]).toBe("KB-001");

View File

@@ -178,6 +178,7 @@ export interface TriageProcessorOptions {
*/
export class TriageProcessor {
private running = false;
private polling = false;
private pollInterval: ReturnType<typeof setInterval> | null = null;
/** The interval (ms) of the currently active `setInterval` timer. */
private activePollMs: number | null = null;
@@ -229,6 +230,8 @@ export class TriageProcessor {
private async poll(): Promise<void> {
if (!this.running) return;
if (this.polling) return;
this.polling = true;
try {
const settings = await this.store.getSettings();
@@ -254,6 +257,8 @@ export class TriageProcessor {
}
} catch (err) {
triageLog.error("Poll error:", err);
} finally {
this.polling = false;
}
}
@@ -265,13 +270,15 @@ export class TriageProcessor {
this.options.onSpecifyStart?.(task);
try {
// Set status inside try so ENOENT (task deleted between poll and specify) is caught
await this.store.updateTask(task.id, { status: "specifying" });
const detail = await this.store.getTask(task.id);
const settings = await this.store.getSettings();
const promptPath = `.kb/tasks/${task.id}/PROMPT.md`;
const agentWork = async () => {
// Set status only after the semaphore slot has been acquired, so
// tasks waiting in the queue don't appear as "specifying".
await this.store.updateTask(task.id, { status: "specifying" });
const agentLogger = new AgentLogger({
store: this.store,
taskId: task.id,