1112 lines
48 KiB
Go
1112 lines
48 KiB
Go
package openai
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/prometheus/client_golang/prometheus/testutil"
|
|
|
|
"iop/packages/go/config"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
// principalTokenCfg builds an EdgeOpenAIConf whose bearer token maps to a fixed
|
|
// principal so usage-metric label values are populated.
|
|
func principalTokenCfg(rawToken, adapter string) config.EdgeOpenAIConf {
|
|
return config.EdgeOpenAIConf{
|
|
Adapter: adapter,
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
|
|
},
|
|
}
|
|
}
|
|
|
|
func requestTokenValue(t *testing.T, l usageLabels, tokenType string) float64 {
|
|
t.Helper()
|
|
return testutil.ToFloat64(openAIUsageTokensTotal.WithLabelValues(
|
|
l.edgeID, l.principalRef, l.principalAlias, l.tokenRef,
|
|
l.modelGroup, l.endpoint, l.responseMode, tokenType,
|
|
))
|
|
}
|
|
|
|
// TestOpenAIUsageMetricsLabelsExcludeSecretsAndHighCardinality asserts the label
|
|
// allowlists never carry raw tokens, prompt/response text, or per-request
|
|
// high-cardinality identifiers (SDD S08).
|
|
func TestOpenAIUsageMetricsLabelsExcludeSecretsAndHighCardinality(t *testing.T) {
|
|
// Substrings that would indicate a per-request identifier, a raw secret, or
|
|
// raw payload text. response_mode / model_group are low-cardinality enums and
|
|
// are intentionally not caught here.
|
|
forbidden := []string{
|
|
"request_id", "session_id", "run_id", "token_hash", "bearer",
|
|
"authorization", "api_key", "prompt", "secret", "raw",
|
|
}
|
|
// token_ref is an allowlisted stable alias, not a raw token; guard against a
|
|
// bare "token" label that could carry the secret itself.
|
|
for _, set := range [][]string{usageRequestLabelNames, usageTokenLabelNames, usageReasoningLabelNames} {
|
|
for _, name := range set {
|
|
if name == "token" {
|
|
t.Fatalf("label %q may carry a raw token value; use token_ref", name)
|
|
}
|
|
for _, bad := range forbidden {
|
|
if strings.Contains(name, bad) {
|
|
t.Fatalf("label %q contains forbidden fragment %q", name, bad)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// The allowlist must still carry the intended identity/attribution labels.
|
|
for _, want := range []string{"edge_id", "principal_ref", "principal_alias", "token_ref", "model_group", "endpoint", "response_mode"} {
|
|
if !containsString(usageRequestLabelNames, want) {
|
|
t.Fatalf("request labels missing %q", want)
|
|
}
|
|
}
|
|
if !containsString(usageRequestLabelNames, "usage_source") || !containsString(usageRequestLabelNames, "status") {
|
|
t.Fatalf("request labels must include status and usage_source: %v", usageRequestLabelNames)
|
|
}
|
|
if !containsString(usageTokenLabelNames, "token_type") {
|
|
t.Fatalf("token labels must include token_type: %v", usageTokenLabelNames)
|
|
}
|
|
}
|
|
|
|
// TestOpenAIUsageMetricsCountTokenTypes drives a normalized chat completion whose
|
|
// run reports input/output/reasoning/cached usage and asserts each token type is
|
|
// counted under a provider_reported request (SDD S06).
|
|
func TestOpenAIUsageMetricsCountTokenTypes(t *testing.T) {
|
|
const rawToken = "sk-count-raw-token"
|
|
const edgeID = "edge-count-token-types"
|
|
const model = "count-model"
|
|
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 2)}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "hi"}
|
|
fake.events <- &iop.RunEvent{Type: "complete", Usage: &iop.Usage{
|
|
InputTokens: 10, OutputTokens: 5, ReasoningTokens: 3, CachedInputTokens: 2,
|
|
}}
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: "normalized",
|
|
}
|
|
before := map[string]float64{
|
|
tokenTypeInput: requestTokenValue(t, labels, tokenTypeInput),
|
|
tokenTypeOutput: requestTokenValue(t, labels, tokenTypeOutput),
|
|
tokenTypeReasoning: requestTokenValue(t, labels, tokenTypeReasoning),
|
|
tokenTypeCachedInput: requestTokenValue(t, labels, tokenTypeCachedInput),
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, "normalized", usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"count-model",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 10, tokenTypeOutput: 5, tokenTypeReasoning: 3, tokenTypeCachedInput: 2,
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType) - before[tokenType]
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got delta %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, "normalized", usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
}
|
|
|
|
// TestOpenAIReasoningObservedDoesNotEstimateTokens verifies that reasoning text
|
|
// observed without a provider-reported reasoning token count increments only the
|
|
// auxiliary reasoning counters and never the reasoning token counter (SDD S07).
|
|
func TestOpenAIReasoningObservedDoesNotEstimateTokens(t *testing.T) {
|
|
const rawToken = "sk-reasoning-raw-token"
|
|
const edgeID = "edge-reasoning-observed"
|
|
const model = "reasoning-model"
|
|
const reasoning = "abcde" // 5 chars
|
|
|
|
fake := &fakeRunService{events: make(chan *iop.RunEvent, 3)}
|
|
fake.events <- &iop.RunEvent{Type: "reasoning_delta", Delta: reasoning}
|
|
fake.events <- &iop.RunEvent{Type: "delta", Delta: "answer"}
|
|
fake.events <- &iop.RunEvent{Type: "complete", Usage: &iop.Usage{InputTokens: 8, OutputTokens: 4}}
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: "normalized",
|
|
}
|
|
reasoningBefore := requestTokenValue(t, labels, tokenTypeReasoning)
|
|
observedBefore := testutil.ToFloat64(openAIReasoningObservedTotal.WithLabelValues(labels.reasoningValues()...))
|
|
charsBefore := testutil.ToFloat64(openAIReasoningCharsTotal.WithLabelValues(labels.reasoningValues()...))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"reasoning-model",
|
|
"messages":[{"role":"user","content":"hi"}]
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
if got := requestTokenValue(t, labels, tokenTypeReasoning) - reasoningBefore; got != 0 {
|
|
t.Fatalf("reasoning token counter must not move when provider did not report reasoning tokens: got delta %v", got)
|
|
}
|
|
if got := testutil.ToFloat64(openAIReasoningObservedTotal.WithLabelValues(labels.reasoningValues()...)) - observedBefore; got != 1 {
|
|
t.Fatalf("reasoning_observed_total: got delta %v, want 1", got)
|
|
}
|
|
if got := testutil.ToFloat64(openAIReasoningCharsTotal.WithLabelValues(labels.reasoningValues()...)) - charsBefore; got != float64(len(reasoning)) {
|
|
t.Fatalf("reasoning_chars_total: got delta %v, want %d", got, len(reasoning))
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelPassthroughPreservesBodyAndObservesUsage verifies that a pure
|
|
// passthrough provider response is relayed byte-for-byte while its usage is
|
|
// observed into metrics (SDD S05/S06).
|
|
func TestProviderTunnelPassthroughPreservesBodyAndObservesUsage(t *testing.T) {
|
|
const rawToken = "sk-passthrough-raw-token"
|
|
const edgeID = "edge-passthrough-usage"
|
|
const model = "ornith:35b"
|
|
providerBody := `{"choices":[{"message":{"role":"assistant","content":"hi there"}}],"usage":{"prompt_tokens":12,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":3},"completion_tokens_details":{"reasoning_tokens":4}}}`
|
|
|
|
fake := &fakeRunService{tunnelFrames: staticProviderTunnelFrames(providerBody)}
|
|
srv := NewServer(config.EdgeOpenAIConf{
|
|
PrincipalTokens: []config.OpenAIPrincipalTokenConf{
|
|
{TokenRef: "iop-tok-alice", TokenHashSHA256: sha256Hex(rawToken), PrincipalRef: "user:alice", PrincipalAlias: "alice"},
|
|
},
|
|
}, fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-vllm": "Ornith-1.0-35B"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthrough,
|
|
}
|
|
before := map[string]float64{
|
|
tokenTypeInput: requestTokenValue(t, labels, tokenTypeInput),
|
|
tokenTypeOutput: requestTokenValue(t, labels, tokenTypeOutput),
|
|
tokenTypeReasoning: requestTokenValue(t, labels, tokenTypeReasoning),
|
|
tokenTypeCachedInput: requestTokenValue(t, labels, tokenTypeCachedInput),
|
|
}
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"ornith:35b",
|
|
"messages":[{"role":"user","content":"hello"}]
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
if w.Body.String() != providerBody {
|
|
t.Fatalf("passthrough body must be preserved byte-for-byte.\n got: %s\nwant: %s", w.Body.String(), providerBody)
|
|
}
|
|
if len(fake.tunnelReqsSnapshot()) != 1 {
|
|
t.Fatalf("expected exactly one provider tunnel dispatch, got %d", len(fake.tunnelReqsSnapshot()))
|
|
}
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 12, tokenTypeOutput: 7, tokenTypeReasoning: 4, tokenTypeCachedInput: 3,
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType) - before[tokenType]
|
|
if got != want {
|
|
t.Fatalf("passthrough token_type %s: got delta %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestOpenAIScopeExcludesA2AAndNonOpenAISurfaces documents that usage metering is
|
|
// scoped to the OpenAI-compatible input surface only (SDD S04). The endpoint
|
|
// label allowlist carries the OpenAI routes and no A2A/CLI endpoint value, and
|
|
// this plan does not modify the A2A surface.
|
|
func TestOpenAIScopeExcludesA2AAndNonOpenAISurfaces(t *testing.T) {
|
|
endpoints := map[string]struct{}{
|
|
usageEndpointChatCompletions: {},
|
|
usageEndpointResponses: {},
|
|
}
|
|
for _, forbidden := range []string{"a2a", "cli", "ollama", "message/send", "tasks"} {
|
|
if _, ok := endpoints[forbidden]; ok {
|
|
t.Fatalf("usage endpoint label must not include non-OpenAI surface %q", forbidden)
|
|
}
|
|
}
|
|
if !containsString(usageTokenLabelNames, "endpoint") {
|
|
t.Fatalf("token metrics must carry an endpoint label to keep OpenAI scope explicit")
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandStreamingEmitsUsageMetrics verifies that a
|
|
// passthrough+sideband streaming response increments the request counter
|
|
// with status=success and usage_source=provider_reported (REVIEW_USAGE_METRIC-1).
|
|
func TestProviderTunnelSidebandStreamingEmitsUsageMetrics(t *testing.T) {
|
|
const rawToken = "sk-sideband-stream-token"
|
|
const edgeID = "edge-sideband-stream-metrics"
|
|
const model = "pool-model"
|
|
providerBody := `{"choices":[{"delta":{"content":"hi"}}]}
|
|
|
|
data: [DONE]
|
|
|
|
`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "text/event-stream"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE,
|
|
Usage: &iop.Usage{InputTokens: 5, OutputTokens: 3},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandNonStreamingEmitsUsageMetrics verifies that a
|
|
// passthrough+sideband non-streaming response increments the request counter
|
|
// with status=success and usage_source=provider_reported (REVIEW_USAGE_METRIC-1).
|
|
func TestProviderTunnelSidebandNonStreamingEmitsUsageMetrics(t *testing.T) {
|
|
const rawToken = "sk-sideband-nonstream-token"
|
|
const edgeID = "edge-sideband-nonstream-metrics"
|
|
const model = "pool-model"
|
|
providerBody := `{"id":"cmpl-1","choices":[{"message":{"role":"assistant","content":"hi"}}]}`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "application/json"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE,
|
|
Usage: &iop.Usage{InputTokens: 7, OutputTokens: 4},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthroughSideband,
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
// Verify token counters also incremented.
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 7,
|
|
tokenTypeOutput: 4,
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType)
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandStreamCallerCancelEmitsCancelMetric verifies that a
|
|
// passthrough+sideband streaming caller disconnect (pre-cancelled context) emits
|
|
// the cancel metric (REVIEW_USAGE_METRIC-1).
|
|
func TestProviderTunnelSidebandStreamCallerCancelEmitsCancelMetric(t *testing.T) {
|
|
const rawToken = "sk-sideband-cancel-token"
|
|
const edgeID = "edge-sideband-cancel"
|
|
const model = "pool-model"
|
|
|
|
// Create a cancelled context so the sideband stream detects caller disconnect.
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
// Frame channel never closed to simulate an in-flight tunnel.
|
|
frames := make(chan *iop.ProviderTunnelFrame)
|
|
|
|
fake := &fakeRunService{tunnelFrames: frames, tunnelWaitTimeout: 3 * time.Second}
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusCancel, usageSourceUnavailable,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
req = req.WithContext(ctx)
|
|
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusCancel, usageSourceUnavailable,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total cancel: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
}
|
|
|
|
// TestChatCompletionsDispatchFailureEmitsErrorMetric verifies that a tunnel
|
|
// dispatch failure increments the request counter with status=error
|
|
// (REVIEW_USAGE_METRIC-2).
|
|
func TestChatCompletionsDispatchFailureEmitsErrorMetric(t *testing.T) {
|
|
const rawToken = "sk-dispatch-fail-token"
|
|
const edgeID = "edge-dispatch-fail"
|
|
const model = "pool-model"
|
|
|
|
fake := &fakeRunService{tunnelErr: fmt.Errorf("node unavailable")}
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthrough, usageStatusError, usageSourceUnavailable,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"messages":[{"role":"user","content":"hello"}]
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusBadGateway {
|
|
t.Fatalf("expected 502, got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
if len(fake.tunnelReqsSnapshot()) == 0 {
|
|
t.Fatal("expected tunnel dispatch attempt, got none")
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthrough, usageStatusError, usageSourceUnavailable,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total error: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
}
|
|
|
|
// TestToolValidationRetryDispatchFailureEmitsErrorMetric verifies that a
|
|
// non-stream Chat tool-validation retry SubmitRun failure emits the error
|
|
// request metric exactly once and returns tool_validation_retry_error (REVIEW_USAGE_METRIC_RETRY-1).
|
|
func TestToolValidationRetryDispatchFailureEmitsErrorMetric(t *testing.T) {
|
|
const rawToken = "sk-retry-dispatch-fail-token"
|
|
const edgeID = "edge-retry-dispatch-fail"
|
|
const model = "retry-model"
|
|
|
|
fake := &fakeRunService{
|
|
eventRuns: []chan *iop.RunEvent{
|
|
bufferedRunEvents(&iop.RunEvent{
|
|
Type: "complete",
|
|
Metadata: map[string]string{
|
|
"finish_reason": "tool_calls",
|
|
runtimeMetadataOpenAIToolCalls: `[{"id":"call_bad","type":"function","function":{"name":"cmd","arguments":"{\"unexpected\":true}"}}]`,
|
|
},
|
|
}),
|
|
},
|
|
runIDs: []string{"run-bad"},
|
|
submitErrAfter: 2,
|
|
}
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, "normalized", usageStatusError, usageSourceUnavailable,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"retry-model",
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"tools":[{"type":"function","function":{"name":"cmd","parameters":{"type":"object","properties":{"commands":{"type":"array","items":{"type":"string"}}},"required":["commands"]}}}]}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if !strings.Contains(w.Body.String(), `"type":"tool_validation_retry_error"`) {
|
|
t.Fatalf("expected tool_validation_retry_error, got: %s", w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, "normalized", usageStatusError, usageSourceUnavailable,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total error: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
}
|
|
|
|
// TestBufferedStreamToolValidationRetryDispatchFailureEmitsErrorMetric verifies
|
|
// that a buffered stream Chat tool-validation retry SubmitRun failure emits
|
|
// the error request metric exactly once (REVIEW_USAGE_METRIC_RETRY-1).
|
|
func TestBufferedStreamToolValidationRetryDispatchFailureEmitsErrorMetric(t *testing.T) {
|
|
const rawToken = "sk-buf-retry-dispatch-fail-token"
|
|
const edgeID = "edge-buf-retry-dispatch"
|
|
const model = "buf-retry-model"
|
|
|
|
fake := &fakeRunService{
|
|
eventRuns: []chan *iop.RunEvent{
|
|
bufferedRunEvents(&iop.RunEvent{
|
|
Type: "complete",
|
|
Metadata: map[string]string{
|
|
"finish_reason": "tool_calls",
|
|
runtimeMetadataOpenAIToolCalls: `[{"id":"call_bad","type":"function","function":{"name":"cmd","arguments":"{\"unexpected\":true}"}}]`,
|
|
},
|
|
}),
|
|
},
|
|
runIDs: []string{"run-bad"},
|
|
submitErrAfter: 2,
|
|
}
|
|
srv := NewServer(principalTokenCfg(rawToken, ""), fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, "normalized", usageStatusError, usageSourceUnavailable,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"buf-retry-model",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"tools":[{"type":"function","function":{"name":"cmd","parameters":{"type":"object","properties":{"commands":{"type":"array","items":{"type":"string"}}},"required":["commands"]}}}]}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
body := w.Body.String()
|
|
if !strings.Contains(body, `tool_validation_retry_error`) {
|
|
t.Fatalf("expected tool_validation_retry_error in stream, got: %s", body)
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, "normalized", usageStatusError, usageSourceUnavailable,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total error: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandStreamBodyWriteFailureEmitsCancelMetric verifies
|
|
// that a passthrough+sideband streaming body write failure to the caller
|
|
// propagates cancel upstream and emits status=cancel (REVIEW_USAGE_METRIC_RETRY-2).
|
|
func TestProviderTunnelSidebandStreamBodyWriteFailureEmitsCancelMetric(t *testing.T) {
|
|
const rawToken = "sk-sideband-bodyfail-cancel"
|
|
const edgeID = "edge-sideband-bodyfail-cancel"
|
|
const model = "pool-bodyfail"
|
|
|
|
// Use a tunnel that emits a single body frame, then the response writer
|
|
// will fail on the second body frame write.
|
|
frames := make(chan *iop.ProviderTunnelFrame, 3)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: http.StatusOK,
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
|
|
Body: []byte(`{"ok":true}`),
|
|
}
|
|
// Second body frame triggers write failure.
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY,
|
|
Body: []byte(`more`),
|
|
}
|
|
|
|
// Use a writer that fails on the second write.
|
|
fake := &fakeRunService{tunnelFrames: frames, tunnelWaitTimeout: 3 * time.Second}
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), fake, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served"}}})
|
|
|
|
cancelBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusCancel, usageSourceUnavailable,
|
|
))
|
|
errorBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusError, usageSourceUnavailable,
|
|
))
|
|
|
|
// Wrap the response writer so the second Write call fails.
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-bodyfail",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
|
|
w := &failingWriteRecorder{
|
|
ResponseRecorder: httptest.NewRecorder(),
|
|
failAfter: 2,
|
|
}
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
cancelAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusCancel, usageSourceUnavailable,
|
|
))
|
|
errorAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusError, usageSourceUnavailable,
|
|
))
|
|
|
|
if cancelAfter-cancelBefore != 1 {
|
|
t.Fatalf("requests_total cancel: got delta %v, want 1", cancelAfter-cancelBefore)
|
|
}
|
|
if errorAfter-errorBefore != 0 {
|
|
t.Fatalf("requests_total error: got delta %v, want 0", errorAfter-errorBefore)
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandStreamingEmitsUsageMetricsBodyOnly verifies that a
|
|
// passthrough+sideband streaming response that contains usage ONLY in the body
|
|
// (no separate USAGE frame) still emits metrics with usage_source=provider_reported.
|
|
// This is a regression test for REVIEW_USAGE_METRIC-1: body-parsed usage must not
|
|
// be ignored when no proto USAGE frame is present.
|
|
func TestProviderTunnelSidebandStreamingEmitsUsageMetricsBodyOnly(t *testing.T) {
|
|
const rawToken = "sk-body-only-stream-token"
|
|
const edgeID = "edge-body-only-stream-metrics"
|
|
const model = "pool-model"
|
|
// BODY is SSE-formatted (matching actual provider tunnel stream): delta chunk
|
|
// followed by a usage chunk on the next line.
|
|
providerBody := `data: {"choices":[{"delta":{"content":"hi"}}]}
|
|
|
|
data: {"usage":{"prompt_tokens":10,"completion_tokens":5,"prompt_tokens_details":{"cached_tokens":2},"completion_tokens_details":{"reasoning_tokens":3},"total_tokens":20}}
|
|
|
|
data: [DONE]
|
|
|
|
`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "text/event-stream"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
// NOTE: No USAGE frame — usage is only in the body JSON.
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthroughSideband,
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
// Verify token counters were populated from body-only usage including
|
|
// reasoning and cached_input (REVIEW_REVIEW_REVIEW_USAGE_METRIC_RETRY).
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 10,
|
|
tokenTypeOutput: 5,
|
|
tokenTypeReasoning: 3,
|
|
tokenTypeCachedInput: 2,
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType)
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandNonStreamingEmitsUsageMetricsBodyOnly verifies that a
|
|
// passthrough+sideband non-streaming response that contains usage ONLY in the body
|
|
// (no separate USAGE frame) still emits metrics with usage_source=provider_reported.
|
|
func TestProviderTunnelSidebandNonStreamingEmitsUsageMetricsBodyOnly(t *testing.T) {
|
|
const rawToken = "sk-body-only-nonstream-token"
|
|
const edgeID = "edge-body-only-nonstream-metrics"
|
|
const model = "pool-model"
|
|
// BODY contains OpenAI-compatible usage JSON including detail objects;
|
|
// no separate USAGE frame.
|
|
providerBody := `{"id":"cmpl-1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":12,"completion_tokens":8,"prompt_tokens_details":{"cached_tokens":3},"completion_tokens_details":{"reasoning_tokens":4},"total_tokens":27}}`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "application/json"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
// NOTE: No USAGE frame — usage is only in the body JSON.
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthroughSideband,
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
// Verify token counters were populated from body-only usage including
|
|
// reasoning and cached_input (REVIEW_REVIEW_REVIEW_USAGE_METRIC_RETRY).
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 12,
|
|
tokenTypeOutput: 8,
|
|
tokenTypeReasoning: 4,
|
|
tokenTypeCachedInput: 3,
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType)
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandStreamProtoOnlyReasoningCachedInput verifies that when
|
|
// only a proto USAGE frame is present (no body usage), the metric path preserves
|
|
// reasoning and cached_input from the proto (REVIEW_REVIEW_REVIEW_USAGE_METRIC_RETRY).
|
|
func TestProviderTunnelSidebandStreamProtoOnlyReasoningCachedInput(t *testing.T) {
|
|
const rawToken = "sk-proto-only-stream-token"
|
|
const edgeID = "edge-proto-only-stream"
|
|
const model = "pool-model"
|
|
// BODY contains no usage object.
|
|
providerBody := `data: {"choices":[{"delta":{"content":"hi"}}]}
|
|
|
|
data: [DONE]
|
|
|
|
`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "text/event-stream"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
// Proto USAGE frame carries full token breakdown including reasoning/cached_input.
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE,
|
|
Usage: &iop.Usage{InputTokens: 10, OutputTokens: 5, ReasoningTokens: 3, CachedInputTokens: 2},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthroughSideband,
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
// All four token types must be present from proto-only.
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 10,
|
|
tokenTypeOutput: 5,
|
|
tokenTypeReasoning: 3,
|
|
tokenTypeCachedInput: 2,
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType)
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandResponseProtoOnlyReasoningCachedInput verifies that
|
|
// for non-streaming sideband, proto-only USAGE frame preserves reasoning and
|
|
// cached_input in metrics (REVIEW_REVIEW_REVIEW_USAGE_METRIC_RETRY).
|
|
func TestProviderTunnelSidebandResponseProtoOnlyReasoningCachedInput(t *testing.T) {
|
|
const rawToken = "sk-proto-only-nonstream-token"
|
|
const edgeID = "edge-proto-only-nonstream"
|
|
const model = "pool-model"
|
|
providerBody := `{"id":"cmpl-1","choices":[{"message":{"role":"assistant","content":"hi"}}]}`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "application/json"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE,
|
|
Usage: &iop.Usage{InputTokens: 7, OutputTokens: 4, ReasoningTokens: 2, CachedInputTokens: 1},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthroughSideband,
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 7,
|
|
tokenTypeOutput: 4,
|
|
tokenTypeReasoning: 2,
|
|
tokenTypeCachedInput: 1,
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType)
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandStreamBodyAndProtoNoDoubleCount verifies that when
|
|
// both body usage and proto USAGE frame are present, the merge rule prevents
|
|
// double-counting input/output while preserving proto-only reasoning/cached_input
|
|
// (REVIEW_REVIEW_REVIEW_USAGE_METRIC_RETRY).
|
|
func TestProviderTunnelSidebandStreamBodyAndProtoNoDoubleCount(t *testing.T) {
|
|
const rawToken = "sk-body-and-proto-stream-token"
|
|
const edgeID = "edge-body-and-proto-stream"
|
|
const model = "pool-model"
|
|
// Body reports input/output but no reasoning/cached_input details.
|
|
providerBody := `data: {"choices":[{"delta":{"content":"hi"}}]}
|
|
|
|
data: {"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}
|
|
|
|
data: [DONE]
|
|
|
|
`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "text/event-stream"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
// Proto frame carries same input/output plus additional reasoning/cached_input.
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE,
|
|
Usage: &iop.Usage{InputTokens: 10, OutputTokens: 5, ReasoningTokens: 3, CachedInputTokens: 2},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthroughSideband,
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"stream":true,
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
// Input/output should NOT be double-counted (merge rule: body wins).
|
|
// Reasoning/cached_input should come from proto (body has 0).
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 10, // from body, NOT 10+10=20
|
|
tokenTypeOutput: 5, // from body, NOT 5+5=10
|
|
tokenTypeReasoning: 3, // from proto
|
|
tokenTypeCachedInput: 2, // from proto
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType)
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestProviderTunnelSidebandResponseBodyAndProtoNoDoubleCount verifies the same
|
|
// merge rule for non-streaming sideband responses (REVIEW_REVIEW_REVIEW_USAGE_METRIC_RETRY).
|
|
func TestProviderTunnelSidebandResponseBodyAndProtoNoDoubleCount(t *testing.T) {
|
|
const rawToken = "sk-body-and-proto-nonstream-token"
|
|
const edgeID = "edge-body-and-proto-nonstream"
|
|
const model = "pool-model"
|
|
// Body reports input/output only.
|
|
providerBody := `{"id":"cmpl-1","choices":[{"message":{"role":"assistant","content":"hi"}}],"usage":{"prompt_tokens":12,"completion_tokens":8,"total_tokens":20}}`
|
|
|
|
frames := make(chan *iop.ProviderTunnelFrame, 5)
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START,
|
|
StatusCode: 200,
|
|
Headers: map[string]string{"Content-Type": "application/json"},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY, Body: []byte(providerBody)}
|
|
// Proto frame carries same input/output plus reasoning/cached_input.
|
|
frames <- &iop.ProviderTunnelFrame{
|
|
Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE,
|
|
Usage: &iop.Usage{InputTokens: 12, OutputTokens: 8, ReasoningTokens: 5, CachedInputTokens: 4},
|
|
}
|
|
frames <- &iop.ProviderTunnelFrame{Kind: iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END, End: true}
|
|
close(frames)
|
|
|
|
srv := NewServer(principalTokenCfg(rawToken, "ollama"), &fakeRunService{tunnelFrames: frames}, nil)
|
|
srv.SetEdgeID(edgeID)
|
|
srv.SetModelCatalog([]config.ModelCatalogEntry{{ID: model, Providers: map[string]string{"prov-1": "served-model"}}})
|
|
|
|
labels := usageLabels{
|
|
edgeID: edgeID, principalRef: "user:alice", principalAlias: "alice", tokenRef: "iop-tok-alice",
|
|
modelGroup: model, endpoint: usageEndpointChatCompletions, responseMode: responseModePassthroughSideband,
|
|
}
|
|
reqBefore := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
|
|
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{
|
|
"model":"pool-model",
|
|
"messages":[{"role":"user","content":"hello"}],
|
|
"metadata":{"iop_response_mode":"passthrough+sideband"}
|
|
}`))
|
|
req.Header.Set("Authorization", "Bearer "+rawToken)
|
|
w := httptest.NewRecorder()
|
|
srv.routes().ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status: got %d body=%s", w.Code, w.Body.String())
|
|
}
|
|
|
|
reqAfter := testutil.ToFloat64(openAIRequestsTotal.WithLabelValues(
|
|
edgeID, "user:alice", "alice", "iop-tok-alice", model,
|
|
usageEndpointChatCompletions, responseModePassthroughSideband, usageStatusSuccess, usageSourceProviderReported,
|
|
))
|
|
if reqAfter-reqBefore != 1 {
|
|
t.Fatalf("requests_total success/provider_reported: got delta %v, want 1", reqAfter-reqBefore)
|
|
}
|
|
// Input/output from body only; reasoning/cached_input from proto.
|
|
for tokenType, want := range map[string]float64{
|
|
tokenTypeInput: 12, // from body, NOT 12+12=24
|
|
tokenTypeOutput: 8, // from body, NOT 8+8=16
|
|
tokenTypeReasoning: 5, // from proto
|
|
tokenTypeCachedInput: 4, // from proto
|
|
} {
|
|
got := requestTokenValue(t, labels, tokenType)
|
|
if got != want {
|
|
t.Fatalf("token_type %s: got %v, want %v", tokenType, got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
// failingWriteRecorder wraps httptest.ResponseRecorder and fails after N writes.
|
|
type failingWriteRecorder struct {
|
|
*httptest.ResponseRecorder
|
|
failAfter int
|
|
writes int
|
|
}
|
|
|
|
func (fw *failingWriteRecorder) Write(p []byte) (int, error) {
|
|
fw.writes++
|
|
if fw.writes > fw.failAfter {
|
|
return 0, fmt.Errorf("simulated write failure")
|
|
}
|
|
return fw.ResponseRecorder.Write(p)
|
|
}
|