package api import (	"context"	"encoding/json"	"errors"	"io"	"net/http"	"net/http/httptest"	"os"	"path/filepath"	"reflect"	"strings"	"sync"	"sync/atomic"	"testing"	"time" 	gin "github.com/gin-gonic/gin"	managementHandlers "github.com/router-for-me/CLIProxyAPI/v7/internal/api/handlers/management"	claudemodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/claude/models"	codexmodels "github.com/router-for-me/CLIProxyAPI/v7/internal/client/codex/models"	proxyconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"	"github.com/router-for-me/CLIProxyAPI/v7/internal/home"	internallogging "github.com/router-for-me/CLIProxyAPI/v7/internal/logging"	"github.com/router-for-me/CLIProxyAPI/v7/internal/pluginhost"	"github.com/router-for-me/CLIProxyAPI/v7/internal/redisqueue"	"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"	"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor"	runtimehelps "github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor/helps"	sdkaccess "github.com/router-for-me/CLIProxyAPI/v7/sdk/access"	"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"	log "github.com/sirupsen/logrus"	logtest "github.com/sirupsen/logrus/hooks/test"	"gopkg.in/yaml.v3") type codexSearchCaptureExecutor struct {	request      *http.Request	body         []byte	authIDs      []string	prepareErr   error	httpErr      error	responseBody io.ReadCloser	statuses     []int	refreshCalls int	httpCalls    int	beforeReturn func()} func (e *codexSearchCaptureExecutor) Identifier() string { return "codex" } func (e *codexSearchCaptureExecutor) Execute(context.Context, *auth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {	return coreexecutor.Response{}, nil} func (e *codexSearchCaptureExecutor) ExecuteStream(context.Context, *auth.Auth, coreexecutor.Request, coreexecutor.Options) (*coreexecutor.StreamResult, error) {	return nil, nil} func (e *codexSearchCaptureExecutor) Refresh(_ context.Context, a *auth.Auth) (*auth.Auth, error) {	e.refreshCalls++	updated := a.Clone()	if updated.Metadata == nil {		updated.Metadata = make(map[string]any)	}	updated.Metadata["access_token"] = "refreshed-home-search-token"	return updated, nil} func (e *codexSearchCaptureExecutor) CountTokens(context.Context, *auth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {	return coreexecutor.Response{}, nil} func (e *codexSearchCaptureExecutor) PrepareRequest(req *http.Request, a *auth.Auth) error {	if e.prepareErr != nil {		return e.prepareErr	}	token, _ := a.Metadata["access_token"].(string)	if strings.TrimSpace(token) == "" && a.Attributes != nil {		token = a.Attributes[auth.AttributeAPIKey]	}	if strings.TrimSpace(token) != "" {		req.Header.Set("Authorization", "Bearer "+token)	}	return nil} type codexSearchGinContextSelector struct {	ginContext *gin.Context} func (s *codexSearchGinContextSelector) Pick(ctx context.Context, _ string, _ string, _ coreexecutor.Options, auths []*auth.Auth) (*auth.Auth, error) {	s.ginContext, _ = ctx.Value("gin").(*gin.Context)	if len(auths) == 0 {		return nil, nil	}	return auths[0], nil} type codexSearchAPIKeyFirstSelector struct{} type codexSearchModelRouter struct {	response pluginapi.ModelRouteResponse	handled  bool	requests []pluginapi.ModelRouteRequest} func (r *codexSearchModelRouter) RouteModel(_ context.Context, req pluginapi.ModelRouteRequest) (pluginapi.ModelRouteResponse, bool) {	r.requests = append(r.requests, req)	return r.response, r.handled} func (s *codexSearchAPIKeyFirstSelector) Pick(_ context.Context, _ string, _ string, _ coreexecutor.Options, auths []*auth.Auth) (*auth.Auth, error) {	for _, candidate := range auths {		if candidate.AuthKind() == auth.AuthKindAPIKey {			return candidate, nil		}	}	if len(auths) == 0 {		return nil, nil	}	return auths[0], nil} func (e *codexSearchCaptureExecutor) HttpRequest(_ context.Context, selected *auth.Auth, req *http.Request) (*http.Response, error) {	if e.httpErr != nil {		return nil, e.httpErr	}	e.request = req.Clone(req.Context())	e.authIDs = append(e.authIDs, selected.ID)	e.httpCalls++	body, err := io.ReadAll(req.Body)	if err != nil {		return nil, err	}	e.body = body	responseBody := e.responseBody	if responseBody == nil {		responseBody = io.NopCloser(strings.NewReader(`{"results":[{"url":"https://example.com"}]}`))	}	statusCode := http.StatusOK	if e.httpCalls <= len(e.statuses) && e.statuses[e.httpCalls-1] > 0 {		statusCode = e.statuses[e.httpCalls-1]	}	if e.beforeReturn != nil {		e.beforeReturn()	}	return &http.Response{		StatusCode: statusCode,		Header:     http.Header{"Content-Type": []string{"application/json"}},		Body:       responseBody,	}, nil} type codexSearchHomeDispatcher struct {	authID string	calls  atomic.Int32	policy atomic.Value} type homeUnauthorizedUsageCapture struct {	authID  string	records chan coreusage.Record} func (p *homeUnauthorizedUsageCapture) HandleUsage(_ context.Context, record coreusage.Record) {	if p == nil || record.ExecutorType != "home-result" || record.AuthID != p.authID {		return	}	select {	case p.records <- record:	default:	}} func (p *homeUnauthorizedUsageCapture) wait(t *testing.T) coreusage.Record {	t.Helper()	select {	case record := <-p.records:		return record	case <-time.After(time.Second):		t.Fatal("timed out waiting for Home unauthorized usage record")		return coreusage.Record{}	}} type noopHomeUnauthorizedUsagePlugin struct{} func (noopHomeUnauthorizedUsagePlugin) HandleUsage(context.Context, coreusage.Record) {} func registerHomeUnauthorizedUsageCapture(t *testing.T, name, authID string) *homeUnauthorizedUsageCapture {	t.Helper()	capture := &homeUnauthorizedUsageCapture{authID: authID, records: make(chan coreusage.Record, 1)}	coreusage.RegisterNamedPlugin(name, capture)	t.Cleanup(func() {		coreusage.RegisterNamedPlugin(name, noopHomeUnauthorizedUsagePlugin{})	})	return capture} func (*codexSearchHomeDispatcher) HeartbeatOK() bool { return true } func (d *codexSearchHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) {	d.calls.Add(1)	authID := d.authID	if authID == "" {		authID = "home-codex-search"	}	return json.Marshal(map[string]any{		"model":      model,		"auth_index": authID,		"auth": map[string]any{			"id":       authID,			"provider": "codex",			"status":   "active",			"metadata": map[string]any{"access_token": "home-search-token"},		},		"concurrency": map[string]any{			"accounted":     true,			"credential_id": authID,			"model":         model,		},	})} func (d *codexSearchHomeDispatcher) RPopAuthWithPolicy(ctx context.Context, model string, sessionID string, headers http.Header, count int, policy string) ([]byte, error) {	d.policy.Store(policy)	return d.RPopAuth(ctx, model, sessionID, headers, count)} func (*codexSearchHomeDispatcher) AbortAmbiguousDispatch() {} type codexSearchBusyHomeDispatcher struct{} func (*codexSearchBusyHomeDispatcher) HeartbeatOK() bool { return true }func (*codexSearchBusyHomeDispatcher) RPopAuth(context.Context, string, string, http.Header, int) ([]byte, error) {	return []byte(`{"error":{"type":"credential_concurrency_exceeded","message":"busy","retry_after_ms":750}}`), nil}func (d *codexSearchBusyHomeDispatcher) RPopAuthWithPolicy(ctx context.Context, model string, sessionID string, headers http.Header, count int, _ string) ([]byte, error) {	return d.RPopAuth(ctx, model, sessionID, headers, count)}func (*codexSearchBusyHomeDispatcher) AbortAmbiguousDispatch() {} type trackedSearchResponseBody struct {	io.Reader	closed atomic.Bool} func (b *trackedSearchResponseBody) Close() error {	b.closed.Store(true)	return nil} type errorSearchResponseBody struct {	payload []byte	read    atomic.Bool	closed  atomic.Bool} func (b *errorSearchResponseBody) Read(p []byte) (int, error) {	if !b.read.CompareAndSwap(false, true) {		return 0, io.EOF	}	return copy(p, b.payload), io.ErrUnexpectedEOF} func (b *errorSearchResponseBody) Close() error {	b.closed.Store(true)	return nil} type drainAwareSearchResponseBody struct {	started   chan struct{}	closed    chan struct{}	startOnce sync.Once	closeOnce sync.Once} func newDrainAwareSearchResponseBody() *drainAwareSearchResponseBody {	return &drainAwareSearchResponseBody{started: make(chan struct{}), closed: make(chan struct{})}} func (b *drainAwareSearchResponseBody) Read([]byte) (int, error) {	b.startOnce.Do(func() { close(b.started) })	<-b.closed	return 0, io.EOF} func (b *drainAwareSearchResponseBody) Close() error {	b.closeOnce.Do(func() { close(b.closed) })	return nil} func TestAuditHomeBusyNormalAndStream429Headers(t *testing.T) {	for _, stream := range []bool{false, true} {		t.Run(map[bool]string{false: "normal", true: "stream"}[stream], func(t *testing.T) {			server := newTestServer(t)			server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}})			server.handlers.AuthManager.PublishHomeDispatch(&codexSearchBusyHomeDispatcher{}, executionregistry.New(), 1) 			body := `{"model":"gpt-5-codex","input":[]}`			if stream {				body = `{"model":"gpt-5-codex","input":[],"stream":true}`			}			rr := httptest.NewRecorder()			req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body))			req.Header.Set("Authorization", "Bearer test-key")			server.engine.ServeHTTP(rr, req) 			if rr.Code != http.StatusTooManyRequests {				t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusTooManyRequests, rr.Body.String())			}			if got := rr.Header().Get("Retry-After"); got != "1" {				t.Fatalf("Retry-After = %q, want 1", got)			}		})	}} func TestAuditHomeCodexSearchBusyReturnsTrustedRetryAfter(t *testing.T) {	server := newTestServer(t)	server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}})	server.handlers.AuthManager.PublishHomeDispatch(&codexSearchBusyHomeDispatcher{}, executionregistry.New(), 1) 	rr := httptest.NewRecorder()	req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5-codex","query":"test"}`))	req.Header.Set("Authorization", "Bearer test-key")	server.engine.ServeHTTP(rr, req) 	if rr.Code != http.StatusTooManyRequests {		t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusTooManyRequests, rr.Body.String())	}	if got := rr.Header().Get("Retry-After"); got != "1" {		t.Fatalf("Retry-After = %q, want 1", got)	}	if !strings.Contains(rr.Body.String(), "busy") {		t.Fatalf("body = %q, want busy error", rr.Body.String())	}} func TestAuditHomeCodexSearchBodyCloseBeforeRelease(t *testing.T) {	server := newTestServer(t)	dispatcher := &codexSearchHomeDispatcher{}	registry := executionregistry.New()	body := newDrainAwareSearchResponseBody()	var releaseAfterBodyClose atomic.Bool	var releaseCount atomic.Int32	registry.SetReleaseSink(func(group executionregistry.ReleaseGroup, _ int64) {		if group != (executionregistry.ReleaseGroup{CredentialID: "home-codex-search", Model: "gpt-5-codex"}) {			t.Errorf("release group = %#v", group)		}		select {		case <-body.closed:			releaseAfterBodyClose.Store(true)		default:		}		releaseCount.Add(1)	})	server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}})	server.handlers.AuthManager.PublishHomeDispatch(dispatcher, registry, 1)	executor := &codexSearchCaptureExecutor{responseBody: body}	server.handlers.AuthManager.RegisterExecutor(executor) 	rr := httptest.NewRecorder()	req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"id":"home-search-drain","model":"gpt-5-codex","query":"test"}`))	req.Header.Set("Authorization", "Bearer test-key")	handlerDone := make(chan struct{})	go func() {		server.engine.ServeHTTP(rr, req)		close(handlerDone)	}() 	select {	case <-body.started:	case <-time.After(time.Second):		t.Fatal("search handler did not start reading the response body")	}	if errDrain := registry.Drain(context.Background()); errDrain != nil {		t.Fatalf("Drain() error = %v", errDrain)	}	if got := releaseCount.Load(); got != 1 {		t.Fatalf("accounted releases = %d, want 1", got)	}	if !releaseAfterBodyClose.Load() {		t.Fatal("accounted Home selection released before the search response body closed")	}	select {	case <-handlerDone:	case <-time.After(time.Second):		t.Fatal("search handler remained blocked after Home drain")	}	if rr.Code != http.StatusOK {		t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String())	}} func TestHomeCodexAlphaSearchForwardsUnauthorizedResponseWithoutRefresh(t *testing.T) {	const upstreamError = `{"error":{"message":"access token expired"}}`	server := newTestServer(t)	server.cfg.RequestLog = true	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},		responseBody: io.NopCloser(strings.NewReader(upstreamError)),	}	server.handlers.AuthManager.RegisterExecutor(executor) 	req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"id":"home-search-refresh","model":"gpt-5-codex","query":"test"}`))	req.Header.Set("Authorization", "Bearer test-key")	rr := httptest.NewRecorder()	c, _ := gin.CreateTestContext(rr)	c.Request = req	server.codexAlphaSearch(c) 	if rr.Code != http.StatusUnauthorized {		t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusUnauthorized, rr.Body.String())	}	if got := rr.Body.String(); got != upstreamError {		t.Fatalf("body = %q, want original upstream error %q", got, upstreamError)	}	if executor.refreshCalls != 0 || executor.httpCalls != 1 {		t.Fatalf("refresh/http calls = %d/%d, want 0/1", executor.refreshCalls, executor.httpCalls)	}	if got := executor.request.Header.Get("Authorization"); got != "Bearer home-search-token" {		t.Fatalf("Authorization = %q, want original Home token", got)	}	if got := dispatcher.calls.Load(); got != 1 {		t.Fatalf("Home RPOP calls = %d, want 1", got)	}	rawAPIResponse, okResponse := c.Get("API_RESPONSE")	if !okResponse {		t.Fatal("API_RESPONSE was not captured")	}	apiResponse, _ := rawAPIResponse.([]byte)	if !strings.Contains(string(apiResponse), "Status: 401") || !strings.Contains(string(apiResponse), upstreamError) {		t.Fatalf("API_RESPONSE = %q, want original upstream 401", apiResponse)	}} func TestHomeCodexAlphaSearchReportsUnauthorizedBeforeEarlyReturn(t *testing.T) {	const upstreamError = `{"error":{"message":"access token expired"}}`	tests := []struct {		name         string		responseBody func() io.ReadCloser		beforeReturn func(*executionregistry.Registry)		wantStatus   int		wantFailBody string	}{		{			name: "response bind failure",			responseBody: func() io.ReadCloser {				return &trackedSearchResponseBody{Reader: strings.NewReader(upstreamError)}			},			beforeReturn: func(registry *executionregistry.Registry) {				_ = registry.Close()			},			wantStatus:   http.StatusServiceUnavailable,			wantFailBody: "upstream unauthorized",		},		{			name: "response read failure",			responseBody: func() io.ReadCloser {				return &errorSearchResponseBody{payload: []byte(upstreamError)}			},			wantStatus:   http.StatusBadGateway,			wantFailBody: upstreamError,		},	} 	for _, test := range tests {		t.Run(test.name, func(t *testing.T) {			server := newTestServer(t)			registry := executionregistry.New()			server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}})			testAuthID := "home-codex-search-" + strings.ReplaceAll(test.name, " ", "-")			server.handlers.AuthManager.PublishHomeDispatch(&codexSearchHomeDispatcher{authID: testAuthID}, registry, 1)			executor := &codexSearchCaptureExecutor{				statuses:     []int{http.StatusUnauthorized},				responseBody: test.responseBody(),			}			if test.beforeReturn != nil {				executor.beforeReturn = func() { test.beforeReturn(registry) }			}			server.handlers.AuthManager.RegisterExecutor(executor)			usageCapture := registerHomeUnauthorizedUsageCapture(t, t.Name(), testAuthID) 			recorder := httptest.NewRecorder()			request := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5-codex","query":"test"}`))			request.Header.Set("Authorization", "Bearer test-key")			server.engine.ServeHTTP(recorder, request) 			if recorder.Code != test.wantStatus {

Lines 1–500 of 3098. Load more