toki
04879f2b43
요청 종료 계수와 실제 provider 시도 사용량을 분리하고, 직접·pool·retry 경로의 attribution을 보존한다. 관련 계약·스펙과 완료된 task archive 정리도 함께 반영한다.
249 lines
11 KiB
Go
249 lines
11 KiB
Go
package openai
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/prometheus/client_golang/prometheus/testutil"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
"iop/packages/go/config"
|
|
"iop/packages/go/streamgate"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
// TestProviderSwitchEmitsUsagePerActualAttemptAndOneRequest drives the real
|
|
// provider-pool request runtime through a tunnel-to-normalized recovery. The
|
|
// replaced attempt and the committed attempt report distinct provider usage;
|
|
// both must be retained while the HTTP request terminal advances only once.
|
|
func TestProviderSwitchEmitsUsagePerActualAttemptAndOneRequest(t *testing.T) {
|
|
const (
|
|
edgeID = "edge-provider-switch-usage"
|
|
routeModel = "client-model"
|
|
)
|
|
initialFrames := bufferedTunnelFrames(
|
|
&iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START, StatusCode: http.StatusOK, Headers: map[string]string{"Content-Type": "application/json"}},
|
|
&iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(`{"choices":[{"message":{"content":"discarded"}}],"usage":{"prompt_tokens":11,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":1},"completion_tokens_details":{"reasoning_tokens":2}}}`)},
|
|
&iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true},
|
|
)
|
|
service := newScriptedPoolRunService(
|
|
scriptedPoolAttempt{
|
|
path: string(edgeservice.ProviderPoolPathTunnel), runID: "usage-attempt-a",
|
|
provider: "provider-a", target: "served-a", frames: initialFrames,
|
|
},
|
|
scriptedPoolAttempt{
|
|
path: string(edgeservice.ProviderPoolPathNormalized), runID: "usage-attempt-b",
|
|
provider: "provider-b", target: "served-b",
|
|
runEvents: bufferedRunEvents(
|
|
&iop.RunEvent{Type: "delta", Delta: "accepted"},
|
|
&iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 17, OutputTokens: 5, ReasoningTokens: 4, CachedInputTokens: 6}},
|
|
),
|
|
},
|
|
)
|
|
rawBody := []byte(`{"model":"client-model","messages":[{"role":"user","content":"hi"}]}`)
|
|
srv, dc := buildStreamGatePoolChatFixture(t, service, rawBody, 3)
|
|
srv.SetEdgeID(edgeID)
|
|
dc.usage = &openAIUsageRecorder{
|
|
request: usageRequestLabels{edgeID: edgeID, routeModel: routeModel, endpoint: usageEndpointChatCompletions},
|
|
attempts: make(map[string]usageDispatchBinding),
|
|
}
|
|
|
|
labelsA := usageAttemptLabels{
|
|
usageRequestLabels: dc.usage.request, usageAttribution: config.UsageAttributionProvider,
|
|
providerID: "provider-a", servedModel: "served-a", responseMode: responseModePassthrough,
|
|
}
|
|
labelsB := usageAttemptLabels{
|
|
usageRequestLabels: dc.usage.request, usageAttribution: config.UsageAttributionProvider,
|
|
providerID: "provider-b", servedModel: "served-b", responseMode: responseModeNormalized,
|
|
}
|
|
aInputBefore := requestTokenValue(t, labelsA, tokenTypeInput)
|
|
aOutputBefore := requestTokenValue(t, labelsA, tokenTypeOutput)
|
|
aReasoningBefore := requestTokenValue(t, labelsA, tokenTypeReasoning)
|
|
aCachedBefore := requestTokenValue(t, labelsA, tokenTypeCachedInput)
|
|
bInputBefore := requestTokenValue(t, labelsB, tokenTypeInput)
|
|
bOutputBefore := requestTokenValue(t, labelsB, tokenTypeOutput)
|
|
bReasoningBefore := requestTokenValue(t, labelsB, tokenTypeReasoning)
|
|
bCachedBefore := requestTokenValue(t, labelsB, tokenTypeCachedInput)
|
|
requestBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "", "", "", routeModel, usageEndpointChatCompletions,
|
|
responseModeNormalized, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
w, sink, runErr := runPoolChatStreamGate(t, srv, dc, newInjectedViolationRegistration(t, dc, "test.usage_switch", 1))
|
|
if runErr != nil {
|
|
t.Fatalf("runtime run: %v", runErr)
|
|
}
|
|
committed, success := sink.terminalStatus()
|
|
if !committed || !success || !strings.Contains(w.body.String(), "accepted") {
|
|
t.Fatalf("replacement response did not commit successfully: committed=%v success=%v body=%q", committed, success, w.body.String())
|
|
}
|
|
dc.finishUsageRequest(usageStatusSuccess, responseModeNormalized)
|
|
|
|
for name, gotWant := range map[string][2]float64{
|
|
"provider-a input": {requestTokenValue(t, labelsA, tokenTypeInput) - aInputBefore, 11},
|
|
"provider-a output": {requestTokenValue(t, labelsA, tokenTypeOutput) - aOutputBefore, 3},
|
|
"provider-a reasoning": {requestTokenValue(t, labelsA, tokenTypeReasoning) - aReasoningBefore, 2},
|
|
"provider-a cached": {requestTokenValue(t, labelsA, tokenTypeCachedInput) - aCachedBefore, 1},
|
|
"provider-b input": {requestTokenValue(t, labelsB, tokenTypeInput) - bInputBefore, 17},
|
|
"provider-b output": {requestTokenValue(t, labelsB, tokenTypeOutput) - bOutputBefore, 5},
|
|
"provider-b reasoning": {requestTokenValue(t, labelsB, tokenTypeReasoning) - bReasoningBefore, 4},
|
|
"provider-b cached": {requestTokenValue(t, labelsB, tokenTypeCachedInput) - bCachedBefore, 6},
|
|
} {
|
|
if gotWant[0] != gotWant[1] {
|
|
t.Fatalf("%s tokens = %v, want %v", name, gotWant[0], gotWant[1])
|
|
}
|
|
}
|
|
requestAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "", "", "", routeModel, usageEndpointChatCompletions,
|
|
responseModeNormalized, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if got := requestAfter - requestBefore; got != 1 {
|
|
t.Fatalf("request terminal count = %v, want 1", got)
|
|
}
|
|
}
|
|
|
|
func TestNonStreamingTunnelAttemptFinalizesBodyUsage(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
endpoint string
|
|
status int
|
|
contentType string
|
|
body string
|
|
endKind iop.ProviderTunnelFrameKind
|
|
wantInput int
|
|
wantOutput int
|
|
wantReasoning int
|
|
wantCached int
|
|
wantProviderReported bool
|
|
wantHolderSet bool
|
|
}{
|
|
{
|
|
name: "chat completions success body",
|
|
endpoint: openAIRebuildEndpointChat,
|
|
status: http.StatusOK,
|
|
contentType: "application/json",
|
|
body: `{"choices":[{"message":{"content":"hello"}}],"usage":{"prompt_tokens":9,"completion_tokens":4,"prompt_tokens_details":{"cached_tokens":1},"completion_tokens_details":{"reasoning_tokens":2}}}`,
|
|
endKind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
|
|
wantInput: 9,
|
|
wantOutput: 4,
|
|
wantReasoning: 2,
|
|
wantCached: 1,
|
|
wantProviderReported: true,
|
|
wantHolderSet: true,
|
|
},
|
|
{
|
|
name: "responses API success body",
|
|
endpoint: openAIRebuildEndpointResponses,
|
|
status: http.StatusOK,
|
|
contentType: "application/json",
|
|
body: `{"output_text":"hello","usage":{"input_tokens":14,"output_tokens":7,"input_tokens_details":{"cached_tokens":3},"output_tokens_details":{"reasoning_tokens":5}}}`,
|
|
endKind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
|
|
wantInput: 14,
|
|
wantOutput: 7,
|
|
wantReasoning: 5,
|
|
wantCached: 3,
|
|
wantProviderReported: true,
|
|
wantHolderSet: true,
|
|
},
|
|
{
|
|
name: "chat completions provider error boundary",
|
|
endpoint: openAIRebuildEndpointChat,
|
|
status: http.StatusOK,
|
|
contentType: "application/json",
|
|
body: `{"error":{"message":"internal error"},"usage":{"prompt_tokens":5,"completion_tokens":2,"prompt_tokens_details":{"cached_tokens":0},"completion_tokens_details":{"reasoning_tokens":1}}}`,
|
|
endKind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
|
|
wantInput: 5,
|
|
wantOutput: 2,
|
|
wantReasoning: 1,
|
|
wantCached: 0,
|
|
wantProviderReported: true,
|
|
wantHolderSet: false,
|
|
},
|
|
{
|
|
name: "responses API provider error boundary",
|
|
endpoint: openAIRebuildEndpointResponses,
|
|
status: http.StatusOK,
|
|
contentType: "application/json",
|
|
body: `{"error":{"message":"internal error"},"usage":{"input_tokens":8,"output_tokens":3,"input_tokens_details":{"cached_tokens":2},"output_tokens_details":{"reasoning_tokens":0}}}`,
|
|
endKind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR,
|
|
wantInput: 8,
|
|
wantOutput: 3,
|
|
wantReasoning: 0,
|
|
wantCached: 2,
|
|
wantProviderReported: true,
|
|
wantHolderSet: false,
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
frames := bufferedTunnelFrames(
|
|
&iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: int32(tc.status),
|
|
Headers: map[string]string{"Content-Type": tc.contentType},
|
|
},
|
|
&iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
|
|
Body: []byte(tc.body),
|
|
},
|
|
&iop.ProviderTunnelFrame{
|
|
Kind: tc.endKind,
|
|
End: tc.endKind == iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END,
|
|
},
|
|
)
|
|
|
|
stream := edgeservice.ProviderTunnelStream{Frames: frames}
|
|
assembler := &providerChatAssembler{streaming: false}
|
|
rewriter := newProviderModelRewriter(false, "request-model")
|
|
state := &openAITunnelCodecState{}
|
|
state.bindEndpoint(tc.endpoint)
|
|
src := newOpenAITunnelEndpointEventSource(stream, 5*time.Second, rewriter, assembler, tc.endpoint, state)
|
|
|
|
holder := &openAIStreamGateUsageHolder{}
|
|
attempt := &openAIAttemptUsage{}
|
|
tracking := &openAIStreamGateUsageTrackingTunnelSource{
|
|
openAITunnelEventSource: src,
|
|
usage: holder,
|
|
attempt: attempt,
|
|
}
|
|
|
|
ctx := context.Background()
|
|
for {
|
|
ev, err := tracking.NextEvent(ctx)
|
|
if err != nil || ev.Kind() == streamgate.EventKindTerminal || ev.Kind() == streamgate.EventKindProviderError {
|
|
break
|
|
}
|
|
}
|
|
|
|
attObs := attempt.snapshot()
|
|
if attObs.providerReported != tc.wantProviderReported {
|
|
t.Fatalf("attempt providerReported = %v, want %v", attObs.providerReported, tc.wantProviderReported)
|
|
}
|
|
if attObs.inputTokens != tc.wantInput || attObs.outputTokens != tc.wantOutput ||
|
|
attObs.reasoningTokens != tc.wantReasoning || attObs.cachedInputTokens != tc.wantCached {
|
|
t.Fatalf("attempt usage = %+v, want input=%d output=%d reasoning=%d cached=%d",
|
|
attObs, tc.wantInput, tc.wantOutput, tc.wantReasoning, tc.wantCached)
|
|
}
|
|
|
|
hObs := holder.get()
|
|
if tc.wantHolderSet {
|
|
if !hObs.providerReported {
|
|
t.Fatalf("holder providerReported = false, want true")
|
|
}
|
|
if hObs.inputTokens != tc.wantInput || hObs.outputTokens != tc.wantOutput ||
|
|
hObs.reasoningTokens != tc.wantReasoning || hObs.cachedInputTokens != tc.wantCached {
|
|
t.Fatalf("holder usage = %+v, want input=%d output=%d reasoning=%d cached=%d",
|
|
hObs, tc.wantInput, tc.wantOutput, tc.wantReasoning, tc.wantCached)
|
|
}
|
|
} else {
|
|
if hObs.providerReported {
|
|
t.Fatalf("holder providerReported = true, want false for non-terminal error")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|