diff --git a/bun.lock b/bun.lock index 691d101..d2975be 100644 --- a/bun.lock +++ b/bun.lock @@ -8,7 +8,7 @@ "@openauthjs/openauth": "^0.4.3", }, "devDependencies": { - "@opencode-ai/plugin": "^1.1.25", + "@opencode-ai/plugin": "^1.1.48", "@types/bun": "latest", }, "peerDependencies": { diff --git a/src/plugin.ts b/src/plugin.ts index 911b338..8295277 100644 --- a/src/plugin.ts +++ b/src/plugin.ts @@ -19,6 +19,8 @@ import { } from "./plugin/request"; import { refreshAccessToken } from "./plugin/token"; import { startOAuthListener, type OAuthListener } from "./plugin/server"; +import { setGlobalState } from "./plugin/state"; +import { geminiQuota } from "./plugin/tools"; import type { GetAuth, LoaderResult, @@ -36,9 +38,13 @@ import type { export const GeminiCLIOAuthPlugin = async ( { client }: PluginContext, ): Promise => ({ + tool: { + "gemini-quota": geminiQuota, + }, auth: { provider: GEMINI_PROVIDER_ID, loader: async (getAuth: GetAuth, provider: Provider): Promise => { + setGlobalState(getAuth, provider, client); const auth = await getAuth(); if (!isOAuthAuth(auth)) { return null; diff --git a/src/plugin/state.ts b/src/plugin/state.ts new file mode 100644 index 0000000..a2452d6 --- /dev/null +++ b/src/plugin/state.ts @@ -0,0 +1,17 @@ +import type { GetAuth, PluginClient, Provider } from "./types"; + +let currentGetAuth: GetAuth | undefined; +let currentProvider: Provider | undefined; +let currentClient: PluginClient | undefined; + +export const setGlobalState = (getAuth: GetAuth, provider: Provider, client: PluginClient) => { + currentGetAuth = getAuth; + currentProvider = provider; + currentClient = client; +}; + +export const getGlobalState = () => ({ + getAuth: currentGetAuth, + provider: currentProvider, + client: currentClient, +}); diff --git a/src/plugin/tools.ts b/src/plugin/tools.ts new file mode 100644 index 0000000..a1b9771 --- /dev/null +++ b/src/plugin/tools.ts @@ -0,0 +1,57 @@ +import { tool } from "@opencode-ai/plugin"; +import { getGlobalState } from "./state"; +import { isOAuthAuth, accessTokenExpired } from "./auth"; +import { refreshAccessToken } from "./token"; +import { ensureProjectContext } from "./project"; +import { GEMINI_CODE_ASSIST_ENDPOINT } from "../constants"; +import type { RetrieveUserQuotaResponse } from "./types"; + +export const geminiQuota = tool({ + description: "Retrieves the current user's quota usage for Gemini models.", + args: {}, + execute: async (_args, _ctx) => { + const { getAuth, client } = getGlobalState(); + + if (!getAuth || !client) { + throw new Error("Gemini plugin not initialized. Please ensure the provider is configured."); + } + + let auth = await getAuth(); + if (!isOAuthAuth(auth)) { + throw new Error("Quota retrieval is only available for OAuth authentication."); + } + + if (accessTokenExpired(auth)) { + const refreshed = await refreshAccessToken(auth, client); + if (refreshed) { + auth = refreshed; + } else { + throw new Error("Failed to refresh access token."); + } + } + + const projectContext = await ensureProjectContext(auth, client); + const projectId = projectContext.effectiveProjectId; + + const url = `${GEMINI_CODE_ASSIST_ENDPOINT}/v1internal:retrieveUserQuota`; + + const response = await fetch(url, { + method: "POST", + headers: { + "Authorization": `Bearer ${auth.access}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + project: projectId + }) + }); + + if (!response.ok) { + const text = await response.text(); + throw new Error(`Failed to retrieve quota: ${response.status} ${response.statusText} - ${text}`); + } + + const data = await response.json() as RetrieveUserQuotaResponse; + return JSON.stringify(data, null, 2); + }, +}); diff --git a/src/plugin/types.ts b/src/plugin/types.ts index 3e37b40..cdd378e 100644 --- a/src/plugin/types.ts +++ b/src/plugin/types.ts @@ -1,5 +1,17 @@ +import type { ToolDefinition } from "@opencode-ai/plugin"; import type { GeminiTokenExchangeResult } from "../gemini/oauth"; +export interface QuotaBucket { + resetTime: string; + tokenType: string; + modelId: string; + remainingFraction: number; +} + +export interface RetrieveUserQuotaResponse { + buckets: QuotaBucket[]; +} + export interface OAuthAuthDetails { type: "oauth"; refresh: string; @@ -62,6 +74,7 @@ export interface PluginResult { loader: (getAuth: GetAuth, provider: Provider) => Promise; methods: AuthMethod[]; }; + tool?: Record; } export interface RefreshParts {