From f02af4a2c5c0f09b7fd65f6f86cbe51b1934bcf3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B4=AA=E6=B3=BD=E9=91=AB?= Date: Fri, 21 Aug 2026 12:00:12 +0800 Subject: [PATCH 1/2] feat(company-codex): add self-service weekly quota increases --- packages/core/api/client.ts | 10 ++ packages/core/api/schemas.test.ts | 18 +++ packages/core/api/schemas.ts | 4 + packages/core/types/api.ts | 4 + packages/views/locales/en/settings.json | 14 ++ packages/views/locales/ja/settings.json | 14 ++ packages/views/locales/ko/settings.json | 14 ++ packages/views/locales/zh-Hans/settings.json | 14 ++ .../components/company-codex-section.test.tsx | 54 +++++++ .../components/company-codex-section.tsx | 125 +++++++++++++++-- server/cmd/server/router.go | 1 + server/internal/handler/company_codex.go | 112 +++++++++++++-- server/internal/handler/company_codex_test.go | 132 ++++++++++++++++++ 13 files changed, 495 insertions(+), 21 deletions(-) diff --git a/packages/core/api/client.ts b/packages/core/api/client.ts index df00f85454b..ae12a727ed6 100644 --- a/packages/core/api/client.ts +++ b/packages/core/api/client.ts @@ -2363,6 +2363,16 @@ export class ApiClient { return created; } + async increaseCompanyCodexQuota(additionalTokens: number): Promise { + const raw = await this.fetch("/api/company-codex/key", { + method: "PATCH", + body: JSON.stringify({ additional_tokens: additionalTokens }), + }); + return parseWithFallback(raw, CompanyCodexKeyStatusSchema, EMPTY_COMPANY_CODEX_KEY_STATUS, { + endpoint: "PATCH /api/company-codex/key", + }); + } + async revokeCompanyCodexKey(): Promise { await this.fetch("/api/company-codex/key", { method: "DELETE" }); } diff --git a/packages/core/api/schemas.test.ts b/packages/core/api/schemas.test.ts index 86ab48b782d..bef7ab7d2e2 100644 --- a/packages/core/api/schemas.test.ts +++ b/packages/core/api/schemas.test.ts @@ -81,6 +81,24 @@ describe("company Codex schemas", () => { expect(CreateCompanyCodexKeyResponseSchema.safeParse({ active: true }).success).toBe(false); }); + it("keeps valid company Codex quota fields and rejects invalid percentages", () => { + expect(CompanyCodexKeyStatusSchema.parse({ + active: true, + weekly_token_limit: 25_000_000, + current_tokens: 5_000_000, + reset_at: "2026-08-28T00:00:00Z", + upstream_remaining_percent: 74.25, + })).toMatchObject({ + weekly_token_limit: 25_000_000, + current_tokens: 5_000_000, + upstream_remaining_percent: 74.25, + }); + expect(CompanyCodexKeyStatusSchema.safeParse({ + active: true, + upstream_remaining_percent: 101, + }).success).toBe(false); + }); + it("keeps future status values and rejects malformed session list items", () => { expect(CompanyCodexSessionListSchema.parse([{ ...session, status: "archived" }])[0]?.status) .toBe("archived"); diff --git a/packages/core/api/schemas.ts b/packages/core/api/schemas.ts index b57b58b640b..0b84a71dd85 100644 --- a/packages/core/api/schemas.ts +++ b/packages/core/api/schemas.ts @@ -70,6 +70,10 @@ export const CompanyCodexKeyStatusSchema = z.object({ active: z.boolean().default(false), key_prefix: z.string().optional(), created_at: z.string().optional(), + weekly_token_limit: z.number().nonnegative().optional(), + current_tokens: z.number().nonnegative().optional(), + reset_at: z.string().optional(), + upstream_remaining_percent: z.number().min(0).max(100).nullable().optional(), }).loose(); export const EMPTY_COMPANY_CODEX_KEY_STATUS: CompanyCodexKeyStatus = { diff --git a/packages/core/types/api.ts b/packages/core/types/api.ts index 7704cff0daa..db2b0574741 100644 --- a/packages/core/types/api.ts +++ b/packages/core/types/api.ts @@ -518,6 +518,10 @@ export interface CompanyCodexKeyStatus { active: boolean; key_prefix?: string; created_at?: string; + weekly_token_limit?: number; + current_tokens?: number; + reset_at?: string; + upstream_remaining_percent?: number | null; } export interface CreateCompanyCodexKeyResponse extends CompanyCodexKeyStatus { diff --git a/packages/views/locales/en/settings.json b/packages/views/locales/en/settings.json index 06c1baec0ed..4f32efcac50 100644 --- a/packages/views/locales/en/settings.json +++ b/packages/views/locales/en/settings.json @@ -842,6 +842,20 @@ "enabled": "Active", "not_enabled": "Not configured", "key_metadata": "{{prefix}} · issued {{date}}", + "quota_summary": "{{used}} / {{limit}} tokens used this week", + "quota_unavailable": "Quota status is temporarily unavailable. Refresh and try again.", + "quota_resets": "resets {{date}}", + "upstream_remaining": "Shared upstream pool remaining: {{percent}}%", + "increase_quota": "Increase quota", + "increase_quota_title": "Increase weekly quota", + "additional_tokens": "Additional tokens", + "quota_preview": "{{current}} + {{additional}} = {{next}} tokens per week", + "quota_invalid": "Enter a positive whole number of tokens.", + "quota_warning": "This updates your existing key without rotating it or resetting usage. It does not add capacity to the shared upstream Pro account pool.", + "quota_cancel": "Cancel", + "confirm_increase": "Confirm increase", + "quota_increased": "Weekly quota increased", + "quota_increase_failed": "Failed to increase weekly quota", "create": "Create access", "rotate": "Rotate key", "revoke": "Revoke", diff --git a/packages/views/locales/ja/settings.json b/packages/views/locales/ja/settings.json index 31c46a09108..17f021ec74d 100644 --- a/packages/views/locales/ja/settings.json +++ b/packages/views/locales/ja/settings.json @@ -841,6 +841,20 @@ "enabled": "有効", "not_enabled": "未設定", "key_metadata": "{{prefix}} · {{date}} に発行", + "quota_summary": "今週 {{used}} / {{limit}} Token を使用", + "quota_unavailable": "クォータの状態を一時的に取得できません。更新してから再試行してください。", + "quota_resets": "{{date}} にリセット", + "upstream_remaining": "共有アップストリームプール残量: {{percent}}%", + "increase_quota": "上限を増やす", + "increase_quota_title": "週間上限を増やす", + "additional_tokens": "追加する Token 数", + "quota_preview": "{{current}} + {{additional}} = 週 {{next}} Token", + "quota_invalid": "正の整数の Token 数を入力してください。", + "quota_warning": "既存のキーを維持したまま上限を増やし、使用量はリセットしません。共有アップストリーム Pro アカウントプールの容量は増えません。", + "quota_cancel": "キャンセル", + "confirm_increase": "増加を確認", + "quota_increased": "週間上限を増やしました", + "quota_increase_failed": "週間上限を増やせませんでした", "create": "アクセスを作成", "rotate": "キーをローテーション", "revoke": "無効化", diff --git a/packages/views/locales/ko/settings.json b/packages/views/locales/ko/settings.json index c508cb09a12..02b123dc787 100644 --- a/packages/views/locales/ko/settings.json +++ b/packages/views/locales/ko/settings.json @@ -607,6 +607,20 @@ "enabled": "활성", "not_enabled": "미설정", "key_metadata": "{{prefix}} · {{date}} 발급", + "quota_summary": "이번 주 {{used}} / {{limit}} Token 사용", + "quota_unavailable": "할당량 상태를 일시적으로 불러올 수 없습니다. 새로 고친 후 다시 시도하세요.", + "quota_resets": "{{date}} 초기화", + "upstream_remaining": "공유 업스트림 풀 잔여량: {{percent}}%", + "increase_quota": "한도 늘리기", + "increase_quota_title": "주간 한도 늘리기", + "additional_tokens": "추가 Token 수", + "quota_preview": "{{current}} + {{additional}} = 주당 {{next}} Token", + "quota_invalid": "양의 정수 Token 수를 입력하세요.", + "quota_warning": "기존 키를 교체하거나 사용량을 초기화하지 않고 한도만 늘립니다. 공유 업스트림 Pro 계정 풀의 용량은 늘어나지 않습니다.", + "quota_cancel": "취소", + "confirm_increase": "증가 확인", + "quota_increased": "주간 한도를 늘렸습니다", + "quota_increase_failed": "주간 한도를 늘리지 못했습니다", "create": "액세스 생성", "rotate": "키 교체", "revoke": "해지", diff --git a/packages/views/locales/zh-Hans/settings.json b/packages/views/locales/zh-Hans/settings.json index 4e86683c5ff..5461922a9a3 100644 --- a/packages/views/locales/zh-Hans/settings.json +++ b/packages/views/locales/zh-Hans/settings.json @@ -841,6 +841,20 @@ "enabled": "已启用", "not_enabled": "未配置", "key_metadata": "{{prefix}} · 签发于 {{date}}", + "quota_summary": "本周已用 {{used}} / {{limit}} Token", + "quota_unavailable": "额度状态暂时不可用，请刷新后重试。", + "quota_resets": "{{date}} 重置", + "upstream_remaining": "上游共享号池剩余：{{percent}}%", + "increase_quota": "增加额度", + "increase_quota_title": "增加每周额度", + "additional_tokens": "追加 Token 数", + "quota_preview": "{{current}} + {{additional}} = 每周 {{next}} Token", + "quota_invalid": "请输入正整数 Token 数。", + "quota_warning": "额度会直接加到现有密钥，不会轮换密钥或清零用量。增加个人额度不会增加上游共享 Pro 号池容量。", + "quota_cancel": "取消", + "confirm_increase": "确认增加", + "quota_increased": "每周额度已增加", + "quota_increase_failed": "增加每周额度失败", "create": "创建访问", "rotate": "轮换密钥", "revoke": "撤销", diff --git a/packages/views/settings/components/company-codex-section.test.tsx b/packages/views/settings/components/company-codex-section.test.tsx index b63c44e3a23..c002058e0ef 100644 --- a/packages/views/settings/components/company-codex-section.test.tsx +++ b/packages/views/settings/components/company-codex-section.test.tsx @@ -7,6 +7,7 @@ import { renderWithI18n } from "../../test/i18n"; const apiMock = vi.hoisted(() => ({ createCompanyCodexKey: vi.fn(), getCompanyCodexKey: vi.fn(), + increaseCompanyCodexQuota: vi.fn(), listCompanyCodexSessions: vi.fn(), revokeCompanyCodexKey: vi.fn(), })); @@ -57,4 +58,57 @@ describe("CompanyCodexSection", () => { expect(screen.getByRole("link", { name: "Download CC Switch (official releases)" })) .toHaveAttribute("href", "https://github.com/farion1231/cc-switch/releases"); }); + + it("increases the existing key quota without rotating it", async () => { + apiMock.getCompanyCodexKey.mockResolvedValue({ + active: true, + key_prefix: "sk-clb-test", + created_at: "2026-08-21T00:00:00Z", + weekly_token_limit: 25_000_000, + current_tokens: 5_000_000, + reset_at: "2026-08-28T00:00:00Z", + upstream_remaining_percent: 74.25, + }); + apiMock.increaseCompanyCodexQuota.mockResolvedValue({ + active: true, + key_prefix: "sk-clb-test", + created_at: "2026-08-21T00:00:00Z", + weekly_token_limit: 32_000_000, + current_tokens: 5_000_000, + reset_at: "2026-08-28T00:00:00Z", + upstream_remaining_percent: 74.25, + }); + + renderWithI18n(); + + await waitFor(() => expect(screen.getByRole("button", { name: "Increase quota" })).toBeEnabled()); + expect(screen.getByText(/5,000,000 \/ 25,000,000 tokens used this week/)).toBeInTheDocument(); + expect(screen.getByText(/Shared upstream pool remaining: 74.25%/)).toBeInTheDocument(); + + fireEvent.click(screen.getByRole("button", { name: "Increase quota" })); + fireEvent.change(screen.getByLabelText("Additional tokens"), { target: { value: "7000000" } }); + expect(screen.getByText("25,000,000 + 7,000,000 = 32,000,000 tokens per week")) + .toBeInTheDocument(); + fireEvent.click(screen.getByRole("button", { name: "Confirm increase" })); + + await waitFor(() => expect(apiMock.increaseCompanyCodexQuota).toHaveBeenCalledWith(7_000_000)); + expect(apiMock.createCompanyCodexKey).not.toHaveBeenCalled(); + await waitFor(() => expect(screen.getByText(/5,000,000 \/ 32,000,000 tokens used this week/)) + .toBeInTheDocument()); + }); + + it("disables quota increases when the gateway quota is unavailable", async () => { + apiMock.getCompanyCodexKey.mockResolvedValue({ + active: true, + key_prefix: "sk-clb-test", + created_at: "2026-08-21T00:00:00Z", + }); + + renderWithI18n(); + + await waitFor(() => expect(screen.getByRole("button", { name: "Increase quota" })) + .toBeDisabled()); + expect(screen.getByText("Quota status is temporarily unavailable. Refresh and try again.")) + .toBeInTheDocument(); + }); }); diff --git a/packages/views/settings/components/company-codex-section.tsx b/packages/views/settings/components/company-codex-section.tsx index 98fe2fb6f08..e0e3cef96d1 100644 --- a/packages/views/settings/components/company-codex-section.tsx +++ b/packages/views/settings/components/company-codex-section.tsx @@ -1,7 +1,7 @@ "use client"; import { useCallback, useEffect, useState } from "react"; -import { Check, Copy, ExternalLink, KeyRound, RefreshCw, ShieldCheck, Trash2 } from "lucide-react"; +import { Check, Copy, ExternalLink, KeyRound, Plus, RefreshCw, ShieldCheck, Trash2 } from "lucide-react"; import { api } from "@multica/core/api"; import type { CompanyCodexKeyStatus, @@ -20,6 +20,7 @@ import { DialogHeader, DialogTitle, } from "@multica/ui/components/ui/dialog"; +import { Input } from "@multica/ui/components/ui/input"; import { Skeleton } from "@multica/ui/components/ui/skeleton"; import { copyText } from "@multica/ui/lib/clipboard"; import { toast } from "sonner"; @@ -37,6 +38,9 @@ export function CompanyCodexSection() { const [loading, setLoading] = useState(true); const [creating, setCreating] = useState(false); const [revoking, setRevoking] = useState(false); + const [quotaOpen, setQuotaOpen] = useState(false); + const [quotaIncrease, setQuotaIncrease] = useState("5000000"); + const [increasingQuota, setIncreasingQuota] = useState(false); const [issued, setIssued] = useState(null); const [copied, setCopied] = useState(null); const [sessionDetail, setSessionDetail] = useState(null); @@ -88,6 +92,25 @@ export function CompanyCodexSection() { } }; + const additionalTokens = Number(quotaIncrease); + const quotaIncreaseValid = Number.isSafeInteger(additionalTokens) && additionalTokens > 0; + + const increaseQuota = async () => { + if (!quotaIncreaseValid) return; + setIncreasingQuota(true); + try { + const result = await api.increaseCompanyCodexQuota(additionalTokens); + setKeyStatus(result); + setQuotaOpen(false); + setQuotaIncrease("5000000"); + toast.success(t(($) => $.company_codex.quota_increased)); + } catch (error) { + toast.error(error instanceof Error ? error.message : t(($) => $.company_codex.quota_increase_failed)); + } finally { + setIncreasingQuota(false); + } + }; + const copy = async (field: Exclude, value: string) => { if (await copyText(value)) { setCopied(field); @@ -133,17 +156,50 @@ export function CompanyCodexSection() { {keyStatus?.active ? ( -

- {t(($) => $.company_codex.key_metadata, { - prefix: keyStatus.key_prefix ?? "", - date: keyStatus.created_at - ? new Date(keyStatus.created_at).toLocaleDateString() - : "", - })} -

+
+

+ {t(($) => $.company_codex.key_metadata, { + prefix: keyStatus.key_prefix ?? "", + date: keyStatus.created_at + ? new Date(keyStatus.created_at).toLocaleDateString() + : "", + })} +

+ {keyStatus.weekly_token_limit ? ( +

+ {t(($) => $.company_codex.quota_summary, { + used: (keyStatus.current_tokens ?? 0).toLocaleString(), + limit: keyStatus.weekly_token_limit.toLocaleString(), + })} + {keyStatus.reset_at + ? ` · ${t(($) => $.company_codex.quota_resets, { + date: new Date(keyStatus.reset_at).toLocaleString(), + })}` + : ""} +

+ ) : ( +

{t(($) => $.company_codex.quota_unavailable)}

+ )} + {keyStatus.upstream_remaining_percent !== null && + keyStatus.upstream_remaining_percent !== undefined ? ( +

{t(($) => $.company_codex.upstream_remaining, { + percent: keyStatus.upstream_remaining_percent.toLocaleString(), + })}

+ ) : null} +
) : null} -
+
+ {keyStatus?.active ? ( + + ) : null} {keyStatus?.active ? ( + + + + + { if (!open) setIssued(null); }}> diff --git a/server/cmd/server/router.go b/server/cmd/server/router.go index 2fe0c8d5967..42e9d9c1ddf 100644 --- a/server/cmd/server/router.go +++ b/server/cmd/server/router.go @@ -1334,6 +1334,7 @@ func NewRouterWithOptions(pool *pgxpool.Pool, hub *realtime.Hub, bus *events.Bus r.Use(handler.RequireHumanActor) r.Get("/key", h.GetCompanyCodexKey) r.Post("/key", h.CreateCompanyCodexKey) + r.Patch("/key", h.IncreaseCompanyCodexKeyQuota) r.Delete("/key", h.RevokeCompanyCodexKey) r.Get("/sessions", h.ListCompanyCodexSessions) r.Get("/sessions/{id}", h.GetCompanyCodexSession) diff --git a/server/internal/handler/company_codex.go b/server/internal/handler/company_codex.go index 5637ba4ed1c..af366819972 100644 --- a/server/internal/handler/company_codex.go +++ b/server/internal/handler/company_codex.go @@ -19,19 +19,36 @@ import ( ) const companyCodexBodyLimit = 2 << 20 +const companyCodexMaxSafeInteger = int64(9007199254740991) type companyCodexKeyStatus struct { - Active bool `json:"active"` - KeyPrefix string `json:"key_prefix,omitempty"` - CreatedAt string `json:"created_at,omitempty"` + Active bool `json:"active"` + KeyPrefix string `json:"key_prefix,omitempty"` + CreatedAt string `json:"created_at,omitempty"` + WeeklyTokenLimit int64 `json:"weekly_token_limit,omitempty"` + CurrentTokens int64 `json:"current_tokens,omitempty"` + ResetAt string `json:"reset_at,omitempty"` + UpstreamRemainingPercent *float64 `json:"upstream_remaining_percent,omitempty"` } type companyCodexBrokerCreateResponse struct { GatewayKeyID string `json:"gateway_key_id"` GatewayKeyPrefix string `json:"gateway_key_prefix"` + WeeklyTokenLimit int64 `json:"weekly_token_limit"` Credential string `json:"credential"` } +type companyCodexQuotaResponse struct { + WeeklyTokenLimit int64 `json:"weekly_token_limit"` + CurrentTokens int64 `json:"current_tokens"` + ResetAt string `json:"reset_at"` + UpstreamRemainingPercent *float64 `json:"upstream_remaining_percent"` +} + +type companyCodexQuotaIncreaseRequest struct { + AdditionalTokens int64 `json:"additional_tokens"` +} + type companyCodexCreateResponse struct { companyCodexKeyStatus Credential string `json:"credential"` @@ -190,13 +207,13 @@ func (h *Handler) GetCompanyCodexKey(w http.ResponseWriter, r *http.Request) { return } workspaceID := ctxWorkspaceID(r.Context()) - var prefix string + var gatewayKeyID, prefix string var createdAt time.Time err := h.DB.QueryRow(r.Context(), ` - SELECT gateway_key_prefix, created_at + SELECT gateway_key_id, gateway_key_prefix, created_at FROM company_codex_key WHERE workspace_id = $1 AND user_id = $2 AND status = 'active' - `, workspaceID, userID).Scan(&prefix, &createdAt) + `, workspaceID, userID).Scan(&gatewayKeyID, &prefix, &createdAt) if errors.Is(err, pgx.ErrNoRows) { writeJSON(w, http.StatusOK, companyCodexKeyStatus{Active: false}) return @@ -205,11 +222,22 @@ func (h *Handler) GetCompanyCodexKey(w http.ResponseWriter, r *http.Request) { writeError(w, http.StatusInternalServerError, "failed to load company Codex access") return } - writeJSON(w, http.StatusOK, companyCodexKeyStatus{ + status := companyCodexKeyStatus{ Active: true, KeyPrefix: prefix, CreatedAt: createdAt.UTC().Format(time.RFC3339), - }) + } + var quota companyCodexQuotaResponse + if err := h.companyCodexBrokerRequest(r.Context(), http.MethodGet, + "/internal/company-codex/keys/"+url.PathEscape(gatewayKeyID), nil, "a); err != nil { + slog.Warn("company Codex quota lookup failed", "error", err, "gateway_key_id", gatewayKeyID) + } else { + status.WeeklyTokenLimit = quota.WeeklyTokenLimit + status.CurrentTokens = quota.CurrentTokens + status.ResetAt = quota.ResetAt + status.UpstreamRemainingPercent = quota.UpstreamRemainingPercent + } + writeJSON(w, http.StatusOK, status) } func (h *Handler) CreateCompanyCodexKey(w http.ResponseWriter, r *http.Request) { @@ -307,9 +335,10 @@ requires_openai_auth = true w.Header().Set("Cache-Control", "no-store") writeJSON(w, http.StatusCreated, companyCodexCreateResponse{ companyCodexKeyStatus: companyCodexKeyStatus{ - Active: true, - KeyPrefix: created.GatewayKeyPrefix, - CreatedAt: storedCreatedAt.UTC().Format(time.RFC3339), + Active: true, + KeyPrefix: created.GatewayKeyPrefix, + CreatedAt: storedCreatedAt.UTC().Format(time.RFC3339), + WeeklyTokenLimit: created.WeeklyTokenLimit, }, Credential: created.Credential, CCSwitchURL: deepLink.String(), @@ -318,6 +347,67 @@ requires_openai_auth = true }) } +func (h *Handler) IncreaseCompanyCodexKeyQuota(w http.ResponseWriter, r *http.Request) { + userID, ok := requireUserID(w, r) + if !ok { + return + } + workspaceID := ctxWorkspaceID(r.Context()) + var requested companyCodexQuotaIncreaseRequest + if !decodeCompanyCodexJSON(w, r, &requested) { + return + } + if requested.AdditionalTokens <= 0 || requested.AdditionalTokens > companyCodexMaxSafeInteger { + writeError(w, http.StatusBadRequest, "additional_tokens must be a positive safe integer") + return + } + + var gatewayKeyID, prefix string + var createdAt time.Time + err := h.DB.QueryRow(r.Context(), ` + SELECT gateway_key_id, gateway_key_prefix, created_at + FROM company_codex_key + WHERE workspace_id = $1 AND user_id = $2 AND status = 'active' + `, workspaceID, userID).Scan(&gatewayKeyID, &prefix, &createdAt) + if errors.Is(err, pgx.ErrNoRows) { + writeError(w, http.StatusNotFound, "active company Codex access required") + return + } + if err != nil { + writeError(w, http.StatusInternalServerError, "failed to load company Codex access") + return + } + + var quota companyCodexQuotaResponse + err = h.companyCodexBrokerRequest(r.Context(), http.MethodPatch, + "/internal/company-codex/keys/"+url.PathEscape(gatewayKeyID), map[string]any{ + "workspace_id": workspaceID, + "user_id": userID, + "additional_tokens": requested.AdditionalTokens, + }, "a) + if err != nil { + slog.Warn("company Codex quota increase failed", "error", err, "gateway_key_id", gatewayKeyID) + writeError(w, http.StatusBadGateway, "company Codex quota is temporarily unavailable") + return + } + slog.Info("company Codex quota increased", + "user_id", userID, + "workspace_id", workspaceID, + "gateway_key_id", gatewayKeyID, + "additional_tokens", requested.AdditionalTokens, + "weekly_token_limit", quota.WeeklyTokenLimit, + ) + writeJSON(w, http.StatusOK, companyCodexKeyStatus{ + Active: true, + KeyPrefix: prefix, + CreatedAt: createdAt.UTC().Format(time.RFC3339), + WeeklyTokenLimit: quota.WeeklyTokenLimit, + CurrentTokens: quota.CurrentTokens, + ResetAt: quota.ResetAt, + UpstreamRemainingPercent: quota.UpstreamRemainingPercent, + }) +} + func (h *Handler) RevokeCompanyCodexKey(w http.ResponseWriter, r *http.Request) { userID, ok := requireUserID(w, r) if !ok { diff --git a/server/internal/handler/company_codex_test.go b/server/internal/handler/company_codex_test.go index fd88f62044a..e6a052766da 100644 --- a/server/internal/handler/company_codex_test.go +++ b/server/internal/handler/company_codex_test.go @@ -1,10 +1,17 @@ package handler import ( + "context" + "encoding/json" + "net/http" "net/http/httptest" "strings" "testing" "unicode/utf8" + + "github.com/multica-ai/multica/server/internal/middleware" + "github.com/multica-ai/multica/server/internal/util" + db "github.com/multica-ai/multica/server/pkg/db/generated" ) func TestTruncateCompanyCodexTextPreservesUTF8(t *testing.T) { @@ -42,3 +49,128 @@ func TestRequireCompanyCodexInternal(t *testing.T) { }) } } + +func TestCompanyCodexQuotaStatusAndIncrease(t *testing.T) { + if testHandler == nil || testPool == nil { + t.Skip("database not available") + } + + const gatewayKeyID = "gateway-company-codex-quota-test" + if _, err := testPool.Exec(context.Background(), ` + DELETE FROM company_codex_key WHERE workspace_id = $1 AND user_id = $2 + `, testWorkspaceID, testUserID); err != nil { + t.Fatalf("clear company Codex key: %v", err) + } + if _, err := testPool.Exec(context.Background(), ` + INSERT INTO company_codex_key ( + workspace_id, user_id, gateway_key_id, gateway_key_prefix, status + ) VALUES ($1, $2, $3, 'sk-clb-quota', 'active') + `, testWorkspaceID, testUserID, gatewayKeyID); err != nil { + t.Fatalf("insert company Codex key: %v", err) + } + t.Cleanup(func() { + _, _ = testPool.Exec(context.Background(), ` + DELETE FROM company_codex_key WHERE workspace_id = $1 AND user_id = $2 + `, testWorkspaceID, testUserID) + }) + + weeklyLimit := int64(25_000_000) + broker := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("X-Company-Codex-Secret"); got != "test-broker-secret" { + http.Error(w, "missing broker secret", http.StatusUnauthorized) + return + } + if r.URL.Path != "/internal/company-codex/keys/"+gatewayKeyID { + http.NotFound(w, r) + return + } + if r.Method == http.MethodPatch { + var body struct { + WorkspaceID string `json:"workspace_id"` + UserID string `json:"user_id"` + AdditionalTokens int64 `json:"additional_tokens"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + if body.WorkspaceID != testWorkspaceID || body.UserID != testUserID { + http.Error(w, "wrong member identity", http.StatusForbidden) + return + } + weeklyLimit += body.AdditionalTokens + } else if r.Method != http.MethodGet { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + writeJSON(w, http.StatusOK, companyCodexQuotaResponse{ + WeeklyTokenLimit: weeklyLimit, + CurrentTokens: 5_000_000, + ResetAt: "2026-08-28T00:00:00Z", + UpstreamRemainingPercent: float64Pointer(74.25), + }) + })) + defer broker.Close() + + h := *testHandler + h.cfg.CompanyCodexBrokerURL = broker.URL + h.cfg.CompanyCodexBrokerSecret = "test-broker-secret" + memberRow, err := h.Queries.GetMemberByUserAndWorkspace(context.Background(), db.GetMemberByUserAndWorkspaceParams{ + UserID: util.MustParseUUID(testUserID), + WorkspaceID: util.MustParseUUID(testWorkspaceID), + }) + if err != nil { + t.Fatalf("load test member: %v", err) + } + withWorkspace := func(req *http.Request) *http.Request { + return req.WithContext(middleware.SetMemberContext(req.Context(), testWorkspaceID, memberRow)) + } + + getResponse := httptest.NewRecorder() + h.GetCompanyCodexKey(getResponse, withWorkspace(newRequest(http.MethodGet, "/api/company-codex/key", nil))) + if getResponse.Code != http.StatusOK { + t.Fatalf("GetCompanyCodexKey status = %d: %s", getResponse.Code, getResponse.Body.String()) + } + var status companyCodexKeyStatus + if err := json.NewDecoder(getResponse.Body).Decode(&status); err != nil { + t.Fatalf("decode quota status: %v", err) + } + if status.WeeklyTokenLimit != 25_000_000 || status.CurrentTokens != 5_000_000 { + t.Fatalf("quota status = %+v", status) + } + + patchResponse := httptest.NewRecorder() + h.IncreaseCompanyCodexKeyQuota(patchResponse, withWorkspace(newRequest( + http.MethodPatch, + "/api/company-codex/key", + map[string]any{"additional_tokens": 7_000_000}, + ))) + if patchResponse.Code != http.StatusOK { + t.Fatalf("IncreaseCompanyCodexKeyQuota status = %d: %s", patchResponse.Code, patchResponse.Body.String()) + } + if err := json.NewDecoder(patchResponse.Body).Decode(&status); err != nil { + t.Fatalf("decode increased quota: %v", err) + } + if status.WeeklyTokenLimit != 32_000_000 || status.KeyPrefix != "sk-clb-quota" { + t.Fatalf("increased quota status = %+v", status) + } +} + +func TestIncreaseCompanyCodexKeyQuotaRejectsNonPositiveAmount(t *testing.T) { + h := &Handler{} + for _, additionalTokens := range []int64{0, -1} { + response := httptest.NewRecorder() + h.IncreaseCompanyCodexKeyQuota(response, newRequest( + http.MethodPatch, + "/api/company-codex/key", + map[string]any{"additional_tokens": additionalTokens}, + )) + if response.Code != http.StatusBadRequest { + t.Fatalf("additional_tokens=%d status = %d, want 400", additionalTokens, response.Code) + } + } +} + +func float64Pointer(value float64) *float64 { + return &value +} From 4cc1bdfbe9550cefc4e138719d75ac5809be0c44 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B4=AA=E6=B3=BD=E9=91=AB?= Date: Fri, 21 Aug 2026 12:14:51 +0800 Subject: [PATCH 2/2] fix(company-codex): avoid logging quota request input --- server/internal/handler/company_codex.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/server/internal/handler/company_codex.go b/server/internal/handler/company_codex.go index af366819972..a3eed69d2a3 100644 --- a/server/internal/handler/company_codex.go +++ b/server/internal/handler/company_codex.go @@ -391,10 +391,7 @@ func (h *Handler) IncreaseCompanyCodexKeyQuota(w http.ResponseWriter, r *http.Re return } slog.Info("company Codex quota increased", - "user_id", userID, - "workspace_id", workspaceID, "gateway_key_id", gatewayKeyID, - "additional_tokens", requested.AdditionalTokens, "weekly_token_limit", quota.WeeklyTokenLimit, ) writeJSON(w, http.StatusOK, companyCodexKeyStatus{