55 lines
2.2 KiB
TypeScript
55 lines
2.2 KiB
TypeScript
import { Plugin } from "@opencode/plugin";
|
|
import { QuotaTracker } from "./quota.js";
|
|
import { QuotaRpc } from "./quota-rpc.js";
|
|
|
|
export default Plugin.define({
|
|
id: "stanp-opencode.quota",
|
|
async setup(context) {
|
|
const controller = new AbortController();
|
|
const quota = new QuotaTracker(context, async () => {
|
|
await registration.events.emit("updated", {});
|
|
});
|
|
const registration = await context.rpc.register(QuotaRpc, {
|
|
status: async () => quota.status() ?? null,
|
|
});
|
|
await quota.init();
|
|
const updates = (async () => {
|
|
for await (const event of context.event.subscribe({ signal: controller.signal })) {
|
|
if (
|
|
event.location &&
|
|
(event.location.directory !== context.location.directory ||
|
|
("workspaceID" in event.location
|
|
? event.location.workspaceID
|
|
: undefined) !== context.location.workspaceID)
|
|
)
|
|
continue;
|
|
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) => {
|
|
if (!controller.signal.aborted)
|
|
console.error("stanp-opencode quota update failed", error);
|
|
});
|
|
await context.session.hook("model.request", async (event) => quota.start(event.sessionID));
|
|
return async () => {
|
|
controller.abort();
|
|
await updates;
|
|
await quota.dispose();
|
|
await registration.dispose();
|
|
};
|
|
},
|
|
});
|