255 lines
11 KiB
Go
255 lines
11 KiB
Go
package plugin
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
managedaccess "cpa-ext/internal/access"
|
|
)
|
|
|
|
func TestManagedKeyAuthenticationAndLifecycle(t *testing.T) {
|
|
app := NewApp()
|
|
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
auth := func(secret string) FrontendAuthResponse {
|
|
raw, _ := json.Marshal(FrontendAuthRequest{Headers: http.Header{"Authorization": []string{"Bearer " + secret}}})
|
|
response, err := app.HandleMethod(MethodFrontendAuthenticate, raw)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var envelope Envelope
|
|
_ = json.Unmarshal(response, &envelope)
|
|
var result FrontendAuthResponse
|
|
_ = json.Unmarshal(envelope.Result, &result)
|
|
return result
|
|
}
|
|
if got := auth("000000"); !got.Authenticated || got.Principal != "key_default" {
|
|
t.Fatalf("default auth = %+v", got)
|
|
}
|
|
if got := auth("key_default"); got.Authenticated {
|
|
t.Fatalf("internal Key ID authenticated as a credential: %+v", got)
|
|
}
|
|
xAPIKeyRaw, _ := json.Marshal(FrontendAuthRequest{Headers: http.Header{"X-Api-Key": []string{"000000"}}})
|
|
xAPIKeyResponse, _ := app.HandleMethod(MethodFrontendAuthenticate, xAPIKeyRaw)
|
|
var xAPIKeyEnvelope Envelope
|
|
_ = json.Unmarshal(xAPIKeyResponse, &xAPIKeyEnvelope)
|
|
var xAPIKeyResult FrontendAuthResponse
|
|
_ = json.Unmarshal(xAPIKeyEnvelope.Result, &xAPIKeyResult)
|
|
if !xAPIKeyResult.Authenticated {
|
|
t.Fatal("x-api-key did not authenticate")
|
|
}
|
|
store, _ := app.currentStore()
|
|
key, _ := store.ManagedKeyByID(context.Background(), "key_default")
|
|
key.Status = managedaccess.StatusDisabled
|
|
if err := store.UpdateManagedKey(context.Background(), key); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got := auth("000000"); got.Authenticated {
|
|
t.Fatalf("disabled key authenticated: %+v", got)
|
|
}
|
|
key.Status = managedaccess.StatusActive
|
|
if err := store.UpdateManagedKey(context.Background(), key); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if got := auth("000000"); !got.Authenticated {
|
|
t.Fatalf("restored key did not authenticate: %+v", got)
|
|
}
|
|
}
|
|
|
|
func TestModelAllowlistAndStrictUpstreamRouting(t *testing.T) {
|
|
app := NewApp()
|
|
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
price := []byte(`{"model":"gpt-5.6-sol","base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"},"fast_pricing_enabled":false,"fast_multiplier":"2.5"}`)
|
|
if response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, price); response.StatusCode != http.StatusOK {
|
|
t.Fatalf("put price=%d %s", response.StatusCode, response.Body)
|
|
}
|
|
store, _ := app.currentStore()
|
|
if err := store.SyncUpstreamAccounts(context.Background(), []managedaccess.UpstreamAccount{{
|
|
CPAAuthID: "oauth-1", CPAAuthIndex: "index-1", Provider: "codex", DisplayName: "OAuth One",
|
|
Disabled: true, Unavailable: true, LastSeenAt: time.Now(),
|
|
}}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
accounts, _ := store.ListUpstreamAccounts(context.Background())
|
|
key, _ := store.ManagedKeyByID(context.Background(), "key_default")
|
|
key.RouteMode = managedaccess.RouteStrict
|
|
key.UpstreamAccountID = accounts[0].ID
|
|
key.AllModels = false
|
|
key.Models = []string{"gpt-5.6-sol"}
|
|
if err := store.UpdateManagedKey(context.Background(), key); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
interceptRaw, _ := json.Marshal(RequestInterceptRequest{
|
|
RequestID: "req-1", RequestedModel: "other-model",
|
|
Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)},
|
|
})
|
|
response, err := app.HandleMethod(MethodRequestBefore, interceptRaw)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var envelope Envelope
|
|
_ = json.Unmarshal(response, &envelope)
|
|
var intercept RequestInterceptResponse
|
|
_ = json.Unmarshal(envelope.Result, &intercept)
|
|
if !intercept.Terminate || intercept.StatusCode != http.StatusForbidden {
|
|
t.Fatalf("model denial = %+v", intercept)
|
|
}
|
|
|
|
pickRaw, _ := json.Marshal(SchedulerPickRequest{
|
|
Model: "gpt-5.6-sol", Options: SchedulerOptions{Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}},
|
|
Candidates: []SchedulerAuthCandidate{{ID: "oauth-1", Provider: "codex"}},
|
|
})
|
|
response, err = app.HandleMethod(MethodSchedulerPick, pickRaw)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_ = json.Unmarshal(response, &envelope)
|
|
var pick SchedulerPickResponse
|
|
_ = json.Unmarshal(envelope.Result, &pick)
|
|
if !pick.Handled || pick.AuthID != "oauth-1" {
|
|
t.Fatalf("scheduler pick = %+v", pick)
|
|
}
|
|
afterRaw, _ := json.Marshal(RequestInterceptRequest{
|
|
RequestID: "req-1", Model: "gpt-5.6-sol", RequestedModel: "gpt-5.6-sol",
|
|
Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID), selectedAuthMetadata: "oauth-1"},
|
|
})
|
|
afterResponse, _ := app.HandleMethod(MethodRequestAfter, afterRaw)
|
|
_ = json.Unmarshal(afterResponse, &envelope)
|
|
_ = json.Unmarshal(envelope.Result, &intercept)
|
|
if intercept.Terminate {
|
|
t.Fatalf("selected strict target was rejected: %+v", intercept)
|
|
}
|
|
|
|
pickRaw, _ = json.Marshal(SchedulerPickRequest{
|
|
Model: "gpt-5.6-sol", Options: SchedulerOptions{Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}},
|
|
Candidates: []SchedulerAuthCandidate{{ID: "oauth-2", Provider: "codex"}},
|
|
})
|
|
response, err = app.HandleMethod(MethodSchedulerPick, pickRaw)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_ = json.Unmarshal(response, &envelope)
|
|
if envelope.OK || envelope.Error == nil || envelope.Error.Code != "bound_upstream_unavailable" {
|
|
t.Fatalf("missing strict target did not fail closed: %s", response)
|
|
}
|
|
}
|
|
|
|
func TestModelAllowlistSupportsFullNameWildcards(t *testing.T) {
|
|
key := managedaccess.ManagedKey{AllModels: false, Models: []string{"deepseek-*", "gpt-5.6-sol"}}
|
|
for _, model := range []string{"deepseek-v4-flash", "DEEPSEEK-SOURCE-B", "gpt-5.6-sol"} {
|
|
if !managedKeyAllowsModel(key, model) {
|
|
t.Fatalf("expected %q to match", model)
|
|
}
|
|
}
|
|
for _, model := range []string{"x-deepseek-v4-flash", "gpt-5.6-terra", ""} {
|
|
if managedKeyAllowsModel(key, model) {
|
|
t.Fatalf("expected %q to be denied", model)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCreateKeyCopiesDefaultRuleSnapshot(t *testing.T) {
|
|
app := NewApp()
|
|
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\n"))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
store, _ := app.currentStore()
|
|
template, err := store.ManagedKeyByID(context.Background(), "key_default")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
template.Name = "renamed-default"
|
|
template.AllModels = false
|
|
template.Models = []string{"deepseek-*"}
|
|
if err := store.UpdateManagedKey(context.Background(), template); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
response := managementCallBody(t, app, http.MethodPost, managementBase+routeKeys, []byte(`{"name":"alice","secret":"alice-000000"}`))
|
|
if response.StatusCode != http.StatusCreated {
|
|
t.Fatalf("create status=%d body=%s", response.StatusCode, response.Body)
|
|
}
|
|
var created managedaccess.ManagedKey
|
|
if err := json.Unmarshal(response.Body, &created); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if created.Secret != "alice-000000" || created.RouteMode != managedaccess.RouteAuto || created.AllModels || len(created.Models) != 1 || created.Models[0] != "deepseek-*" {
|
|
t.Fatalf("created = %+v", created)
|
|
}
|
|
}
|
|
|
|
func TestBillingAdmissionConcurrencyAndUsageSettlement(t *testing.T) {
|
|
app := NewApp()
|
|
if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
price := []byte(`{"model":"deepseek-v4-flash","base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"},"fast_pricing_enabled":false,"fast_multiplier":"2.5"}`)
|
|
if response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, price); response.StatusCode != http.StatusOK {
|
|
t.Fatalf("put price=%d %s", response.StatusCode, response.Body)
|
|
}
|
|
key, _ := app.store.ManagedKeyByID(context.Background(), "key_default")
|
|
key.AllModels = false
|
|
key.Models = []string{"deepseek-*"}
|
|
if err := app.store.UpdateManagedKey(context.Background(), key); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
intercept := func(requestID string) RequestInterceptResponse {
|
|
raw, _ := json.Marshal(RequestInterceptRequest{RequestID: requestID, Model: "deepseek-v4-flash", RequestedModel: "deepseek-source-a", Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}})
|
|
response, err := app.HandleMethod(MethodRequestBefore, raw)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var envelope Envelope
|
|
_ = json.Unmarshal(response, &envelope)
|
|
var result RequestInterceptResponse
|
|
_ = json.Unmarshal(envelope.Result, &result)
|
|
return result
|
|
}
|
|
if denied := intercept("zero"); !denied.Terminate || denied.StatusCode != http.StatusTooManyRequests || !strings.Contains(string(denied.ResponseBody), "billing_quota_exhausted") {
|
|
t.Fatalf("zero quota response=%+v body=%s", denied, denied.ResponseBody)
|
|
}
|
|
patch := []byte(`{"id":"key_default","billing":{"quota_usd":"1","reset_period":"none","max_concurrency":1}}`)
|
|
if response := managementCallBody(t, app, http.MethodPatch, managementBase+routeKeys, patch); response.StatusCode != http.StatusOK {
|
|
t.Fatalf("update billing=%d %s", response.StatusCode, response.Body)
|
|
}
|
|
if allowed := intercept("one"); allowed.Terminate {
|
|
t.Fatalf("funded request denied: %+v", allowed)
|
|
}
|
|
afterRaw, _ := json.Marshal(RequestInterceptRequest{RequestID: "one", Model: "unpriced-resolved-model", RequestedModel: "deepseek-source-a", Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}})
|
|
afterResponse, _ := app.HandleMethod(MethodRequestAfter, afterRaw)
|
|
var afterEnvelope Envelope
|
|
_ = json.Unmarshal(afterResponse, &afterEnvelope)
|
|
var afterResult RequestInterceptResponse
|
|
_ = json.Unmarshal(afterEnvelope.Result, &afterResult)
|
|
if !afterResult.Terminate || afterResult.StatusCode != http.StatusServiceUnavailable || !strings.Contains(string(afterResult.ResponseBody), "billing_price_unavailable") {
|
|
t.Fatalf("missing price response=%+v body=%s", afterResult, afterResult.ResponseBody)
|
|
}
|
|
if denied := intercept("two"); !denied.Terminate || !strings.Contains(string(denied.ResponseBody), "billing_concurrency_exceeded") {
|
|
t.Fatalf("concurrency response=%+v body=%s", denied, denied.ResponseBody)
|
|
}
|
|
completionRaw, _ := json.Marshal(RequestCompletion{RequestID: "one", StartedAt: time.Now().Add(-time.Second), CompletedAt: time.Now(), Outcome: RequestCompletionSucceeded, Metadata: map[string]any{callerScopeMetadata: managedaccess.CallerScope(key.ID)}})
|
|
if _, err := app.HandleMethod(MethodRequestComplete, completionRaw); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
usage := UsageRecord{APIKey: "000000", Model: "deepseek-v4-flash", RequestedAt: time.Now(), Detail: UsageDetail{OutputTokens: 10_000, TotalTokens: 10_000}}
|
|
raw, _ := json.Marshal(usage)
|
|
if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
state, err := app.store.BillingState(context.Background(), key.ID, time.Now())
|
|
if err != nil || state.SpentMicros != 150_000 || state.BalanceMicros != 850_000 || state.ActiveRequests != 0 {
|
|
t.Fatalf("settled state=%+v err=%v", state, err)
|
|
}
|
|
}
|