Merge pull request Stack-Cairn#270 from inkdust2021/codex/provider-usage-query

feat(providers): add provider usage queries

Author: su-fen

Cargo.lock

@@ -39,6 +39,7 @@ func vetAgentRequest(sm session.AgentView, env *gatewayv2.GatewayEnvelope) error
 		*gatewayv2.GatewayEnvelope_HistoryWorkdirs,
 		*gatewayv2.GatewayEnvelope_HistoryBranch,
 		*gatewayv2.GatewayEnvelope_ProviderList,
+		*gatewayv2.GatewayEnvelope_ProviderUsage,
 		*gatewayv2.GatewayEnvelope_ProviderModels,
 		*gatewayv2.GatewayEnvelope_SettingsGet,
 		*gatewayv2.GatewayEnvelope_SettingsUpdate,

crates/agent-gateway/internal/proto/v2/gateway.pb.go

@@ -0,0 +1,23 @@
+package pbws
+
+import (
+	"testing"
+
+	gatewayv2 "github.com/liveagent/agent-gateway/internal/proto/v2"
+	"github.com/liveagent/agent-gateway/internal/session"
+)
+
+func TestVetAgentRequestAllowsProviderUsage(t *testing.T) {
+	env := &gatewayv2.GatewayEnvelope{
+		Payload: &gatewayv2.GatewayEnvelope_ProviderUsage{
+			ProviderUsage: &gatewayv2.ProviderUsageRequest{
+				ProviderId: "provider-1",
+				Refresh:    true,
+			},
+		},
+	}
+
+	if err := vetAgentRequest(session.AgentView{}, env); err != nil {
+		t.Fatalf("vetAgentRequest() error = %v", err)
+	}
+}

crates/agent-gateway/internal/protocol/pbws/guard.go

@@ -58,6 +58,7 @@ message GatewayEnvelope {
     WorkspaceWatchRequest workspace_watch = 90;
     ManagedProcessRequest managed_process_request = 91;
     HistoryBranchRequest history_branch = 92;
+    ProviderUsageRequest provider_usage = 93;
   }
 
   // Legacy tunnel control/frame payloads (pre-rewrite protocol) and the
@@ -126,6 +127,7 @@ message AgentEnvelope {
     ManagedProcessResponse managed_process_response = 91;
     ManagedProcessSnapshot managed_process_snapshot = 92;
     HistoryBranchResponse history_branch_resp = 93;
+    ProviderUsageResponse provider_usage_resp = 94;
     ErrorResponse error = 99;
   }
 
@@ -1134,3 +1136,12 @@ message ProviderModelsRequest {
 message ProviderModelsResponse {
   string models_json = 1;
 }
+
+message ProviderUsageRequest {
+  string provider_id = 1;
+  bool refresh = 2;
+}
+
+message ProviderUsageResponse {
+  string result_json = 1;
+}

crates/agent-gateway/internal/protocol/pbws/guard_test.go

@@ -0,0 +1,48 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { createWebModuleLoader } from "../helpers/load-web-module.mjs";
+
+const loader = createWebModuleLoader();
+const forms = loader.loadModule("src/pages/settings/providerUtils.ts");
+
+const usageQuery = {
+  enabled: true,
+  mode: "newapi",
+  script: "",
+  baseUrl: "https://usage.example.test",
+  accessToken: "",
+  accessTokenConfigured: true,
+  userId: "user-1",
+  accessKeyId: "key-1",
+  secretAccessKey: "",
+  secretAccessKeyConfigured: true,
+};
+
+test("WebUI usage query draft preserves configured redacted secrets when saved", () => {
+  const draft = forms.createUsageQueryDraft(usageQuery, true);
+
+  assert.notEqual(draft.accessToken, "");
+  assert.notEqual(draft.secretAccessKey, "");
+  assert.deepEqual(forms.serializeUsageQueryDraft(draft, true), usageQuery);
+});
+
+test("WebUI usage test action accepts only a persisted provider id", () => {
+  assert.equal(forms.getPersistedUsageQueryProviderId(undefined), null);
+  assert.equal(forms.getPersistedUsageQueryProviderId({ id: "" }), null);
+  assert.equal(forms.getPersistedUsageQueryProviderId({ id: "provider-a" }), "provider-a");
+});
+
+test("WebUI custom usage query needs confirmation before its first enabled save", () => {
+  assert.equal(
+    forms.requiresCustomUsageQueryConfirmation({ ...usageQuery, mode: "custom" }, false),
+    true,
+  );
+  assert.equal(
+    forms.requiresCustomUsageQueryConfirmation({ ...usageQuery, mode: "custom" }, true),
+    false,
+  );
+  assert.equal(
+    forms.requiresCustomUsageQueryConfirmation({ ...usageQuery, mode: "custom", enabled: true }, true),
+    false,
+  );
+});

crates/agent-gateway/proto/v2/gateway.proto

@@ -0,0 +1,106 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { createGatewayV2Codec } from "../helpers/gateway-v2.mjs";
+import { createWebModuleLoader } from "../helpers/load-web-module.mjs";
+
+const requestCalls = [];
+const loader = createWebModuleLoader({
+  mocks: {
+    "@/lib/gatewaySocket": {
+      getGatewayWebSocketClient() {
+        return {
+          providerUsageQuery(providerId, refresh) {
+            requestCalls.push({ providerId, refresh });
+            return Promise.resolve({
+              entries: [{ label: "Credits", value: "11" }],
+              queriedAt: 123,
+              error: null,
+              isStale: false,
+            });
+          },
+        };
+      },
+    },
+    "@/lib/storage": { loadToken: () => "gateway-token" },
+  },
+});
+const usage = loader.loadModule("src/lib/providers/usageQuery.ts");
+const adapters = loader.loadModule("src/lib/gatewaySocketV2/adapters.ts");
+const codec = createGatewayV2Codec(loader);
+
+test("WebUI query client refreshes one provider through the Gateway", async () => {
+  requestCalls.length = 0;
+
+  const result = await usage.queryProviderUsage("provider-a", true);
+
+  assert.equal(result.entries[0].value, "11");
+  assert.deepEqual(requestCalls, [{ providerId: "provider-a", refresh: true }]);
+});
+
+test("WebUI protobuf encodes usage request and decodes JSON response", () => {
+  const request = codec.decodeClientFrame(
+    adapters.encodeRequestFrame(
+      "request-1",
+      "provider.usage.query",
+      { provider_id: "provider-a", refresh: true },
+      "desktop-agent",
+    ),
+  );
+
+  assert.equal(request.case, "agentRequest");
+  assert.deepEqual(request.json.agent_request.provider_usage, {
+    provider_id: "provider-a",
+    refresh: true,
+  });
+
+  const frame = codec.encodeServerFrame({
+    request_id: "request-1",
+    agent_id: "desktop-agent",
+    agent_response: {
+      provider_usage_resp: {
+        result_json: JSON.stringify({
+          entries: [{ label: "Credits", value: "11" }],
+          queriedAt: 123,
+          error: null,
+          isStale: false,
+        }),
+      },
+    },
+  });
+  const decoded = adapters.decodeServerFrame(adapters.decodeServerFrameBinary(frame), {
+    agentOnline: true,
+  });
+
+  assert.deepEqual(decoded, {
+    kind: "response",
+    requestId: "request-1",
+    agentId: "desktop-agent",
+    payload: {
+      entries: [{ label: "Credits", value: "11" }],
+      queriedAt: 123,
+      error: null,
+      isStale: false,
+    },
+  });
+});
+
+test("WebUI protobuf rejects malformed usage response JSON", () => {
+  const frame = codec.encodeServerFrame({
+    request_id: "request-2",
+    agent_id: "desktop-agent",
+    agent_response: {
+      provider_usage_resp: { result_json: "{not-json" },
+    },
+  });
+
+  const decoded = adapters.decodeServerFrame(adapters.decodeServerFrameBinary(frame), {
+    agentOnline: true,
+  });
+
+  assert.deepEqual(decoded, {
+    kind: "error",
+    requestId: "request-2",
+    agentId: "desktop-agent",
+    message: "provider usage response is not valid JSON",
+  });
+});

crates/agent-gateway/test/webui/provider-usage-query-form.test.mjs

@@ -1387,6 +1387,34 @@ export const translations: Record<Locale, Record<string, string>> = {
     "settings.promptCacheRetention": "缓存保留",
     "settings.promptCacheRetentionShort": "5 分钟（默认）",
     "settings.promptCacheRetentionLong": "1 小时（仅官方 API）",
+    "settings.providerUsageQuery": "用量查询",
+    "settings.providerUsageEnabled": "启用用量查询",
+    "settings.providerUsageMode": "查询方式",
+    "settings.providerUsageMode.codingPlan": "Coding Plan",
+    "settings.providerUsageMode.general": "通用模板",
+    "settings.providerUsageMode.newapi": "NewAPI 模板",
+    "settings.providerUsageMode.custom": "自定义脚本",
+    "settings.providerUsageTemplate.general": "使用内置通用余额模板和当前供应商 API Key。",
+    "settings.providerUsageTemplate.newapi": "使用内置 NewAPI 模板；可选择单独配置查询地址与凭据。",
+    "settings.providerUsageScript": "自定义查询脚本",
+    "settings.providerUsageScriptPlaceholder": "返回 { request, extractor } 的 JavaScript 表达式",
+    "settings.providerUsageBaseUrl": "查询 Base URL（可选）",
+    "settings.providerUsageAccessToken": "Access Token",
+    "settings.providerUsageUserId": "User ID",
+    "settings.providerUsageAccessKeyId": "AccessKey ID",
+    "settings.providerUsageSecretAccessKey": "SecretAccessKey",
+    "settings.providerUsageAutoRefreshFixedHint": "启用后每 5 分钟自动刷新当前使用的供应商。",
+    "settings.providerUsageTest": "测试查询",
+    "settings.providerUsageTestSavedHint": "请先保存供应商，才能测试查询。",
+    "settings.providerUsageTestRunning": "正在查询…",
+    "settings.providerUsageTestFailed": "查询失败",
+    "settings.providerUsageTestEmpty": "查询成功，但没有返回可展示的数据。",
+    "settings.providerUsageCustomConfirmTitle": "启用自定义用量脚本？",
+    "settings.providerUsageCustomConfirmDescription":
+      "自定义脚本可以向其定义的目标发起网络请求。请仅使用可信脚本和端点。",
+    "settings.providerUsageCustomConfirmDetail":
+      "脚本仅会在保存后由本机服务执行，查询凭据不会发送到 WebUI。",
+    "settings.providerUsageCustomConfirmAction": "确认启用",
     "settings.nativeWebSearch": "联网搜索",
     "settings.thinking": "Thinking",
     "settings.thinkingEnabled": "启用 Thinking",
@@ -3480,6 +3508,38 @@ export const translations: Record<Locale, Record<string, string>> = {
     "settings.promptCacheRetention": "Cache retention",
     "settings.promptCacheRetentionShort": "5 minutes (default)",
     "settings.promptCacheRetentionLong": "1 hour (official API only)",
+    "settings.providerUsageQuery": "Usage query",
+    "settings.providerUsageEnabled": "Enable usage query",
+    "settings.providerUsageMode": "Query mode",
+    "settings.providerUsageMode.codingPlan": "Coding Plan",
+    "settings.providerUsageMode.general": "General template",
+    "settings.providerUsageMode.newapi": "NewAPI template",
+    "settings.providerUsageMode.custom": "Custom script",
+    "settings.providerUsageTemplate.general":
+      "Uses the built-in general balance template with this provider's API key.",
+    "settings.providerUsageTemplate.newapi":
+      "Uses the built-in NewAPI template; optionally configure a separate query URL and credentials.",
+    "settings.providerUsageScript": "Custom query script",
+    "settings.providerUsageScriptPlaceholder":
+      "A JavaScript expression returning { request, extractor }",
+    "settings.providerUsageBaseUrl": "Query Base URL (optional)",
+    "settings.providerUsageAccessToken": "Access Token",
+    "settings.providerUsageUserId": "User ID",
+    "settings.providerUsageAccessKeyId": "AccessKey ID",
+    "settings.providerUsageSecretAccessKey": "SecretAccessKey",
+    "settings.providerUsageAutoRefreshFixedHint":
+      "When enabled, the active provider refreshes automatically every 5 minutes.",
+    "settings.providerUsageTest": "Test query",
+    "settings.providerUsageTestSavedHint": "Save the provider before testing its query.",
+    "settings.providerUsageTestRunning": "Querying…",
+    "settings.providerUsageTestFailed": "Query failed",
+    "settings.providerUsageTestEmpty": "The query succeeded but returned no displayable data.",
+    "settings.providerUsageCustomConfirmTitle": "Enable custom usage script?",
+    "settings.providerUsageCustomConfirmDescription":
+      "A custom script can make network requests to the destinations it defines. Use only trusted scripts and endpoints.",
+    "settings.providerUsageCustomConfirmDetail":
+      "The script runs only after saving through the local service; query credentials are never sent to the WebUI.",
+    "settings.providerUsageCustomConfirmAction": "Enable query",
     "settings.nativeWebSearch": "Web Search",
     "settings.thinking": "Thinking",
     "settings.thinkingEnabled": "Enable Thinking",

crates/agent-gateway/test/webui/provider-usage-query.test.mjs

@@ -2703,6 +2703,13 @@ export class GatewayWebSocketClient {
     });
   }
 
+  async providerUsageQuery<T = unknown>(providerId: string, refresh: boolean): Promise<T> {
+    return this.requestWithRecovery<T>("provider.usage.query", {
+      provider_id: providerId,
+      refresh,
+    });
+  }
+
   dispose() {
     this.disposed = true;
     this.terminalStream.dispose();
@@ -3810,6 +3817,7 @@ export type GatewayWebSocketClientLike = {
     apiKey: string,
     useSystemProxy?: boolean,
   ): Promise<unknown>;
+  providerUsageQuery<T = unknown>(providerId: string, refresh: boolean): Promise<T>;
   dispose(): void;
 };
 

crates/agent-gateway/web/src/i18n/config.ts

@@ -66,6 +66,7 @@ import {
   MemoryManageRequestSchema,
   ProviderListRequestSchema,
   ProviderModelsRequestSchema,
+  ProviderUsageRequestSchema,
   SettingsGetRequestSchema,
   SettingsResetSshKnownHostRequestSchema,
   SettingsUpdateRequestSchema,
@@ -573,6 +574,14 @@ function agentRequestPayload(type: string, body: J): GatewayEnvelope["payload"]
           useSystemProxy: bool(body.use_system_proxy),
         }),
       };
+    case "provider.usage.query":
+      return {
+        case: "providerUsage",
+        value: create(ProviderUsageRequestSchema, {
+          providerId: trimStr(body.provider_id),
+          refresh: bool(body.refresh),
+        }),
+      };
     case "settings.get":
       return { case: "settingsGet", value: create(SettingsGetRequestSchema, {}) };
     case "settings.update":
@@ -1020,6 +1029,13 @@ function decodeAgentResponse(envelope: AgentEnvelope, options: { agentOnline: bo
         frameError("provider model response is not valid JSON");
       }
       break;
+    case "providerUsageResp":
+      try {
+        return parseJson(payload.value.resultJson);
+      } catch {
+        frameError("provider usage response is not valid JSON");
+      }
+      break;
     case "settingsGetResp": {
       const raw = payload.value.settingsJson.trim();
       if (!raw) return {};