mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-03 06:35:00 +08:00
chore: exclude test changes from Home fixes
This commit is contained in:
@@ -27,21 +27,10 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi"
|
||||
)
|
||||
|
||||
type apiUsageCapturePlugin struct {
|
||||
records chan coreusage.Record
|
||||
}
|
||||
|
||||
func (p apiUsageCapturePlugin) HandleUsage(_ context.Context, record coreusage.Record) {
|
||||
if p.records != nil {
|
||||
p.records <- record
|
||||
}
|
||||
}
|
||||
|
||||
type codexSearchCaptureExecutor struct {
|
||||
request *http.Request
|
||||
body []byte
|
||||
@@ -362,53 +351,6 @@ func TestHomeCodexAlphaSearchRefreshesUnauthorizedSelectionOnce(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeCodexAlphaSearchReportsEveryUnauthorizedAttempt(t *testing.T) {
|
||||
records := make(chan coreusage.Record, 8)
|
||||
const pluginName = "api-home-search-unauthorized-test"
|
||||
coreusage.RegisterNamedPlugin(pluginName, apiUsageCapturePlugin{records: records})
|
||||
t.Cleanup(func() {
|
||||
coreusage.RegisterNamedPlugin(pluginName, apiUsageCapturePlugin{})
|
||||
})
|
||||
|
||||
server := newTestServer(t)
|
||||
dispatcher := &codexSearchHomeDispatcher{}
|
||||
server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}})
|
||||
server.handlers.AuthManager.PublishHomeDispatch(dispatcher, executionregistry.New(), 1)
|
||||
executor := &codexSearchCaptureExecutor{statuses: []int{http.StatusUnauthorized, http.StatusUnauthorized}}
|
||||
server.handlers.AuthManager.RegisterExecutor(executor)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"id":"home-search-unauthorized","model":"gpt-5-codex","query":"test"}`))
|
||||
req.Header.Set("Authorization", "Bearer test-key")
|
||||
rr := httptest.NewRecorder()
|
||||
server.engine.ServeHTTP(rr, req)
|
||||
|
||||
if rr.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusUnauthorized, rr.Body.String())
|
||||
}
|
||||
if executor.refreshCalls != 1 || executor.httpCalls != 2 {
|
||||
t.Fatalf("refresh/http calls = %d/%d, want 1/2", executor.refreshCalls, executor.httpCalls)
|
||||
}
|
||||
wantHashes := map[string]bool{
|
||||
auth.AccessTokenSHA256(&auth.Auth{Metadata: map[string]any{"access_token": "home-search-token"}}): false,
|
||||
auth.AccessTokenSHA256(&auth.Auth{Metadata: map[string]any{"access_token": "refreshed-home-search-token"}}): false,
|
||||
}
|
||||
deadline := time.After(time.Second)
|
||||
for remaining := len(wantHashes); remaining > 0; {
|
||||
select {
|
||||
case record := <-records:
|
||||
if record.AuthID != "home-codex-search" || record.Fail.StatusCode != http.StatusUnauthorized {
|
||||
continue
|
||||
}
|
||||
if seen, ok := wantHashes[record.AccessTokenSHA256]; ok && !seen {
|
||||
wantHashes[record.AccessTokenSHA256] = true
|
||||
remaining--
|
||||
}
|
||||
case <-deadline:
|
||||
t.Fatalf("unauthorized attempt fingerprints = %#v", wantHashes)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeCodexAlphaSearchEndsSelectionAcrossDirectHTTPPaths(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -18,42 +18,8 @@ import (
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
)
|
||||
|
||||
type liveUsageCapturePlugin struct {
|
||||
records chan coreusage.Record
|
||||
}
|
||||
|
||||
func (p liveUsageCapturePlugin) HandleUsage(_ context.Context, record coreusage.Record) {
|
||||
if p.records != nil {
|
||||
p.records <- record
|
||||
}
|
||||
}
|
||||
|
||||
func waitForLiveUnauthorizedHashes(t *testing.T, records <-chan coreusage.Record, authID string, hashes ...string) {
|
||||
t.Helper()
|
||||
want := make(map[string]bool, len(hashes))
|
||||
for _, hash := range hashes {
|
||||
want[hash] = false
|
||||
}
|
||||
deadline := time.After(time.Second)
|
||||
for remaining := len(want); remaining > 0; {
|
||||
select {
|
||||
case record := <-records:
|
||||
if record.AuthID != authID || record.Fail.StatusCode != http.StatusUnauthorized {
|
||||
continue
|
||||
}
|
||||
if seen, ok := want[record.AccessTokenSHA256]; ok && !seen {
|
||||
want[record.AccessTokenSHA256] = true
|
||||
remaining--
|
||||
}
|
||||
case <-deadline:
|
||||
t.Fatalf("unauthorized attempt fingerprints = %#v", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type apiKeyFirstSelector struct{}
|
||||
|
||||
func (*apiKeyFirstSelector) Pick(_ context.Context, _ string, _ string, _ coreexecutor.Options, auths []*auth.Auth) (*auth.Auth, error) {
|
||||
@@ -674,44 +640,6 @@ func TestHandlerRefreshesUnauthorizedHomeSelectionOnce(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerReportsEveryUnauthorizedHomeAttempt(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
records := make(chan coreusage.Record, 8)
|
||||
const pluginName = "live-home-unauthorized-test"
|
||||
coreusage.RegisterNamedPlugin(pluginName, liveUsageCapturePlugin{records: records})
|
||||
t.Cleanup(func() {
|
||||
coreusage.RegisterNamedPlugin(pluginName, liveUsageCapturePlugin{})
|
||||
})
|
||||
|
||||
manager := auth.NewManager(nil, nil, nil)
|
||||
manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}})
|
||||
manager.PublishHomeDispatch(&homeDispatcher{}, executionregistry.New(), 1)
|
||||
executor := &captureExecutor{
|
||||
statuses: []int{http.StatusUnauthorized, http.StatusUnauthorized},
|
||||
responseBody: io.NopCloser(strings.NewReader("unauthorized")),
|
||||
}
|
||||
manager.RegisterExecutor(executor)
|
||||
handler := NewHandler(manager, nil)
|
||||
router := gin.New()
|
||||
router.POST("/v1/live", handler.Handle)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/live", strings.NewReader(`{"model":"gpt-live-1-codex","sdp":"v=0"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
if recorder.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("status = %d, want %d; body=%s", recorder.Code, http.StatusUnauthorized, recorder.Body.String())
|
||||
}
|
||||
if executor.refreshCalls.Load() != 1 || executor.httpCalls.Load() != 2 {
|
||||
t.Fatalf("refresh/http calls = %d/%d, want 1/2", executor.refreshCalls.Load(), executor.httpCalls.Load())
|
||||
}
|
||||
waitForLiveUnauthorizedHashes(t, records, "home-codex-live",
|
||||
auth.AccessTokenSHA256(&auth.Auth{Metadata: map[string]any{"access_token": "home-live-token"}}),
|
||||
auth.AccessTokenSHA256(&auth.Auth{Metadata: map[string]any{"access_token": "refreshed-home-live-token"}}),
|
||||
)
|
||||
}
|
||||
|
||||
func TestHandlerUsesLiveModelForHomeDispatch(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -964,62 +892,6 @@ func TestHandleSidebandRefreshesUnauthorizedHomeHandshakeOnce(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleSidebandReportsEveryUnauthorizedHomeHandshake(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
records := make(chan coreusage.Record, 8)
|
||||
const pluginName = "live-sideband-home-unauthorized-test"
|
||||
coreusage.RegisterNamedPlugin(pluginName, liveUsageCapturePlugin{records: records})
|
||||
t.Cleanup(func() {
|
||||
coreusage.RegisterNamedPlugin(pluginName, liveUsageCapturePlugin{})
|
||||
})
|
||||
|
||||
var upstreamCalls atomic.Int32
|
||||
upstreamServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) {
|
||||
upstreamCalls.Add(1)
|
||||
writer.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer upstreamServer.Close()
|
||||
|
||||
manager := auth.NewManager(nil, nil, nil)
|
||||
manager.SetConfig(&config.Config{Home: config.HomeConfig{Enabled: true}})
|
||||
manager.PublishHomeDispatch(&homeDispatcher{}, executionregistry.New(), 1)
|
||||
executor := &captureExecutor{}
|
||||
manager.RegisterExecutor(executor)
|
||||
selection, errSelect := manager.SelectHomeAuthByKind(context.Background(), "codex", defaultLiveModel, auth.AuthKindOAuth, coreexecutor.Options{})
|
||||
if errSelect != nil {
|
||||
t.Fatalf("SelectHomeAuthByKind() error = %v", errSelect)
|
||||
}
|
||||
selection.Retain()
|
||||
defer selection.End("test_complete")
|
||||
|
||||
handler := NewHandler(manager, nil)
|
||||
handler.sidebandAPIBaseURL = "ws" + strings.TrimPrefix(upstreamServer.URL, "http") + "/v1"
|
||||
handler.sessions.put("call-home-unauthorized", liveSession{authID: "home-codex-live", model: defaultLiveModel, homeSelection: selection})
|
||||
router := gin.New()
|
||||
router.GET("/v1/live/:call_id", handler.HandleSideband)
|
||||
downstreamServer := httptest.NewServer(router)
|
||||
defer downstreamServer.Close()
|
||||
|
||||
wsURL := "ws" + strings.TrimPrefix(downstreamServer.URL, "http") + "/v1/live/call-home-unauthorized"
|
||||
client, response, errDial := websocket.DefaultDialer.Dial(wsURL, nil)
|
||||
if client != nil {
|
||||
_ = client.Close()
|
||||
}
|
||||
if response != nil && response.Body != nil {
|
||||
defer func() { _ = response.Body.Close() }()
|
||||
}
|
||||
if errDial == nil || response == nil || response.StatusCode != http.StatusUnauthorized {
|
||||
t.Fatalf("sideband dial = response %#v error %v, want 401", response, errDial)
|
||||
}
|
||||
if executor.refreshCalls.Load() != 1 || upstreamCalls.Load() != 2 {
|
||||
t.Fatalf("refresh/upstream calls = %d/%d, want 1/2", executor.refreshCalls.Load(), upstreamCalls.Load())
|
||||
}
|
||||
waitForLiveUnauthorizedHashes(t, records, "home-codex-live",
|
||||
auth.AccessTokenSHA256(&auth.Auth{Metadata: map[string]any{"access_token": "home-live-token"}}),
|
||||
auth.AccessTokenSHA256(&auth.Auth{Metadata: map[string]any{"access_token": "refreshed-home-live-token"}}),
|
||||
)
|
||||
}
|
||||
|
||||
func TestPrepareCallRequestRewritesMultipart(t *testing.T) {
|
||||
const boundary = "live-model-boundary"
|
||||
body := multipartBody(boundary, "v=0-offer", `{"model":"future-live-model","instructions":"hi"}`)
|
||||
|
||||
@@ -112,37 +112,6 @@ func TestAuthAccessTokenSHA256SupportsKnownMetadataShapes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshAuthViaHomeRejectsDisabledAuthEnvelope(t *testing.T) {
|
||||
raw, errMarshal := json.Marshal(homeRefreshAuthEnvelope{
|
||||
Auth: cliproxyauth.Auth{
|
||||
ID: "disabled-home-auth",
|
||||
Provider: "codex",
|
||||
Status: cliproxyauth.StatusDisabled,
|
||||
Disabled: true,
|
||||
Metadata: map[string]any{"access_token": "disabled-access-token"},
|
||||
},
|
||||
AuthIndex: "disabled-home-auth",
|
||||
})
|
||||
if errMarshal != nil {
|
||||
t.Fatalf("marshal home envelope: %v", errMarshal)
|
||||
}
|
||||
client := &fakeHomeRefreshClient{raw: raw}
|
||||
oldCurrentHomeRefreshClient := currentHomeRefreshClient
|
||||
currentHomeRefreshClient = func() homeRefreshClient { return client }
|
||||
t.Cleanup(func() { currentHomeRefreshClient = oldCurrentHomeRefreshClient })
|
||||
|
||||
cfg := &config.Config{Home: config.HomeConfig{Enabled: true}}
|
||||
auth := &cliproxyauth.Auth{ID: "disabled-home-auth", Index: "disabled-home-auth", Provider: "codex"}
|
||||
updated, handled, errRefresh := RefreshAuthViaHome(context.Background(), cfg, auth)
|
||||
if updated != nil {
|
||||
t.Fatalf("RefreshAuthViaHome() auth = %#v, want nil", updated)
|
||||
}
|
||||
statusErr, okStatus := errRefresh.(interface{ StatusCode() int })
|
||||
if !handled || !okStatus || statusErr.StatusCode() != http.StatusUnauthorized {
|
||||
t.Fatalf("RefreshAuthViaHome() = handled %v err %v, want unauthorized", handled, errRefresh)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshAuthViaHomeAcceptsAuthEnvelope(t *testing.T) {
|
||||
raw, errMarshal := json.Marshal(struct {
|
||||
Auth cliproxyauth.Auth `json:"auth"`
|
||||
|
||||
@@ -8,7 +8,6 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
)
|
||||
|
||||
@@ -483,28 +482,6 @@ func TestUsageReporterBuildRecordIncludesLatency(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterUpdatesAccessTokenFingerprint(t *testing.T) {
|
||||
initial := &cliproxyauth.Auth{
|
||||
ID: "usage-auth",
|
||||
Index: "usage-auth",
|
||||
Provider: "antigravity",
|
||||
Metadata: map[string]any{"access_token": "initial-token"},
|
||||
}
|
||||
updated := initial.Clone()
|
||||
updated.Metadata["access_token"] = "refreshed-token"
|
||||
reporter := NewUsageReporter(context.Background(), "antigravity", "gemini-3-pro", initial)
|
||||
|
||||
reporter.UpdateAccessTokenFingerprint(updated)
|
||||
record := reporter.buildRecord(usage.Detail{TotalTokens: 3}, false)
|
||||
want := authAccessTokenSHA256(updated)
|
||||
if record.AccessTokenSHA256 != want {
|
||||
t.Fatalf("access token fingerprint = %q, want %q", record.AccessTokenSHA256, want)
|
||||
}
|
||||
if record.AccessTokenSHA256 == authAccessTokenSHA256(initial) {
|
||||
t.Fatal("usage reporter retained the pre-refresh token fingerprint")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUsageReporterTrackHTTPClientStartsTTFTBeforeRoundTrip(t *testing.T) {
|
||||
delay := 40 * time.Millisecond
|
||||
reporter := NewUsageReporter(context.Background(), "openai", "gpt-5.4", nil)
|
||||
|
||||
@@ -1,100 +0,0 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
)
|
||||
|
||||
type homeResultCapturePlugin struct {
|
||||
records chan coreusage.Record
|
||||
}
|
||||
|
||||
func (p homeResultCapturePlugin) HandleUsage(_ context.Context, record coreusage.Record) {
|
||||
if p.records != nil {
|
||||
p.records <- record
|
||||
}
|
||||
}
|
||||
|
||||
func TestReportHomeUnauthorizedPublishesTokenVersionedFailure(t *testing.T) {
|
||||
records := make(chan coreusage.Record, 8)
|
||||
const pluginName = "auth-home-result-test"
|
||||
coreusage.RegisterNamedPlugin(pluginName, homeResultCapturePlugin{records: records})
|
||||
t.Cleanup(func() {
|
||||
coreusage.RegisterNamedPlugin(pluginName, homeResultCapturePlugin{})
|
||||
})
|
||||
|
||||
auth := &Auth{
|
||||
ID: "home-result-auth",
|
||||
Index: "home-result-index",
|
||||
Provider: "codex",
|
||||
Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth},
|
||||
Metadata: map[string]any{
|
||||
"token": map[string]any{"accessToken": " current-access-token "},
|
||||
},
|
||||
}
|
||||
ctx := coreusage.WithRequestedModelAlias(context.Background(), "client-model")
|
||||
NewManager(nil, nil, nil).ReportHomeUnauthorized(ctx, auth, "codex", "upstream-model")
|
||||
|
||||
deadline := time.After(time.Second)
|
||||
for {
|
||||
select {
|
||||
case record := <-records:
|
||||
if record.AuthID != auth.ID {
|
||||
continue
|
||||
}
|
||||
if !record.Failed || record.Fail.StatusCode != http.StatusUnauthorized {
|
||||
t.Fatalf("failure = %#v, want 401", record.Fail)
|
||||
}
|
||||
if record.AuthIndex != auth.Index {
|
||||
t.Fatalf("auth index = %q, want %q", record.AuthIndex, auth.Index)
|
||||
}
|
||||
if record.AccessTokenSHA256 != AccessTokenSHA256(auth) || record.AccessTokenSHA256 == "" {
|
||||
t.Fatalf("access token fingerprint = %q", record.AccessTokenSHA256)
|
||||
}
|
||||
if record.Model != "upstream-model" || record.Alias != "client-model" {
|
||||
t.Fatalf("model/alias = %q/%q", record.Model, record.Alias)
|
||||
}
|
||||
if coreusage.GenerateEnabled(record.Generate) {
|
||||
t.Fatal("result-only unauthorized record was marked as generation")
|
||||
}
|
||||
if record.Detail.TotalTokens != 0 {
|
||||
t.Fatalf("result-only tokens = %d, want 0", record.Detail.TotalTokens)
|
||||
}
|
||||
return
|
||||
case <-deadline:
|
||||
t.Fatal("timed out waiting for Home unauthorized usage record")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReportHomeUnauthorizedRequiresTokenFingerprint(t *testing.T) {
|
||||
records := make(chan coreusage.Record, 1)
|
||||
const pluginName = "auth-home-result-empty-token-test"
|
||||
coreusage.RegisterNamedPlugin(pluginName, homeResultCapturePlugin{records: records})
|
||||
t.Cleanup(func() {
|
||||
coreusage.RegisterNamedPlugin(pluginName, homeResultCapturePlugin{})
|
||||
})
|
||||
|
||||
NewManager(nil, nil, nil).ReportHomeUnauthorized(context.Background(), &Auth{
|
||||
ID: "home-result-no-token",
|
||||
Index: "home-result-no-token",
|
||||
Provider: "codex",
|
||||
}, "codex", "model")
|
||||
|
||||
timer := time.NewTimer(50 * time.Millisecond)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case record := <-records:
|
||||
if record.AuthID == "home-result-no-token" {
|
||||
t.Fatalf("unexpected usage record without token fingerprint: %#v", record)
|
||||
}
|
||||
case <-timer.C:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3,17 +3,13 @@ package auth
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
internalconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executionregistry"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
coreusage "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/usage"
|
||||
)
|
||||
|
||||
const homeUnauthorizedRefreshProvider = "home-unauthorized-refresh"
|
||||
@@ -43,22 +39,14 @@ func (d *homeUnauthorizedRefreshDispatcher) RPopAuth(context.Context, string, st
|
||||
func (*homeUnauthorizedRefreshDispatcher) AbortAmbiguousDispatch() {}
|
||||
|
||||
type homeUnauthorizedRefreshExecutor struct {
|
||||
streamMode string
|
||||
refreshErr error
|
||||
keepStale bool
|
||||
retainSelection bool
|
||||
requirePrepared bool
|
||||
alwaysUnauthorized bool
|
||||
countAccessTokens []string
|
||||
nilRetryStream bool
|
||||
nilRetryChunks bool
|
||||
executeCalls atomic.Int32
|
||||
countCalls atomic.Int32
|
||||
streamCalls atomic.Int32
|
||||
refreshCalls atomic.Int32
|
||||
prepareCalls atomic.Int32
|
||||
refreshInputsMu sync.Mutex
|
||||
refreshInputs []string
|
||||
streamMode string
|
||||
refreshErr error
|
||||
keepStale bool
|
||||
retainSelection bool
|
||||
executeCalls atomic.Int32
|
||||
countCalls atomic.Int32
|
||||
streamCalls atomic.Int32
|
||||
refreshCalls atomic.Int32
|
||||
}
|
||||
|
||||
func (*homeUnauthorizedRefreshExecutor) Identifier() string { return homeUnauthorizedRefreshProvider }
|
||||
@@ -73,9 +61,6 @@ func (e *homeUnauthorizedRefreshExecutor) Execute(_ context.Context, auth *Auth,
|
||||
if authAccessToken(auth) == "stale-access-token" {
|
||||
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}
|
||||
}
|
||||
if e.requirePrepared && auth.Metadata["project_id"] != "prepared-project" {
|
||||
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusBadRequest, Message: "missing prepared auth"}
|
||||
}
|
||||
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
||||
}
|
||||
|
||||
@@ -98,15 +83,6 @@ func (e *homeUnauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth
|
||||
return nil, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}
|
||||
}
|
||||
}
|
||||
if e.requirePrepared && auth.Metadata["project_id"] != "prepared-project" {
|
||||
return nil, &Error{HTTPStatus: http.StatusBadRequest, Message: "missing prepared auth"}
|
||||
}
|
||||
if e.nilRetryStream {
|
||||
return nil, nil
|
||||
}
|
||||
if e.nilRetryChunks {
|
||||
return &cliproxyexecutor.StreamResult{}, nil
|
||||
}
|
||||
chunks := make(chan cliproxyexecutor.StreamChunk, 1)
|
||||
chunks <- cliproxyexecutor.StreamChunk{Payload: []byte("ok")}
|
||||
close(chunks)
|
||||
@@ -115,9 +91,6 @@ func (e *homeUnauthorizedRefreshExecutor) ExecuteStream(_ context.Context, auth
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) {
|
||||
e.refreshCalls.Add(1)
|
||||
e.refreshInputsMu.Lock()
|
||||
e.refreshInputs = append(e.refreshInputs, authAccessToken(auth))
|
||||
e.refreshInputsMu.Unlock()
|
||||
if e.refreshErr != nil {
|
||||
return nil, e.refreshErr
|
||||
}
|
||||
@@ -129,42 +102,14 @@ func (e *homeUnauthorizedRefreshExecutor) Refresh(_ context.Context, auth *Auth)
|
||||
updated.Metadata = make(map[string]any)
|
||||
}
|
||||
updated.Metadata["access_token"] = "fresh-access-token"
|
||||
if e.requirePrepared {
|
||||
delete(updated.Metadata, "project_id")
|
||||
}
|
||||
return updated, nil
|
||||
}
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) ShouldPrepareRequestAuth(auth *Auth) bool {
|
||||
return e.requirePrepared && auth != nil && auth.Metadata["project_id"] != "prepared-project"
|
||||
}
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) PrepareRequestAuth(_ context.Context, auth *Auth) (*Auth, error) {
|
||||
e.prepareCalls.Add(1)
|
||||
updated := auth.Clone()
|
||||
if updated.Metadata == nil {
|
||||
updated.Metadata = make(map[string]any)
|
||||
}
|
||||
updated.Metadata["project_id"] = "prepared-project"
|
||||
return updated, nil
|
||||
}
|
||||
|
||||
func (e *homeUnauthorizedRefreshExecutor) CountTokens(ctx context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
call := int(e.countCalls.Add(1))
|
||||
if call <= len(e.countAccessTokens) {
|
||||
effective := auth.Clone()
|
||||
if effective.Metadata == nil {
|
||||
effective.Metadata = make(map[string]any)
|
||||
}
|
||||
effective.Metadata["access_token"] = e.countAccessTokens[call-1]
|
||||
NotifyAccessTokenFingerprint(ctx, effective)
|
||||
}
|
||||
if e.alwaysUnauthorized || authAccessToken(auth) == "stale-access-token" {
|
||||
func (e *homeUnauthorizedRefreshExecutor) CountTokens(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
e.countCalls.Add(1)
|
||||
if authAccessToken(auth) == "stale-access-token" {
|
||||
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusUnauthorized, Message: "expired access token"}
|
||||
}
|
||||
if e.requirePrepared && auth.Metadata["project_id"] != "prepared-project" {
|
||||
return cliproxyexecutor.Response{}, &Error{HTTPStatus: http.StatusBadRequest, Message: "missing prepared auth"}
|
||||
}
|
||||
return cliproxyexecutor.Response{Payload: []byte("ok")}, nil
|
||||
}
|
||||
|
||||
@@ -224,59 +169,6 @@ func TestHomeUnauthorizedRefreshesSameSelectionBeforeRedispatch(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedRefreshRepreparesAuthBeforeRetry(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
run func(*Manager) error
|
||||
}{
|
||||
{
|
||||
name: "execute",
|
||||
run: func(manager *Manager) error {
|
||||
_, errExecute := manager.Execute(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
||||
return errExecute
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "count_tokens",
|
||||
run: func(manager *Manager) error {
|
||||
_, errCount := manager.ExecuteCount(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
||||
return errCount
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "stream",
|
||||
run: func(manager *Manager) error {
|
||||
result, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
||||
if errStream != nil {
|
||||
return errStream
|
||||
}
|
||||
for chunk := range result.Chunks {
|
||||
if chunk.Err != nil {
|
||||
return chunk.Err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{requirePrepared: true}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
if errRun := test.run(manager); errRun != nil {
|
||||
t.Fatalf("execution error = %v", errRun)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want 1", got)
|
||||
}
|
||||
if got := executor.prepareCalls.Load(); got != 2 {
|
||||
t.Fatalf("prepare calls = %d, want initial preparation and refreshed preparation", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedRefreshUpdatesRetainedSelection(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{retainSelection: true}
|
||||
@@ -304,7 +196,7 @@ func TestHomeUnauthorizedRefreshUpdatesRetainedSelection(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRefreshHomeSelectionReusesConcurrentNewerToken(t *testing.T) {
|
||||
executor := &homeUnauthorizedRefreshExecutor{requirePrepared: true}
|
||||
executor := &homeUnauthorizedRefreshExecutor{}
|
||||
selection := &HomeDispatchSelection{
|
||||
Auth: &Auth{ID: "home-refresh-auth", Provider: homeUnauthorizedRefreshProvider, Attributes: map[string]string{AttributeAuthKind: AuthKindOAuth}, Metadata: map[string]any{"access_token": "fresh-access-token"}},
|
||||
Executor: executor,
|
||||
@@ -320,12 +212,6 @@ func TestRefreshHomeSelectionReusesConcurrentNewerToken(t *testing.T) {
|
||||
if got := executor.refreshCalls.Load(); got != 0 {
|
||||
t.Fatalf("refresh calls = %d, want 0 when selection already has a newer token", got)
|
||||
}
|
||||
if got := executor.prepareCalls.Load(); got != 1 {
|
||||
t.Fatalf("prepare calls = %d, want reused token prepared once", got)
|
||||
}
|
||||
if updated.Metadata["project_id"] != "prepared-project" {
|
||||
t.Fatalf("reused auth metadata = %#v, want prepared project", updated.Metadata)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedRefreshIsAttemptedAtMostOnce(t *testing.T) {
|
||||
@@ -345,75 +231,6 @@ func TestHomeUnauthorizedRefreshIsAttemptedAtMostOnce(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeCountTokensReportsEveryUnauthorizedAttempt(t *testing.T) {
|
||||
records := make(chan coreusage.Record, 8)
|
||||
const pluginName = "auth-home-count-unauthorized-test"
|
||||
coreusage.RegisterNamedPlugin(pluginName, homeResultCapturePlugin{records: records})
|
||||
t.Cleanup(func() {
|
||||
coreusage.RegisterNamedPlugin(pluginName, homeResultCapturePlugin{})
|
||||
})
|
||||
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{
|
||||
alwaysUnauthorized: true,
|
||||
countAccessTokens: []string{"executor-internal-token", "retry-internal-token"},
|
||||
}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
_, errCount := manager.ExecuteCount(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{})
|
||||
if statusCodeFromError(errCount) != http.StatusUnauthorized {
|
||||
t.Fatalf("ExecuteCount() error = %v, want final 401", errCount)
|
||||
}
|
||||
if got := executor.refreshCalls.Load(); got != 1 {
|
||||
t.Fatalf("refresh calls = %d, want 1", got)
|
||||
}
|
||||
executor.refreshInputsMu.Lock()
|
||||
refreshInputs := append([]string(nil), executor.refreshInputs...)
|
||||
executor.refreshInputsMu.Unlock()
|
||||
if len(refreshInputs) != 1 || refreshInputs[0] != "executor-internal-token" {
|
||||
t.Fatalf("refresh input tokens = %#v, want internally refreshed token", refreshInputs)
|
||||
}
|
||||
if got := executor.countCalls.Load(); got != 2 {
|
||||
t.Fatalf("CountTokens calls = %d, want 2", got)
|
||||
}
|
||||
|
||||
wantHashes := map[string]bool{
|
||||
AccessTokenSHA256(&Auth{Metadata: map[string]any{"access_token": "executor-internal-token"}}): false,
|
||||
AccessTokenSHA256(&Auth{Metadata: map[string]any{"access_token": "retry-internal-token"}}): false,
|
||||
}
|
||||
matchedRecords := 0
|
||||
deadline := time.After(time.Second)
|
||||
for remaining := len(wantHashes); remaining > 0; {
|
||||
select {
|
||||
case record := <-records:
|
||||
if record.AuthID != "home-refresh-auth" || record.Fail.StatusCode != http.StatusUnauthorized {
|
||||
continue
|
||||
}
|
||||
matchedRecords++
|
||||
if seen, ok := wantHashes[record.AccessTokenSHA256]; ok && !seen {
|
||||
wantHashes[record.AccessTokenSHA256] = true
|
||||
remaining--
|
||||
}
|
||||
case <-deadline:
|
||||
t.Fatalf("unauthorized attempt fingerprints = %#v", wantHashes)
|
||||
}
|
||||
}
|
||||
timer := time.NewTimer(50 * time.Millisecond)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case record := <-records:
|
||||
if record.AuthID == "home-refresh-auth" && record.Fail.StatusCode == http.StatusUnauthorized {
|
||||
matchedRecords++
|
||||
}
|
||||
case <-timer.C:
|
||||
if matchedRecords != 2 {
|
||||
t.Fatalf("unauthorized usage records = %d, want exactly 2", matchedRecords)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeNoCandidateAfterRefreshFailurePreservesRefreshError(t *testing.T) {
|
||||
refreshErr := &Error{Code: "refresh_temporarily_unavailable", HTTPStatus: http.StatusServiceUnavailable, Message: "refresh unavailable"}
|
||||
noCandidate := &Error{Code: "auth_not_found", HTTPStatus: http.StatusServiceUnavailable, Message: "no auth available"}
|
||||
@@ -458,45 +275,6 @@ func TestHomeUnauthorizedStreamRefreshesAtMostOnceAcrossRedispatch(t *testing.T)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedBootstrapRetryRejectsEmptyStream(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
nilRetryStream bool
|
||||
nilRetryChunks bool
|
||||
}{
|
||||
{name: "nil result", nilRetryStream: true},
|
||||
{name: "nil chunks", nilRetryChunks: true},
|
||||
} {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{
|
||||
streamMode: "bootstrap",
|
||||
nilRetryStream: test.nilRetryStream,
|
||||
nilRetryChunks: test.nilRetryChunks,
|
||||
}
|
||||
manager := newHomeUnauthorizedRefreshManager(dispatcher, executor)
|
||||
|
||||
result, errStream := manager.ExecuteStream(context.Background(), []string{homeUnauthorizedRefreshProvider}, cliproxyexecutor.Request{Model: "model-a"}, cliproxyexecutor.Options{Stream: true})
|
||||
if errStream != nil {
|
||||
t.Fatalf("ExecuteStream() error = %v", errStream)
|
||||
}
|
||||
var streamErr error
|
||||
for chunk := range result.Chunks {
|
||||
if chunk.Err != nil {
|
||||
streamErr = chunk.Err
|
||||
}
|
||||
}
|
||||
var authErr *Error
|
||||
if !errors.As(streamErr, &authErr) || authErr.Code != "empty_stream" {
|
||||
t.Fatalf("stream error = %#v, want empty_stream", streamErr)
|
||||
}
|
||||
if got := executor.streamCalls.Load(); got != 2 {
|
||||
t.Fatalf("stream calls = %d, want initial attempt and one retry", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHomeUnauthorizedStartedStreamDoesNotReplay(t *testing.T) {
|
||||
dispatcher := &homeUnauthorizedRefreshDispatcher{}
|
||||
executor := &homeUnauthorizedRefreshExecutor{streamMode: "started"}
|
||||
|
||||
Reference in New Issue
Block a user