feat: add quota tracking

src/index.ts

@@ -1,9 +1,12 @@
 import { Agent, Model, Plugin, Provider } from "@opencode/plugin";
+import type { RpcRegistration } from "@opencode/plugin/promise/rpc";
 import type { ModelInfo } from "@opencode/client";
 import { z } from "zod";
 import { gptAgents, settings, workerNames } from "./config.js";
 import { discover } from "./provider.js";
 import { descriptions, routingGuide } from "./routing.js";
+import { QuotaTracker } from "./quota.js";
+import { QuotaRpc } from "./quota-rpc.js";
 
 // Permission actions use wildcards, not filesystem globs.
 function matches(value: string, pattern: string) {
@@ -50,6 +53,14 @@ export default Plugin.define({
         let codegraphTools: string[] = [];
         let agentTransformReady = false;
         const controller = new AbortController();
+        let quotaRegistration: RpcRegistration<typeof QuotaRpc> | undefined;
+        const quota = new QuotaTracker(context, async () => {
+            await quotaRegistration?.events.emit("updated", {});
+        });
+        quotaRegistration = await context.rpc.register(QuotaRpc, {
+            status: async ({ sessionID }) => (await quota.status(sessionID)) ?? null,
+        });
+        await quota.init();
         // Subscribe before the first read: configuration can publish its catalog while
         // asynchronous setup is running.
         const updates = (async () => {
@@ -69,6 +80,20 @@ export default Plugin.define({
                     inherited = permissions;
                     permissionKey = key;
                     await context.agent.reload();
+                } else if (
+                    event.type === "session.execution.succeeded" ||
+                    event.type === "session.execution.failed" ||
+                    event.type === "session.execution.interrupted" ||
+                    event.type === "session.idle"
+                ) {
+                    await quota.complete(event.data.sessionID);
+                } else if (event.type === "credential.updated") {
+                    await quota.authChanged();
+                } else if (
+                    event.type === "credential.switched" &&
+                    String(event.data.integrationID) === "openai"
+                ) {
+                    await quota.authChanged(true);
                 }
             }
         })().catch((error) => {
@@ -251,6 +276,7 @@ export default Plugin.define({
                 properties: { ...properties, agent: { type: "string", enum: workers } },
             };
         });
+        await context.session.hook("model.request", async (event) => quota.start(event.sessionID));
         await context.tool.hook("execute.before", (event) => {
             if (event.tool !== "subagent") return;
             const { agent } = delegation.parse(event.input);
@@ -278,6 +304,8 @@ export default Plugin.define({
         return async () => {
             controller.abort();
             await updates;
+            await quota.dispose();
+            await quotaRegistration?.dispose();
         };
     },
 });