mirror of
https://github.com/zgs225/cliproxy-plugin-commandcode.git
synced 2026-09-26 03:32:50 +08:00
feat: initial implementation of CLIProxyAPI Command Code plugin (stage 1)
This commit is contained in:
+166
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
+280
@@ -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"`
|
||||
}
|
||||
+259
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user