diff options
| author | Adam Malczewski <[email protected]> | 2026-06-28 15:31:49 +0900 |
|---|---|---|
| committer | Adam Malczewski <[email protected]> | 2026-06-28 15:31:49 +0900 |
| commit | b60586285863f8bb82242a9df49c4d47e1235755 (patch) | |
| tree | dd38669dbd8092987bc50d16dcf523c68d43c460 /packages/provider-concurrency/src | |
| parent | fb4a9217b55dd3ba11670104ac23536416d36940 (diff) | |
| parent | 076edf7d1dfc4dc818f173f751dcb1e57b5baaeb (diff) | |
| download | dispatch-b60586285863f8bb82242a9df49c4d47e1235755.tar.gz dispatch-b60586285863f8bb82242a9df49c4d47e1235755.zip | |
Merge branch 'feature/workspace-star' into predev
# Conflicts:
# packages/provider-concurrency/src/concurrency-manager.ts
# packages/provider-concurrency/src/extension.ts
Diffstat (limited to 'packages/provider-concurrency/src')
5 files changed, 478 insertions, 49 deletions
diff --git a/packages/provider-concurrency/src/concurrency-manager.test.ts b/packages/provider-concurrency/src/concurrency-manager.test.ts index 185c6c2..357a5d1 100644 --- a/packages/provider-concurrency/src/concurrency-manager.test.ts +++ b/packages/provider-concurrency/src/concurrency-manager.test.ts @@ -90,7 +90,7 @@ function createManager(opts?: { describe("createConcurrencyManager", () => { it("returns no-op release for providers with no configured limit", async () => { const { manager } = createManager(); - const release = await manager.acquire("unknown", "conv1", 0); + const release = await manager.acquire("unknown", "conv1", "default", 0); expect(typeof release).toBe("function"); // No state → release is a no-op, no error. release(); @@ -101,7 +101,7 @@ describe("createConcurrencyManager", () => { const { manager } = createManager(); manager.setLimit("umans", 4); - const release1 = await manager.acquire("umans", "conv1", 0); + const release1 = await manager.acquire("umans", "conv1", "default", 0); const status = manager.getStatus("umans"); expect(status).toEqual({ providerId: "umans", @@ -120,11 +120,11 @@ describe("createConcurrencyManager", () => { const { manager } = createManager(); manager.setLimit("umans", 1); - const release1 = await manager.acquire("umans", "conv1", 100); + const release1 = await manager.acquire("umans", "conv1", "default", 100); // Second request should block (at limit). let resolved = false; - const promise2 = manager.acquire("umans", "conv2", 200).then((r) => { + const promise2 = manager.acquire("umans", "conv2", "default", 200).then((r) => { resolved = true; return r; }); @@ -150,13 +150,13 @@ describe("createConcurrencyManager", () => { manager.setLimit("umans", 1); // Hold the single slot. - const release0 = await manager.acquire("umans", "holder", 0); + const release0 = await manager.acquire("umans", "holder", "default", 0); // Three agents queue with different prompt start times. // Agent C started latest (t=300), Agent A started earliest (t=100). const results: string[] = []; const acquireAndRecord = (conv: string, promptAt: number) => - manager.acquire("umans", conv, promptAt).then((r) => { + manager.acquire("umans", conv, "default", promptAt).then((r) => { results.push(conv); return r; }); @@ -193,7 +193,7 @@ describe("createConcurrencyManager", () => { const { manager, timers } = createManager(); manager.setLimit("umans", 1); - const release1 = await manager.acquire("umans", "conv1", 0); + const release1 = await manager.acquire("umans", "conv1", "default", 0); release1(); // Simulate a 429 → queue pauses. @@ -204,7 +204,7 @@ describe("createConcurrencyManager", () => { // A new acquire should block (paused, even though under limit). let resolved = false; - const promise = manager.acquire("umans", "conv2", 0).then((r) => { + const promise = manager.acquire("umans", "conv2", "default", 0).then((r) => { resolved = true; return r; }); @@ -233,7 +233,7 @@ describe("createConcurrencyManager", () => { const { manager, timers } = createManager(); manager.setLimit("umans", 1); - const release = await manager.acquire("umans", "conv1", 0); + const release = await manager.acquire("umans", "conv1", "default", 0); expect(manager.getStatus("umans")?.inFlight).toBe(1); // Advance past the slot timeout (5000ms) and fire the watchdog. @@ -253,11 +253,11 @@ describe("createConcurrencyManager", () => { manager.setLimit("umans", 1); // Hold the slot. - await manager.acquire("umans", "holder", 0); + await manager.acquire("umans", "holder", "default", 0); // Queue a waiter. let resolved = false; - const promise = manager.acquire("umans", "waiter", 10).then((r) => { + const promise = manager.acquire("umans", "waiter", "default", 10).then((r) => { resolved = true; return r; }); @@ -280,11 +280,11 @@ describe("createConcurrencyManager", () => { const { manager } = createManager(); manager.setLimit("umans", 1); - const release1 = await manager.acquire("umans", "conv1", 0); + const release1 = await manager.acquire("umans", "conv1", "default", 0); // Queue a waiter. let resolved = false; - const promise = manager.acquire("umans", "conv2", 100).then((r) => { + const promise = manager.acquire("umans", "conv2", "default", 100).then((r) => { resolved = true; return r; }); @@ -307,11 +307,11 @@ describe("createConcurrencyManager", () => { const { manager } = createManager(); manager.setLimit("umans", 1); - const release1 = await manager.acquire("umans", "conv1", 0); + const release1 = await manager.acquire("umans", "conv1", "default", 0); // Queue two waiters. - const p2 = manager.acquire("umans", "conv2", 100); - const p3 = manager.acquire("umans", "conv3", 200); + const p2 = manager.acquire("umans", "conv2", "default", 100); + const p3 = manager.acquire("umans", "conv3", "default", 200); await Promise.resolve(); await Promise.resolve(); @@ -369,7 +369,7 @@ describe("createConcurrencyManager", () => { const { manager } = createManager(); manager.setLimit("umans", 2); - const release = await manager.acquire("umans", "conv1", 0); + const release = await manager.acquire("umans", "conv1", "default", 0); expect(manager.getStatus("umans")?.inFlight).toBe(1); release(); @@ -385,9 +385,9 @@ describe("createConcurrencyManager", () => { manager.setLimit("umans", 3); const releases = await Promise.all([ - manager.acquire("umans", "conv1", 0), - manager.acquire("umans", "conv2", 0), - manager.acquire("umans", "conv3", 0), + manager.acquire("umans", "conv1", "default", 0), + manager.acquire("umans", "conv2", "default", 0), + manager.acquire("umans", "conv3", "default", 0), ]); expect(manager.getStatus("umans")?.inFlight).toBe(3); @@ -402,12 +402,12 @@ describe("createConcurrencyManager", () => { const { manager, timers } = createManager({ releaseCooldownMs: 200 }); manager.setLimit("umans", 1); - const release1 = await manager.acquire("umans", "conv1", 0); + const release1 = await manager.acquire("umans", "conv1", "default", 0); expect(manager.getStatus("umans")?.inFlight).toBe(1); // Queue a waiter. let resolved = false; - const promise2 = manager.acquire("umans", "conv2", 100).then((r) => { + const promise2 = manager.acquire("umans", "conv2", "default", 100).then((r) => { resolved = true; return r; }); @@ -436,7 +436,7 @@ describe("createConcurrencyManager", () => { const { manager, timers } = createManager({ releaseCooldownMs: 200 }); manager.setLimit("umans", 2); - const release = await manager.acquire("umans", "conv1", 0); + const release = await manager.acquire("umans", "conv1", "default", 0); expect(manager.getStatus("umans")?.inFlight).toBe(1); release(); @@ -454,7 +454,7 @@ describe("createConcurrencyManager", () => { const { manager } = createManager({ releaseCooldownMs: 200 }); manager.setLimit("umans", 1); // Acquire + release to schedule a cooldown timer. - manager.acquire("umans", "conv1", 0).then((release) => { + manager.acquire("umans", "conv1", "default", 0).then((release) => { release(); // Now there's a pending cooldown timer — destroy should clean it up. expect(() => manager.destroy()).not.toThrow(); @@ -466,11 +466,11 @@ describe("createConcurrencyManager", () => { manager.setLimit("umans", 1); // Hold the single slot. - const release1 = await manager.acquire("umans", "conv1", 0); + const release1 = await manager.acquire("umans", "conv1", "default", 0); // Second request should trigger onQueued. let queuedCalled = false; - const promise = manager.acquire("umans", "conv2", 100, () => { + const promise = manager.acquire("umans", "conv2", "default", 100, () => { queuedCalled = true; }); await Promise.resolve(); @@ -490,7 +490,7 @@ describe("createConcurrencyManager", () => { manager.setLimit("umans", 2); let queuedCalled = false; - const release = await manager.acquire("umans", "conv1", 0, () => { + const release = await manager.acquire("umans", "conv1", "default", 0, () => { queuedCalled = true; }); @@ -966,3 +966,295 @@ describe("createConcurrencyManager", () => { expect(status?.notice).toBeUndefined(); }); }); + +// ─── Starred-workspace priority tests ─────────────────────────────────────── + +describe("starred-workspace priority", () => { + it("starred-workspace agents are admitted before non-starred (regardless of promptStartedAt)", async () => { + // Use a callback backed by a Set so we can star/unstar at runtime. + const starred = new Set<string>(); + const timers = createFakeTimers(); + const manager = createConcurrencyManager({ + now: timers.now, + slotTimeoutMs: 5000, + watchdogIntervalMs: 1000, + defaultPauseMs: 30000, + isWorkspaceStarred: (wsId: string) => starred.has(wsId), + setTimeout: timers.setTimeout, + clearTimeout: timers.clearTimeout, + setInterval: timers.setInterval, + clearInterval: timers.clearInterval, + }); + manager.setLimit("umans", 1); + + // Hold the single slot. + const release0 = await manager.acquire("umans", "holder", "default", 0); + + // Three agents queue: + // - convA (workspace "ws-normal", promptAt=100) — non-starred, earliest + // - convB (workspace "ws-starred", promptAt=200) — starred, later + // - convC (workspace "ws-normal", promptAt=300) — non-starred, latest + starred.add("ws-starred"); + + const results: string[] = []; + const acquireAndRecord = (conv: string, wsId: string, promptAt: number) => + manager.acquire("umans", conv, wsId, promptAt).then((r) => { + results.push(conv); + return r; + }); + + const pA = acquireAndRecord("convA", "ws-normal", 100); + const pB = acquireAndRecord("convB", "ws-starred", 200); + const pC = acquireAndRecord("convC", "ws-normal", 300); + + await Promise.resolve(); + await Promise.resolve(); + expect(results).toEqual([]); // none resolved yet. + + // Release the holder. The starred agent (convB, t=200) should get the + // slot FIRST, even though convA (t=100) started earlier. + release0(); + + const rB = await pB; + expect(results).toEqual(["convB"]); + rB(); + + // Now the oldest non-starred (convA, t=100) should be next. + const rA = await pA; + expect(results).toEqual(["convB", "convA"]); + rA(); + + // Then convC (t=300). + const rC = await pC; + expect(results).toEqual(["convB", "convA", "convC"]); + rC(); + + manager.destroy(); + }); + + it("within the starred group, oldest-agent-first is preserved", async () => { + const starred = new Set<string>(["ws-starred"]); + const timers = createFakeTimers(); + const manager = createConcurrencyManager({ + now: timers.now, + slotTimeoutMs: 5000, + watchdogIntervalMs: 1000, + defaultPauseMs: 30000, + isWorkspaceStarred: (wsId: string) => starred.has(wsId), + setTimeout: timers.setTimeout, + clearTimeout: timers.clearTimeout, + setInterval: timers.setInterval, + clearInterval: timers.clearInterval, + }); + manager.setLimit("umans", 1); + + const release0 = await manager.acquire("umans", "holder", "default", 0); + + const results: string[] = []; + const acquireAndRecord = (conv: string, wsId: string, promptAt: number) => + manager.acquire("umans", conv, wsId, promptAt).then((r) => { + results.push(conv); + return r; + }); + + // Two starred agents: convLate (t=300) queues first, convEarly (t=100) second. + const pLate = acquireAndRecord("convLate", "ws-starred", 300); + const pEarly = acquireAndRecord("convEarly", "ws-starred", 100); + + await Promise.resolve(); + await Promise.resolve(); + + release0(); + + // convEarly (t=100) should win within the starred group (oldest-first). + const rEarly = await pEarly; + expect(results).toEqual(["convEarly"]); + rEarly(); + + const rLate = await pLate; + expect(results).toEqual(["convEarly", "convLate"]); + rLate(); + + manager.destroy(); + }); + + it("starring a workspace while agents are queued re-prioritizes them immediately", async () => { + const timers = createFakeTimers(); + const manager = createConcurrencyManager({ + now: timers.now, + slotTimeoutMs: 5000, + watchdogIntervalMs: 1000, + defaultPauseMs: 30000, + setTimeout: timers.setTimeout, + clearTimeout: timers.clearTimeout, + setInterval: timers.setInterval, + clearInterval: timers.clearInterval, + }); + manager.setLimit("umans", 1); + + const release0 = await manager.acquire("umans", "holder", "default", 0); + + // convA (non-starred, t=100) queues first. + let resolvedA = false; + const pA = manager.acquire("umans", "convA", "ws-normal", 100).then((r) => { + resolvedA = true; + return r; + }); + // convB (non-starred, t=200) queues second. + let resolvedB = false; + const pB = manager.acquire("umans", "convB", "ws-to-star", 200).then((r) => { + resolvedB = true; + return r; + }); + + await Promise.resolve(); + await Promise.resolve(); + expect(resolvedA).toBe(false); + expect(resolvedB).toBe(false); + + // Now star convB's workspace AFTER it's queued. notifyWorkspaceStarred + // updates the internal cache + re-sorts + tries to grant. + manager.notifyWorkspaceStarred("ws-to-star", true); + + // Release the holder — convB (now starred) should jump ahead of convA. + release0(); + + const rB = await pB; + expect(resolvedB).toBe(true); + expect(resolvedA).toBe(false); + rB(); + + // Now convA gets the next slot. + const rA = await pA; + expect(resolvedA).toBe(true); + rA(); + + manager.destroy(); + }); + + it("unstar a workspace demotes its queued agents", async () => { + const timers = createFakeTimers(); + const manager = createConcurrencyManager({ + now: timers.now, + slotTimeoutMs: 5000, + watchdogIntervalMs: 1000, + defaultPauseMs: 30000, + setTimeout: timers.setTimeout, + clearTimeout: timers.clearTimeout, + setInterval: timers.setInterval, + clearInterval: timers.clearInterval, + }); + manager.setLimit("umans", 1); + + // Initially star "ws-starred" via the internal cache. + manager.notifyWorkspaceStarred("ws-starred", true); + + const release0 = await manager.acquire("umans", "holder", "default", 0); + + // convA (starred, t=200) queues first. + const pA = manager.acquire("umans", "convA", "ws-starred", 200); + // convB (non-starred, t=100) queues second but is older. + const pB = manager.acquire("umans", "convB", "ws-normal", 100); + + await Promise.resolve(); + await Promise.resolve(); + + // Unstar convA's workspace — it should now be behind convB (which is older). + manager.notifyWorkspaceStarred("ws-starred", false); + + release0(); + + // convB (t=100, now non-starred but oldest) should win. + const rB = await pB; + expect(rB).toBeDefined(); + rB(); + + const rA = await pA; + rA(); + + manager.destroy(); + }); + + it("notifyWorkspaceStarred re-sorts queues and tries to grant when capacity is free", async () => { + const timers = createFakeTimers(); + const manager = createConcurrencyManager({ + now: timers.now, + slotTimeoutMs: 5000, + watchdogIntervalMs: 1000, + defaultPauseMs: 30000, + setTimeout: timers.setTimeout, + clearTimeout: timers.clearTimeout, + setInterval: timers.setInterval, + clearInterval: timers.clearInterval, + }); + manager.setLimit("umans", 1); + + // Slot is held. Two agents queued (both non-starred). + const release0 = await manager.acquire("umans", "holder", "default", 0); + const pA = manager.acquire("umans", "convA", "ws-normal", 100); + const pB = manager.acquire("umans", "convB", "ws-to-star", 200); + + await Promise.resolve(); + await Promise.resolve(); + + // Star convB's workspace — notifyWorkspaceStarred re-sorts + tries to + // grant. But the slot is still held, so no one is granted yet. + manager.notifyWorkspaceStarred("ws-to-star", true); + + // Release the slot — convB (now starred) should get it. + release0(); + + const rB = await pB; + expect(rB).toBeDefined(); + rB(); + + const rA = await pA; + rA(); + + manager.destroy(); + }); + + it("without isWorkspaceStarred callback, all agents are non-starred (backward compatible)", async () => { + const timers = createFakeTimers(); + const manager = createConcurrencyManager({ + now: timers.now, + slotTimeoutMs: 5000, + watchdogIntervalMs: 1000, + defaultPauseMs: 30000, + setTimeout: timers.setTimeout, + clearTimeout: timers.clearTimeout, + setInterval: timers.setInterval, + clearInterval: timers.clearInterval, + }); + manager.setLimit("umans", 1); + + const release0 = await manager.acquire("umans", "holder", "default", 0); + + const results: string[] = []; + const acquireAndRecord = (conv: string, wsId: string, promptAt: number) => + manager.acquire("umans", conv, wsId, promptAt).then((r) => { + results.push(conv); + return r; + }); + + // Queue in non-sorted order: B (t=200), A (t=100). + const pB = acquireAndRecord("convB", "ws-any", 200); + const pA = acquireAndRecord("convA", "ws-any", 100); + + await Promise.resolve(); + await Promise.resolve(); + + release0(); + + // Without a callback, oldest-first ordering applies (no starred priority). + const rA = await pA; + expect(results).toEqual(["convA"]); + rA(); + + const rB = await pB; + expect(results).toEqual(["convA", "convB"]); + rB(); + + manager.destroy(); + }); +}); diff --git a/packages/provider-concurrency/src/concurrency-manager.ts b/packages/provider-concurrency/src/concurrency-manager.ts index 1d05bb0..ea66a49 100644 --- a/packages/provider-concurrency/src/concurrency-manager.ts +++ b/packages/provider-concurrency/src/concurrency-manager.ts @@ -77,10 +77,19 @@ export interface ProviderConcurrencyStatus { export interface ConcurrencyLimiter { /** * Acquire a concurrency slot for `providerId`. Resolves immediately when a - * slot is available; otherwise blocks (queued by oldest-agent-first) until - * one frees up. The returned function MUST be called when the response - * stream completes (in a `finally` block). For providers with no configured - * limit, resolves instantly with a no-op release. + * slot is available; otherwise blocks (queued by starred-workspace-first, + * then oldest-agent-first) until one frees up. The returned function MUST be + * called when the response stream completes (in a `finally` block). For + * providers with no configured limit, resolves instantly with a no-op + * release. + * + * **Priority:** agents from **starred workspaces** are always admitted before + * agents from non-starred workspaces (regardless of `promptStartedAt`). + * Within each group (starred vs non-starred), oldest-agent-first ordering is + * preserved. The starred status is looked up via the injected + * `isWorkspaceStarred` callback at sort time, so starring a workspace while + * agents are queued takes effect on the next sort (new acquire or slot + * release). * * If `onQueued` is provided and the request cannot be granted immediately * (at limit or paused), it is called synchronously BEFORE the Promise is @@ -89,14 +98,18 @@ export interface ConcurrencyLimiter { * * @param providerId The provider to limit (e.g. "umans", "openai-compat"). * @param conversationId The agent requesting the slot. + * @param workspaceId The workspace the agent belongs to (for starred + * priority scheduling). Defaults to `"default"`. * @param promptStartedAt When the agent's current prompt (turn) started - * (epoch-ms). Used for oldest-agent-first scheduling. + * (epoch-ms). Used for oldest-agent-first scheduling + * within each starred group. * @param onQueued Called synchronously when the request is enqueued * (not granted immediately). Optional. */ acquire( providerId: string, conversationId: string, + workspaceId: string, promptStartedAt: number, onQueued?: () => void, ): Promise<() => void>; @@ -146,6 +159,13 @@ export interface ConcurrencyService extends ConcurrencyLimiter { getStatus(providerId: string): ProviderConcurrencyStatus | undefined; /** Status for every provider with a configured limit. */ getStatusAll(): readonly ProviderConcurrencyStatus[]; + /** + * Notify the limiter that a workspace's starred state changed. Updates the + * in-memory starred cache so subsequent queue sorts re-evaluate priority + * (a newly-starred workspace's already-queued agents jump ahead). Called by + * the transport layer after persisting the starred toggle. + */ + notifyWorkspaceStarred(workspaceId: string, starred: boolean): void; /** Stop the watchdog + clear all timers. */ destroy(): void; } @@ -161,6 +181,7 @@ interface Slot { interface QueuedWaiter { readonly conversationId: string; + readonly workspaceId: string; readonly promptStartedAt: number; readonly resolve: (release: () => void) => void; } @@ -232,6 +253,14 @@ export interface ConcurrencyManagerOpts { * poll never becomes an unhandled rejection. */ readonly onUsagePollError?: (providerId: string, err: unknown) => void; + /** + * Injected callback: returns whether a workspace is starred (for priority + * scheduling). When provided, agents from starred workspaces jump ahead of + * non-starred agents in the queue. When omitted (or returns `false`), all + * agents are treated as non-starred (backward-compatible). This is an I/O + * effect injected so the manager stays pure + unit-testable with a fake. + */ + readonly isWorkspaceStarred?: (workspaceId: string) => boolean; } /** Min interval between usage-gate fallback repolls (ms). The release trigger is immediate. */ @@ -254,6 +283,16 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre const setInterval = opts.setInterval ?? globalThis.setInterval.bind(globalThis); const clearInterval = opts.clearInterval ?? globalThis.clearInterval.bind(globalThis); + // In-memory cache of starred workspace IDs. Populated by the extension on + // activation (from the conversation store) + updated via + // `notifyWorkspaceStarred`. The `isWorkspaceStarred` callback reads this + // synchronously so the queue sort comparator (sync) can re-evaluate priority + // on every sort — a newly-starred workspace's already-queued agents jump + // ahead on the next sort (new acquire or slot release). + const starredWorkspaces = new Set<string>(); + const isWorkspaceStarred = + opts.isWorkspaceStarred ?? ((wsId: string) => starredWorkspaces.has(wsId)); + const states = new Map<string, ProviderState>(); const cooldownOverrides = new Map<string, number>(); const cooldownTimers = new Set<ReturnType<typeof setTimeout>>(); @@ -343,6 +382,25 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre * (`inFlight < limit`). Synchronous. */ function grantLoop(state: ProviderState, providerId: string): void { + * Priority comparator for queued waiters: starred-workspace agents first, + * then oldest-agent-first (ascending `promptStartedAt`) within each group. + * Called at sort time (both on insert and before granting) so a workspace + * starred AFTER an agent queued is re-evaluated on the next sort. + */ + function compareWaiters(a: QueuedWaiter, b: QueuedWaiter): number { + const aStarred = isWorkspaceStarred(a.workspaceId); + const bStarred = isWorkspaceStarred(b.workspaceId); + if (aStarred !== bStarred) return aStarred ? -1 : 1; // starred first + return a.promptStartedAt - b.promptStartedAt; // oldest first within group + } + + function tryGrantNext(providerId: string): void { + const state = states.get(providerId); + if (state === undefined) return; + if (state.paused) return; + // Re-sort before granting: a workspace may have been starred/unstarred + // since the waiters were enqueued, so priority may have changed. + state.queue.sort(compareWaiters); while (state.queue.length > 0 && state.inFlight < state.limit) { const waiter = state.queue[0]; if (waiter === undefined) break; @@ -533,7 +591,7 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre // ── Public API ───────────────────────────────────────────────────────────── const manager: ConcurrencyService = { - acquire(providerId, conversationId, promptStartedAt, onQueued) { + acquire(providerId, conversationId, workspaceId, promptStartedAt, onQueued) { const state = states.get(providerId); if (state === undefined) { // No limit configured → unlimited. @@ -561,7 +619,7 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre // "queued" status signal while we're still synchronous. onQueued?.(); - // Queue (oldest-agent-first by promptStartedAt). + // Queue (starred-workspace-first, then oldest-agent-first). return new Promise<() => void>((resolve) => { state.queue.push({ conversationId, promptStartedAt, resolve }); // Keep sorted ascending by promptStartedAt (oldest first). @@ -572,6 +630,10 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre if (fetchUsage !== undefined) { armGateRepoll(providerId, state); } + state.queue.push({ conversationId, workspaceId, promptStartedAt, resolve }); + // Keep sorted by priority (starred first, then oldest-agent-first). + // The queue is typically tiny (<20), so a simple sort is fine. + state.queue.sort(compareWaiters); }); }, @@ -751,6 +813,22 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre .filter((s): s is ProviderConcurrencyStatus => s !== undefined); }, + notifyWorkspaceStarred(workspaceId, starred) { + if (starred) { + starredWorkspaces.add(workspaceId); + } else { + starredWorkspaces.delete(workspaceId); + } + // Re-sort all queues so a newly-starred workspace's already-queued + // agents jump ahead immediately (no need to wait for the next acquire). + for (const [providerId, state] of states) { + if (state.queue.length > 0) { + state.queue.sort(compareWaiters); + tryGrantNext(providerId); + } + } + }, + destroy() { clearInterval(watchdogTimer); for (const timer of cooldownTimers) { diff --git a/packages/provider-concurrency/src/extension.ts b/packages/provider-concurrency/src/extension.ts index 378655b..48c019b 100644 --- a/packages/provider-concurrency/src/extension.ts +++ b/packages/provider-concurrency/src/extension.ts @@ -1,3 +1,4 @@ +import { conversationStoreHandle } from "@dispatch/conversation-store"; import type { Extension, HostAPI, Logger, Manifest, StorageNamespace } from "@dispatch/kernel"; import type { ConcurrencyManagerOpts, ConcurrencyService } from "./concurrency-manager.js"; import { createConcurrencyManager } from "./concurrency-manager.js"; @@ -11,6 +12,7 @@ export const manifest: Manifest = { trust: "bundled", activation: "eager", capabilities: { db: true }, + dependsOn: ["conversation-store"], contributes: { services: ["provider-concurrency/service"] }, }; @@ -109,6 +111,7 @@ function createPersistedService( getCooldowns: inner.getCooldowns.bind(inner), getStatus: inner.getStatus.bind(inner), getStatusAll: inner.getStatusAll.bind(inner), + notifyWorkspaceStarred: inner.notifyWorkspaceStarred.bind(inner), destroy: inner.destroy.bind(inner), }; } @@ -215,6 +218,16 @@ export async function activate(host: HostAPI): Promise<void> { return provider.getUsage(); }; + // Resolve the conversation store to seed the in-memory starred-workspace + // cache. The `isWorkspaceStarred` callback reads this cache synchronously + // (the queue sort comparator is sync), so we must populate it before the + // manager handles its first acquire. `dependsOn: ["conversation-store"]` + // in the manifest guarantees the store is registered before we activate. + const conversationStore = host.getService(conversationStoreHandle); + + // The manager owns the in-memory `starredWorkspaces` set internally (the + // default `isWorkspaceStarred` callback checks it). We seed it by calling + // `notifyWorkspaceStarred` for each starred workspace found in the store. const managerOpts: ConcurrencyManagerOpts = { now: () => Date.now(), slotTimeoutMs: SLOT_TIMEOUT_MS, @@ -274,6 +287,27 @@ export async function activate(host: HostAPI): Promise<void> { await loadAutoReduce(storage, inner, logger); await loadCooldowns(storage, inner, logger); + // Seed the in-memory starred cache from the conversation store so the + // priority scheduling is correct on a fresh server start (previously-starred + // workspaces are respected without requiring the user to re-star them). + try { + const workspaces = await conversationStore.listWorkspaces(); + for (const ws of workspaces) { + if (ws.starred) { + inner.notifyWorkspaceStarred(ws.id, true); + } + } + if (workspaces.some((w) => w.starred)) { + logger.info("provider-concurrency: restored starred workspaces", { + count: workspaces.filter((w) => w.starred).length, + }); + } + } catch (err) { + logger.warn("provider-concurrency: failed to load starred workspaces", { + err: err instanceof Error ? err.message : String(err), + }); + } + const service = createPersistedService(inner, storage, logger); host.provideService(concurrencyServiceHandle, service); logger.info("provider-concurrency: registered"); diff --git a/packages/provider-concurrency/src/provider-wrapper.test.ts b/packages/provider-concurrency/src/provider-wrapper.test.ts index e59ab39..7554e64 100644 --- a/packages/provider-concurrency/src/provider-wrapper.test.ts +++ b/packages/provider-concurrency/src/provider-wrapper.test.ts @@ -17,12 +17,21 @@ function fakeProvider(events: ProviderEvent[]): ProviderContract { /** A fake limiter that records acquire/release calls. */ function recordingLimiter(): ConcurrencyLimiter & { - acquireCalls: { providerId: string; conversationId: string; promptStartedAt: number }[]; + acquireCalls: { + providerId: string; + conversationId: string; + workspaceId: string; + promptStartedAt: number; + }[]; releaseCalls: number; rateLimitReports: string[]; } { - const acquireCalls: { providerId: string; conversationId: string; promptStartedAt: number }[] = - []; + const acquireCalls: { + providerId: string; + conversationId: string; + workspaceId: string; + promptStartedAt: number; + }[] = []; const releaseCalls: { count: number } = { count: 0 }; const rateLimitReports: string[] = []; @@ -32,8 +41,8 @@ function recordingLimiter(): ConcurrencyLimiter & { return releaseCalls.count; }, rateLimitReports, - acquire(providerId, conversationId, promptStartedAt) { - acquireCalls.push({ providerId, conversationId, promptStartedAt }); + acquire(providerId, conversationId, workspaceId, promptStartedAt) { + acquireCalls.push({ providerId, conversationId, workspaceId, promptStartedAt }); return Promise.resolve(() => { releaseCalls.count++; }); @@ -52,7 +61,7 @@ describe("wrapProviderWithConcurrency", () => { ]); const limiter = recordingLimiter(); - const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", 12345); + const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 12345); const events: ProviderEvent[] = []; for await (const e of wrapped.stream([], [])) { @@ -61,7 +70,12 @@ describe("wrapProviderWithConcurrency", () => { // Slot acquired before stream, released after. expect(limiter.acquireCalls).toEqual([ - { providerId: "test-provider", conversationId: "conv1", promptStartedAt: 12345 }, + { + providerId: "test-provider", + conversationId: "conv1", + workspaceId: "default", + promptStartedAt: 12345, + }, ]); expect(limiter.releaseCalls).toBe(1); expect(events).toEqual([ @@ -79,7 +93,7 @@ describe("wrapProviderWithConcurrency", () => { }, }; const limiter = recordingLimiter(); - const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", 0); + const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0); await expect(async () => { for await (const _e of wrapped.stream([], [])) { @@ -95,7 +109,7 @@ describe("wrapProviderWithConcurrency", () => { { type: "error", message: "Too many requests", code: "429", retryable: true }, ]); const limiter = recordingLimiter(); - const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", 0); + const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0); const events: ProviderEvent[] = []; for await (const e of wrapped.stream([], [])) { @@ -113,7 +127,7 @@ describe("wrapProviderWithConcurrency", () => { { type: "error", message: "Internal error", code: "500", retryable: true }, ]); const limiter = recordingLimiter(); - const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", 0); + const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0); for await (const _e of wrapped.stream([], [])) { // consume @@ -131,7 +145,7 @@ describe("wrapProviderWithConcurrency", () => { listModels: async () => [{ id: "model-1" }], }; const limiter = recordingLimiter(); - const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", 0); + const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0); expect(wrapped.id).toBe("my-provider"); expect(wrapped.listModels).toBeDefined(); @@ -144,7 +158,7 @@ describe("wrapProviderWithConcurrency", () => { let acquiredCalled = false; const blockingLimiter: ConcurrencyLimiter = { - acquire(_providerId, _convId, _promptAt, onQueued) { + acquire(_providerId, _convId, _wsId, _promptAt, onQueued) { // Simulate a queued request: call onQueued, then resolve on next tick. onQueued?.(); return new Promise((resolve) => { @@ -161,6 +175,7 @@ describe("wrapProviderWithConcurrency", () => { provider, blockingLimiter, "conv1", + "default", 0, () => { queuedCalled = true; @@ -183,7 +198,7 @@ describe("wrapProviderWithConcurrency", () => { let acquiredCalled = false; const immediateLimiter: ConcurrencyLimiter = { - acquire(_providerId, _convId, _promptAt, _onQueued) { + acquire(_providerId, _convId, _wsId, _promptAt, _onQueued) { // Grant immediately — do NOT call onQueued. return Promise.resolve(() => {}); }, @@ -195,6 +210,7 @@ describe("wrapProviderWithConcurrency", () => { provider, immediateLimiter, "conv1", + "default", 0, () => { queuedCalled = true; @@ -229,7 +245,7 @@ describe("wrapProviderWithConcurrency", () => { }, }; const limiter = recordingLimiter(); - const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", 0); + const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0); const messages = [{ role: "user" as const, chunks: [{ type: "text" as const, text: "hi" }] }]; const tools = [{ name: "test_tool", description: "test", parameters: {} }]; diff --git a/packages/provider-concurrency/src/provider-wrapper.ts b/packages/provider-concurrency/src/provider-wrapper.ts index aa08e5b..1e3f2c0 100644 --- a/packages/provider-concurrency/src/provider-wrapper.ts +++ b/packages/provider-concurrency/src/provider-wrapper.ts @@ -23,6 +23,8 @@ import type { ConcurrencyLimiter } from "./concurrency-manager.js"; * @param provider The underlying provider to wrap. * @param limiter The concurrency limiter (acquire/release/reportRateLimit). * @param conversationId The agent requesting the stream (for slot attribution). + * @param workspaceId The workspace the agent belongs to (for starred + * priority scheduling in the limiter queue). * @param promptStartedAt When the agent's current prompt (turn) started * (epoch-ms, for oldest-agent-first scheduling). * @param onQueued Called synchronously when `acquire()` decides to @@ -36,6 +38,7 @@ export function wrapProviderWithConcurrency( provider: ProviderContract, limiter: ConcurrencyLimiter, conversationId: string, + workspaceId: string, promptStartedAt: number, onQueued?: () => void, onAcquired?: () => void, @@ -50,7 +53,13 @@ export function wrapProviderWithConcurrency( tools: readonly ToolContract[], opts?: ProviderStreamOptions, ): AsyncIterable<ProviderEvent> { - const release = await limiter.acquire(providerId, conversationId, promptStartedAt, onQueued); + const release = await limiter.acquire( + providerId, + conversationId, + workspaceId, + promptStartedAt, + onQueued, + ); onAcquired?.(); try { for await (const event of innerStream(messages, tools, opts)) { |
