package api import (	"context"	"encoding/json"	"errors"	"io"	"net/http"	"net/http/httptest"	"os"	"path/filepath"	"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"	proxyconfig "github.com/router-for-me/CLIProxyAPI/v7/internal/config"	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"	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"	sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"	"github.com/router-for-me/CLIProxyAPI/v7/sdk/pluginapi") type codexSearchCaptureExecutor struct {	request      *http.Request	body         []byte	authIDs      []string	prepareErr   error	httpErr      error	responseBody io.ReadCloser	statuses     []int	refreshCalls int	httpCalls    int} 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]	}	return &http.Response{		StatusCode: statusCode,		Header:     http.Header{"Content-Type": []string{"application/json"}},		Body:       responseBody,	}, nil} type codexSearchHomeDispatcher struct {	calls  atomic.Int32	policy atomic.Value} func (*codexSearchHomeDispatcher) HeartbeatOK() bool { return true } func (d *codexSearchHomeDispatcher) RPopAuth(_ context.Context, model string, _ string, _ http.Header, _ int) ([]byte, error) {	d.calls.Add(1)	return json.Marshal(map[string]any{		"model":      model,		"auth_index": "home-codex-search",		"auth": map[string]any{			"id":       "home-codex-search",			"provider": "codex",			"status":   "active",			"metadata": map[string]any{"access_token": "home-search-token"},		},		"concurrency": map[string]any{			"accounted":     true,			"credential_id": "home-codex-search",			"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 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 TestHomeCodexAlphaSearchRefreshesUnauthorizedSelectionOnce(t *testing.T) {	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.StatusOK}}	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()	server.engine.ServeHTTP(rr, req) 	if rr.Code != http.StatusOK {		t.Fatalf("status = %d, want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String())	}	if executor.refreshCalls != 1 || executor.httpCalls != 2 {		t.Fatalf("refresh/http calls = %d/%d, want 1/2", executor.refreshCalls, executor.httpCalls)	}	if got := executor.request.Header.Get("Authorization"); got != "Bearer refreshed-home-search-token" {		t.Fatalf("retry Authorization = %q, want refreshed token", got)	}	if got := dispatcher.calls.Load(); got != 1 {		t.Fatalf("Home RPOP calls = %d, want 1", got)	}} func TestHomeCodexAlphaSearchEndsSelectionAcrossDirectHTTPPaths(t *testing.T) {	tests := []struct {		name       string		configure  func(*codexSearchCaptureExecutor, *trackedSearchResponseBody)		wantStatus int		wantClosed bool	}{		{			name: "request build failure",			configure: func(executor *codexSearchCaptureExecutor, _ *trackedSearchResponseBody) {				executor.prepareErr = errors.New("request preparation failed")			},			wantStatus: http.StatusBadGateway,		},		{			name: "HTTP error",			configure: func(executor *codexSearchCaptureExecutor, _ *trackedSearchResponseBody) {				executor.httpErr = errors.New("upstream unavailable")			},			wantStatus: http.StatusBadGateway,		},		{			name: "response body close",			configure: func(executor *codexSearchCaptureExecutor, body *trackedSearchResponseBody) {				executor.responseBody = body			},			wantStatus: http.StatusOK,			wantClosed: true,		},	} 	for _, test := range tests {		t.Run(test.name, func(t *testing.T) {			server := newTestServer(t)			dispatcher := &codexSearchHomeDispatcher{}			registry := executionregistry.New()			server.handlers.AuthManager.SetConfig(&proxyconfig.Config{Home: proxyconfig.HomeConfig{Enabled: true}})			server.handlers.AuthManager.PublishHomeDispatch(dispatcher, registry, 1)			body := &trackedSearchResponseBody{Reader: strings.NewReader(`{"results":[]}`)}			executor := &codexSearchCaptureExecutor{}			test.configure(executor, body)			server.handlers.AuthManager.RegisterExecutor(executor) 			req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"id":"home-search-session","model":"gpt-5-codex","query":"test"}`))			req.Header.Set("Authorization", "Bearer test-key")			rr := httptest.NewRecorder()			server.engine.ServeHTTP(rr, req)			if rr.Code != test.wantStatus {				t.Fatalf("status = %d, want %d; body=%s", rr.Code, test.wantStatus, rr.Body.String())			}			if got := dispatcher.calls.Load(); got != 1 {				t.Fatalf("Home RPOP calls = %d, want 1", got)			}			if got, _ := dispatcher.policy.Load().(string); got != auth.CredentialPolicyCodexAlphaSearchV1 {				t.Fatalf("Home credential policy = %q, want %q", got, auth.CredentialPolicyCodexAlphaSearchV1)			}			if got := body.closed.Load(); got != test.wantClosed {				t.Fatalf("response body closed = %t, want %t", got, test.wantClosed)			}			if errDrain := registry.Drain(context.Background()); errDrain != nil {				t.Fatalf("Drain() error = %v", errDrain)			}		})	}} func newTestServer(t *testing.T) *Server {	t.Helper()	return newTestServerWithOptions(t)} func newTestServerWithOptions(t *testing.T, opts ...ServerOption) *Server {	t.Helper() 	gin.SetMode(gin.TestMode) 	tmpDir := t.TempDir()	authDir := filepath.Join(tmpDir, "auth")	if err := os.MkdirAll(authDir, 0o700); err != nil {		t.Fatalf("failed to create auth dir: %v", err)	} 	cfg := &proxyconfig.Config{		SDKConfig: sdkconfig.SDKConfig{			APIKeys: []string{"test-key"},		},		Port:                   0,		AuthDir:                authDir,		Debug:                  true,		LoggingToFile:          false,		UsageStatisticsEnabled: false,	} 	authManager := auth.NewManager(nil, nil, nil)	accessManager := sdkaccess.NewManager() 	configPath := filepath.Join(tmpDir, "config.yaml")	return NewServer(cfg, authManager, accessManager, configPath, opts...)} func TestHealthz(t *testing.T) {	server := newTestServer(t) 	t.Run("GET", func(t *testing.T) {		req := httptest.NewRequest(http.MethodGet, "/healthz", nil)		rr := httptest.NewRecorder()		server.engine.ServeHTTP(rr, req) 		if rr.Code != http.StatusOK {			t.Fatalf("unexpected status code: got %d want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String())		} 		var resp struct {			Status string `json:"status"`		}		if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {			t.Fatalf("failed to parse response JSON: %v; body=%s", err, rr.Body.String())		}		if resp.Status != "ok" {			t.Fatalf("unexpected response status: got %q want %q", resp.Status, "ok")		}	}) 	t.Run("HEAD", func(t *testing.T) {		req := httptest.NewRequest(http.MethodHead, "/healthz", nil)		rr := httptest.NewRecorder()		server.engine.ServeHTTP(rr, req) 		if rr.Code != http.StatusOK {			t.Fatalf("unexpected status code: got %d want %d; body=%s", rr.Code, http.StatusOK, rr.Body.String())		}		if rr.Body.Len() != 0 {			t.Fatalf("expected empty body for HEAD request, got %q", rr.Body.String())		}	})} func TestCodexLiveRoutesRequireAuthAndAreRegistered(t *testing.T) {	server := newTestServer(t) 	for _, path := range []string{"/v1/live", "/v1/realtime/calls"} {		unauthorized := httptest.NewRequest(http.MethodPost, path, nil)		unauthorizedRecorder := httptest.NewRecorder()		server.engine.ServeHTTP(unauthorizedRecorder, unauthorized)		if unauthorizedRecorder.Code != http.StatusUnauthorized {			t.Fatalf("%s unauthorized status = %d, want %d", path, unauthorizedRecorder.Code, http.StatusUnauthorized)		}

Lines 1–500 of 2269. Load more