cpa-plugin/internal/plugin/management_test.go

package plugin

import (
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"testing"
	"time"
)

func managementCall(t *testing.T, app *App, method, path string) ManagementResponse {
	return managementCallBody(t, app, method, path, nil)
}

func managementCallBody(t *testing.T, app *App, method, path string, body []byte) ManagementResponse {
	t.Helper()
	request, err := json.Marshal(ManagementRequest{Method: method, Path: path, Body: body})
	if err != nil {
		t.Fatal(err)
	}
	raw, err := app.HandleMethod(MethodManagementHandle, request)
	if err != nil {
		t.Fatal(err)
	}
	var envelope Envelope
	if err := json.Unmarshal(raw, &envelope); err != nil {
		t.Fatal(err)
	}
	var response ManagementResponse
	if err := json.Unmarshal(envelope.Result, &response); err != nil {
		t.Fatal(err)
	}
	return response
}

func TestManagementRegistrationDeclaresUsageAPIAndUI(t *testing.T) {
	raw, err := NewApp().HandleMethod(MethodManagementRegister, nil)
	if err != nil {
		t.Fatal(err)
	}
	var envelope Envelope
	if err := json.Unmarshal(raw, &envelope); err != nil {
		t.Fatal(err)
	}
	var registration ManagementRegistrationResponse
	if err := json.Unmarshal(envelope.Result, &registration); err != nil {
		t.Fatal(err)
	}
	if len(registration.Routes) != 4 || registration.Routes[0].Path != managementBase+routeUsage || registration.Routes[1].Path != managementBase+routePrices {
		t.Fatalf("unexpected management routes: %+v", registration.Routes)
	}
	if len(registration.Resources) != 1 || registration.Resources[0].Path != resourceBase+resourceUI {
		t.Fatalf("unexpected resource routes: %+v", registration.Resources)
	}
}

func TestPriceManagementCalculatesUsageWithLongContextAndFastPolicy(t *testing.T) {
	app := NewApp()
	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil {
		t.Fatal(err)
	}
	priceBody := []byte(`{
        "model":"gpt-5.6-sol",
        "base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"},
        "long_context":{"threshold_input_tokens":272000,"comparison":"gt","input_per_1m":"5","cache_read_per_1m":"0.5","cache_write_per_1m":"6.25","output_per_1m":"22.5"},
        "fast_pricing_enabled":true,
        "fast_multiplier":"2.5"
    }`)
	response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, priceBody)
	if response.StatusCode != http.StatusOK {
		t.Fatalf("put price status = %d, body = %s", response.StatusCode, response.Body)
	}

	record := UsageRecord{
		Provider: "codex", Model: "gpt-5.6-sol", ServiceTier: "priority", RequestedAt: time.Now(),
		Detail: UsageDetail{InputTokens: 272_001, OutputTokens: 10_000, TotalTokens: 282_001},
	}
	raw, _ := json.Marshal(record)
	if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
		t.Fatal(err)
	}
	usageResponse := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
	var payload usageListResponse
	if err := json.Unmarshal(usageResponse.Body, &payload); err != nil {
		t.Fatal(err)
	}
	got := payload.Records[0]
	if !got.CostAvailable || got.CostUSD == nil || *got.CostUSD != 3.962513 || got.PriceTier != "long_context" || !got.FastRequested || !got.FastPricingApplied {
		t.Fatalf("unexpected priced usage: %+v", got)
	}

	pricesResponse := managementCall(t, app, http.MethodGet, managementBase+routePrices)
	if pricesResponse.StatusCode != http.StatusOK || !strings.Contains(string(pricesResponse.Body), `"fast_multiplier":"2.5"`) {
		t.Fatalf("unexpected prices response: %d %s", pricesResponse.StatusCode, pricesResponse.Body)
	}
}

func TestFastRequestCanUseStandardPricing(t *testing.T) {
	app := NewApp()
	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil {
		t.Fatal(err)
	}
	priceBody := []byte(`{"model":"gpt-5.6-sol","base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"},"fast_pricing_enabled":false,"fast_multiplier":"2.5"}`)
	if response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, priceBody); response.StatusCode != http.StatusOK {
		t.Fatalf("put price status = %d, body = %s", response.StatusCode, response.Body)
	}
	record := UsageRecord{Provider: "codex", Model: "gpt-5.6-sol", Speed: "fast", RequestedAt: time.Now(), Detail: UsageDetail{InputTokens: 100_000, TotalTokens: 100_000}}
	raw, _ := json.Marshal(record)
	if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
		t.Fatal(err)
	}
	response := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
	var payload usageListResponse
	if err := json.Unmarshal(response.Body, &payload); err != nil {
		t.Fatal(err)
	}
	got := payload.Records[0]
	if got.CostUSD == nil || *got.CostUSD != 0.25 || !got.FastRequested || got.FastPricingApplied {
		t.Fatalf("unexpected standard-priced Fast usage: %+v", got)
	}
}

func TestCanceledRequestWithoutUsageRemainsVisible(t *testing.T) {
	app := NewApp()
	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
		t.Fatal(err)
	}
	startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
	completion := RequestCompletion{
		RequestID: "request-canceled", TraceID: "trace-canceled", SourceFormat: "openai-response",
		Model: "deepseek-v4-flash", RequestedModel: "deepseek-flash", Stream: true,
		Outcome: RequestCompletionCanceled, StatusCode: 499, Error: "context canceled",
		StartedAt: startedAt, CompletedAt: startedAt.Add(2 * time.Second),
		Metadata: map[string]any{"request_path": "/v1/responses"},
	}
	raw, _ := json.Marshal(completion)
	if _, err := app.HandleMethod(MethodRequestComplete, raw); err != nil {
		t.Fatal(err)
	}
	response := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
	var payload usageListResponse
	if err := json.Unmarshal(response.Body, &payload); err != nil {
		t.Fatal(err)
	}
	if len(payload.Records) != 1 {
		t.Fatalf("records = %d, want 1", len(payload.Records))
	}
	got := payload.Records[0]
	if got.RequestID != "request-canceled" || got.TraceID != "trace-canceled" || got.Outcome != "canceled" || !got.Failed || got.StatusCode != 499 || got.RequestType != "SSE" || got.CostAvailable {
		t.Fatalf("unexpected canceled request: %+v", got)
	}
}

func TestCompactUsageIsPricedAndUsesJSONMetrics(t *testing.T) {
	app := NewApp()
	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
		t.Fatal(err)
	}
	priceBody := []byte(`{"model":"deepseek-v4-flash","base":{"input_per_1m":"2.5","cache_read_per_1m":"0.25","cache_write_per_1m":"3.125","output_per_1m":"15"},"fast_pricing_enabled":true,"fast_multiplier":"2.5"}`)
	if response := managementCallBody(t, app, http.MethodPut, managementBase+routePrices, priceBody); response.StatusCode != http.StatusOK {
		t.Fatalf("put price status = %d, body = %s", response.StatusCode, response.Body)
	}
	record := UsageRecord{
		RequestID: "compact-success", ExecutionID: "compact-attempt", Provider: "openai",
		Model: "deepseek-v4-flash", Endpoint: "POST /v1/responses/compact", RequestedAt: time.Now(),
		Latency: time.Second, TTFT: 200 * time.Millisecond,
		Detail: UsageDetail{InputTokens: 1_000, OutputTokens: 100, TotalTokens: 1_100},
	}
	raw, _ := json.Marshal(record)
	if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
		t.Fatal(err)
	}
	response := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
	var payload usageListResponse
	if err := json.Unmarshal(response.Body, &payload); err != nil {
		t.Fatal(err)
	}
	if len(payload.Records) != 1 {
		t.Fatalf("records = %d, want 1", len(payload.Records))
	}
	got := payload.Records[0]
	if got.RequestType != "JSON" || got.Endpoint != "POST /v1/responses/compact" || got.TTFTMilliseconds != 0 || got.SpeedTPS != nil {
		t.Fatalf("unexpected compact metadata: %+v", got)
	}
	if !got.CostAvailable || got.CostUSD == nil || *got.CostUSD != 0.004 {
		t.Fatalf("unexpected compact cost: %+v", got)
	}
}

func TestCompactFailureAndCancellationRemainVisible(t *testing.T) {
	for _, test := range []struct {
		name       string
		outcome    RequestCompletionOutcome
		statusCode int
	}{
		{name: "failed", outcome: RequestCompletionFailed, statusCode: 500},
		{name: "canceled", outcome: RequestCompletionCanceled, statusCode: 499},
	} {
		t.Run(test.name, func(t *testing.T) {
			app := NewApp()
			if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
				t.Fatal(err)
			}
			startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
			completion := RequestCompletion{
				RequestID: "compact-" + test.name, Model: "deepseek-v4-flash", Stream: false,
				Outcome: test.outcome, StatusCode: test.statusCode, Error: test.name,
				StartedAt: startedAt, CompletedAt: startedAt.Add(time.Second),
				Metadata: map[string]any{"request_path": "/v1/responses/compact"},
			}
			raw, _ := json.Marshal(completion)
			if _, err := app.HandleMethod(MethodRequestComplete, raw); err != nil {
				t.Fatal(err)
			}
			response := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
			var payload usageListResponse
			if err := json.Unmarshal(response.Body, &payload); err != nil {
				t.Fatal(err)
			}
			if len(payload.Records) != 1 {
				t.Fatalf("records = %d, want 1", len(payload.Records))
			}
			got := payload.Records[0]
			if got.Outcome != string(test.outcome) || !got.Failed || got.StatusCode != test.statusCode || got.RequestType != "JSON" || got.Endpoint != "/v1/responses/compact" || got.SpeedTPS != nil || got.CostAvailable {
				t.Fatalf("unexpected compact terminal record: %+v", got)
			}
		})
	}
}

func TestUsageAndLifecycleMergeInEitherOrderAndDeduplicateExecution(t *testing.T) {
	for _, usageFirst := range []bool{true, false} {
		app := NewApp()
		if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
			t.Fatal(err)
		}
		startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
		usage := UsageRecord{
			RequestID: "request-1", ExecutionID: "attempt-1", TraceID: "trace-1",
			Provider: "openai", Model: "deepseek-v4-flash", RequestedAt: startedAt,
			Endpoint: "POST /v1/responses", Detail: UsageDetail{InputTokens: 10, TotalTokens: 10},
		}
		completion := RequestCompletion{
			RequestID: "request-1", TraceID: "trace-1", Model: "deepseek-v4-flash",
			Outcome: RequestCompletionFailed, StatusCode: 500, Error: "upstream failed",
			StartedAt: startedAt, CompletedAt: startedAt.Add(time.Second),
		}
		usageRaw, _ := json.Marshal(usage)
		completionRaw, _ := json.Marshal(completion)
		calls := []struct {
			method string
			raw    []byte
		}{{MethodUsageHandle, usageRaw}, {MethodRequestComplete, completionRaw}}
		if !usageFirst {
			calls[0], calls[1] = calls[1], calls[0]
		}
		for _, call := range calls {
			if _, err := app.HandleMethod(call.method, call.raw); err != nil {
				t.Fatal(err)
			}
		}
		// 重复的 Usage 回调不能生成第二条执行记录。
		if _, err := app.HandleMethod(MethodUsageHandle, usageRaw); err != nil {
			t.Fatal(err)
		}
		response := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
		var payload usageListResponse
		if err := json.Unmarshal(response.Body, &payload); err != nil {
			t.Fatal(err)
		}
		if len(payload.Records) != 1 || payload.Records[0].ExecutionID != "attempt-1" || payload.Records[0].Outcome != "failed" || payload.Records[0].StatusCode != 500 {
			t.Fatalf("usageFirst=%v records=%+v", usageFirst, payload.Records)
		}
	}
}

func TestRetriedExecutionsRemainSeparate(t *testing.T) {
	app := NewApp()
	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: false\n"))); err != nil {
		t.Fatal(err)
	}
	startedAt := time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC)
	for index, failed := range []bool{true, false} {
		usage := UsageRecord{
			RequestID: "request-retry", ExecutionID: fmt.Sprintf("attempt-%d", index+1),
			Provider: "openai", Model: "deepseek-v4-flash", RequestedAt: startedAt.Add(time.Duration(index) * time.Second),
			Failed: failed, Detail: UsageDetail{InputTokens: int64(10 + index), TotalTokens: int64(10 + index)},
		}
		raw, _ := json.Marshal(usage)
		if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
			t.Fatal(err)
		}
	}
	completion := RequestCompletion{RequestID: "request-retry", Outcome: RequestCompletionSucceeded, StartedAt: startedAt, CompletedAt: startedAt.Add(2 * time.Second)}
	raw, _ := json.Marshal(completion)
	if _, err := app.HandleMethod(MethodRequestComplete, raw); err != nil {
		t.Fatal(err)
	}
	response := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
	var payload usageListResponse
	if err := json.Unmarshal(response.Body, &payload); err != nil {
		t.Fatal(err)
	}
	if len(payload.Records) != 2 || payload.Records[0].ExecutionID != "attempt-2" || payload.Records[0].Failed || payload.Records[1].ExecutionID != "attempt-1" || !payload.Records[1].Failed {
		t.Fatalf("unexpected retry records: %+v", payload.Records)
	}
}

func TestUsageManagementResponseContainsDisplayFields(t *testing.T) {
	app := NewApp()
	if _, err := app.HandleMethod(MethodPluginRegister, lifecycleRequest(t, SchemaVersion, testConfig(t, "enabled: true\ncodex_only: true\n"))); err != nil {
		t.Fatal(err)
	}
	record := UsageRecord{
		Provider:        "codex",
		APIKey:          "test-key",
		Model:           "gpt-5.5",
		ReasoningEffort: "high",
		ServiceTier:     "priority",
		ExecutorType:    "CodexExecutor",
		Endpoint:        "POST /v1/responses",
		ClientIP:        "192.0.2.10",
		RequestedAt:     time.Date(2026, 8, 14, 12, 0, 0, 0, time.UTC),
		Latency:         1500 * time.Millisecond,
		TTFT:            250 * time.Millisecond,
		Detail: UsageDetail{
			InputTokens:         10,
			OutputTokens:        5,
			TotalTokens:         15,
			CachedTokens:        4,
			CacheReadTokens:     4,
			CacheCreationTokens: 1,
			ReasoningTokens:     2,
		},
	}
	raw, _ := json.Marshal(record)
	if _, err := app.HandleMethod(MethodUsageHandle, raw); err != nil {
		t.Fatal(err)
	}

	response := managementCall(t, app, http.MethodGet, managementBase+routeUsage)
	if response.StatusCode != http.StatusOK {
		t.Fatalf("status = %d, body = %s", response.StatusCode, response.Body)
	}
	var payload usageListResponse
	if err := json.Unmarshal(response.Body, &payload); err != nil {
		t.Fatal(err)
	}
	if len(payload.Records) != 1 {
		t.Fatalf("records = %d, want 1", len(payload.Records))
	}
	got := payload.Records[0]
	if got.APIKey != "test-key" || got.Model != "gpt-5.5" || got.ReasoningEffort != "high" || got.ServiceTier != "priority" || got.ExecutorType != "CodexExecutor" {
		t.Fatalf("unexpected usage identity fields: %+v", got)
	}
	if got.RequestType != "SSE" || got.Endpoint != "POST /v1/responses" || got.ClientIP != "192.0.2.10" {
		t.Fatalf("unexpected request metadata: %+v", got)
	}
	if got.TotalTokens != 15 || got.TTFTMilliseconds != 250 || got.CacheReadTokens != 4 || got.CacheWriteTokens != 1 {
		t.Fatalf("unexpected usage counters: %+v", got)
	}
	if got.SpeedTPS == nil || *got.SpeedTPS != 4 || got.CacheRate == nil || *got.CacheRate != 40 {
		t.Fatalf("unexpected derived display fields: %+v", got)
	}
	if got.KeyAlias != "" || got.CostUSD != nil || got.CostAvailable {
		t.Fatalf("unexpected usage item: %+v", got)
	}
}

func TestUsageResourceServesTablePage(t *testing.T) {
	response := managementCall(t, NewApp(), http.MethodGet, resourceBase+resourceUI)
	page := string(response.Body)
	if response.StatusCode != http.StatusOK || !strings.Contains(page, "最近用量记录") {
		t.Fatalf("unexpected UI response: status=%d", response.StatusCode)
	}
	for _, column := range []string{"Key / 别名", "推理强度", "生成速度", "缓存写入", "总成本", "客户端 IP"} {
		if !strings.Contains(page, column) {
			t.Fatalf("UI does not contain column %q", column)
		}
	}
	if !strings.Contains(page, `return "compact"`) || !strings.Contains(page, "isCompactEndpoint(record.endpoint)") {
		t.Fatal("UI does not contain compact display rules")
	}
}