Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/openai_ws_passthrough_turn_pricing_test.go

@@ -0,0 +1,235 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+
+	coderws "github.com/coder/websocket"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+	"github.com/tidwall/gjson"
+)
+
+// startPassthroughHookRecordingServer 与 startPassthroughLifecycleServer 相同，
+// 但把一组会记录调用的 hooks 传给 ingress，用于观察透传路径的 turn 回调。
+func startPassthroughHookRecordingServer(
+	t *testing.T,
+	controlCtx context.Context,
+	svc *OpenAIGatewayService,
+	account *Account,
+	hooks *OpenAIWSIngressHooks,
+) (*httptest.Server, <-chan error) {
+	t.Helper()
+	serverErr := make(chan error, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
+		if err != nil {
+			serverErr <- err
+			return
+		}
+		defer func() { _ = conn.CloseNow() }()
+
+		msgType, firstMessage, err := ReadOpenAIWSClientMessage(
+			controlCtx,
+			conn,
+			3*time.Second,
+			coderws.StatusPolicyViolation,
+			"missing first response.create message",
+		)
+		if err != nil {
+			serverErr <- err
+			return
+		}
+		if msgType != coderws.MessageText {
+			serverErr <- errors.New("first message was not text")
+			return
+		}
+
+		recorder := httptest.NewRecorder()
+		ginCtx, _ := gin.CreateTestContext(recorder)
+		req := r.Clone(controlCtx)
+		req.Header = req.Header.Clone()
+		ginCtx.Request = req
+		serverErr <- svc.ProxyResponsesWebSocketFromClient(controlCtx, ginCtx, conn, account, "sk-test", firstMessage, hooks)
+	}))
+	return server, serverErr
+}
+
+// TestPassthroughIngressReportsTurnStartedBeforeAfterTurnWithoutBeforeTurn 钉死
+// ws_v2 透传 ingress 与 handler 侧 turn 定价的耦合：透传 relay 不触发
+// BeforeTurn，但会在每个 AfterTurn 前通过 TurnStarted 报告同一 turn 的开始时刻。
+//
+// handler 的 recordTurnStart 保存该时刻，AfterTurn 再用 currentOr(turnStart)
+// 作为计费 PricingAt；不触发 BeforeTurn 也意味着透传仍没有 turn 级利润复核。
+func TestPassthroughIngressReportsTurnStartedBeforeAfterTurnWithoutBeforeTurn(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+	controlCtx, cancelControl := context.WithCancelCause(context.Background())
+	defer cancelControl(context.Canceled)
+
+	upstream := newStagedPassthroughConn()
+	upstream.Send(`{"type":"response.completed","response":{"id":"resp_pricing","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
+
+	var hooksMu sync.Mutex
+	beforeTurnCalls := 0
+	expectedTurnStartedAt := time.Date(2026, time.August, 17, 9, 59, 59, 0, time.UTC)
+	type hookEvent struct {
+		name      string
+		turn      int
+		startedAt time.Time
+	}
+	var hookEvents []hookEvent
+	hooks := &OpenAIWSIngressHooks{
+		InitialTurnStartedAt: expectedTurnStartedAt,
+		TurnStarted: func(turn int, startedAt time.Time) {
+			hooksMu.Lock()
+			hookEvents = append(hookEvents, hookEvent{name: "TurnStarted", turn: turn, startedAt: startedAt})
+			hooksMu.Unlock()
+		},
+		BeforeTurn: func(int) error {
+			hooksMu.Lock()
+			beforeTurnCalls++
+			hooksMu.Unlock()
+			return nil
+		},
+		AfterTurn: func(turn int, _ *OpenAIForwardResult, _ error) {
+			hooksMu.Lock()
+			hookEvents = append(hookEvents, hookEvent{name: "AfterTurn", turn: turn})
+			hooksMu.Unlock()
+		},
+	}
+
+	server, serverErr := startPassthroughHookRecordingServer(
+		t,
+		controlCtx,
+		newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream),
+		passthroughLifecycleAccount(),
+		hooks,
+	)
+	defer server.Close()
+	clientConn := dialPassthroughLifecycleClient(t, server)
+	defer func() { _ = clientConn.CloseNow() }()
+
+	event, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
+	require.NoError(t, err)
+	require.Equal(t, "response.completed", gjson.GetBytes(event, "type").String())
+
+	// 等待连接自然结束（inter-turn idle 超时），确保 AfterTurn 已提交。
+	_, _ = readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
+	select {
+	case <-serverErr:
+	case <-time.After(3 * time.Second):
+		t.Fatal("passthrough ingress did not exit")
+	}
+
+	hooksMu.Lock()
+	gotBefore := beforeTurnCalls
+	gotEvents := append([]hookEvent(nil), hookEvents...)
+	hooksMu.Unlock()
+
+	require.Zero(t, gotBefore, "透传 ingress 不应调用 BeforeTurn")
+	require.GreaterOrEqual(t, len(gotEvents), 2, "透传 ingress 应报告 TurnStarted 和 AfterTurn")
+	require.Equal(t, "TurnStarted", gotEvents[0].name)
+	require.Equal(t, expectedTurnStartedAt, gotEvents[0].startedAt, "TurnStarted 必须携带入口冻结的首轮开始时刻")
+	require.Equal(t, "AfterTurn", gotEvents[1].name)
+	require.Equal(t, gotEvents[0].turn, gotEvents[1].turn, "TurnStarted 后应提交同一 turn 的 AfterTurn")
+}
+
+func TestPassthroughIngressFreezesSubsequentTurnBeforeRequestPolicy(t *testing.T) {
+	testPassthroughIngressFreezesSubsequentTurnBeforeRequestPolicy(t, coderws.MessageText)
+}
+
+func TestPassthroughIngressFreezesBinarySubsequentTurnBeforeRequestPolicy(t *testing.T) {
+	testPassthroughIngressFreezesSubsequentTurnBeforeRequestPolicy(t, coderws.MessageBinary)
+}
+
+func testPassthroughIngressFreezesSubsequentTurnBeforeRequestPolicy(t *testing.T, secondMessageType coderws.MessageType) {
+	t.Helper()
+	gin.SetMode(gin.TestMode)
+	controlCtx, cancelControl := context.WithCancelCause(context.Background())
+	defer cancelControl(context.Canceled)
+
+	upstream := newStagedPassthroughConn()
+	upstream.Send(`{"type":"response.completed","response":{"id":"resp_first","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
+
+	type turnStart struct {
+		turn      int
+		startedAt time.Time
+	}
+	turnStarts := make(chan turnStart, 2)
+	beforeRequestEntered := make(chan time.Time, 1)
+	releaseBeforeRequest := make(chan struct{})
+	hooks := &OpenAIWSIngressHooks{
+		InitialTurnStartedAt: time.Now(),
+		TurnStarted: func(turn int, startedAt time.Time) {
+			turnStarts <- turnStart{turn: turn, startedAt: startedAt}
+		},
+		BeforeRequest: func(turn int, _ []byte, _ string) error {
+			if turn == 2 {
+				beforeRequestEntered <- time.Now()
+				<-releaseBeforeRequest
+			}
+			return nil
+		},
+	}
+
+	server, serverErr := startPassthroughHookRecordingServer(
+		t,
+		controlCtx,
+		newPassthroughLifecycleService(passthroughLifecycleConfig(), upstream),
+		passthroughLifecycleAccount(),
+		hooks,
+	)
+	defer server.Close()
+	clientConn := dialPassthroughLifecycleClient(t, server)
+	defer func() { _ = clientConn.CloseNow() }()
+
+	require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, 3*time.Second), "type").String())
+	firstCompleted, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
+	require.NoError(t, err)
+	require.Equal(t, "resp_first", gjson.GetBytes(firstCompleted, "response.id").String())
+	select {
+	case first := <-turnStarts:
+		require.Equal(t, 1, first.turn)
+	case <-time.After(time.Second):
+		t.Fatal("first turn start was not reported")
+	}
+
+	writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
+	err = clientConn.Write(writeCtx, secondMessageType, []byte(`{"type":"response.create","model":"gpt-5.1","previous_response_id":"resp_first"}`))
+	cancelWrite()
+	require.NoError(t, err)
+
+	var policyEnteredAt time.Time
+	select {
+	case policyEnteredAt = <-beforeRequestEntered:
+	case <-time.After(time.Second):
+		t.Fatal("second turn did not enter BeforeRequest")
+	}
+	close(releaseBeforeRequest)
+	require.Equal(t, "response.create", gjson.GetBytes(requirePassthroughUpstreamWrite(t, upstream, 3*time.Second), "type").String())
+	upstream.Send(`{"type":"response.completed","response":{"id":"resp_second","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`)
+	secondCompleted, err := readPassthroughLifecycleFrame(t, clientConn, 3*time.Second)
+	require.NoError(t, err)
+	require.Equal(t, "resp_second", gjson.GetBytes(secondCompleted, "response.id").String())
+
+	select {
+	case second := <-turnStarts:
+		require.Equal(t, 2, second.turn)
+		require.False(t, second.startedAt.After(policyEnteredAt), "第二轮开始时刻必须在 BeforeRequest 策略执行前冻结")
+	case <-time.After(time.Second):
+		t.Fatal("second turn start was not reported")
+	}
+
+	_ = clientConn.CloseNow()
+	cancelControl(context.Canceled)
+	select {
+	case <-serverErr:
+	case <-time.After(3 * time.Second):
+		t.Fatal("passthrough ingress did not exit")
+	}
+}