604 lines
19 KiB
Go
604 lines
19 KiB
Go
|
package service
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"encoding/json"
|
|||
|
|
"errors"
|
|||
|
|
"fmt"
|
|||
|
|
"net/http"
|
|||
|
|
"net/http/httptest"
|
|||
|
|
"strings"
|
|||
|
|
"sync"
|
|||
|
|
"testing"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|||
|
|
coderws "github.com/coder/websocket"
|
|||
|
|
"github.com/stretchr/testify/require"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
type liveTestFrame struct {
|
|||
|
|
messageType coderws.MessageType
|
|||
|
|
payload []byte
|
|||
|
|
err error
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type liveTestFrameConn struct {
|
|||
|
|
reads chan liveTestFrame
|
|||
|
|
writes chan liveTestFrame
|
|||
|
|
closed chan struct{}
|
|||
|
|
closeOnce sync.Once
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func newLiveTestFrameConn() *liveTestFrameConn {
|
|||
|
|
return &liveTestFrameConn{
|
|||
|
|
reads: make(chan liveTestFrame, 8),
|
|||
|
|
writes: make(chan liveTestFrame, 8),
|
|||
|
|
closed: make(chan struct{}),
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestFrameConn) ReadFrame(ctx context.Context) (coderws.MessageType, []byte, error) {
|
|||
|
|
select {
|
|||
|
|
case frame := <-c.reads:
|
|||
|
|
return frame.messageType, frame.payload, frame.err
|
|||
|
|
case <-c.closed:
|
|||
|
|
return coderws.MessageText, nil, coderws.CloseError{Code: coderws.StatusNormalClosure}
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
return coderws.MessageText, nil, context.Cause(ctx)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestFrameConn) WriteFrame(ctx context.Context, messageType coderws.MessageType, payload []byte) error {
|
|||
|
|
frame := liveTestFrame{messageType: messageType, payload: append([]byte(nil), payload...)}
|
|||
|
|
select {
|
|||
|
|
case c.writes <- frame:
|
|||
|
|
return nil
|
|||
|
|
case <-c.closed:
|
|||
|
|
return errors.New("connection closed")
|
|||
|
|
case <-ctx.Done():
|
|||
|
|
return context.Cause(ctx)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestFrameConn) WriteJSON(ctx context.Context, value any) error {
|
|||
|
|
payload, err := json.Marshal(value)
|
|||
|
|
if err != nil {
|
|||
|
|
return err
|
|||
|
|
}
|
|||
|
|
return c.WriteFrame(ctx, coderws.MessageText, payload)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestFrameConn) ReadMessage(ctx context.Context) ([]byte, error) {
|
|||
|
|
_, payload, err := c.ReadFrame(ctx)
|
|||
|
|
return payload, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestFrameConn) Ping(context.Context) error { return nil }
|
|||
|
|
|
|||
|
|
func (c *liveTestFrameConn) Close() error {
|
|||
|
|
c.closeOnce.Do(func() { close(c.closed) })
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type liveTestDialer struct {
|
|||
|
|
conn *liveTestFrameConn
|
|||
|
|
url string
|
|||
|
|
headers http.Header
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (d *liveTestDialer) Dial(
|
|||
|
|
_ context.Context,
|
|||
|
|
wsURL string,
|
|||
|
|
headers http.Header,
|
|||
|
|
_ string,
|
|||
|
|
) (openAIWSClientConn, int, http.Header, error) {
|
|||
|
|
d.url = wsURL
|
|||
|
|
d.headers = headers.Clone()
|
|||
|
|
return d.conn, http.StatusSwitchingProtocols, nil, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type liveTestAccountRepo struct {
|
|||
|
|
AccountRepository
|
|||
|
|
account *Account
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *liveTestAccountRepo) GetByID(context.Context, int64) (*Account, error) {
|
|||
|
|
return r.account, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type liveTestStore struct {
|
|||
|
|
GatewayCache
|
|||
|
|
mu sync.Mutex
|
|||
|
|
record *LiveCallRecord
|
|||
|
|
// 注入 store 故障（模拟 Redis 抖动），区别于 ErrLiveCallNotFound。
|
|||
|
|
claimErr error
|
|||
|
|
getCallErr error
|
|||
|
|
getControllerErr error
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *liveTestStore) SaveLiveCall(_ context.Context, record *LiveCallRecord, _ time.Duration) error {
|
|||
|
|
s.mu.Lock()
|
|||
|
|
defer s.mu.Unlock()
|
|||
|
|
copy := *record
|
|||
|
|
s.record = ©
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *liveTestStore) GetLiveCall(_ context.Context, callHash string) (*LiveCallRecord, error) {
|
|||
|
|
s.mu.Lock()
|
|||
|
|
defer s.mu.Unlock()
|
|||
|
|
if s.getCallErr != nil {
|
|||
|
|
return nil, s.getCallErr
|
|||
|
|
}
|
|||
|
|
if s.record == nil || s.record.CallHash != callHash {
|
|||
|
|
return nil, ErrLiveCallNotFound
|
|||
|
|
}
|
|||
|
|
copy := *s.record
|
|||
|
|
return ©, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *liveTestStore) ClaimLiveController(_ context.Context, callHash, controller, owner string) (bool, error) {
|
|||
|
|
s.mu.Lock()
|
|||
|
|
defer s.mu.Unlock()
|
|||
|
|
if s.claimErr != nil {
|
|||
|
|
return false, s.claimErr
|
|||
|
|
}
|
|||
|
|
if s.record == nil || s.record.CallHash != callHash || s.record.Controller == LiveControllerClosed {
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
if controller == LiveControllerObserver && s.record.Controller != LiveControllerPending {
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
if controller == LiveControllerProxy && s.record.Controller != LiveControllerPending && s.record.Controller != LiveControllerObserver {
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
s.record.Controller = controller
|
|||
|
|
s.record.ControllerOwner = owner
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *liveTestStore) ReleaseLiveController(_ context.Context, callHash, owner string) (bool, error) {
|
|||
|
|
s.mu.Lock()
|
|||
|
|
defer s.mu.Unlock()
|
|||
|
|
if s.record == nil || s.record.CallHash != callHash || s.record.ControllerOwner != owner {
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
s.record.Controller = LiveControllerPending
|
|||
|
|
s.record.ControllerOwner = ""
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *liveTestStore) GetLiveController(_ context.Context, callHash string) (string, error) {
|
|||
|
|
s.mu.Lock()
|
|||
|
|
defer s.mu.Unlock()
|
|||
|
|
if s.getControllerErr != nil {
|
|||
|
|
return "", s.getControllerErr
|
|||
|
|
}
|
|||
|
|
if s.record == nil || s.record.CallHash != callHash {
|
|||
|
|
return "", ErrLiveCallNotFound
|
|||
|
|
}
|
|||
|
|
return s.record.Controller, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (s *liveTestStore) MarkLiveCallClosed(_ context.Context, callHash string, _ time.Duration) (bool, error) {
|
|||
|
|
s.mu.Lock()
|
|||
|
|
defer s.mu.Unlock()
|
|||
|
|
if s.record == nil || s.record.CallHash != callHash || s.record.Controller == LiveControllerClosed {
|
|||
|
|
return false, nil
|
|||
|
|
}
|
|||
|
|
s.record.Controller = LiveControllerClosed
|
|||
|
|
s.record.ControllerOwner = ""
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type liveTestConcurrencyCache struct {
|
|||
|
|
ConcurrencyCache
|
|||
|
|
mu sync.Mutex
|
|||
|
|
releases int
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestConcurrencyCache) AcquireLiveLease(
|
|||
|
|
context.Context,
|
|||
|
|
int64,
|
|||
|
|
int,
|
|||
|
|
int64,
|
|||
|
|
int,
|
|||
|
|
int64,
|
|||
|
|
string,
|
|||
|
|
bool,
|
|||
|
|
) (bool, error) {
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestConcurrencyCache) RefreshLiveLease(
|
|||
|
|
context.Context,
|
|||
|
|
int64,
|
|||
|
|
int64,
|
|||
|
|
int64,
|
|||
|
|
string,
|
|||
|
|
) (bool, error) {
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *liveTestConcurrencyCache) ReleaseLiveLease(
|
|||
|
|
context.Context,
|
|||
|
|
int64,
|
|||
|
|
int64,
|
|||
|
|
int64,
|
|||
|
|
string,
|
|||
|
|
) error {
|
|||
|
|
c.mu.Lock()
|
|||
|
|
c.releases++
|
|||
|
|
c.mu.Unlock()
|
|||
|
|
return nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type liveTestUsageRepo struct {
|
|||
|
|
UsageLogRepository
|
|||
|
|
mu sync.Mutex
|
|||
|
|
logs []*UsageLog
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *liveTestUsageRepo) Create(_ context.Context, log *UsageLog) (bool, error) {
|
|||
|
|
r.mu.Lock()
|
|||
|
|
defer r.mu.Unlock()
|
|||
|
|
copy := *log
|
|||
|
|
r.logs = append(r.logs, ©)
|
|||
|
|
return true, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestRunLiveControllerClosesExpiredSession(t *testing.T) {
|
|||
|
|
upstream := newLiveTestFrameConn()
|
|||
|
|
record := &LiveCallRecord{ExpiresAt: time.Now().Add(20 * time.Millisecond)}
|
|||
|
|
service := &OpenAIGatewayService{}
|
|||
|
|
|
|||
|
|
err := service.runLiveController(context.Background(), record, upstream, make(chan error))
|
|||
|
|
require.ErrorIs(t, err, context.DeadlineExceeded)
|
|||
|
|
|
|||
|
|
select {
|
|||
|
|
case frame := <-upstream.writes:
|
|||
|
|
require.Equal(t, coderws.MessageText, frame.messageType)
|
|||
|
|
require.JSONEq(t, `{"type":"session.close"}`, string(frame.payload))
|
|||
|
|
case <-time.After(time.Second):
|
|||
|
|
t.Fatal("没有向上游发送 session.close")
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestFinalizeLiveCallIsIdempotentAndWritesZeroUsage(t *testing.T) {
|
|||
|
|
record := &LiveCallRecord{
|
|||
|
|
CallID: "call_secret",
|
|||
|
|
CallHash: hashLiveCallID("call_secret"),
|
|||
|
|
AccountID: 11,
|
|||
|
|
APIKeyID: 22,
|
|||
|
|
UserID: 33,
|
|||
|
|
GroupID: 44,
|
|||
|
|
LeaseID: "lease-1",
|
|||
|
|
Model: "gpt-live-test",
|
|||
|
|
CreatedAt: time.Now().Add(-time.Second),
|
|||
|
|
ExpiresAt: time.Now().Add(time.Hour),
|
|||
|
|
Controller: LiveControllerPending,
|
|||
|
|
InboundEndpoint: "/v1/live",
|
|||
|
|
}
|
|||
|
|
store := &liveTestStore{}
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), record, time.Hour))
|
|||
|
|
concurrencyCache := &liveTestConcurrencyCache{}
|
|||
|
|
usageRepo := &liveTestUsageRepo{}
|
|||
|
|
service := &OpenAIGatewayService{
|
|||
|
|
cache: store,
|
|||
|
|
concurrencyService: NewConcurrencyService(concurrencyCache),
|
|||
|
|
usageLogRepo: usageRepo,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
service.finalizeLiveCall(record)
|
|||
|
|
service.finalizeLiveCall(record)
|
|||
|
|
|
|||
|
|
concurrencyCache.mu.Lock()
|
|||
|
|
require.Equal(t, 1, concurrencyCache.releases)
|
|||
|
|
concurrencyCache.mu.Unlock()
|
|||
|
|
usageRepo.mu.Lock()
|
|||
|
|
require.Len(t, usageRepo.logs, 1)
|
|||
|
|
log := usageRepo.logs[0]
|
|||
|
|
usageRepo.mu.Unlock()
|
|||
|
|
require.Equal(t, RequestTypeLive, log.RequestType)
|
|||
|
|
require.Equal(t, record.CallHash, log.RequestID)
|
|||
|
|
require.NotEqual(t, record.CallID, log.RequestID)
|
|||
|
|
require.NotNil(t, log.DurationMs)
|
|||
|
|
require.Zero(t, log.InputTokens)
|
|||
|
|
require.Zero(t, log.OutputTokens)
|
|||
|
|
require.Zero(t, log.TotalCost)
|
|||
|
|
require.Zero(t, log.ActualCost)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestGetLiveCallForIdentityRejectsMismatchedCaller(t *testing.T) {
|
|||
|
|
groupID := int64(44)
|
|||
|
|
record := &LiveCallRecord{
|
|||
|
|
CallID: "call_identity",
|
|||
|
|
CallHash: hashLiveCallID("call_identity"),
|
|||
|
|
APIKeyID: 22,
|
|||
|
|
UserID: 33,
|
|||
|
|
GroupID: groupID,
|
|||
|
|
Controller: LiveControllerPending,
|
|||
|
|
}
|
|||
|
|
store := &liveTestStore{}
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), record, time.Hour))
|
|||
|
|
service := &OpenAIGatewayService{cache: store}
|
|||
|
|
|
|||
|
|
_, err := service.GetLiveCallForIdentity(context.Background(), record.CallID, LiveCallIdentity{
|
|||
|
|
APIKeyID: 99,
|
|||
|
|
UserID: record.UserID,
|
|||
|
|
GroupID: &groupID,
|
|||
|
|
})
|
|||
|
|
require.ErrorIs(t, err, ErrLiveIdentityMismatch)
|
|||
|
|
|
|||
|
|
loaded, err := service.GetLiveCallForIdentity(context.Background(), record.CallID, LiveCallIdentity{
|
|||
|
|
APIKeyID: record.APIKeyID,
|
|||
|
|
UserID: record.UserID,
|
|||
|
|
GroupID: &groupID,
|
|||
|
|
})
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
require.Equal(t, record.AccountID, loaded.AccountID)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func TestProxyLiveSidebandForwardsTextAndBinary(t *testing.T) {
|
|||
|
|
account := &Account{
|
|||
|
|
ID: 11,
|
|||
|
|
Platform: PlatformOpenAI,
|
|||
|
|
Type: AccountTypeOAuth,
|
|||
|
|
Concurrency: 2,
|
|||
|
|
Credentials: map[string]any{
|
|||
|
|
"access_token": "test-access-token",
|
|||
|
|
"chatgpt_account_id": "acct_test",
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
record := &LiveCallRecord{
|
|||
|
|
CallID: "call_proxy",
|
|||
|
|
CallHash: hashLiveCallID("call_proxy"),
|
|||
|
|
AccountID: account.ID,
|
|||
|
|
APIKeyID: 22,
|
|||
|
|
UserID: 33,
|
|||
|
|
LeaseID: "lease-1",
|
|||
|
|
CreatedAt: time.Now(),
|
|||
|
|
ExpiresAt: time.Now().Add(time.Minute),
|
|||
|
|
Controller: LiveControllerPending,
|
|||
|
|
}
|
|||
|
|
attestationCipher := newLiveAttestationCipher(&config.Config{
|
|||
|
|
JWT: config.JWTConfig{Secret: "live-sideband-test-secret"},
|
|||
|
|
})
|
|||
|
|
var err error
|
|||
|
|
record.AttestationCiphertext, err = attestationCipher.Encrypt(`{"v":1,"s":0,"t":"v1.sideband"}`)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
store := &liveTestStore{}
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), record, time.Hour))
|
|||
|
|
upstream := newLiveTestFrameConn()
|
|||
|
|
dialer := &liveTestDialer{conn: upstream}
|
|||
|
|
service := &OpenAIGatewayService{
|
|||
|
|
accountRepo: &liveTestAccountRepo{account: account},
|
|||
|
|
cache: store,
|
|||
|
|
openaiWSPassthroughDialer: dialer,
|
|||
|
|
liveAttestationCipher: attestationCipher,
|
|||
|
|
}
|
|||
|
|
proxyResult := make(chan error, 1)
|
|||
|
|
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
|||
|
|
downstream, err := coderws.Accept(writer, request, nil)
|
|||
|
|
if err != nil {
|
|||
|
|
proxyResult <- err
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
defer func() { _ = downstream.CloseNow() }()
|
|||
|
|
proxyResult <- service.ProxyLiveSideband(request.Context(), record, downstream)
|
|||
|
|
}))
|
|||
|
|
defer server.Close()
|
|||
|
|
|
|||
|
|
client, _, err := coderws.Dial(
|
|||
|
|
context.Background(),
|
|||
|
|
"ws"+strings.TrimPrefix(server.URL, "http"),
|
|||
|
|
nil,
|
|||
|
|
)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
defer func() { _ = client.CloseNow() }()
|
|||
|
|
|
|||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|||
|
|
defer cancel()
|
|||
|
|
require.NoError(t, client.Write(ctx, coderws.MessageText, []byte(`{"type":"client.text"}`)))
|
|||
|
|
clientText := <-upstream.writes
|
|||
|
|
require.Equal(t, coderws.MessageText, clientText.messageType)
|
|||
|
|
require.JSONEq(t, `{"type":"client.text"}`, string(clientText.payload))
|
|||
|
|
|
|||
|
|
require.NoError(t, client.Write(ctx, coderws.MessageBinary, []byte{1, 2, 3}))
|
|||
|
|
clientBinary := <-upstream.writes
|
|||
|
|
require.Equal(t, coderws.MessageBinary, clientBinary.messageType)
|
|||
|
|
require.Equal(t, []byte{1, 2, 3}, clientBinary.payload)
|
|||
|
|
|
|||
|
|
upstream.reads <- liveTestFrame{messageType: coderws.MessageText, payload: []byte(`{"type":"server.text"}`)}
|
|||
|
|
messageType, payload, err := client.Read(ctx)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
require.Equal(t, coderws.MessageText, messageType)
|
|||
|
|
require.JSONEq(t, `{"type":"server.text"}`, string(payload))
|
|||
|
|
|
|||
|
|
upstream.reads <- liveTestFrame{messageType: coderws.MessageBinary, payload: []byte{4, 5, 6}}
|
|||
|
|
messageType, payload, err = client.Read(ctx)
|
|||
|
|
require.NoError(t, err)
|
|||
|
|
require.Equal(t, coderws.MessageBinary, messageType)
|
|||
|
|
require.Equal(t, []byte{4, 5, 6}, payload)
|
|||
|
|
|
|||
|
|
require.Equal(t, "wss://chatgpt.com/backend-api/codex/call_proxy", dialer.url)
|
|||
|
|
require.Equal(t, "Bearer test-access-token", dialer.headers.Get("Authorization"))
|
|||
|
|
require.Equal(t, "acct_test", dialer.headers.Get("Chatgpt-Account-Id"))
|
|||
|
|
require.Equal(t, `{"v":1,"s":0,"t":"v1.sideband"}`, dialer.headers.Get(liveAttestationHeader))
|
|||
|
|
upstream.reads <- liveTestFrame{err: coderws.CloseError{Code: coderws.StatusNormalClosure}}
|
|||
|
|
require.ErrorIs(t, <-proxyResult, ErrLiveCallNotFound)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestLiveSessionEndedTreatsLeaseLossAsTerminal 锁定：租约续租失败（ErrLiveUnavailable）
|
|||
|
|
// 必须判为会话终结。RefreshLiveLease 的 Lua 在 leaseID 被 GC 后不会重新写入，若把它
|
|||
|
|
// 当临时错误交给 observer 重连，会话会空转到 ExpiresAt 且不计入任何并发限制。
|
|||
|
|
func TestLiveSessionEndedTreatsLeaseLossAsTerminal(t *testing.T) {
|
|||
|
|
cases := []struct {
|
|||
|
|
name string
|
|||
|
|
err error
|
|||
|
|
want bool
|
|||
|
|
}{
|
|||
|
|
{"租约丢失", ErrLiveUnavailable, true},
|
|||
|
|
{"租约丢失（被包装）", fmt.Errorf("refresh live lease: %w", ErrLiveUnavailable), true},
|
|||
|
|
{"上游报告会话已关闭", ErrLiveCallNotFound, true},
|
|||
|
|
{"到达会话时长上限", context.DeadlineExceeded, true},
|
|||
|
|
{"控制权被他人接管", ErrLiveControllerChanged, false},
|
|||
|
|
{"临时读错误", errors.New("unexpected EOF"), false},
|
|||
|
|
{"无错误", nil, false},
|
|||
|
|
}
|
|||
|
|
for _, tc := range cases {
|
|||
|
|
t.Run(tc.name, func(t *testing.T) {
|
|||
|
|
require.Equal(t, tc.want, liveSessionEnded(tc.err))
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestWaitForLiveObserverRetryLeavesExpiryToLoopFinalize 锁定：已过期但控制权仍在
|
|||
|
|
// observer 手上时返回 true，让调用方回到 observeLiveCall 循环顶部的过期分支去
|
|||
|
|
// finalize（写 usage log + 释放租约）。在此处直接返回 false 会让会话静默结束、不留记录。
|
|||
|
|
func TestWaitForLiveObserverRetryLeavesExpiryToLoopFinalize(t *testing.T) {
|
|||
|
|
record := &LiveCallRecord{
|
|||
|
|
CallID: "call_expired",
|
|||
|
|
CallHash: hashLiveCallID("call_expired"),
|
|||
|
|
Controller: LiveControllerObserver,
|
|||
|
|
ExpiresAt: time.Now().Add(-time.Minute),
|
|||
|
|
}
|
|||
|
|
store := &liveTestStore{}
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), record, time.Hour))
|
|||
|
|
svc := &OpenAIGatewayService{cache: store}
|
|||
|
|
|
|||
|
|
require.True(t, svc.waitForLiveObserverRetry(record),
|
|||
|
|
"过期判定必须留给循环顶部，否则不会写 usage log")
|
|||
|
|
|
|||
|
|
// 控制权已被他人接管时仍必须停止重试，避免与新控制者抢同一个 call。
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), &LiveCallRecord{
|
|||
|
|
CallID: record.CallID,
|
|||
|
|
CallHash: record.CallHash,
|
|||
|
|
Controller: LiveControllerProxy,
|
|||
|
|
ExpiresAt: time.Now().Add(time.Hour),
|
|||
|
|
}, time.Hour))
|
|||
|
|
require.False(t, svc.waitForLiveObserverRetry(record))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestWaitForLiveObserverRetryTreatsStoreErrorAsRetryable 锁定：store 报错（Redis
|
|||
|
|
// 抖动）不等于控制权被接管，必须返回 true 交回 observeLiveCall 循环顶部，由它做
|
|||
|
|
// 有限次重试与 ExpiresAt 兜底 finalize；记录确实不存在时才停止重试。
|
|||
|
|
func TestWaitForLiveObserverRetryTreatsStoreErrorAsRetryable(t *testing.T) {
|
|||
|
|
record := &LiveCallRecord{
|
|||
|
|
CallID: "call_flaky_store",
|
|||
|
|
CallHash: hashLiveCallID("call_flaky_store"),
|
|||
|
|
Controller: LiveControllerObserver,
|
|||
|
|
ExpiresAt: time.Now().Add(time.Hour),
|
|||
|
|
}
|
|||
|
|
store := &liveTestStore{getControllerErr: errors.New("redis: connection refused")}
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), record, time.Hour))
|
|||
|
|
svc := &OpenAIGatewayService{cache: store}
|
|||
|
|
|
|||
|
|
require.True(t, svc.waitForLiveObserverRetry(record),
|
|||
|
|
"store 报错必须继续重试，否则 Redis 抖动会让会话静默结束、不留记录")
|
|||
|
|
|
|||
|
|
// 记录已被清理（ErrLiveCallNotFound）不是故障，应停止重试。
|
|||
|
|
require.False(t, (&OpenAIGatewayService{cache: &liveTestStore{}}).waitForLiveObserverRetry(record))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestObserveLiveCallStoreOutageFallsBackToExpiryFinalize 锁定：observer 遇到持续
|
|||
|
|
// store 报错时不能静默退出，必须按 record.ExpiresAt 兜底 finalize（写 usage log +
|
|||
|
|
// 释放租约）。
|
|||
|
|
func TestObserveLiveCallStoreOutageFallsBackToExpiryFinalize(t *testing.T) {
|
|||
|
|
restore := liveObserverStoreRetryInterval
|
|||
|
|
liveObserverStoreRetryInterval = time.Millisecond
|
|||
|
|
t.Cleanup(func() { liveObserverStoreRetryInterval = restore })
|
|||
|
|
|
|||
|
|
cases := []struct {
|
|||
|
|
name string
|
|||
|
|
inject func(*liveTestStore)
|
|||
|
|
}{
|
|||
|
|
{"GetLiveCall 持续报错", func(s *liveTestStore) { s.getCallErr = errors.New("redis: i/o timeout") }},
|
|||
|
|
{"ClaimLiveController 报错", func(s *liveTestStore) { s.claimErr = errors.New("redis: i/o timeout") }},
|
|||
|
|
}
|
|||
|
|
for _, tc := range cases {
|
|||
|
|
t.Run(tc.name, func(t *testing.T) {
|
|||
|
|
record := &LiveCallRecord{
|
|||
|
|
CallID: "call_store_outage",
|
|||
|
|
CallHash: hashLiveCallID("call_store_outage"),
|
|||
|
|
AccountID: 11,
|
|||
|
|
APIKeyID: 22,
|
|||
|
|
UserID: 33,
|
|||
|
|
LeaseID: "lease-1",
|
|||
|
|
Model: "gpt-live-test",
|
|||
|
|
CreatedAt: time.Now().Add(-time.Minute),
|
|||
|
|
ExpiresAt: time.Now().Add(-time.Second), // 已到期：兜底无需等待
|
|||
|
|
Controller: LiveControllerPending,
|
|||
|
|
}
|
|||
|
|
store := &liveTestStore{}
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), record, time.Hour))
|
|||
|
|
tc.inject(store)
|
|||
|
|
concurrencyCache := &liveTestConcurrencyCache{}
|
|||
|
|
usageRepo := &liveTestUsageRepo{}
|
|||
|
|
svc := &OpenAIGatewayService{
|
|||
|
|
cache: store,
|
|||
|
|
concurrencyService: NewConcurrencyService(concurrencyCache),
|
|||
|
|
usageLogRepo: usageRepo,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
svc.observeLiveCall(record)
|
|||
|
|
|
|||
|
|
concurrencyCache.mu.Lock()
|
|||
|
|
require.Equal(t, 1, concurrencyCache.releases, "store 故障时租约释放不能丢")
|
|||
|
|
concurrencyCache.mu.Unlock()
|
|||
|
|
usageRepo.mu.Lock()
|
|||
|
|
require.Len(t, usageRepo.logs, 1, "store 故障时 usage log 不能丢")
|
|||
|
|
require.Equal(t, RequestTypeLive, usageRepo.logs[0].RequestType)
|
|||
|
|
usageRepo.mu.Unlock()
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type liveTestBestEffortUsageRepo struct {
|
|||
|
|
liveTestUsageRepo
|
|||
|
|
bestEffortErr error
|
|||
|
|
bestEffortCalls int
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (r *liveTestBestEffortUsageRepo) CreateBestEffort(_ context.Context, _ *UsageLog) error {
|
|||
|
|
r.mu.Lock()
|
|||
|
|
defer r.mu.Unlock()
|
|||
|
|
r.bestEffortCalls++
|
|||
|
|
return r.bestEffortErr
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// TestFinalizeLiveCallUsageLogFallsBackToSyncCreate 锁定：finalize 是该会话唯一一次
|
|||
|
|
// 落库机会（MarkLiveCallClosed 已标记 first），best-effort 写入失败必须走同步 Create
|
|||
|
|
// 兜底，而不是丢弃错误。
|
|||
|
|
func TestFinalizeLiveCallUsageLogFallsBackToSyncCreate(t *testing.T) {
|
|||
|
|
record := &LiveCallRecord{
|
|||
|
|
CallID: "call_usage_fallback",
|
|||
|
|
CallHash: hashLiveCallID("call_usage_fallback"),
|
|||
|
|
AccountID: 11,
|
|||
|
|
APIKeyID: 22,
|
|||
|
|
UserID: 33,
|
|||
|
|
LeaseID: "lease-1",
|
|||
|
|
Model: "gpt-live-test",
|
|||
|
|
CreatedAt: time.Now().Add(-time.Second),
|
|||
|
|
ExpiresAt: time.Now().Add(time.Hour),
|
|||
|
|
Controller: LiveControllerPending,
|
|||
|
|
}
|
|||
|
|
store := &liveTestStore{}
|
|||
|
|
require.NoError(t, store.SaveLiveCall(context.Background(), record, time.Hour))
|
|||
|
|
usageRepo := &liveTestBestEffortUsageRepo{bestEffortErr: errors.New("usage log queue dropped")}
|
|||
|
|
svc := &OpenAIGatewayService{
|
|||
|
|
cache: store,
|
|||
|
|
concurrencyService: NewConcurrencyService(&liveTestConcurrencyCache{}),
|
|||
|
|
usageLogRepo: usageRepo,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
svc.finalizeLiveCall(record)
|
|||
|
|
|
|||
|
|
usageRepo.mu.Lock()
|
|||
|
|
defer usageRepo.mu.Unlock()
|
|||
|
|
require.Equal(t, 1, usageRepo.bestEffortCalls)
|
|||
|
|
require.Len(t, usageRepo.logs, 1, "best-effort 失败后必须同步兜底落库")
|
|||
|
|
require.Equal(t, record.CallHash, usageRepo.logs[0].RequestID)
|
|||
|
|
}
|