AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告

M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 11:45:54 +08:00
commit 5759c1862e
807 changed files with 114727 additions and 0 deletions
@@ -0,0 +1,141 @@
+package gateway
+
+import (
+	"context"
+	"io"
+	"log/slog"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"aigateway.local/core/internal/audit"
+	"aigateway.local/core/internal/contentpolicy"
+	"aigateway.local/core/internal/pricing"
+	provideropenai "aigateway.local/core/internal/provider/openai"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+func TestContentPolicyAndPricingIntegration(t *testing.T) {
+	databaseURL := os.Getenv("CONTENT_PRICING_TEST_DATABASE_URL")
+	if databaseURL == "" {
+		t.Skip("CONTENT_PRICING_TEST_DATABASE_URL is not set")
+	}
+	ctx := context.Background()
+	pool, err := pgxpool.New(ctx, databaseURL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer pool.Close()
+	const blockID = "00000000-0000-4000-8000-000000001601"
+	const priceID = "00000000-0000-4000-8000-000000001602"
+	_, err = pool.Exec(ctx, `INSERT INTO gateway.content_policies(id,name,action,priority,rules,enabled) VALUES($1,'integration block','block',2000,'[{"name":"forbidden","pattern":"forbidden","replacement":"[BLOCKED]"}]',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, blockID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = pool.Exec(ctx, `INSERT INTO gateway.model_prices(id,provider_code,model_pattern,input_microunits_per_million,output_microunits_per_million,currency,effective_from,enabled) VALUES($1,'environment','gpt-test',2000000,8000000,'USD',clock_timestamp()-interval '1 hour',true) ON CONFLICT(id) DO UPDATE SET enabled=true`, priceID)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`)
+		_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.content_policies WHERE id=$1`, blockID)
+		_, _ = pool.Exec(context.Background(), `DELETE FROM gateway.model_prices WHERE id=$1`, priceID)
+	}()
+
+	var upstreamCalls atomic.Int64
+	bodySeen := make(chan string, 1)
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		upstreamCalls.Add(1)
+		body, _ := io.ReadAll(r.Body)
+		bodySeen <- string(body)
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = io.WriteString(w, `{"id":"ok","usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`)
+	}))
+	defer upstream.Close()
+	adapter, _ := provideropenai.New(upstream.URL, "upstream-key")
+	engine := contentpolicy.NewEngine(pool, time.Minute, slog.Default())
+	if err := engine.Reload(ctx); err != nil {
+		t.Fatal(err)
+	}
+	prices := pricing.NewService(pool, time.Minute, slog.Default())
+	if err := prices.Reload(ctx); err != nil {
+		t.Fatal(err)
+	}
+	recorder := audit.NewRecorder(pool, slog.Default(), 100, 1, 10*time.Millisecond)
+	recordCtx, cancel := context.WithCancel(ctx)
+	stopped := make(chan struct{})
+	go func() { recorder.Run(recordCtx); close(stopped) }()
+	defer func() { cancel(); <-stopped }()
+	proxy := NewProxy(adapter, "test-key", 1<<20, slog.Default())
+	proxy.SetAuditRecorder(recorder)
+	proxy.SetContentPolicyEngine(engine)
+	proxy.SetPricingService(prices)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		proxy.ServeHTTP(w, r.WithContext(WithRequestID(r.Context(), r.Header.Get("X-Request-ID"))))
+	}))
+	defer server.Close()
+
+	req, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/chat/completions", strings.NewReader(`{"model":"gpt-test","messages":[{"role":"user","content":"token sk-1234567890123456"}]}`))
+	req.Header.Set("Authorization", "Bearer test-key")
+	req.Header.Set("X-Request-ID", "m3-content-price-redact")
+	response, err := http.DefaultClient.Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, _ = io.ReadAll(response.Body)
+	_ = response.Body.Close()
+	if response.StatusCode != 200 || response.Header.Get("X-Gateway-Content-Redacted") != "true" {
+		t.Fatalf("unexpected response %d headers=%v", response.StatusCode, response.Header)
+	}
+	select {
+	case body := <-bodySeen:
+		if strings.Contains(body, "sk-1234567890123456") || !strings.Contains(body, "[REDACTED_API_KEY]") {
+			t.Fatalf("upstream saw unsafe body %s", body)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("upstream did not receive request")
+	}
+
+	blocked, _ := http.NewRequest(http.MethodPost, server.URL+"/v1/responses", strings.NewReader(`{"model":"gpt-test","input":"forbidden"}`))
+	blocked.Header.Set("Authorization", "Bearer test-key")
+	blocked.Header.Set("X-Request-ID", "m3-content-price-block")
+	blockedResponse, err := http.DefaultClient.Do(blocked)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, _ = io.ReadAll(blockedResponse.Body)
+	_ = blockedResponse.Body.Close()
+	if blockedResponse.StatusCode != http.StatusUnprocessableEntity || upstreamCalls.Load() != 1 {
+		t.Fatalf("block failed status=%d calls=%d", blockedResponse.StatusCode, upstreamCalls.Load())
+	}
+
+	deadline := time.Now().Add(2 * time.Second)
+	for {
+		var cost *int64
+		var labels string
+		err = pool.QueryRow(ctx, `SELECT cost_microunits,labels::text FROM gateway.audit_events WHERE request_id='m3-content-price-redact' ORDER BY recorded_at DESC LIMIT 1`).Scan(&cost, &labels)
+		if err == nil && cost != nil && *cost == 6000 && strings.Contains(labels, "content_redacted") {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("audit cost/redaction missing cost=%v labels=%s err=%v", cost, labels, err)
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+	deadline = time.Now().Add(2 * time.Second)
+	for {
+		var count int
+		err = pool.QueryRow(ctx, `SELECT count(*) FROM gateway.outbox_events WHERE event_type='content_policy.matched' AND aggregate_id IN ('m3-content-price-redact','m3-content-price-block')`).Scan(&count)
+		if err == nil && count == 2 {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatalf("content policy notification outbox missing count=%d err=%v", count, err)
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+}