1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
|
import type {
ChatMessage,
ProviderContract,
ProviderEvent,
ProviderStreamOptions,
ToolContract,
} from "@dispatch/kernel";
import type { ConcurrencyLimiter } from "./concurrency-manager.js";
/**
* Wrap a provider's `stream` method with concurrency limiting.
*
* A slot is acquired BEFORE the first event is yielded (before the HTTP
* request is sent — the `await limiter.acquire()` runs before the generator
* body starts iterating the inner stream). The slot is released in a `finally`
* block AFTER the inner stream completes (the full response stream, not just
* HTTP headers — matching the Umans concurrency model where a slot is held
* only while tokens are actually generating).
*
* 429 detection: if the provider yields an `error` event with `code: "429"`,
* the limiter is notified so it can pause the queue for that provider.
*
* @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
* queue the request (cannot grant immediately).
* Lets the caller emit a "queued" status signal.
* @param onAcquired Called when `acquire()` resolves (slot granted,
* whether immediately or after queueing). Lets the
* caller emit an "active" status signal.
*/
export function wrapProviderWithConcurrency(
provider: ProviderContract,
limiter: ConcurrencyLimiter,
conversationId: string,
workspaceId: string,
promptStartedAt: number,
onQueued?: () => void,
onAcquired?: () => void,
): ProviderContract {
const innerStream = provider.stream;
const providerId = provider.id;
return {
id: provider.id,
stream: async function* (
messages: readonly ChatMessage[],
tools: readonly ToolContract[],
opts?: ProviderStreamOptions,
): AsyncIterable<ProviderEvent> {
const release = await limiter.acquire(
providerId,
conversationId,
workspaceId,
promptStartedAt,
onQueued,
);
onAcquired?.();
try {
for await (const event of innerStream(messages, tools, opts)) {
if (event.type === "error" && event.code === "429") {
limiter.reportRateLimit(providerId);
}
yield event;
}
} finally {
release();
}
},
...(provider.listModels !== undefined ? { listModels: provider.listModels } : {}),
};
}
|
