1 files changed, 77 insertions, 0 deletions
diff --git a/packages/provider-concurrency/src/provider-wrapper.ts b/packages/provider-concurrency/src/provider-wrapper.ts
new file mode 100644
index 0000000..1e3f2c0
--- /dev/null
+++ b/packages/provider-concurrency/src/provider-wrapper.ts
@@ -0,0 +1,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 } : {}),
+  };
+}