From 4ce221c22aedd8d5c24f1318c72cd83bb851e636 Mon Sep 17 00:00:00 2001 From: yuez Date: Fri, 4 Sep 2026 10:39:04 +0800 Subject: [PATCH] feat: initial implementation of CLIProxyAPI Command Code plugin (stage 1) --- .gitignore | 17 + LICENSE | 21 + Makefile | 24 + README.md | 264 ++++++++++ assets/logo.svg | 14 + go.mod | 5 + go.sum | 4 + main.go | 155 ++++++ plugin/auth.go | 166 ++++++ plugin/auth_test.go | 166 ++++++ plugin/management.go | 208 ++++++++ plugin/management_test.go | 152 ++++++ plugin/plugin.go | 285 ++++++++++ plugin/plugin_test.go | 139 +++++ plugin/quota_page.go | 1028 +++++++++++++++++++++++++++++++++++++ plugin/types.go | 280 ++++++++++ plugin/usage.go | 259 ++++++++++ plugin/usage_test.go | 222 ++++++++ 18 files changed, 3409 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 Makefile create mode 100644 README.md create mode 100644 assets/logo.svg create mode 100644 go.mod create mode 100644 go.sum create mode 100644 main.go create mode 100644 plugin/auth.go create mode 100644 plugin/auth_test.go create mode 100644 plugin/management.go create mode 100644 plugin/management_test.go create mode 100644 plugin/plugin.go create mode 100644 plugin/plugin_test.go create mode 100644 plugin/quota_page.go create mode 100644 plugin/types.go create mode 100644 plugin/usage.go create mode 100644 plugin/usage_test.go diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..50a1a62 --- /dev/null +++ b/.gitignore @@ -0,0 +1,17 @@ +# Binaries & Shared Libraries +*.dylib +*.so +*.dll +*.h +*.exe +*.test + +# Go workspace & coverage +coverage.txt +*.out + +# OS & IDE +.DS_Store +.idea/ +.vscode/ +*.swp diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..f1ddef3 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 yuez (zgs225) + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..62f8247 --- /dev/null +++ b/Makefile @@ -0,0 +1,24 @@ +UNAME_S := $(shell uname -s) +ifeq ($(UNAME_S),Darwin) + TARGET := commandcode.dylib +else ifeq ($(OS),Windows_NT) + TARGET := commandcode.dll +else + TARGET := commandcode.so +endif + +.PHONY: all build test clean lint + +all: build + +build: + CGO_ENABLED=1 go build -buildmode=c-shared -o $(TARGET) main.go + +test: + go test -v -race ./... + +clean: + rm -f commandcode.dylib commandcode.so commandcode.dll commandcode.h + +lint: + go vet ./... diff --git a/README.md b/README.md new file mode 100644 index 0000000..6aac73c --- /dev/null +++ b/README.md @@ -0,0 +1,264 @@ +# CLIProxyAPI Command Code Plugin (`commandcode`) + +[![Go Version](https://img.shields.io/badge/Go-1.22%2B-blue.svg)](https://golang.org) +[![CLIProxyAPI Plugin ABI](https://img.shields.io/badge/C%20ABI-v1-emerald.svg)](https://help.router-for.me/plugin/development.html) +[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE) + +[CLIProxyAPI](https://github.com/router-for-me/CLIProxyAPI) 动态 C ABI 插件，用于提供 **Command Code** 凭据认证、上游配额与窗口限额查询、以及嵌入式配额监控仪表盘卡片（QuotaCard）。 + +--- + +## 目录 + +- [功能特性](#功能特性) +- [系统架构](#系统架构) +- [快速开始](#快速开始) + - [构建插件](#构建插件) + - [安装与目录结构](#安装与目录结构) + - [宿主配置 (`config.yaml`)](#宿主配置-configyaml) + - [凭据文件配置](#凭据文件配置) +- [管理端点与资源页](#管理端点与资源页) + - [1. 浏览器资源页 (`QuotaCard`)](#1-浏览器资源页-quotacard) + - [2. 管理 API: 查询用量 (`GET`)](#2-管理-api-查询用量-get) + - [3. 管理 API: 测试用量 (`POST`)](#3-管理-api-测试用量-post) +- [用量数据结构说明](#用量数据结构说明) +- [开发与测试](#开发与测试) +- [许可证](#许可证) + +--- + +## 功能特性 + +1. **标准 C ABI 兼容**： + - 导出 `cliproxy_plugin_init`、`cliproxyPluginCall`、`cliproxyPluginFree`、`cliproxyPluginShutdown`。 + - 遵照 CLIProxyAPI 官方 JSON Envelope 规范（`ok`, `result`, `error`）。 +2. **双核心能力声明**： + - `auth_provider`: 参与凭据识别、加载、解析与刷新。 + - `management_api`: 注册插件自有的管理端点与浏览器资源页面。 +3. **凭据自动解析 (`auth.parse`)**： + - 自动识别 `commandcode-*.json` 凭据文件、`type: "commandcode"` 配置或包含 `session_token` / Cookie 的凭据。 + - 提取并规范化 `__Secure-commandcode_prod_.session_token`，存入宿主持久化凭据库。 +4. **精确用量与双滑动窗口限额解析**： + - 上游接口：`GET https://api.commandcode.ai/internal/billing/credits`。 + - 请求优先走宿主提供的 `host.http.do` 回调（复用宿主代理、日志与鉴权管道），离线或未注入宿主时自动无缝降级至 Go 标准 `net/http`。 + - 全面解析 `credits`（月度基础额度、开源奖励额度、总可用额度）与 `windowLimits`（5小时短期滑动窗口、周度窗口限额，计算已用量、上限、剩余量、使用百分比及重置时间）。 +5. **嵌入式纯单文件 QuotaCard 资源页**： + - 页面挂载于 `/v0/resource/plugins/commandcode/quota`。 + - 零外部 CDN 依赖，纯内置 HTML + CSS + JS，深色/浅色模式自适应。 + - 具有进度条颜色变化、5小时/周限额卡片、秒级动态重置倒计时、同源 `localStorage` 鉴权与一键刷新。 + +--- + +## 系统架构 + +``` +┌────────────────────────────────────────────────────────┐ +│ CLIProxyAPI │ +│ │ +│ ┌─────────────────────────┐ ┌─────────────────────┐ │ +│ │ Auth Management │ │ Management Center │ │ +│ │ (reads auths/*.json) │ │ (/v0/management) │ │ +│ └───────────┬─────────────┘ └──────────┬──────────┘ │ +│ │ C ABI │ C ABI │ +│ ▼ ▼ │ +│ ┌──────────────────────────────────────────────────┐ │ +│ │ cliproxy-plugin-commandcode.dylib/.so │ │ +│ │ │ │ +│ │ • auth.identifier / auth.parse │ │ +│ │ • management.register / management.handle │ │ +│ │ • Usage Parser & Window Limits Formatter │ │ +│ │ • Embedded Single-file HTML/CSS/JS QuotaCard │ │ +│ └───────────────────────────┬──────────────────────┘ │ +│ │ │ +│ │ host.http.do │ +│ ▼ │ +│ ┌──────────────────────────────────────────────────┐ │ +│ │ Host Transport / Proxy Pipeline │ │ +│ └───────────────────────────┬──────────────────────┘ │ +└──────────────────────────────┼─────────────────────────┘ + │ Upstream HTTPS + ▼ + https://api.commandcode.ai/internal/billing/credits +``` + +--- + +## 快速开始 + +### 构建插件 + +项目提供标准的 `Makefile`，可直接编译与操作系统相对应的 C 共享动态库： + +```bash +# 自动编译出 commandcode.dylib (macOS) 或 commandcode.so (Linux) +make build + +# 运行完整单元测试与竞态检测 +make test + +# 清理构建产物 +make clean +``` + +### 安装与目录结构 + +将编译出的动态库放入 CLIProxyAPI 的插件目录中： + +```bash +# macOS +mkdir -p plugins/darwin/arm64 +cp commandcode.dylib plugins/darwin/arm64/commandcode.dylib + +# Linux +mkdir -p plugins/linux/amd64 +cp commandcode.so plugins/linux/amd64/commandcode.so +``` + +### 宿主配置 (`config.yaml`) + +在 CLIProxyAPI 的 `config.yaml` 中启用插件并配置默认参数： + +```yaml +plugins: + enabled: true + dir: "plugins" + configs: + commandcode: + enabled: true + priority: 1 + session_token: "YOUR_COMMANDCODE_SESSION_TOKEN" + api_base: "https://api.commandcode.ai" # 可选，默认为官方接口 +``` + +### 凭据文件配置 + +除了在 `config.yaml` 中全局配置，你也可以在 CLIProxyAPI 的 `auths/` 凭据目录下创建凭据文件（如 `auths/commandcode-main.json`）： + +```json +{ + "type": "commandcode", + "session_token": "YOUR_COMMANDCODE_SESSION_TOKEN", + "email": "user@example.com", + "label": "Command Code Pro" +} +``` + +或者直接放入浏览器 Cookie： + +```json +{ + "type": "commandcode", + "cookie": "__Secure-commandcode_prod_.session_token=YOUR_COMMANDCODE_SESSION_TOKEN; Path=/;" +} +``` + +插件的 `auth.parse` 会自动拦截并完成凭据加载。 + +--- + +## 管理端点与资源页 + +### 1. 浏览器资源页 (`QuotaCard`) + +- **访问路径**：`GET http://:8317/v0/resource/plugins/commandcode/quota` +- **菜单名**：`Command Code 配额` +- **说明**： + - 资源请求本身无需经过管理认证，可在浏览器中直接打开或嵌入仪表盘。 + - 在同源模式下，页面 JavaScript 会自动读取 `localStorage` 中的管理密钥向 `/v0/management/plugins/commandcode/usage` 请求数据。 + - 若在独立或跨域测试环境下打开，页面提供内置的诊断面板，可手动输入 Management Key 或测试 Session Token。 + +### 2. 管理 API: 查询用量 (`GET`) + +- **端点**：`GET /v0/management/plugins/commandcode/usage` +- **认证**：需要管理密钥 (`Authorization: Bearer ` 或 `X-Management-Key: `) +- **可选查询参数**： + - `session_token`: 临时覆盖查询的 token + - `api_base`: 临时覆盖的上游基础 URL +- **响应示例**： + +```json +{ + "ok": true, + "credits": { + "monthly_credits": 1000.0, + "opensource_monthly_credits": 500.0, + "total_credits": 1500.0, + "details": { + "monthlyCredits": 1000, + "opensourceMonthlyCredits": 500 + } + }, + "window_limits": { + "five_hour": { + "used": 12.5, + "cap": 100.0, + "remaining": 87.5, + "percentage": 12.5, + "exceeded": false, + "reset_at": "2025-03-04T16:30:00Z", + "reset_in_seconds": 7200 + }, + "weekly": { + "used": 150.0, + "cap": 1000.0, + "remaining": 850.0, + "percentage": 15.0, + "exceeded": false, + "reset_at": "2025-03-10T00:00:00Z", + "reset_in_seconds": 475200 + } + }, + "updated_at": "2025-03-04T14:30:00Z" +} +``` + +### 3. 管理 API: 测试用量 (`POST`) + +- **端点**：`POST /v0/management/plugins/commandcode/usage` +- **请求体**： + +```json +{ + "session_token": "YOUR_TEMPORARY_TOKEN", + "api_base": "https://api.commandcode.ai" +} +``` + +--- + +## 用量数据结构说明 + +| 字段 | 类型 | 说明 | +| :--- | :--- | :--- | +| `credits.monthly_credits` | `float64` | 当前账单周期的月度基础 Credits 额度 | +| `credits.opensource_monthly_credits` | `float64` | 开源项目贡献者获得的奖励额度 | +| `credits.total_credits` | `float64` | 可用 Credits 总计 (`monthly + opensource`) | +| `window_limits.five_hour.used` | `float64` | 5小时滑动窗口内已消耗的量 | +| `window_limits.five_hour.cap` | `float64` | 5小时滑动窗口上限 | +| `window_limits.five_hour.remaining` | `float64` | 5小时滑动窗口剩余可用量 | +| `window_limits.five_hour.percentage` | `float64` | 5小时窗口使用百分比（0-100%） | +| `window_limits.five_hour.exceeded` | `bool` | 是否已触发 5 小时限额熔断 | +| `window_limits.five_hour.reset_at` | `string` | 5小时窗口重置时间的 RFC3339 字符串 | +| `window_limits.five_hour.reset_in_seconds`| `int64` | 距离 5 小时窗口重置的剩余秒数 | +| `window_limits.weekly.*` | - | 每周限额对应指标（结构同 5 小时窗口） | + +--- + +## 开发与测试 + +```bash +# 运行单元测试 +go test -v ./... + +# 运行代码规范检查 +go vet ./... + +# 运行竞态检查测试 +go test -race -v ./... +``` + +--- + +## 许可证 + +本项目基于 [MIT License](LICENSE) 开源。 diff --git a/assets/logo.svg b/assets/logo.svg new file mode 100644 index 0000000..5b35629 --- /dev/null +++ b/assets/logo.svg @@ -0,0 +1,14 @@ + + + + + + + + + + + + + + diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..e7a7f4f --- /dev/null +++ b/go.mod @@ -0,0 +1,5 @@ +module github.com/zgs225/cliproxy-plugin-commandcode + +go 1.22 + +require gopkg.in/yaml.v3 v3.0.1 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..a62c313 --- /dev/null +++ b/go.sum @@ -0,0 +1,4 @@ +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/main.go b/main.go new file mode 100644 index 0000000..f1a75f7 --- /dev/null +++ b/main.go @@ -0,0 +1,155 @@ +package main + +/* +#include +#include + +typedef struct { + void* ptr; + size_t len; +} cliproxy_buffer; + +typedef int (*cliproxy_host_call_fn)(void*, const char*, const uint8_t*, size_t, cliproxy_buffer*); +typedef void (*cliproxy_host_free_fn)(void*, size_t); + +typedef struct { + uint32_t abi_version; + void* host_ctx; + cliproxy_host_call_fn call; + cliproxy_host_free_fn free_buffer; +} cliproxy_host_api; + +typedef int (*cliproxy_plugin_call_fn)(char*, uint8_t*, size_t, cliproxy_buffer*); +typedef void (*cliproxy_plugin_free_fn)(void*, size_t); +typedef void (*cliproxy_plugin_shutdown_fn)(void); + +typedef struct { + uint32_t abi_version; + cliproxy_plugin_call_fn call; + cliproxy_plugin_free_fn free_buffer; + cliproxy_plugin_shutdown_fn shutdown; +} cliproxy_plugin_api; + +extern int cliproxyPluginCall(char*, uint8_t*, size_t, cliproxy_buffer*); +extern void cliproxyPluginFree(void*, size_t); +extern void cliproxyPluginShutdown(void); + +static const cliproxy_host_api* stored_host; + +static void store_host_api(const cliproxy_host_api* host) { + stored_host = host; +} + +static int call_host_api(const char* method, const uint8_t* request, size_t request_len, cliproxy_buffer* response) { + if (stored_host == NULL || stored_host->call == NULL) { + return 1; + } + return stored_host->call(stored_host->host_ctx, method, request, request_len, response); +} + +static void free_host_buffer(void* ptr, size_t len) { + if (stored_host != NULL && stored_host->free_buffer != NULL && ptr != NULL) { + stored_host->free_buffer(ptr, len); + } +} +*/ +import "C" + +import ( + "fmt" + "unsafe" + + "github.com/zgs225/cliproxy-plugin-commandcode/plugin" +) + +const abiVersion uint32 = 1 + +func main() {} + +func init() { + plugin.SetHostCaller(callHost) +} + +//export cliproxy_plugin_init +func cliproxy_plugin_init(host *C.cliproxy_host_api, pluginAPI *C.cliproxy_plugin_api) C.int { + if pluginAPI == nil { + return 1 + } + C.store_host_api(host) + pluginAPI.abi_version = C.uint32_t(abiVersion) + pluginAPI.call = C.cliproxy_plugin_call_fn(C.cliproxyPluginCall) + pluginAPI.free_buffer = C.cliproxy_plugin_free_fn(C.cliproxyPluginFree) + pluginAPI.shutdown = C.cliproxy_plugin_shutdown_fn(C.cliproxyPluginShutdown) + return 0 +} + +//export cliproxyPluginCall +func cliproxyPluginCall(method *C.char, request *C.uint8_t, requestLen C.size_t, response *C.cliproxy_buffer) C.int { + if response != nil { + response.ptr = nil + response.len = 0 + } + if method == nil { + writeResponse(response, plugin.ErrorEnvelope("invalid_method", "method is required")) + return 1 + } + + var requestBytes []byte + if request != nil && requestLen > 0 { + requestBytes = C.GoBytes(unsafe.Pointer(request), C.int(requestLen)) + } + + raw, errHandle := plugin.DefaultPlugin().HandleMethod(C.GoString(method), requestBytes) + if errHandle != nil { + writeResponse(response, plugin.ErrorEnvelope("plugin_error", errHandle.Error())) + return 1 + } + + writeResponse(response, raw) + return 0 +} + +//export cliproxyPluginFree +func cliproxyPluginFree(ptr unsafe.Pointer, len C.size_t) { + if ptr != nil { + C.free(ptr) + } + _ = len +} + +//export cliproxyPluginShutdown +func cliproxyPluginShutdown() {} + +func writeResponse(response *C.cliproxy_buffer, raw []byte) { + if response == nil || len(raw) == 0 { + return + } + ptr := C.CBytes(raw) + if ptr == nil { + return + } + response.ptr = ptr + response.len = C.size_t(len(raw)) +} + +func callHost(method string, payload []byte) ([]byte, error) { + cMethod := C.CString(method) + defer C.free(unsafe.Pointer(cMethod)) + + var response C.cliproxy_buffer + var req *C.uint8_t + if len(payload) > 0 { + req = (*C.uint8_t)(C.CBytes(payload)) + defer C.free(unsafe.Pointer(req)) + } + + rc := C.call_host_api(cMethod, req, C.size_t(len(payload)), &response) + if rc != 0 { + return nil, fmt.Errorf("host call %s failed with code %d", method, int(rc)) + } + if response.ptr == nil || response.len == 0 { + return nil, nil + } + defer C.free_host_buffer(response.ptr, response.len) + return C.GoBytes(response.ptr, C.int(response.len)), nil +} diff --git a/plugin/auth.go b/plugin/auth.go new file mode 100644 index 0000000..533db64 --- /dev/null +++ b/plugin/auth.go @@ -0,0 +1,166 @@ +package plugin + +import ( + "encoding/json" + "fmt" + "path/filepath" + "regexp" + "strings" + "time" +) + +var cookieRegex = regexp.MustCompile(`(?:^|;\s*)__Secure-commandcode_prod_\.session_token=([^;]+)`) + +// RawAuthContent represents possible structures inside a commandcode credential JSON file. +type RawAuthContent struct { + Type string `json:"type"` + Provider string `json:"provider"` + ID string `json:"id"` + Label string `json:"label"` + Name string `json:"name"` + Email string `json:"email"` + SessionToken string `json:"session_token"` + CommandCodeSession string `json:"commandcode_session_token"` + Cookie string `json:"cookie"` + Token string `json:"token"` + UpstreamBase string `json:"api_base"` +} + +// ExtractSessionToken extracts the clean session token from a raw string or cookie string. +func ExtractSessionToken(raw string) string { + raw = strings.TrimSpace(raw) + if raw == "" { + return "" + } + if matches := cookieRegex.FindStringSubmatch(raw); len(matches) > 1 { + return strings.TrimSpace(matches[1]) + } + if strings.HasPrefix(raw, "__Secure-commandcode_prod_.session_token=") { + trimmed := strings.TrimPrefix(raw, "__Secure-commandcode_prod_.session_token=") + if idx := strings.Index(trimmed, ";"); idx != -1 { + trimmed = trimmed[:idx] + } + return strings.TrimSpace(trimmed) + } + return raw +} + +// FormatSessionCookie ensures the token is formatted as the upstream Cookie header value. +func FormatSessionCookie(token string) string { + clean := ExtractSessionToken(token) + if clean == "" { + return "" + } + return "__Secure-commandcode_prod_.session_token=" + clean +} + +// ParseAuth handles auth.parse requests for Command Code credentials. +func ParseAuth(req AuthParseRequest) (AuthParseResponse, error) { + lowerFileName := strings.ToLower(req.FileName) + isCommandCodeFile := strings.HasPrefix(lowerFileName, "commandcode") && strings.HasSuffix(lowerFileName, ".json") + isCommandCodeProvider := strings.EqualFold(req.Provider, PluginID) + + var content RawAuthContent + var rawMap map[string]any + if len(req.RawJSON) > 0 { + if err := json.Unmarshal(req.RawJSON, &content); err == nil { + _ = json.Unmarshal(req.RawJSON, &rawMap) + } + } + + isExplicitCommandCode := strings.EqualFold(content.Type, PluginID) || + strings.EqualFold(content.Provider, PluginID) || + content.SessionToken != "" || + content.CommandCodeSession != "" || + strings.Contains(content.Cookie, "__Secure-commandcode_prod_.session_token") + + if !isCommandCodeFile && !isCommandCodeProvider && !isExplicitCommandCode { + return AuthParseResponse{Handled: false}, nil + } + + // Extract session token + token := content.SessionToken + if token == "" { + token = content.CommandCodeSession + } + if token == "" && content.Cookie != "" { + token = ExtractSessionToken(content.Cookie) + } + if token == "" && (isCommandCodeFile || isCommandCodeProvider || isExplicitCommandCode) { + token = content.Token + } + token = ExtractSessionToken(token) + + // Determine ID + authID := content.ID + if authID == "" && req.FileName != "" { + base := filepath.Base(req.FileName) + authID = strings.TrimSuffix(base, filepath.Ext(base)) + } + if authID == "" { + authID = "commandcode-default" + } + + // Determine Label + label := content.Label + if label == "" { + label = content.Name + } + if label == "" && content.Email != "" { + label = fmt.Sprintf("Command Code (%s)", content.Email) + } + if label == "" { + label = fmt.Sprintf("Command Code (%s)", authID) + } + + // Build clean StorageJSON + storageMap := map[string]any{ + "type": PluginID, + "provider": PluginID, + "session_token": token, + } + if content.Email != "" { + storageMap["email"] = content.Email + } + if content.Label != "" { + storageMap["label"] = content.Label + } + if content.UpstreamBase != "" { + storageMap["api_base"] = content.UpstreamBase + } + for k, v := range rawMap { + if _, exists := storageMap[k]; !exists { + storageMap[k] = v + } + } + storageJSON, _ := json.Marshal(storageMap) + + metadata := map[string]any{ + "type": PluginID, + "session_token": token, + } + if content.Email != "" { + metadata["email"] = content.Email + } + + attributes := map[string]string{ + "provider": PluginID, + } + + authData := AuthData{ + Provider: PluginID, + ID: authID, + FileName: req.FileName, + Label: label, + Disabled: false, + StorageJSON: storageJSON, + Metadata: metadata, + Attributes: attributes, + NextRefreshAfter: time.Now().Add(24 * time.Hour).UTC(), + } + + return AuthParseResponse{ + Handled: true, + Auth: authData, + }, nil +} diff --git a/plugin/auth_test.go b/plugin/auth_test.go new file mode 100644 index 0000000..22d1917 --- /dev/null +++ b/plugin/auth_test.go @@ -0,0 +1,166 @@ +package plugin + +import ( + "encoding/json" + "testing" +) + +func TestExtractSessionToken(t *testing.T) { + tests := []struct { + name string + input string + expected string + }{ + { + name: "raw token", + input: "abc123token", + expected: "abc123token", + }, + { + name: "single cookie", + input: "__Secure-commandcode_prod_.session_token=secret_tok_123", + expected: "secret_tok_123", + }, + { + name: "cookie with semicolons and trailing params", + input: "__Secure-commandcode_prod_.session_token=secret_tok_123; Path=/; Secure; HttpOnly", + expected: "secret_tok_123", + }, + { + name: "multi-cookie string", + input: "some_other_cookie=xyz; __Secure-commandcode_prod_.session_token=secret_tok_123; foo=bar", + expected: "secret_tok_123", + }, + { + name: "empty string", + input: " ", + expected: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := ExtractSessionToken(tt.input) + if got != tt.expected { + t.Errorf("ExtractSessionToken(%q) = %q, want %q", tt.input, got, tt.expected) + } + }) + } +} + +func TestFormatSessionCookie(t *testing.T) { + got := FormatSessionCookie("my-token") + want := "__Secure-commandcode_prod_.session_token=my-token" + if got != want { + t.Errorf("FormatSessionCookie() = %q, want %q", got, want) + } + + gotCookie := FormatSessionCookie("__Secure-commandcode_prod_.session_token=my-token; Path=/") + if gotCookie != want { + t.Errorf("FormatSessionCookie() from cookie = %q, want %q", gotCookie, want) + } +} + +func TestParseAuth_ExplicitJSON(t *testing.T) { + raw := []byte(`{ + "type": "commandcode", + "session_token": "test-session-token-xyz", + "email": "user@example.com", + "label": "My Command Code Auth" + }`) + + resp, err := ParseAuth(AuthParseRequest{ + FileName: "custom.json", + RawJSON: raw, + }) + if err != nil { + t.Fatalf("ParseAuth error: %v", err) + } + if !resp.Handled { + t.Fatal("expected Handled=true for explicit commandcode type") + } + + auth := resp.Auth + if auth.Provider != PluginID { + t.Errorf("Provider = %q, want %q", auth.Provider, PluginID) + } + if auth.ID != "custom" { + t.Errorf("ID = %q, want %q", auth.ID, "custom") + } + if auth.Label != "My Command Code Auth" { + t.Errorf("Label = %q, want %q", auth.Label, "My Command Code Auth") + } + + var storage map[string]any + if err := json.Unmarshal(auth.StorageJSON, &storage); err != nil { + t.Fatalf("failed to unmarshal StorageJSON: %v", err) + } + if storage["session_token"] != "test-session-token-xyz" { + t.Errorf("StorageJSON session_token = %v, want test-session-token-xyz", storage["session_token"]) + } + + if auth.Metadata["session_token"] != "test-session-token-xyz" { + t.Errorf("Metadata session_token = %v, want test-session-token-xyz", auth.Metadata["session_token"]) + } +} + +func TestParseAuth_FileNameMatch(t *testing.T) { + raw := []byte(`{ + "token": "tok_987654" + }`) + + resp, err := ParseAuth(AuthParseRequest{ + FileName: "commandcode-work.json", + RawJSON: raw, + }) + if err != nil { + t.Fatalf("ParseAuth error: %v", err) + } + if !resp.Handled { + t.Fatal("expected Handled=true for commandcode-*.json filename") + } + if resp.Auth.ID != "commandcode-work" { + t.Errorf("ID = %q, want commandcode-work", resp.Auth.ID) + } + if resp.Auth.Metadata["session_token"] != "tok_987654" { + t.Errorf("session_token = %v, want tok_987654", resp.Auth.Metadata["session_token"]) + } +} + +func TestParseAuth_CookieFormat(t *testing.T) { + raw := []byte(`{ + "cookie": "__Secure-commandcode_prod_.session_token=cookie_tok_456; Path=/" + }`) + + resp, err := ParseAuth(AuthParseRequest{ + FileName: "any.json", + RawJSON: raw, + }) + if err != nil { + t.Fatalf("ParseAuth error: %v", err) + } + if !resp.Handled { + t.Fatal("expected Handled=true for cookie with __Secure-commandcode_prod_.session_token") + } + if resp.Auth.Metadata["session_token"] != "cookie_tok_456" { + t.Errorf("session_token = %v, want cookie_tok_456", resp.Auth.Metadata["session_token"]) + } +} + +func TestParseAuth_UnrelatedFile(t *testing.T) { + raw := []byte(`{ + "type": "openai", + "api_key": "sk-123456" + }`) + + resp, err := ParseAuth(AuthParseRequest{ + FileName: "openai-test.json", + RawJSON: raw, + }) + if err != nil { + t.Fatalf("ParseAuth error: %v", err) + } + if resp.Handled { + t.Fatal("expected Handled=false for unrelated credential file") + } +} diff --git a/plugin/management.go b/plugin/management.go new file mode 100644 index 0000000..2c7af20 --- /dev/null +++ b/plugin/management.go @@ -0,0 +1,208 @@ +package plugin + +import ( + "context" + "encoding/json" + "net/http" + "strings" + "time" +) + +// RegisterManagement handles management.register method. +func RegisterManagement() (ManagementRegistrationResponse, error) { + return ManagementRegistrationResponse{ + Routes: []ManagementRoute{ + { + Method: http.MethodGet, + Path: "/plugins/commandcode/usage", + Description: "Query Command Code credits and window limits usage", + }, + { + Method: http.MethodPost, + Path: "/plugins/commandcode/usage", + Description: "Query Command Code credits and window limits usage with custom session_token", + }, + }, + Resources: []ResourceRoute{ + { + Path: "/quota", + Menu: "Command Code 配额", + Description: "Command Code 用量与限额卡片", + }, + }, + }, nil +} + +// HandleManagement processes management.handle requests for API routes and resource pages. +func HandleManagement(ctx context.Context, req ManagementRequest, cfg *PluginConfig) (ManagementResponse, error) { + method := strings.ToUpper(strings.TrimSpace(req.Method)) + path := strings.TrimSpace(req.Path) + + // 1. Serve Quota Resource Page + if method == http.MethodGet && (strings.HasSuffix(path, "/quota") || strings.HasSuffix(path, "/quota/")) { + return ManagementResponse{ + StatusCode: http.StatusOK, + Headers: map[string][]string{ + "Content-Type": {"text/html; charset=utf-8"}, + }, + Body: GetQuotaPageHTML(), + }, nil + } + + // 2. Serve Usage API (GET / POST) + if strings.HasSuffix(path, "/plugins/commandcode/usage") || strings.HasSuffix(path, "/usage") { + switch method { + case http.MethodGet: + return handleGetUsage(ctx, req, cfg) + case http.MethodPost: + return handlePostUsage(ctx, req, cfg) + default: + return ManagementResponse{ + StatusCode: http.StatusMethodNotAllowed, + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + }, + Body: []byte(`{"ok":false,"error":"method not allowed"}`), + }, nil + } + } + + // Unknown path + return ManagementResponse{ + StatusCode: http.StatusNotFound, + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + }, + Body: []byte(`{"ok":false,"error":"not found"}`), + }, nil +} + +func handleGetUsage(ctx context.Context, req ManagementRequest, cfg *PluginConfig) (ManagementResponse, error) { + sessionToken := "" + apiBase := "" + + // Check query params + if len(req.Query) > 0 { + if tokens, ok := req.Query["session_token"]; ok && len(tokens) > 0 { + sessionToken = tokens[0] + } else if tokens, ok := req.Query["token"]; ok && len(tokens) > 0 { + sessionToken = tokens[0] + } + if bases, ok := req.Query["api_base"]; ok && len(bases) > 0 { + apiBase = bases[0] + } + } + + // Fallback to plugin config + if sessionToken == "" && cfg != nil { + sessionToken = cfg.GetSessionToken() + } + if apiBase == "" && cfg != nil { + apiBase = cfg.GetAPIBase() + } + + return executeUsageQuery(ctx, apiBase, sessionToken, req.HostCallbackID) +} + +func handlePostUsage(ctx context.Context, req ManagementRequest, cfg *PluginConfig) (ManagementResponse, error) { + var body struct { + SessionToken string `json:"session_token"` + Token string `json:"token"` + APIBase string `json:"api_base"` + } + + if len(req.Body) > 0 { + _ = json.Unmarshal(req.Body, &body) + } + + sessionToken := body.SessionToken + if sessionToken == "" { + sessionToken = body.Token + } + apiBase := body.APIBase + + // Fallback to plugin config if body didn't specify + if sessionToken == "" && cfg != nil { + sessionToken = cfg.GetSessionToken() + } + if apiBase == "" && cfg != nil { + apiBase = cfg.GetAPIBase() + } + + return executeUsageQuery(ctx, apiBase, sessionToken, req.HostCallbackID) +} + +func executeUsageQuery(ctx context.Context, apiBase, sessionToken, hostCallbackID string) (ManagementResponse, error) { + if strings.TrimSpace(sessionToken) == "" { + resBytes, _ := json.Marshal(map[string]any{ + "ok": false, + "error": "session_token is required. Configure session_token in plugin config, provide a credential file, or pass session_token in request", + }) + return ManagementResponse{ + StatusCode: http.StatusBadRequest, + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + }, + Body: resBytes, + }, nil + } + + raw, statusCode, errFetch := FetchCreditsRaw(ctx, apiBase, sessionToken, hostCallbackID) + if errFetch != nil { + resBytes, _ := json.Marshal(map[string]any{ + "ok": false, + "status_code": statusCode, + "error": errFetch.Error(), + }) + if statusCode == 0 || statusCode == http.StatusOK { + statusCode = http.StatusBadGateway + } + return ManagementResponse{ + StatusCode: statusCode, + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + }, + Body: resBytes, + }, nil + } + + if statusCode != http.StatusOK { + resBytes, _ := json.Marshal(map[string]any{ + "ok": false, + "status_code": statusCode, + "error": "upstream returned non-200 status", + "body": string(raw), + }) + return ManagementResponse{ + StatusCode: statusCode, + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + }, + Body: resBytes, + }, nil + } + + usage, errParse := ParseAndFormatUsage(raw, time.Now().UTC()) + if errParse != nil { + resBytes, _ := json.Marshal(map[string]any{ + "ok": false, + "error": "failed to parse upstream usage: " + errParse.Error(), + }) + return ManagementResponse{ + StatusCode: http.StatusBadGateway, + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + }, + Body: resBytes, + }, nil + } + + resBytes, _ := json.Marshal(usage) + return ManagementResponse{ + StatusCode: http.StatusOK, + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + }, + Body: resBytes, + }, nil +} diff --git a/plugin/management_test.go b/plugin/management_test.go new file mode 100644 index 0000000..0efa909 --- /dev/null +++ b/plugin/management_test.go @@ -0,0 +1,152 @@ +package plugin + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestRegisterManagement(t *testing.T) { + resp, err := RegisterManagement() + if err != nil { + t.Fatalf("RegisterManagement error: %v", err) + } + + if len(resp.Routes) != 2 { + t.Fatalf("len(Routes) = %d, want 2", len(resp.Routes)) + } + if resp.Routes[0].Method != http.MethodGet || resp.Routes[0].Path != "/plugins/commandcode/usage" { + t.Errorf("Route 0 mismatch: %+v", resp.Routes[0]) + } + if resp.Routes[1].Method != http.MethodPost || resp.Routes[1].Path != "/plugins/commandcode/usage" { + t.Errorf("Route 1 mismatch: %+v", resp.Routes[1]) + } + + if len(resp.Resources) != 1 { + t.Fatalf("len(Resources) = %d, want 1", len(resp.Resources)) + } + if resp.Resources[0].Path != "/quota" || resp.Resources[0].Menu != "Command Code 配额" { + t.Errorf("Resource 0 mismatch: %+v", resp.Resources[0]) + } +} + +func TestHandleManagement_QuotaResource(t *testing.T) { + paths := []string{ + "/quota", + "/v0/resource/plugins/commandcode/quota", + } + + for _, p := range paths { + req := ManagementRequest{ + Method: http.MethodGet, + Path: p, + } + resp, err := HandleManagement(context.Background(), req, nil) + if err != nil { + t.Fatalf("HandleManagement error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Errorf("StatusCode = %d, want 200", resp.StatusCode) + } + ct := resp.Headers["Content-Type"] + if len(ct) == 0 || !strings.Contains(ct[0], "text/html") { + t.Errorf("Content-Type = %v, want text/html", ct) + } + bodyStr := string(resp.Body) + if !strings.Contains(bodyStr, "Command Code 配额") { + t.Errorf("Body does not contain expected title") + } + } +} + +func TestHandleManagement_GetUsage(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "credits": {"monthlyCredits": 888}, + "windowLimits": {"fiveHour": {"used": 2, "cap": 20}} + }`)) + })) + defer ts.Close() + + SetHostCaller(nil) + SetDefaultHTTPClient(ts.Client()) + + cfg := &PluginConfig{ + SessionToken: "configured-token", + APIBase: ts.URL, + } + + req := ManagementRequest{ + Method: http.MethodGet, + Path: "/v0/management/plugins/commandcode/usage", + } + + resp, err := HandleManagement(context.Background(), req, cfg) + if err != nil { + t.Fatalf("HandleManagement error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("StatusCode = %d, want 200, body=%s", resp.StatusCode, string(resp.Body)) + } + + var usage FormattedUsageResponse + if err := json.Unmarshal(resp.Body, &usage); err != nil { + t.Fatalf("Unmarshal body error: %v", err) + } + if !usage.OK { + t.Fatal("expected OK=true") + } + if usage.Credits.MonthlyCredits != 888 { + t.Errorf("MonthlyCredits = %v, want 888", usage.Credits.MonthlyCredits) + } +} + +func TestHandleManagement_PostUsage(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + cookie := r.Header.Get("Cookie") + if !strings.Contains(cookie, "post-token-999") { + w.WriteHeader(http.StatusUnauthorized) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "credits": {"monthlyCredits": 666}, + "windowLimits": {"fiveHour": {"used": 1, "cap": 10}} + }`)) + })) + defer ts.Close() + + SetHostCaller(nil) + SetDefaultHTTPClient(ts.Client()) + + reqBody, _ := json.Marshal(map[string]string{ + "session_token": "post-token-999", + "api_base": ts.URL, + }) + + req := ManagementRequest{ + Method: http.MethodPost, + Path: "/plugins/commandcode/usage", + Body: reqBody, + } + + resp, err := HandleManagement(context.Background(), req, nil) + if err != nil { + t.Fatalf("HandleManagement error: %v", err) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("StatusCode = %d, want 200, body=%s", resp.StatusCode, string(resp.Body)) + } + + var usage FormattedUsageResponse + if err := json.Unmarshal(resp.Body, &usage); err != nil { + t.Fatalf("Unmarshal error: %v", err) + } + if usage.Credits.MonthlyCredits != 666 { + t.Errorf("MonthlyCredits = %v, want 666", usage.Credits.MonthlyCredits) + } +} diff --git a/plugin/plugin.go b/plugin/plugin.go new file mode 100644 index 0000000..a83350c --- /dev/null +++ b/plugin/plugin.go @@ -0,0 +1,285 @@ +package plugin + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "time" + + "gopkg.in/yaml.v3" +) + +const ( + PluginID = "commandcode" + PluginName = "commandcode" + PluginVersion = "0.1.0" + PluginAuthor = "zgs225" + PluginRepo = "https://github.com/zgs225/cliproxy-plugin-commandcode" + PluginLogo = "https://raw.githubusercontent.com/zgs225/cliproxy-plugin-commandcode/main/assets/logo.svg" + SchemaVersion1 = 1 +) + +// PluginConfig holds the runtime configuration parsed from YAML. +type PluginConfig struct { + mu sync.RWMutex + SessionToken string `yaml:"session_token" json:"session_token"` + APIBase string `yaml:"api_base" json:"api_base"` +} + +// UpdateFromYAML updates the configuration from raw YAML bytes. +func (c *PluginConfig) UpdateFromYAML(raw []byte) error { + if len(raw) == 0 { + return nil + } + var tmp struct { + SessionToken string `yaml:"session_token"` + APIBase string `yaml:"api_base"` + } + if err := yaml.Unmarshal(raw, &tmp); err != nil { + return fmt.Errorf("unmarshal config_yaml: %w", err) + } + + c.mu.Lock() + defer c.mu.Unlock() + + if tmp.SessionToken != "" { + c.SessionToken = ExtractSessionToken(tmp.SessionToken) + } + if tmp.APIBase != "" { + c.APIBase = strings.TrimRight(tmp.APIBase, "/") + } + if c.APIBase == "" { + c.APIBase = DefaultAPIBase + } + return nil +} + +// GetSessionToken safely returns the session token. +func (c *PluginConfig) GetSessionToken() string { + c.mu.RLock() + defer c.mu.RUnlock() + return c.SessionToken +} + +// SetSessionToken safely sets the session token. +func (c *PluginConfig) SetSessionToken(token string) { + c.mu.Lock() + defer c.mu.Unlock() + c.SessionToken = ExtractSessionToken(token) +} + +// GetAPIBase safely returns the API base URL. +func (c *PluginConfig) GetAPIBase() string { + c.mu.RLock() + defer c.mu.RUnlock() + if c.APIBase == "" { + return DefaultAPIBase + } + return c.APIBase +} + +// Plugin encapsulates the Command Code plugin instance. +type Plugin struct { + config *PluginConfig +} + +var ( + defaultPlugin = NewPlugin() +) + +// DefaultPlugin returns the singleton plugin instance. +func DefaultPlugin() *Plugin { + return defaultPlugin +} + +// NewPlugin creates a new Plugin instance. +func NewPlugin() *Plugin { + return &Plugin{ + config: &PluginConfig{ + APIBase: DefaultAPIBase, + }, + } +} + +// HandleMethod dispatches an ABI call to the corresponding handler. +func (p *Plugin) HandleMethod(method string, requestBytes []byte) ([]byte, error) { + switch method { + case "plugin.register": + return p.handleRegister(requestBytes) + case "plugin.reconfigure": + return p.handleReconfigure(requestBytes) + case "plugin.quiesce", "plugin.shutdown": + return OkEnvelope(map[string]any{"shutdown": true}) + + case "auth.identifier": + return OkEnvelope(IdentifierResponse{Identifier: PluginID}) + case "auth.parse": + return p.handleAuthParse(requestBytes) + case "auth.login.start": + return p.handleAuthLoginStart() + case "auth.login.poll": + return p.handleAuthLoginPoll() + case "auth.refresh": + return p.handleAuthRefresh(requestBytes) + + case "management.register": + return p.handleManagementRegister() + case "management.handle": + return p.handleManagementHandle(requestBytes) + + default: + return ErrorEnvelope("unknown_method", "unknown method: "+method), nil + } +} + +func (p *Plugin) handleRegister(raw []byte) ([]byte, error) { + if len(raw) > 0 { + var req LifecycleRequest + if err := json.Unmarshal(raw, &req); err == nil && len(req.ConfigYAML) > 0 { + _ = p.config.UpdateFromYAML(req.ConfigYAML) + } + } + return OkEnvelope(Registration{ + SchemaVersion: SchemaVersion1, + Metadata: Metadata{ + Name: PluginName, + Version: PluginVersion, + Author: PluginAuthor, + GitHubRepository: PluginRepo, + Logo: PluginLogo, + ConfigFields: []ConfigField{ + { + Name: "session_token", + Type: "string", + Description: "Command Code session token (__Secure-commandcode_prod_.session_token cookie value)", + }, + { + Name: "api_base", + Type: "string", + Description: "Command Code API base URL (default: https://api.commandcode.ai)", + }, + }, + }, + Capabilities: RegistrationCapability{ + AuthProvider: true, + ManagementAPI: true, + }, + }) +} + +func (p *Plugin) handleReconfigure(raw []byte) ([]byte, error) { + if len(raw) > 0 { + var req LifecycleRequest + if err := json.Unmarshal(raw, &req); err == nil && len(req.ConfigYAML) > 0 { + _ = p.config.UpdateFromYAML(req.ConfigYAML) + } + } + return p.handleRegister(raw) +} + +func (p *Plugin) handleAuthParse(raw []byte) ([]byte, error) { + var req AuthParseRequest + if len(raw) > 0 { + if err := json.Unmarshal(raw, &req); err != nil { + return ErrorEnvelope("invalid_request", "failed to parse AuthParseRequest: "+err.Error()), nil + } + } + + resp, err := ParseAuth(req) + if err != nil { + return ErrorEnvelope("auth_parse_error", err.Error()), nil + } + + // Cache token in config if config doesn't have one yet + if resp.Handled && resp.Auth.Metadata != nil { + if tok, ok := resp.Auth.Metadata["session_token"].(string); ok && tok != "" { + if p.config.GetSessionToken() == "" { + p.config.SetSessionToken(tok) + } + } + } + + return OkEnvelope(resp) +} + +func (p *Plugin) handleAuthLoginStart() ([]byte, error) { + return OkEnvelope(map[string]any{ + "Provider": PluginID, + "URL": "https://commandcode.ai", + "State": "manual", + "ExpiresAt": time.Now().Add(5 * time.Minute).UTC(), + }) +} + +func (p *Plugin) handleAuthLoginPoll() ([]byte, error) { + return OkEnvelope(map[string]any{ + "Status": "error", + "Message": "Command Code interactive login is not supported; please configure session_token or provide a commandcode-*.json credential file", + }) +} + +func (p *Plugin) handleAuthRefresh(raw []byte) ([]byte, error) { + var req AuthRefreshRequest + if len(raw) > 0 { + _ = json.Unmarshal(raw, &req) + } + + authData := AuthData{ + Provider: PluginID, + ID: req.AuthID, + StorageJSON: req.StorageJSON, + Metadata: req.Metadata, + Attributes: req.Attributes, + NextRefreshAfter: time.Now().Add(24 * time.Hour).UTC(), + } + return OkEnvelope(AuthRefreshResponse{ + Auth: authData, + NextRefreshAfter: authData.NextRefreshAfter, + }) +} + +func (p *Plugin) handleManagementRegister() ([]byte, error) { + resp, err := RegisterManagement() + if err != nil { + return ErrorEnvelope("management_register_error", err.Error()), nil + } + return OkEnvelope(resp) +} + +func (p *Plugin) handleManagementHandle(raw []byte) ([]byte, error) { + var req ManagementRequest + if len(raw) > 0 { + if err := json.Unmarshal(raw, &req); err != nil { + return ErrorEnvelope("invalid_request", "failed to parse ManagementRequest: "+err.Error()), nil + } + } + + resp, err := HandleManagement(context.Background(), req, p.config) + if err != nil { + return ErrorEnvelope("management_handle_error", err.Error()), nil + } + return OkEnvelope(resp) +} + +// OkEnvelope builds a successful Envelope response. +func OkEnvelope(v any) ([]byte, error) { + raw, err := json.Marshal(v) + if err != nil { + return nil, err + } + return json.Marshal(Envelope{OK: true, Result: raw}) +} + +// ErrorEnvelope builds an error Envelope response. +func ErrorEnvelope(code, message string) []byte { + raw, _ := json.Marshal(Envelope{ + OK: false, + Error: &EnvelopeError{ + Code: code, + Message: message, + }, + }) + return raw +} diff --git a/plugin/plugin_test.go b/plugin/plugin_test.go new file mode 100644 index 0000000..e92708a --- /dev/null +++ b/plugin/plugin_test.go @@ -0,0 +1,139 @@ +package plugin + +import ( + "encoding/json" + "testing" +) + +func TestPluginRegister_And_Reconfigure(t *testing.T) { + p := NewPlugin() + + configYAML := []byte(` +session_token: "my-yaml-token" +api_base: "https://custom-api.commandcode.ai" +`) + lifecycleReq, _ := json.Marshal(LifecycleRequest{ConfigYAML: configYAML}) + + // Test plugin.register + regBytes, err := p.HandleMethod("plugin.register", lifecycleReq) + if err != nil { + t.Fatalf("handleMethod(plugin.register) error: %v", err) + } + + var env Envelope + if err := json.Unmarshal(regBytes, &env); err != nil { + t.Fatalf("unmarshal envelope error: %v", err) + } + if !env.OK { + t.Fatalf("expected env.OK=true, got false: %+v", env.Error) + } + + var reg Registration + if err := json.Unmarshal(env.Result, ®); err != nil { + t.Fatalf("unmarshal registration error: %v", err) + } + + if reg.Metadata.Name != PluginName { + t.Errorf("Metadata.Name = %q, want %q", reg.Metadata.Name, PluginName) + } + if reg.Metadata.Version != PluginVersion { + t.Errorf("Metadata.Version = %q, want %q", reg.Metadata.Version, PluginVersion) + } + if !reg.Capabilities.AuthProvider { + t.Errorf("Capabilities.AuthProvider = false, want true") + } + if !reg.Capabilities.ManagementAPI { + t.Errorf("Capabilities.ManagementAPI = false, want true") + } + + // Verify config fields + if len(reg.Metadata.ConfigFields) != 2 { + t.Fatalf("ConfigFields len = %d, want 2", len(reg.Metadata.ConfigFields)) + } + fieldNames := map[string]bool{} + for _, f := range reg.Metadata.ConfigFields { + fieldNames[f.Name] = true + } + if !fieldNames["session_token"] || !fieldNames["api_base"] { + t.Errorf("ConfigFields missing session_token or api_base: %+v", reg.Metadata.ConfigFields) + } + + // Verify config parsed + if p.config.GetSessionToken() != "my-yaml-token" { + t.Errorf("SessionToken = %q, want my-yaml-token", p.config.GetSessionToken()) + } + if p.config.GetAPIBase() != "https://custom-api.commandcode.ai" { + t.Errorf("APIBase = %q, want https://custom-api.commandcode.ai", p.config.GetAPIBase()) + } + + // Test plugin.reconfigure + reconfYAML := []byte(` +session_token: "new-token-abc" +`) + reconfReq, _ := json.Marshal(LifecycleRequest{ConfigYAML: reconfYAML}) + reconfBytes, err := p.HandleMethod("plugin.reconfigure", reconfReq) + if err != nil { + t.Fatalf("handleMethod(plugin.reconfigure) error: %v", err) + } + var reconfEnv Envelope + if err := json.Unmarshal(reconfBytes, &reconfEnv); err != nil || !reconfEnv.OK { + t.Fatalf("reconfigure failed: %+v", reconfEnv) + } + if p.config.GetSessionToken() != "new-token-abc" { + t.Errorf("SessionToken after reconfigure = %q, want new-token-abc", p.config.GetSessionToken()) + } +} + +func TestPluginAuthIdentifier(t *testing.T) { + p := NewPlugin() + raw, err := p.HandleMethod("auth.identifier", nil) + if err != nil { + t.Fatalf("auth.identifier error: %v", err) + } + + var env Envelope + if err := json.Unmarshal(raw, &env); err != nil || !env.OK { + t.Fatalf("envelope error: %+v", env) + } + + var idResp IdentifierResponse + if err := json.Unmarshal(env.Result, &idResp); err != nil { + t.Fatalf("unmarshal idResp error: %v", err) + } + if idResp.Identifier != PluginID { + t.Errorf("Identifier = %q, want %q", idResp.Identifier, PluginID) + } +} + +func TestPluginUnknownMethod(t *testing.T) { + p := NewPlugin() + raw, err := p.HandleMethod("unknown.method.test", nil) + if err != nil { + t.Fatalf("expected no go error, got %v", err) + } + + var env Envelope + if err := json.Unmarshal(raw, &env); err != nil { + t.Fatalf("unmarshal error: %v", err) + } + if env.OK { + t.Fatal("expected env.OK=false for unknown method") + } + if env.Error == nil || env.Error.Code != "unknown_method" { + t.Errorf("Error = %+v, want code=unknown_method", env.Error) + } +} + +func TestEnvelopeError(t *testing.T) { + raw := ErrorEnvelope("test_code", "test error message") + var env Envelope + if err := json.Unmarshal(raw, &env); err != nil { + t.Fatalf("unmarshal error: %v", err) + } + if env.OK { + t.Fatal("expected OK=false") + } + if env.Error.Code != "test_code" || env.Error.Message != "test error message" { + t.Errorf("env.Error = %+v", env.Error) + } +} diff --git a/plugin/quota_page.go b/plugin/quota_page.go new file mode 100644 index 0000000..be5b2af --- /dev/null +++ b/plugin/quota_page.go @@ -0,0 +1,1028 @@ +package plugin + +import ( + "encoding/base64" +) + +// QuotaPageHTML contains the full standalone HTML/CSS/JS single-file QuotaCard UI. +const QuotaPageHTML = ` + + + + + Command Code 配额与用量 - CLIProxyAPI + + + +
+ +
+
+
+ + + + +
+
+
+ Command Code 配额 + v0.1.0 +
+
CLIProxyAPI 实时限额与 Credits 用量监控
+
+
+ +
+
+ + 正在检查... +
+ + + + +
+
+ + +
+ + 加载失败 +
+ + +
+
+ 诊断与访问配置 + +
+
+
+ + +
+
+ + +
+
+
+ +
+
+ + +
+
+
+ 月度额度 (Monthly Credits) +
💰
+
+
-
+
当前周期 Command Code 基础配额
+
+ +
+
+ 开源奖励额度 (Open Source) +
🎁
+
+
-
+
开源贡献者赠送与奖励额度
+
+ +
+
+ 可用额度总计 (Total Credits) +
✨
+
+
-
+
月度额度与奖励额度汇总
+
+
+ + +
+ +
+
+
+ 短期滑动窗口 +
5 小时限制 (5-Hour Window)
+
+
- %
+
+ +
+
+ - + / - +
+
剩余可用: -
+
+ +
+
+
+ +
+
+ + 重置倒计时 +
+
--:--:--
+
+
+ + +
+
+
+ 周度窗口 +
每周限制 (Weekly Window)
+
+
- %
+
+ +
+
+ - + / - +
+
剩余可用: -
+
+ +
+
+
+ +
+
+ + 重置倒计时 +
+
--:--:--
+
+
+
+ + +
+
最后同步时间: -
+
+ Provider: commandcode +
+
+
+ + + + +` + +// GetQuotaPageHTML returns the raw HTML bytes. +func GetQuotaPageHTML() []byte { + return []byte(QuotaPageHTML) +} + +// GetQuotaPageBase64 returns the base64-encoded Quota page HTML. +func GetQuotaPageBase64() string { + return base64.StdEncoding.EncodeToString([]byte(QuotaPageHTML)) +} diff --git a/plugin/types.go b/plugin/types.go new file mode 100644 index 0000000..2bc5961 --- /dev/null +++ b/plugin/types.go @@ -0,0 +1,280 @@ +package plugin + +import ( + "encoding/json" + "fmt" + "strconv" + "strings" + "time" +) + +// Envelope matches the CLIProxyAPI ABI JSON Envelope. +type Envelope struct { + OK bool `json:"ok"` + Result json.RawMessage `json:"result,omitempty"` + Error *EnvelopeError `json:"error,omitempty"` +} + +// EnvelopeError represents an error inside the ABI Envelope. +type EnvelopeError struct { + Code string `json:"code"` + Message string `json:"message"` + Retryable bool `json:"retryable,omitempty"` + HTTPStatus int `json:"http_status,omitempty"` +} + +// LifecycleRequest represents the payload for plugin.register or plugin.reconfigure. +type LifecycleRequest struct { + ConfigYAML []byte `json:"config_yaml"` +} + +// Registration describes the plugin registration response. +type Registration struct { + SchemaVersion uint32 `json:"schema_version"` + Metadata Metadata `json:"metadata"` + Capabilities RegistrationCapability `json:"capabilities"` +} + +// Metadata describes the plugin metadata. +type Metadata struct { + Name string `json:"Name"` + Version string `json:"Version"` + Author string `json:"Author,omitempty"` + GitHubRepository string `json:"GitHubRepository,omitempty"` + Logo string `json:"Logo,omitempty"` + ConfigFields []ConfigField `json:"ConfigFields,omitempty"` +} + +// ConfigField describes one configuration field for the plugin. +type ConfigField struct { + Name string `json:"Name"` + Type string `json:"Type"` + EnumValues []string `json:"EnumValues,omitempty"` + Description string `json:"Description"` +} + +// RegistrationCapability declares the capabilities implemented by this plugin. +type RegistrationCapability struct { + AuthProvider bool `json:"auth_provider"` + ManagementAPI bool `json:"management_api"` +} + +// IdentifierResponse is returned by auth.identifier. +type IdentifierResponse struct { + Identifier string `json:"identifier"` +} + +// AuthData describes a credential record. +type AuthData struct { + Provider string `json:"Provider"` + ID string `json:"ID"` + FileName string `json:"FileName"` + Label string `json:"Label"` + Prefix string `json:"Prefix,omitempty"` + ProxyURL string `json:"ProxyURL,omitempty"` + Disabled bool `json:"Disabled,omitempty"` + StorageJSON []byte `json:"StorageJSON"` + Metadata map[string]any `json:"Metadata,omitempty"` + Attributes map[string]string `json:"Attributes,omitempty"` + NextRefreshAfter time.Time `json:"NextRefreshAfter,omitempty"` +} + +// AuthParseRequest is passed to auth.parse. +type AuthParseRequest struct { + Provider string `json:"Provider"` + Path string `json:"Path"` + FileName string `json:"FileName"` + RawJSON []byte `json:"RawJSON"` + Host map[string]any `json:"Host,omitempty"` +} + +// AuthParseResponse is returned by auth.parse. +type AuthParseResponse struct { + Handled bool `json:"Handled"` + Auth AuthData `json:"Auth"` + Auths []AuthData `json:"Auths,omitempty"` +} + +// AuthRefreshRequest is passed to auth.refresh. +type AuthRefreshRequest struct { + AuthID string `json:"AuthID"` + AuthProvider string `json:"AuthProvider"` + StorageJSON []byte `json:"StorageJSON"` + Metadata map[string]any `json:"Metadata,omitempty"` + Attributes map[string]string `json:"Attributes,omitempty"` +} + +// AuthRefreshResponse is returned by auth.refresh. +type AuthRefreshResponse struct { + Auth AuthData `json:"Auth"` + NextRefreshAfter time.Time `json:"NextRefreshAfter,omitempty"` +} + +// ManagementRegistrationResponse is returned by management.register. +type ManagementRegistrationResponse struct { + Routes []ManagementRoute `json:"routes,omitempty"` + Resources []ResourceRoute `json:"resources,omitempty"` +} + +// ManagementRoute describes one Management API route. +type ManagementRoute struct { + Method string `json:"Method"` + Path string `json:"Path"` + Menu string `json:"Menu,omitempty"` + Description string `json:"Description,omitempty"` +} + +// ResourceRoute describes one browser-navigable resource route. +type ResourceRoute struct { + Path string `json:"Path"` + Menu string `json:"Menu"` + Description string `json:"Description"` +} + +// ManagementRequest is received by management.handle. +type ManagementRequest struct { + Method string `json:"Method"` + Path string `json:"Path"` + Headers map[string][]string `json:"Headers"` + Query map[string][]string `json:"Query"` + Body []byte `json:"Body"` + HostCallbackID string `json:"host_callback_id,omitempty"` +} + +// ManagementResponse is returned by management.handle. +type ManagementResponse struct { + StatusCode int `json:"StatusCode"` + Headers map[string][]string `json:"Headers"` + Body []byte `json:"Body"` +} + +// HostHTTPRequest describes a request dispatched through host.http.do. +type HostHTTPRequest struct { + Method string `json:"method"` + URL string `json:"url"` + Headers map[string][]string `json:"headers,omitempty"` + Body []byte `json:"body,omitempty"` + HostCallbackID string `json:"host_callback_id,omitempty"` +} + +// HostHTTPResponse describes a response received from host.http.do. +type HostHTTPResponse struct { + StatusCode int `json:"StatusCode"` + Headers map[string][]string `json:"Headers"` + Body []byte `json:"Body"` +} + +// FlexibleTime handles parsing timestamps from upstream that may be unix seconds, unix milliseconds, or RFC3339 strings. +type FlexibleTime struct { + time.Time +} + +// UnmarshalJSON parses various timestamp formats. +func (ft *FlexibleTime) UnmarshalJSON(b []byte) error { + s := strings.Trim(string(b), "\" \t\r\n") + if s == "" || s == "null" || s == "0" { + ft.Time = time.Time{} + return nil + } + + // Try parsing as integer / float number (Unix timestamp) + if n, err := strconv.ParseInt(s, 10, 64); err == nil { + if n > 1e11 { + // Milliseconds + ft.Time = time.UnixMilli(n).UTC() + } else { + // Seconds + ft.Time = time.Unix(n, 0).UTC() + } + return nil + } + + if f, err := strconv.ParseFloat(s, 64); err == nil { + sec := int64(f) + if sec > 1e11 { + ft.Time = time.UnixMilli(sec).UTC() + } else { + ft.Time = time.Unix(sec, 0).UTC() + } + return nil + } + + // Try standard RFC3339 / ISO8601 layouts + formats := []string{ + time.RFC3339Nano, + time.RFC3339, + "2006-01-02T15:04:05.999999999", + "2006-01-02T15:04:05", + "2006-01-02 15:04:05", + } + for _, layout := range formats { + if t, err := time.Parse(layout, s); err == nil { + ft.Time = t.UTC() + return nil + } + } + + return fmt.Errorf("cannot parse %q as FlexibleTime", string(b)) +} + +// UpstreamCreditsResponse reflects the payload returned by Command Code's /internal/billing/credits. +type UpstreamCreditsResponse struct { + Credits map[string]any `json:"credits"` + WindowLimits UpstreamWindowLimits `json:"windowLimits"` +} + +// UpstreamWindowLimits carries fiveHour and weekly window metrics. +type UpstreamWindowLimits struct { + FiveHour UpstreamWindowLimit `json:"fiveHour"` + Weekly UpstreamWindowLimit `json:"weekly"` +} + +// UpstreamWindowLimit represents one quota window from upstream. +type UpstreamWindowLimit struct { + Used float64 `json:"used"` + Cap float64 `json:"cap"` + Exceeded bool `json:"exceeded"` + ResetAt FlexibleTime `json:"resetAt"` +} + +// UsageCreditsData is the formatted credits section. +type UsageCreditsData struct { + MonthlyCredits float64 `json:"monthly_credits"` + OpensourceMonthlyCredits float64 `json:"opensource_monthly_credits"` + TotalCredits float64 `json:"total_credits"` + Details map[string]any `json:"details,omitempty"` +} + +// UsageWindowLimitData is the formatted window limit section. +type UsageWindowLimitData struct { + Used float64 `json:"used"` + Cap float64 `json:"cap"` + Remaining float64 `json:"remaining"` + Percentage float64 `json:"percentage"` + Exceeded bool `json:"exceeded"` + ResetAt string `json:"reset_at"` + ResetInSeconds int64 `json:"reset_in_seconds"` +} + +// UsageWindowLimitsData contains both windows. +type UsageWindowLimitsData struct { + FiveHour UsageWindowLimitData `json:"five_hour"` + Weekly UsageWindowLimitData `json:"weekly"` +} + +// FormattedUsageData is the complete formatted usage payload. +type FormattedUsageData struct { + Credits UsageCreditsData `json:"credits"` + WindowLimits UsageWindowLimitsData `json:"window_limits"` + UpdatedAt string `json:"updated_at"` +} + +// FormattedUsageResponse is returned by GET /plugins/commandcode/usage and POST /plugins/commandcode/usage. +type FormattedUsageResponse struct { + OK bool `json:"ok"` + Data FormattedUsageData `json:"data"` + Credits UsageCreditsData `json:"credits"` + WindowLimits UsageWindowLimitsData `json:"window_limits"` + UpdatedAt string `json:"updated_at"` + Error string `json:"error,omitempty"` +} diff --git a/plugin/usage.go b/plugin/usage.go new file mode 100644 index 0000000..f6d7121 --- /dev/null +++ b/plugin/usage.go @@ -0,0 +1,259 @@ +package plugin + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "net/http" + "strings" + "time" +) + +const ( + DefaultAPIBase = "https://api.commandcode.ai" +) + +// HTTPDoer abstracts HTTP requests for testing and fallback. +type HTTPDoer interface { + Do(req *http.Request) (*http.Response, error) +} + +// HostCallerFunc is the signature for calling the host API via C ABI. +type HostCallerFunc func(method string, payload []byte) ([]byte, error) + +var ( + defaultHTTPClient HTTPDoer = &http.Client{Timeout: 15 * time.Second} + hostCaller HostCallerFunc // Set by main if host API is available +) + +// SetHostCaller registers the host API callback runner. +func SetHostCaller(fn HostCallerFunc) { + hostCaller = fn +} + +// SetDefaultHTTPClient overrides the default standard HTTP client (useful for unit tests). +func SetDefaultHTTPClient(client HTTPDoer) { + if client != nil { + defaultHTTPClient = client + } +} + +// FetchCreditsRaw fetches raw upstream credit data via host.http.do or net/http fallback. +func FetchCreditsRaw(ctx context.Context, apiBase, sessionToken string, hostCallbackID string) ([]byte, int, error) { + cleanToken := ExtractSessionToken(sessionToken) + if cleanToken == "" { + return nil, http.StatusBadRequest, errors.New("missing session_token: please provide a valid Command Code session token") + } + + if apiBase == "" { + apiBase = DefaultAPIBase + } + url := fmt.Sprintf("%s/internal/billing/credits", strings.TrimRight(apiBase, "/")) + cookieValue := FormatSessionCookie(cleanToken) + + // 1. Try host.http.do if hostCaller is configured + if hostCaller != nil { + reqPayload := HostHTTPRequest{ + Method: http.MethodGet, + URL: url, + Headers: map[string][]string{ + "Cookie": {cookieValue}, + "Accept": {"application/json"}, + "User-Agent": {"cliproxy-plugin-commandcode/0.1.0"}, + }, + HostCallbackID: hostCallbackID, + } + rawReq, errMarshal := json.Marshal(reqPayload) + if errMarshal == nil { + respBytes, errCall := hostCaller("host.http.do", rawReq) + if errCall == nil && len(respBytes) > 0 { + var env Envelope + if errEnv := json.Unmarshal(respBytes, &env); errEnv == nil { + if !env.OK { + errMsg := "host HTTP request failed" + if env.Error != nil { + errMsg = fmt.Sprintf("%s: %s", env.Error.Code, env.Error.Message) + } + return nil, http.StatusBadGateway, fmt.Errorf("host.http.do error: %s", errMsg) + } + var hostResp HostHTTPResponse + if errResp := json.Unmarshal(env.Result, &hostResp); errResp == nil { + // hostResp.Body is automatically base64-decoded by json.Unmarshal for []byte + status := hostResp.StatusCode + if status == 0 { + status = http.StatusOK + } + return hostResp.Body, status, nil + } + } + } + } + // If hostCaller fails, seamlessly fallback to net/http + } + + // 2. Fallback to Go net/http client + httpReq, errNew := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if errNew != nil { + return nil, http.StatusInternalServerError, fmt.Errorf("create HTTP request: %w", errNew) + } + httpReq.Header.Set("Cookie", cookieValue) + httpReq.Header.Set("Accept", "application/json") + httpReq.Header.Set("User-Agent", "cliproxy-plugin-commandcode/0.1.0") + + res, errDo := defaultHTTPClient.Do(httpReq) + if errDo != nil { + return nil, http.StatusBadGateway, fmt.Errorf("upstream request failed: %w", errDo) + } + defer func() { + _ = res.Body.Close() + }() + + body, errRead := io.ReadAll(res.Body) + if errRead != nil { + return nil, http.StatusBadGateway, fmt.Errorf("read upstream response body: %w", errRead) + } + + return body, res.StatusCode, nil +} + +// ParseAndFormatUsage parses upstream credits JSON into structured usage metrics. +func ParseAndFormatUsage(raw []byte, now time.Time) (*FormattedUsageResponse, error) { + if len(raw) == 0 { + return nil, errors.New("empty response body from upstream") + } + + // First try to parse as UpstreamCreditsResponse + var upstream UpstreamCreditsResponse + if err := json.Unmarshal(raw, &upstream); err != nil { + // Fallback: check if wrapped in { "data": { ... } } + var wrapped struct { + Data UpstreamCreditsResponse `json:"data"` + } + if errWrap := json.Unmarshal(raw, &wrapped); errWrap == nil && (len(wrapped.Data.Credits) > 0 || wrapped.Data.WindowLimits.FiveHour.Cap > 0) { + upstream = wrapped.Data + } else { + return nil, fmt.Errorf("unmarshal upstream response: %w", err) + } + } + + if now.IsZero() { + now = time.Now().UTC() + } + + // Format credits + creditsData := formatCredits(upstream.Credits) + + // Format window limits + windowLimitsData := formatWindowLimits(upstream.WindowLimits, now) + + nowRFC := now.Format(time.RFC3339) + data := FormattedUsageData{ + Credits: creditsData, + WindowLimits: windowLimitsData, + UpdatedAt: nowRFC, + } + + return &FormattedUsageResponse{ + OK: true, + Data: data, + Credits: creditsData, + WindowLimits: windowLimitsData, + UpdatedAt: nowRFC, + }, nil +} + +func formatCredits(credits map[string]any) UsageCreditsData { + data := UsageCreditsData{ + Details: credits, + } + if credits == nil { + return data + } + + data.MonthlyCredits = getFloatFromMap(credits, "monthlyCredits", "monthly_credits") + data.OpensourceMonthlyCredits = getFloatFromMap(credits, "opensourceMonthlyCredits", "opensource_monthly_credits") + data.TotalCredits = data.MonthlyCredits + data.OpensourceMonthlyCredits + + return data +} + +func formatWindowLimits(upstream UpstreamWindowLimits, now time.Time) UsageWindowLimitsData { + return UsageWindowLimitsData{ + FiveHour: formatSingleWindow(upstream.FiveHour, now), + Weekly: formatSingleWindow(upstream.Weekly, now), + } +} + +func formatSingleWindow(w UpstreamWindowLimit, now time.Time) UsageWindowLimitData { + var remaining float64 + var percentage float64 + + if w.Cap > 0 { + remaining = w.Cap - w.Used + if remaining < 0 { + remaining = 0 + } + percentage = (w.Used / w.Cap) * 100.0 + if percentage > 100.0 { + percentage = 100.0 + } + percentage = math.Round(percentage*100) / 100 + } + + var resetAtStr string + var resetInSeconds int64 + if !w.ResetAt.IsZero() { + resetAtStr = w.ResetAt.UTC().Format(time.RFC3339) + diff := w.ResetAt.UTC().Sub(now) + if diff > 0 { + resetInSeconds = int64(diff.Seconds()) + } else { + resetInSeconds = 0 + } + } + + return UsageWindowLimitData{ + Used: w.Used, + Cap: w.Cap, + Remaining: remaining, + Percentage: percentage, + Exceeded: w.Exceeded, + ResetAt: resetAtStr, + ResetInSeconds: resetInSeconds, + } +} + +func getFloatFromMap(m map[string]any, keys ...string) float64 { + for _, key := range keys { + if val, exists := m[key]; exists && val != nil { + switch v := val.(type) { + case float64: + return v + case float32: + return float64(v) + case int: + return float64(v) + case int64: + return float64(v) + case json.Number: + if f, err := v.Float64(); err == nil { + return f + } + } + } + } + return 0 +} + +// DecodeBase64OrRaw tries to decode base64, returning raw if not base64. +func DecodeBase64OrRaw(in []byte) []byte { + decoded, err := base64.StdEncoding.DecodeString(string(in)) + if err == nil && len(decoded) > 0 { + return decoded + } + return in +} diff --git a/plugin/usage_test.go b/plugin/usage_test.go new file mode 100644 index 0000000..0743862 --- /dev/null +++ b/plugin/usage_test.go @@ -0,0 +1,222 @@ +package plugin + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestFlexibleTime(t *testing.T) { + tests := []struct { + name string + input string + wantYear int + }{ + { + name: "unix seconds", + input: `{"resetAt": 1741123456}`, + wantYear: 2025, + }, + { + name: "unix milliseconds", + input: `{"resetAt": 1741123456000}`, + wantYear: 2025, + }, + { + name: "RFC3339 string", + input: `{"resetAt": "2025-06-15T12:00:00Z"}`, + wantYear: 2025, + }, + { + name: "null", + input: `{"resetAt": null}`, + wantYear: 1, // Zero time year + }, + { + name: "empty string", + input: `{"resetAt": ""}`, + wantYear: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var res struct { + ResetAt FlexibleTime `json:"resetAt"` + } + if err := json.Unmarshal([]byte(tt.input), &res); err != nil { + t.Fatalf("Unmarshal error: %v", err) + } + if res.ResetAt.Year() != tt.wantYear { + t.Errorf("Year = %d, want %d", res.ResetAt.Year(), tt.wantYear) + } + }) + } +} + +func TestParseAndFormatUsage(t *testing.T) { + raw := []byte(`{ + "credits": { + "monthlyCredits": 1000.0, + "opensourceMonthlyCredits": 500.0, + "extraBonus": 50.0 + }, + "windowLimits": { + "fiveHour": { + "used": 25.0, + "cap": 100.0, + "exceeded": false, + "resetAt": 1741123456 + }, + "weekly": { + "used": 200.0, + "cap": 1000.0, + "exceeded": false, + "resetAt": "2025-03-10T12:00:00Z" + } + } + }`) + + now := time.Date(2025, 3, 4, 12, 0, 0, 0, time.UTC) + usage, err := ParseAndFormatUsage(raw, now) + if err != nil { + t.Fatalf("ParseAndFormatUsage error: %v", err) + } + + if !usage.OK { + t.Fatal("expected OK=true") + } + + // Verify credits + if usage.Credits.MonthlyCredits != 1000.0 { + t.Errorf("MonthlyCredits = %v, want 1000", usage.Credits.MonthlyCredits) + } + if usage.Credits.OpensourceMonthlyCredits != 500.0 { + t.Errorf("OpensourceMonthlyCredits = %v, want 500", usage.Credits.OpensourceMonthlyCredits) + } + if usage.Credits.TotalCredits != 1500.0 { + t.Errorf("TotalCredits = %v, want 1500", usage.Credits.TotalCredits) + } + + // Verify 5-hour window + fiveHour := usage.WindowLimits.FiveHour + if fiveHour.Used != 25.0 { + t.Errorf("FiveHour Used = %v, want 25", fiveHour.Used) + } + if fiveHour.Cap != 100.0 { + t.Errorf("FiveHour Cap = %v, want 100", fiveHour.Cap) + } + if fiveHour.Remaining != 75.0 { + t.Errorf("FiveHour Remaining = %v, want 75", fiveHour.Remaining) + } + if fiveHour.Percentage != 25.0 { + t.Errorf("FiveHour Percentage = %v, want 25", fiveHour.Percentage) + } + if fiveHour.Exceeded { + t.Errorf("FiveHour Exceeded = true, want false") + } + + // Verify weekly window + weekly := usage.WindowLimits.Weekly + if weekly.Used != 200.0 { + t.Errorf("Weekly Used = %v, want 200", weekly.Used) + } + if weekly.Cap != 1000.0 { + t.Errorf("Weekly Cap = %v, want 1000", weekly.Cap) + } + if weekly.Remaining != 800.0 { + t.Errorf("Weekly Remaining = %v, want 800", weekly.Remaining) + } + if weekly.Percentage != 20.0 { + t.Errorf("Weekly Percentage = %v, want 20", weekly.Percentage) + } + if weekly.ResetAt != "2025-03-10T12:00:00Z" { + t.Errorf("Weekly ResetAt = %v, want 2025-03-10T12:00:00Z", weekly.ResetAt) + } +} + +func TestFetchCreditsRaw_FallbackHTTP(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/internal/billing/credits" { + t.Errorf("unexpected path: %s", r.URL.Path) + http.NotFound(w, r) + return + } + + cookie := r.Header.Get("Cookie") + expectedCookie := "__Secure-commandcode_prod_.session_token=test-session-123" + if cookie != expectedCookie { + t.Errorf("Cookie = %q, want %q", cookie, expectedCookie) + w.WriteHeader(http.StatusUnauthorized) + return + } + + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"credits":{"monthlyCredits":100},"windowLimits":{"fiveHour":{"used":1,"cap":10}}}`)) + })) + defer ts.Close() + + // Ensure hostCaller is nil for fallback test + SetHostCaller(nil) + SetDefaultHTTPClient(ts.Client()) + + body, status, err := FetchCreditsRaw(context.Background(), ts.URL, "test-session-123", "") + if err != nil { + t.Fatalf("FetchCreditsRaw error: %v", err) + } + if status != http.StatusOK { + t.Errorf("status = %d, want 200", status) + } + if len(body) == 0 { + t.Fatal("expected non-empty body") + } + + usage, errParse := ParseAndFormatUsage(body, time.Time{}) + if errParse != nil { + t.Fatalf("ParseAndFormatUsage error: %v", errParse) + } + if usage.Credits.MonthlyCredits != 100 { + t.Errorf("MonthlyCredits = %v, want 100", usage.Credits.MonthlyCredits) + } +} + +func TestFetchCreditsRaw_MissingToken(t *testing.T) { + _, status, err := FetchCreditsRaw(context.Background(), "", "", "") + if err == nil { + t.Fatal("expected error for missing token") + } + if status != http.StatusBadRequest { + t.Errorf("status = %d, want 400", status) + } +} + +func TestFetchCreditsRaw_HostCaller(t *testing.T) { + mockResponsePayload := []byte(`{"credits":{"monthlyCredits":500},"windowLimits":{"fiveHour":{"used":5,"cap":50}}}`) + + SetHostCaller(func(method string, payload []byte) ([]byte, error) { + if method != "host.http.do" { + t.Errorf("method = %s, want host.http.do", method) + } + hostResp := HostHTTPResponse{ + StatusCode: http.StatusOK, + Body: mockResponsePayload, + } + respJSON, _ := json.Marshal(hostResp) + return json.Marshal(Envelope{OK: true, Result: respJSON}) + }) + defer SetHostCaller(nil) + + body, status, err := FetchCreditsRaw(context.Background(), "https://api.commandcode.ai", "mock-token", "cb-123") + if err != nil { + t.Fatalf("FetchCreditsRaw with hostCaller error: %v", err) + } + if status != http.StatusOK { + t.Errorf("status = %d, want 200", status) + } + if string(body) != string(mockResponsePayload) { + t.Errorf("body = %s, want %s", string(body), string(mockResponsePayload)) + } +}