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) }