From 275df644af00d0c14efcdc31d42ab6b5cad970ee Mon Sep 17 00:00:00 2001 From: GeJiaXiang <353358601@qq.com> Date: Sat, 16 May 2026 10:07:06 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=E5=AE=A1?= =?UTF-8?q?=E8=AE=A1=E7=B3=BB=E7=BB=9F=E5=AF=B9=E6=8E=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Dockerfile | 5 +- main.go | 2 + middleware/audit.go | 154 +++++++++++++++++++++++++++++++++ middleware/audit_test.go | 90 ++++++++++++++++++++ relay/mjproxy_handler.go | 25 ++++++ relay/relay_task.go | 18 ++++ router/relay-router.go | 11 ++- router/video-router.go | 7 +- service/audit/audit_test.go | 77 +++++++++++++++++ service/audit/capture.go | 72 ++++++++++++++++ service/audit/context.go | 46 ++++++++++ service/audit/reporter.go | 165 ++++++++++++++++++++++++++++++++++++ service/audit/types.go | 100 ++++++++++++++++++++++ service/text_quota.go | 36 ++++++++ 14 files changed, 800 insertions(+), 8 deletions(-) create mode 100644 middleware/audit.go create mode 100644 middleware/audit_test.go create mode 100644 service/audit/audit_test.go create mode 100644 service/audit/capture.go create mode 100644 service/audit/context.go create mode 100644 service/audit/reporter.go create mode 100644 service/audit/types.go diff --git a/Dockerfile b/Dockerfile index 17c4398d4eb..ba943b18046 100644 --- a/Dockerfile +++ b/Dockerfile @@ -3,7 +3,7 @@ FROM oven/bun:1@sha256:0733e50325078969732ebe3b15ce4c4be5082f18c4ac1a0f0ca4839c2 WORKDIR /build COPY web/default/package.json . COPY web/default/bun.lock . -RUN bun install +RUN bun install --registry https://registry.npmmirror.com COPY ./web/default . COPY ./VERSION . RUN DISABLE_ESLINT_PLUGIN='true' VITE_REACT_APP_VERSION=$(cat VERSION) bun run build @@ -13,7 +13,7 @@ FROM oven/bun:1@sha256:0733e50325078969732ebe3b15ce4c4be5082f18c4ac1a0f0ca4839c2 WORKDIR /build COPY web/classic/package.json . COPY web/classic/bun.lock . -RUN bun install +RUN bun install --registry https://registry.npmmirror.com COPY ./web/classic . COPY ./VERSION . RUN VITE_REACT_APP_VERSION=$(cat VERSION) bun run build @@ -38,6 +38,7 @@ RUN go build -ldflags "-s -w -X 'github.com/QuantumNous/new-api/common.Version=$ FROM debian:bookworm-slim@sha256:f06537653ac770703bc45b4b113475bd402f451e85223f0f2837acbf89ab020a +RUN sed -i 's|http://deb.debian.org/debian-security|http://repo.huaweicloud.com/debian-security|g; s|http://deb.debian.org/debian|http://repo.huaweicloud.com/debian|g' /etc/apt/sources.list.d/debian.sources RUN apt-get update \ && apt-get install -y --no-install-recommends ca-certificates tzdata libasan8 wget \ && rm -rf /var/lib/apt/lists/* \ diff --git a/main.go b/main.go index 3361b8ce933..30cde1692d3 100644 --- a/main.go +++ b/main.go @@ -23,6 +23,7 @@ import ( "github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/router" "github.com/QuantumNous/new-api/service" + auditservice "github.com/QuantumNous/new-api/service/audit" _ "github.com/QuantumNous/new-api/setting/performance_setting" "github.com/QuantumNous/new-api/setting/ratio_setting" @@ -269,6 +270,7 @@ func InitResources() error { common.InitEnv() logger.SetupLogger() + auditservice.InitFromEnv() // Initialize model settings ratio_setting.InitRatioSettings() diff --git a/middleware/audit.go b/middleware/audit.go new file mode 100644 index 00000000000..9486d76e74c --- /dev/null +++ b/middleware/audit.go @@ -0,0 +1,154 @@ +package middleware + +import ( + "io" + "strconv" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + auditservice "github.com/QuantumNous/new-api/service/audit" + "github.com/gin-gonic/gin" +) + +type auditResponseWriter struct { + gin.ResponseWriter + capture *auditservice.CaptureBuffer +} + +func (w *auditResponseWriter) Write(data []byte) (int, error) { + w.capture.Write(data) + return w.ResponseWriter.Write(data) +} + +func (w *auditResponseWriter) WriteString(data string) (int, error) { + w.capture.Write(common.StringToByteSlice(data)) + return w.ResponseWriter.WriteString(data) +} + +func Audit() gin.HandlerFunc { + return func(c *gin.Context) { + if !auditservice.Enabled() { + c.Next() + return + } + + startedAt := time.Now() + responseCapture := auditservice.NewCaptureBuffer(auditservice.MaxBodyBytes()) + originalWriter := c.Writer + c.Writer = &auditResponseWriter{ + ResponseWriter: originalWriter, + capture: responseCapture, + } + + c.Next() + + requestBody := captureRequestBody(c, auditservice.MaxBodyBytes()) + responseBody := responseCapture.Body(c.Writer.Header().Get("Content-Type")) + responseBody.StatusCode = c.Writer.Status() + modelInfo := auditModelInfo(c) + + auditservice.ReportAsync(auditservice.Event{ + Version: auditservice.ProtocolVersion, + Event: auditservice.EventRequestResponse, + RequestID: c.GetString(common.RequestIdKey), + Timestamp: time.Now(), + Node: auditservice.NodeName(), + Route: c.FullPath(), + User: auditservice.UserInfo{ + ID: c.GetInt("id"), + Username: firstNonEmpty( + c.GetString("username"), + common.GetContextKeyString(c, constant.ContextKeyUserName), + ), + }, + Key: auditservice.KeyInfo{ + ID: c.GetInt("token_id"), + Name: c.GetString("token_name"), + }, + Client: auditservice.ClientInfo{ + IP: c.ClientIP(), + Method: c.Request.Method, + Path: c.Request.URL.RequestURI(), + UserAgent: c.Request.UserAgent(), + }, + Model: modelInfo, + Billing: auditservice.BillingInfoFromContext(c), + Request: requestBody, + Response: responseBody, + DurationMS: time.Since(startedAt).Milliseconds(), + Metadata: auditMetadata(c), + }) + + c.Writer = originalWriter + } +} + +func captureRequestBody(c *gin.Context, maxBodyBytes int64) auditservice.Body { + storage, err := common.GetBodyStorage(c) + if err != nil { + common.SysError("audit request body capture failed: " + err.Error()) + return auditservice.Body{ContentType: c.Request.Header.Get("Content-Type")} + } + + data, err := storage.Bytes() + if err != nil { + common.SysError("audit request body read failed: " + err.Error()) + return auditservice.Body{ContentType: c.Request.Header.Get("Content-Type")} + } + if _, err := storage.Seek(0, io.SeekStart); err != nil { + common.SysError("audit request body rewind failed: " + err.Error()) + } + c.Request.Body = io.NopCloser(storage) + + return auditservice.BodyFromBytes(data, c.Request.Header.Get("Content-Type"), maxBodyBytes) +} + +func auditModelInfo(c *gin.Context) auditservice.ModelInfo { + info := auditservice.ModelInfoFromContext(c) + if info.Name == "" { + info.Name = common.GetContextKeyString(c, constant.ContextKeyOriginalModel) + } + if info.OriginName == "" { + info.OriginName = common.GetContextKeyString(c, constant.ContextKeyOriginalModel) + } + return info +} + +func auditMetadata(c *gin.Context) map[string]string { + metadata := make(map[string]string) + putString := func(key string, value string) { + if value != "" { + metadata[key] = value + } + } + putInt := func(key string, value int) { + if value != 0 { + metadata[key] = strconv.Itoa(value) + } + } + + putString("model", common.GetContextKeyString(c, constant.ContextKeyOriginalModel)) + putString("group", common.GetContextKeyString(c, constant.ContextKeyUsingGroup)) + putString("token_group", common.GetContextKeyString(c, constant.ContextKeyTokenGroup)) + putString("channel_name", common.GetContextKeyString(c, constant.ContextKeyChannelName)) + if _, exists := common.GetContextKey(c, constant.ContextKeyIsStream); exists { + putString("is_stream", strconv.FormatBool(common.GetContextKeyBool(c, constant.ContextKeyIsStream))) + } + putInt("channel_id", common.GetContextKeyInt(c, constant.ContextKeyChannelId)) + putInt("channel_type", common.GetContextKeyInt(c, constant.ContextKeyChannelType)) + + if len(metadata) == 0 { + return nil + } + return metadata +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + return "" +} diff --git a/middleware/audit_test.go b/middleware/audit_test.go new file mode 100644 index 00000000000..c6f7134b755 --- /dev/null +++ b/middleware/audit_test.go @@ -0,0 +1,90 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + auditservice "github.com/QuantumNous/new-api/service/audit" + "github.com/gin-gonic/gin" +) + +func TestAuditMiddlewareReportsRequestAndResponseByUserTokenAndIP(t *testing.T) { + gin.SetMode(gin.TestMode) + + events := make(chan auditservice.Event, 1) + auditServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var event auditservice.Event + if err := common.DecodeJson(r.Body, &event); err != nil { + t.Fatalf("failed to decode audit event: %v", err) + } + events <- event + w.WriteHeader(http.StatusAccepted) + })) + defer auditServer.Close() + + os.Setenv("AUDIT_ENDPOINT_URL", auditServer.URL) + os.Setenv("AUDIT_API_KEY", "test-key") + os.Setenv("AUDIT_TIMEOUT_SECONDS", "1") + os.Setenv("AUDIT_MAX_BODY_BYTES", "1024") + auditservice.InitFromEnv() + defer func() { + os.Unsetenv("AUDIT_ENDPOINT_URL") + os.Unsetenv("AUDIT_API_KEY") + os.Unsetenv("AUDIT_TIMEOUT_SECONDS") + os.Unsetenv("AUDIT_MAX_BODY_BYTES") + auditservice.InitFromEnv() + }() + + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set("id", 42) + c.Set("username", "alice") + c.Set("token_id", 12) + c.Set("token_name", "prod-key") + c.Next() + }) + router.Use(Audit()) + router.POST("/v1/chat/completions", func(c *gin.Context) { + c.JSON(http.StatusCreated, gin.H{"message": "ok"}) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"prompt":"hello"}`)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("User-Agent", "audit-test") + request.RemoteAddr = "203.0.113.10:12345" + + router.ServeHTTP(recorder, request) + if recorder.Code != http.StatusCreated { + t.Fatalf("expected handler response status 201, got %d", recorder.Code) + } + + select { + case event := <-events: + if event.User.ID != 42 || event.User.Username != "alice" { + t.Fatalf("unexpected user dimensions: %+v", event.User) + } + if event.Key.ID != 12 || event.Key.Name != "prod-key" { + t.Fatalf("unexpected key dimensions: %+v", event.Key) + } + if event.Client.IP != "203.0.113.10" { + t.Fatalf("unexpected client ip: %q", event.Client.IP) + } + if event.Request.Content != `{"prompt":"hello"}` { + t.Fatalf("unexpected request body: %q", event.Request.Content) + } + if !strings.Contains(event.Response.Content, `"message":"ok"`) { + t.Fatalf("unexpected response body: %q", event.Response.Content) + } + if event.Response.StatusCode != http.StatusCreated { + t.Fatalf("unexpected response status: %d", event.Response.StatusCode) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for audit event") + } +} diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index 5b0750fec43..48361c86617 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -20,8 +20,10 @@ import ( relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/service" + auditservice "github.com/QuantumNous/new-api/service/audit" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" ) @@ -233,6 +235,7 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR tokenName := c.GetString("token_name") logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, constant.MjActionSwapFace) other := service.GenerateMjOtherInfo(info, priceData) + setPerCallAuditBilling(c, info, modelName, priceData, other) model.RecordConsumeLog(c, info.UserId, model.RecordConsumeLogParams{ ChannelId: info.ChannelId, ModelName: modelName, @@ -539,6 +542,7 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt tokenName := c.GetString("token_name") logContent := fmt.Sprintf("模型固定价格 %.2f，分组倍率 %.2f，操作 %s，ID %s", priceData.ModelPrice, priceData.GroupRatioInfo.GroupRatio, midjRequest.Action, midjResponse.Result) other := service.GenerateMjOtherInfo(relayInfo, priceData) + setPerCallAuditBilling(c, relayInfo, modelName, priceData, other) model.RecordConsumeLog(c, relayInfo.UserId, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, ModelName: modelName, @@ -661,6 +665,27 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt return nil } +func setPerCallAuditBilling(c *gin.Context, info *relaycommon.RelayInfo, modelName string, priceData types.PriceData, details map[string]interface{}) { + auditservice.SetModelInfo(c, auditservice.ModelInfo{ + Name: modelName, + OriginName: info.OriginModelName, + UpstreamName: info.UpstreamModelName, + }) + auditservice.SetBillingInfo(c, auditservice.BillingInfo{ + Quota: priceData.Quota, + ModelRatio: priceData.ModelRatio, + GroupRatio: priceData.GroupRatioInfo.GroupRatio, + ModelPrice: priceData.ModelPrice, + UsePrice: priceData.UsePrice, + FreeModel: priceData.FreeModel, + QuotaToPreConsume: priceData.QuotaToPreConsume, + FinalPreConsumedQuota: info.FinalPreConsumedQuota, + QuotaPerUnit: common.QuotaPerUnit, + OtherRatios: priceData.OtherRatios, + Details: details, + }) +} + type taskChangeParams struct { ID string Action string diff --git a/relay/relay_task.go b/relay/relay_task.go index 098e23828b6..2f344b1f232 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -19,6 +19,7 @@ import ( relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/service" + auditservice "github.com/QuantumNous/new-api/service/audit" "github.com/gin-gonic/gin" ) @@ -248,6 +249,23 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe info.PriceData.OtherRatios = adjustedRatios info.PriceData.Quota = finalQuota } + auditservice.SetModelInfo(c, auditservice.ModelInfo{ + Name: modelName, + OriginName: info.OriginModelName, + UpstreamName: info.UpstreamModelName, + }) + auditservice.SetBillingInfo(c, auditservice.BillingInfo{ + Quota: finalQuota, + ModelRatio: info.PriceData.ModelRatio, + GroupRatio: info.PriceData.GroupRatioInfo.GroupRatio, + ModelPrice: info.PriceData.ModelPrice, + UsePrice: info.PriceData.UsePrice, + FreeModel: info.PriceData.FreeModel, + QuotaToPreConsume: info.PriceData.QuotaToPreConsume, + FinalPreConsumedQuota: info.FinalPreConsumedQuota, + QuotaPerUnit: common.QuotaPerUnit, + OtherRatios: info.PriceData.OtherRatios, + }) return &TaskSubmitResult{ UpstreamTaskID: upstreamTaskID, diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd..eaba7ec6f62 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -19,6 +19,7 @@ func SetRelayRouter(router *gin.Engine) { modelsRouter := router.Group("/v1/models") modelsRouter.Use(middleware.RouteTag("relay")) modelsRouter.Use(middleware.TokenAuth()) + modelsRouter.Use(middleware.Audit()) { modelsRouter.GET("", func(c *gin.Context) { switch { @@ -44,6 +45,7 @@ func SetRelayRouter(router *gin.Engine) { geminiRouter := router.Group("/v1beta/models") geminiRouter.Use(middleware.RouteTag("relay")) geminiRouter.Use(middleware.TokenAuth()) + geminiRouter.Use(middleware.Audit()) { geminiRouter.GET("", func(c *gin.Context) { controller.ListModels(c, constant.ChannelTypeGemini) @@ -53,6 +55,7 @@ func SetRelayRouter(router *gin.Engine) { geminiCompatibleRouter := router.Group("/v1beta/openai/models") geminiCompatibleRouter.Use(middleware.RouteTag("relay")) geminiCompatibleRouter.Use(middleware.TokenAuth()) + geminiCompatibleRouter.Use(middleware.Audit()) { geminiCompatibleRouter.GET("", func(c *gin.Context) { controller.ListModels(c, constant.ChannelTypeOpenAI) @@ -62,7 +65,7 @@ func SetRelayRouter(router *gin.Engine) { playgroundRouter := router.Group("/pg") playgroundRouter.Use(middleware.RouteTag("relay")) playgroundRouter.Use(middleware.SystemPerformanceCheck()) - playgroundRouter.Use(middleware.UserAuth(), middleware.Distribute()) + playgroundRouter.Use(middleware.UserAuth(), middleware.Audit(), middleware.Distribute()) { playgroundRouter.POST("/chat/completions", controller.Playground) } @@ -70,6 +73,7 @@ func SetRelayRouter(router *gin.Engine) { relayV1Router.Use(middleware.RouteTag("relay")) relayV1Router.Use(middleware.SystemPerformanceCheck()) relayV1Router.Use(middleware.TokenAuth()) + relayV1Router.Use(middleware.Audit()) relayV1Router.Use(middleware.ModelRequestRateLimit()) { // WebSocket 路由（统一到 Relay） @@ -179,7 +183,7 @@ func SetRelayRouter(router *gin.Engine) { relaySunoRouter := router.Group("/suno") relaySunoRouter.Use(middleware.RouteTag("relay")) relaySunoRouter.Use(middleware.SystemPerformanceCheck()) - relaySunoRouter.Use(middleware.TokenAuth(), middleware.Distribute()) + relaySunoRouter.Use(middleware.TokenAuth(), middleware.Audit(), middleware.Distribute()) { relaySunoRouter.POST("/submit/:action", controller.RelayTask) relaySunoRouter.POST("/fetch", controller.RelayTaskFetch) @@ -190,6 +194,7 @@ func SetRelayRouter(router *gin.Engine) { relayGeminiRouter.Use(middleware.RouteTag("relay")) relayGeminiRouter.Use(middleware.SystemPerformanceCheck()) relayGeminiRouter.Use(middleware.TokenAuth()) + relayGeminiRouter.Use(middleware.Audit()) relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) relayGeminiRouter.Use(middleware.Distribute()) { @@ -202,7 +207,7 @@ func SetRelayRouter(router *gin.Engine) { func registerMjRouterGroup(relayMjRouter *gin.RouterGroup) { relayMjRouter.GET("/image/:id", relay.RelayMidjourneyImage) - relayMjRouter.Use(middleware.TokenAuth(), middleware.Distribute()) + relayMjRouter.Use(middleware.TokenAuth(), middleware.Audit(), middleware.Distribute()) { relayMjRouter.POST("/submit/action", controller.RelayMidjourney) relayMjRouter.POST("/submit/shorten", controller.RelayMidjourney) diff --git a/router/video-router.go b/router/video-router.go index 46145110452..8dadde51d6b 100644 --- a/router/video-router.go +++ b/router/video-router.go @@ -12,13 +12,14 @@ func SetVideoRouter(router *gin.Engine) { videoProxyRouter := router.Group("/v1") videoProxyRouter.Use(middleware.RouteTag("relay")) videoProxyRouter.Use(middleware.TokenOrUserAuth()) + videoProxyRouter.Use(middleware.Audit()) { videoProxyRouter.GET("/videos/:task_id/content", controller.VideoProxy) } videoV1Router := router.Group("/v1") videoV1Router.Use(middleware.RouteTag("relay")) - videoV1Router.Use(middleware.TokenAuth(), middleware.Distribute()) + videoV1Router.Use(middleware.TokenAuth(), middleware.Audit(), middleware.Distribute()) { videoV1Router.POST("/video/generations", controller.RelayTask) videoV1Router.GET("/video/generations/:task_id", controller.RelayTaskFetch) @@ -33,7 +34,7 @@ func SetVideoRouter(router *gin.Engine) { klingV1Router := router.Group("/kling/v1") klingV1Router.Use(middleware.RouteTag("relay")) - klingV1Router.Use(middleware.KlingRequestConvert(), middleware.TokenAuth(), middleware.Distribute()) + klingV1Router.Use(middleware.KlingRequestConvert(), middleware.TokenAuth(), middleware.Audit(), middleware.Distribute()) { klingV1Router.POST("/videos/text2video", controller.RelayTask) klingV1Router.POST("/videos/image2video", controller.RelayTask) @@ -44,7 +45,7 @@ func SetVideoRouter(router *gin.Engine) { // Jimeng official API routes - direct mapping to official API format jimengOfficialGroup := router.Group("jimeng") jimengOfficialGroup.Use(middleware.RouteTag("relay")) - jimengOfficialGroup.Use(middleware.JimengRequestConvert(), middleware.TokenAuth(), middleware.Distribute()) + jimengOfficialGroup.Use(middleware.JimengRequestConvert(), middleware.TokenAuth(), middleware.Audit(), middleware.Distribute()) { // Maps to: /?Action=CVSync2AsyncSubmitTask&Version=2022-08-31 and /?Action=CVSync2AsyncGetResult&Version=2022-08-31 jimengOfficialGroup.POST("/", controller.RelayTask) diff --git a/service/audit/audit_test.go b/service/audit/audit_test.go new file mode 100644 index 00000000000..ee555ec8fc4 --- /dev/null +++ b/service/audit/audit_test.go @@ -0,0 +1,77 @@ +package audit + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" +) + +func TestReporterPostsAuditEventWithBearerKey(t *testing.T) { + var gotAuth string + var gotContentType string + var gotEvent Event + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotContentType = r.Header.Get("Content-Type") + if err := common.DecodeJson(r.Body, &gotEvent); err != nil { + t.Fatalf("failed to decode audit event: %v", err) + } + w.WriteHeader(http.StatusAccepted) + })) + defer server.Close() + + reporter := NewReporter(Config{ + EndpointURL: server.URL, + APIKey: "audit-key", + Timeout: time.Second, + }, server.Client()) + + err := reporter.Report(Event{ + Version: ProtocolVersion, + Event: EventRequestResponse, + RequestID: "req-1", + User: UserInfo{ID: 7, Username: "alice"}, + Key: KeyInfo{ID: 9, Name: "prod-key"}, + Client: ClientInfo{IP: "203.0.113.10"}, + Request: Body{Content: "hello"}, + Response: Body{Content: "world"}, + }) + if err != nil { + t.Fatalf("expected report to succeed, got %v", err) + } + + if gotAuth != "Bearer audit-key" { + t.Fatalf("expected bearer auth header, got %q", gotAuth) + } + if gotContentType != "application/json" { + t.Fatalf("expected json content type, got %q", gotContentType) + } + if gotEvent.User.Username != "alice" || gotEvent.Key.Name != "prod-key" || gotEvent.Client.IP != "203.0.113.10" { + t.Fatalf("unexpected dimensions in event: %+v", gotEvent) + } + if gotEvent.Request.Content != "hello" || gotEvent.Response.Content != "world" { + t.Fatalf("unexpected body content in event: %+v", gotEvent) + } +} + +func TestCaptureBufferTruncatesWithoutDroppingOriginalWrite(t *testing.T) { + buffer := NewCaptureBuffer(5) + + buffer.Write([]byte("hello")) + buffer.Write([]byte(" world")) + + body := buffer.Body("text/plain") + if body.Content != "hello" { + t.Fatalf("expected captured content to be truncated, got %q", body.Content) + } + if !body.Truncated { + t.Fatalf("expected body to be marked truncated") + } + if body.SizeBytes != 11 { + t.Fatalf("expected original size 11, got %d", body.SizeBytes) + } +} diff --git a/service/audit/capture.go b/service/audit/capture.go new file mode 100644 index 00000000000..ffddbb50949 --- /dev/null +++ b/service/audit/capture.go @@ -0,0 +1,72 @@ +package audit + +import "sync" + +type CaptureBuffer struct { + limit int64 + size int64 + truncated bool + data []byte + mu sync.Mutex +} + +func NewCaptureBuffer(limit int64) *CaptureBuffer { + capacity := int64(4096) + if limit > 0 { + capacity = minInt64(limit, 4096) + } + return &CaptureBuffer{ + limit: limit, + data: make([]byte, 0, capacity), + } +} + +func (b *CaptureBuffer) Write(p []byte) { + b.mu.Lock() + defer b.mu.Unlock() + + b.size += int64(len(p)) + if b.limit <= 0 { + b.data = append(b.data, p...) + return + } + if int64(len(b.data)) >= b.limit { + if len(p) > 0 { + b.truncated = true + } + return + } + + remaining := b.limit - int64(len(b.data)) + if int64(len(p)) > remaining { + b.data = append(b.data, p[:remaining]...) + b.truncated = true + return + } + b.data = append(b.data, p...) +} + +func (b *CaptureBuffer) Body(contentType string) Body { + b.mu.Lock() + defer b.mu.Unlock() + + return Body{ + ContentType: contentType, + Content: string(b.data), + SizeBytes: b.size, + Truncated: b.truncated, + } +} + +func BodyFromBytes(data []byte, contentType string, limit int64) Body { + buffer := NewCaptureBuffer(limit) + buffer.Write(data) + return buffer.Body(contentType) +} + +func minInt64(a int64, b int64) int64 { + if a < b { + return a + } + return b +} diff --git a/service/audit/context.go b/service/audit/context.go new file mode 100644 index 00000000000..ff236e2c885 --- /dev/null +++ b/service/audit/context.go @@ -0,0 +1,46 @@ +package audit + +import "github.com/gin-gonic/gin" + +const ( + contextKeyModelInfo = "audit_model_info" + contextKeyBillingInfo = "audit_billing_info" +) + +func SetModelInfo(c *gin.Context, info ModelInfo) { + if c == nil { + return + } + c.Set(contextKeyModelInfo, info) +} + +func ModelInfoFromContext(c *gin.Context) ModelInfo { + if c == nil { + return ModelInfo{} + } + value, exists := c.Get(contextKeyModelInfo) + if !exists { + return ModelInfo{} + } + info, _ := value.(ModelInfo) + return info +} + +func SetBillingInfo(c *gin.Context, info BillingInfo) { + if c == nil { + return + } + c.Set(contextKeyBillingInfo, info) +} + +func BillingInfoFromContext(c *gin.Context) BillingInfo { + if c == nil { + return BillingInfo{} + } + value, exists := c.Get(contextKeyBillingInfo) + if !exists { + return BillingInfo{} + } + info, _ := value.(BillingInfo) + return info +} diff --git a/service/audit/reporter.go b/service/audit/reporter.go new file mode 100644 index 00000000000..4ffdf280ecc --- /dev/null +++ b/service/audit/reporter.go @@ -0,0 +1,165 @@ +package audit + +import ( + "bytes" + "context" + "fmt" + "net/http" + "os" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/QuantumNous/new-api/common" +) + +const ( + envEndpointURL = "AUDIT_ENDPOINT_URL" + envAPIKey = "AUDIT_API_KEY" + envTimeout = "AUDIT_TIMEOUT_SECONDS" + envMaxBodyBytes = "AUDIT_MAX_BODY_BYTES" + defaultTimeout = 3 * time.Second + defaultMaxBody = int64(0) +) + +type Reporter struct { + config Config + client *http.Client +} + +var defaultReporter atomic.Value + +func LoadConfigFromEnv() Config { + timeout := defaultTimeout + if raw := strings.TrimSpace(os.Getenv(envTimeout)); raw != "" { + if seconds, err := strconv.Atoi(raw); err == nil && seconds > 0 { + timeout = time.Duration(seconds) * time.Second + } + } + + maxBodyBytes := defaultMaxBody + if raw := strings.TrimSpace(os.Getenv(envMaxBodyBytes)); raw != "" { + if bytesLimit, err := strconv.ParseInt(raw, 10, 64); err == nil && bytesLimit >= 0 { + maxBodyBytes = bytesLimit + } + } + + return Config{ + EndpointURL: strings.TrimSpace(os.Getenv(envEndpointURL)), + APIKey: strings.TrimSpace(os.Getenv(envAPIKey)), + Timeout: timeout, + MaxBodyBytes: maxBodyBytes, + NodeName: strings.TrimSpace(os.Getenv("NODE_NAME")), + } +} + +func InitFromEnv() { + config := LoadConfigFromEnv() + if !config.Enabled() { + defaultReporter.Store((*Reporter)(nil)) + return + } + defaultReporter.Store(NewReporter(config, nil)) + common.SysLog("audit reporter enabled") +} + +func DefaultReporter() *Reporter { + if reporter, ok := defaultReporter.Load().(*Reporter); ok { + return reporter + } + return nil +} + +func Enabled() bool { + reporter := DefaultReporter() + return reporter != nil && reporter.config.Enabled() +} + +func MaxBodyBytes() int64 { + reporter := DefaultReporter() + if reporter == nil { + return defaultMaxBody + } + return reporter.config.MaxBodyBytes +} + +func NodeName() string { + reporter := DefaultReporter() + if reporter == nil { + return "" + } + return reporter.config.NodeName +} + +func NewReporter(config Config, client *http.Client) *Reporter { + if config.Timeout <= 0 { + config.Timeout = defaultTimeout + } + if config.MaxBodyBytes < 0 { + config.MaxBodyBytes = 0 + } + if client == nil { + client = &http.Client{Timeout: config.Timeout} + } + return &Reporter{ + config: config, + client: client, + } +} + +func (r *Reporter) Report(event Event) error { + if r == nil || !r.config.Enabled() { + return nil + } + if event.Version == "" { + event.Version = ProtocolVersion + } + if event.Event == "" { + event.Event = EventRequestResponse + } + if event.Timestamp.IsZero() { + event.Timestamp = time.Now() + } + if event.Node == "" { + event.Node = r.config.NodeName + } + + payload, err := common.Marshal(event) + if err != nil { + return fmt.Errorf("marshal audit event: %w", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), r.config.Timeout) + defer cancel() + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, r.config.EndpointURL, bytes.NewReader(payload)) + if err != nil { + return fmt.Errorf("create audit request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+r.config.APIKey) + + resp, err := r.client.Do(req) + if err != nil { + return fmt.Errorf("post audit event: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return fmt.Errorf("audit endpoint returned status %d", resp.StatusCode) + } + return nil +} + +func ReportAsync(event Event) { + reporter := DefaultReporter() + if reporter == nil || !reporter.config.Enabled() { + return + } + go func() { + if err := reporter.Report(event); err != nil { + common.SysError("audit report failed: " + err.Error()) + } + }() +} diff --git a/service/audit/types.go b/service/audit/types.go new file mode 100644 index 00000000000..9cf03eaabbd --- /dev/null +++ b/service/audit/types.go @@ -0,0 +1,100 @@ +package audit + +import "time" + +const ( + ProtocolVersion = "new-api.audit.v1" + EventRequestResponse = "request_response" +) + +type Config struct { + EndpointURL string + APIKey string + Timeout time.Duration + MaxBodyBytes int64 + NodeName string +} + +func (c Config) Enabled() bool { + return c.EndpointURL != "" && c.APIKey != "" +} + +type Event struct { + Version string `json:"version"` + Event string `json:"event"` + RequestID string `json:"request_id,omitempty"` + Timestamp time.Time `json:"timestamp"` + Node string `json:"node,omitempty"` + Route string `json:"route,omitempty"` + User UserInfo `json:"user"` + Key KeyInfo `json:"key"` + Client ClientInfo `json:"client"` + Model ModelInfo `json:"model,omitempty"` + Billing BillingInfo `json:"billing,omitempty"` + Request Body `json:"request"` + Response Body `json:"response"` + DurationMS int64 `json:"duration_ms"` + Metadata map[string]string `json:"metadata,omitempty"` +} + +type UserInfo struct { + ID int `json:"id,omitempty"` + Username string `json:"username,omitempty"` +} + +type KeyInfo struct { + ID int `json:"id,omitempty"` + Name string `json:"name,omitempty"` +} + +type ClientInfo struct { + IP string `json:"ip,omitempty"` + Method string `json:"method,omitempty"` + Path string `json:"path,omitempty"` + UserAgent string `json:"user_agent,omitempty"` +} + +type Body struct { + ContentType string `json:"content_type,omitempty"` + Content string `json:"content"` + SizeBytes int64 `json:"size_bytes"` + Truncated bool `json:"truncated"` + StatusCode int `json:"status_code,omitempty"` +} + +type ModelInfo struct { + Name string `json:"name,omitempty"` + OriginName string `json:"origin_name,omitempty"` + UpstreamName string `json:"upstream_name,omitempty"` +} + +type BillingInfo struct { + Quota int `json:"quota,omitempty"` + PromptTokens int `json:"prompt_tokens,omitempty"` + CompletionTokens int `json:"completion_tokens,omitempty"` + TotalTokens int `json:"total_tokens,omitempty"` + CacheTokens int `json:"cache_tokens,omitempty"` + CacheCreationTokens int `json:"cache_creation_tokens,omitempty"` + CacheCreationTokens5m int `json:"cache_creation_tokens_5m,omitempty"` + CacheCreationTokens1h int `json:"cache_creation_tokens_1h,omitempty"` + ImageTokens int `json:"image_tokens,omitempty"` + AudioTokens int `json:"audio_tokens,omitempty"` + UseTimeSeconds int64 `json:"use_time_seconds,omitempty"` + ModelRatio float64 `json:"model_ratio,omitempty"` + GroupRatio float64 `json:"group_ratio,omitempty"` + CompletionRatio float64 `json:"completion_ratio,omitempty"` + CacheRatio float64 `json:"cache_ratio,omitempty"` + CacheCreationRatio float64 `json:"cache_creation_ratio,omitempty"` + CacheCreationRatio5m float64 `json:"cache_creation_ratio_5m,omitempty"` + CacheCreationRatio1h float64 `json:"cache_creation_ratio_1h,omitempty"` + ImageRatio float64 `json:"image_ratio,omitempty"` + ModelPrice float64 `json:"model_price,omitempty"` + UsePrice bool `json:"use_price,omitempty"` + FreeModel bool `json:"free_model,omitempty"` + QuotaToPreConsume int `json:"quota_to_pre_consume,omitempty"` + FinalPreConsumedQuota int `json:"final_pre_consumed_quota,omitempty"` + QuotaPerUnit float64 `json:"quota_per_unit,omitempty"` + UsageSemantic string `json:"usage_semantic,omitempty"` + OtherRatios map[string]float64 `json:"other_ratios,omitempty"` + Details map[string]any `json:"details,omitempty"` +} diff --git a/service/text_quota.go b/service/text_quota.go index 3f344dc3e57..b5cc123c9cd 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -13,6 +13,7 @@ import ( "github.com/QuantumNous/new-api/pkg/billingexpr" perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics" relaycommon "github.com/QuantumNous/new-api/relay/common" + auditservice "github.com/QuantumNous/new-api/service/audit" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/QuantumNous/new-api/types" @@ -458,6 +459,41 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us if tieredBillingApplied { InjectTieredBillingInfo(other, relayInfo, tieredResult) } + auditservice.SetModelInfo(ctx, auditservice.ModelInfo{ + Name: summary.ModelName, + OriginName: relayInfo.OriginModelName, + UpstreamName: relayInfo.UpstreamModelName, + }) + auditservice.SetBillingInfo(ctx, auditservice.BillingInfo{ + Quota: summary.Quota, + PromptTokens: summary.PromptTokens, + CompletionTokens: summary.CompletionTokens, + TotalTokens: summary.TotalTokens, + CacheTokens: summary.CacheTokens, + CacheCreationTokens: summary.CacheCreationTokens, + CacheCreationTokens5m: summary.CacheCreationTokens5m, + CacheCreationTokens1h: summary.CacheCreationTokens1h, + ImageTokens: summary.ImageTokens, + AudioTokens: summary.AudioTokens, + UseTimeSeconds: summary.UseTimeSeconds, + ModelRatio: summary.ModelRatio, + GroupRatio: summary.GroupRatio, + CompletionRatio: summary.CompletionRatio, + CacheRatio: summary.CacheRatio, + CacheCreationRatio: summary.CacheCreationRatio, + CacheCreationRatio5m: summary.CacheCreationRatio5m, + CacheCreationRatio1h: summary.CacheCreationRatio1h, + ImageRatio: summary.ImageRatio, + ModelPrice: summary.ModelPrice, + UsePrice: relayInfo.PriceData.UsePrice, + FreeModel: relayInfo.PriceData.FreeModel, + QuotaToPreConsume: relayInfo.PriceData.QuotaToPreConsume, + FinalPreConsumedQuota: relayInfo.FinalPreConsumedQuota, + QuotaPerUnit: common.QuotaPerUnit, + UsageSemantic: summary.UsageSemantic, + OtherRatios: relayInfo.PriceData.OtherRatios, + Details: other, + }) model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, From 9f3cfaaf29abb62bcc4249b8de835cb9966448d4 Mon Sep 17 00:00:00 2001 From: GeJiaXiang <353358601@qq.com> Date: Sun, 17 May 2026 01:02:53 +0800 Subject: [PATCH 2/2] =?UTF-8?q?feat:=20=E4=B8=8A=E6=8A=A5=E5=AF=B9?= =?UTF-8?q?=E8=AF=9D=E5=BD=92=E9=9B=86=E5=80=99=E9=80=89=E5=AD=97=E6=AE=B5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- middleware/audit.go | 90 ++++++++++++++++++-- middleware/audit_test.go | 174 +++++++++++++++++++++++++++++++++++++++ service/audit/types.go | 35 ++++---- 3 files changed, 278 insertions(+), 21 deletions(-) diff --git a/middleware/audit.go b/middleware/audit.go index 9486d76e74c..958de8680c3 100644 --- a/middleware/audit.go +++ b/middleware/audit.go @@ -2,13 +2,16 @@ package middleware import ( "io" + "regexp" "strconv" + "strings" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" auditservice "github.com/QuantumNous/new-api/service/audit" "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" ) type auditResponseWriter struct { @@ -72,12 +75,13 @@ func Audit() gin.HandlerFunc { Path: c.Request.URL.RequestURI(), UserAgent: c.Request.UserAgent(), }, - Model: modelInfo, - Billing: auditservice.BillingInfoFromContext(c), - Request: requestBody, - Response: responseBody, - DurationMS: time.Since(startedAt).Milliseconds(), - Metadata: auditMetadata(c), + Model: modelInfo, + Billing: auditservice.BillingInfoFromContext(c), + Request: requestBody, + Response: responseBody, + DurationMS: time.Since(startedAt).Milliseconds(), + Conversation: auditConversation(c, requestBody, responseBody), + Metadata: auditMetadata(c), }) c.Writer = originalWriter @@ -104,6 +108,80 @@ func captureRequestBody(c *gin.Context, maxBodyBytes int64) auditservice.Body { return auditservice.BodyFromBytes(data, c.Request.Header.Get("Content-Type"), maxBodyBytes) } +var auditSessionPattern = regexp.MustCompile(`_session_([A-Za-z0-9-]+)$`) + +func auditConversation(c *gin.Context, requestBody auditservice.Body, responseBody auditservice.Body) auditservice.ConversationInfo { + candidates := make(map[string]string) + put := func(key string, value string) { + if value != "" { + candidates[key] = value + } + } + + put("header.session_id", c.GetHeader("Session_id")) + put("header.x_amp_thread_id", c.GetHeader("X-Amp-Thread-Id")) + put("header.x_session_id", c.GetHeader("X-Session-ID")) + put("header.x_client_request_id", c.GetHeader("X-Client-Request-Id")) + + requestContent := requestBody.Content + put("body.conversation_id", gjson.Get(requestContent, "conversation_id").String()) + put("body.session_id", gjson.Get(requestContent, "session_id").String()) + put("body.previous_response_id", gjson.Get(requestContent, "previous_response_id").String()) + put("body.prompt_cache_key", gjson.Get(requestContent, "prompt_cache_key").String()) + put("body.metadata.conversation_id", gjson.Get(requestContent, "metadata.conversation_id").String()) + put("body.metadata.session_id", gjson.Get(requestContent, "metadata.session_id").String()) + userID := gjson.Get(requestContent, "metadata.user_id").String() + put("body.metadata.user_id", userID) + if matches := auditSessionPattern.FindStringSubmatch(userID); len(matches) >= 2 { + put("body.metadata.user_id_session", matches[1]) + } else if len(userID) > 0 && userID[0] == '{' { + put("body.metadata.user_id_session", gjson.Get(userID, "session_id").String()) + } + + responseID, messageID := auditResponseIDs(responseBody.Content) + put("response.id", responseID) + put("response.message_id", messageID) + + if len(candidates) == 0 { + return auditservice.ConversationInfo{} + } + return auditservice.ConversationInfo{Candidates: candidates} +} + +func auditResponseIDs(content string) (string, string) { + if id := gjson.Get(content, "id").String(); id != "" { + return id, "" + } + if id := gjson.Get(content, "response.id").String(); id != "" { + return id, "" + } + if id := gjson.Get(content, "message.id").String(); id != "" { + return "", id + } + + for _, line := range strings.Split(content, "\n") { + line = strings.TrimSpace(line) + if !strings.HasPrefix(line, "data:") { + continue + } + data := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if data == "" || data == "[DONE]" { + continue + } + if id := gjson.Get(data, "response.id").String(); id != "" { + return id, "" + } + if id := gjson.Get(data, "id").String(); id != "" { + return id, "" + } + if id := gjson.Get(data, "message.id").String(); id != "" { + return "", id + } + } + + return "", "" +} + func auditModelInfo(c *gin.Context) auditservice.ModelInfo { info := auditservice.ModelInfoFromContext(c) if info.Name == "" { diff --git a/middleware/audit_test.go b/middleware/audit_test.go index c6f7134b755..eee9ce87dd5 100644 --- a/middleware/audit_test.go +++ b/middleware/audit_test.go @@ -88,3 +88,177 @@ func TestAuditMiddlewareReportsRequestAndResponseByUserTokenAndIP(t *testing.T) t.Fatal("timed out waiting for audit event") } } + +func TestAuditMiddlewareReportsConversationEvidence(t *testing.T) { + gin.SetMode(gin.TestMode) + + events := make(chan auditservice.Event, 1) + auditServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var event auditservice.Event + if err := common.DecodeJson(r.Body, &event); err != nil { + t.Fatalf("failed to decode audit event: %v", err) + } + events <- event + w.WriteHeader(http.StatusAccepted) + })) + defer auditServer.Close() + + os.Setenv("AUDIT_ENDPOINT_URL", auditServer.URL) + os.Setenv("AUDIT_API_KEY", "test-key") + os.Setenv("AUDIT_TIMEOUT_SECONDS", "1") + os.Setenv("AUDIT_MAX_BODY_BYTES", "4096") + auditservice.InitFromEnv() + defer func() { + os.Unsetenv("AUDIT_ENDPOINT_URL") + os.Unsetenv("AUDIT_API_KEY") + os.Unsetenv("AUDIT_TIMEOUT_SECONDS") + os.Unsetenv("AUDIT_MAX_BODY_BYTES") + auditservice.InitFromEnv() + }() + + router := gin.New() + router.Use(Audit()) + router.POST("/v1/responses", func(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"id": "resp-current", "object": "response"}) + }) + + body := `{"conversation_id":"conv-body","session_id":"body-session","previous_response_id":"resp-prev","prompt_cache_key":"cache-key","metadata":{"conversation_id":"meta-conv","user_id":"user_x_account__session_11111111-2222-3333-4444-555555555555","session_id":"meta-session"}}` + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Session_id", "codex-session") + request.Header.Set("X-Amp-Thread-Id", "T-amp-thread") + request.Header.Set("X-Session-ID", "explicit-session") + request.Header.Set("X-Client-Request-Id", "client-request") + recorder := httptest.NewRecorder() + + router.ServeHTTP(recorder, request) + + select { + case event := <-events: + candidates := event.Conversation.Candidates + for key, expected := range map[string]string{ + "header.session_id": "codex-session", + "header.x_amp_thread_id": "T-amp-thread", + "header.x_session_id": "explicit-session", + "header.x_client_request_id": "client-request", + "body.conversation_id": "conv-body", + "body.session_id": "body-session", + "body.previous_response_id": "resp-prev", + "body.prompt_cache_key": "cache-key", + "body.metadata.conversation_id": "meta-conv", + "body.metadata.user_id": "user_x_account__session_11111111-2222-3333-4444-555555555555", + "body.metadata.user_id_session": "11111111-2222-3333-4444-555555555555", + "body.metadata.session_id": "meta-session", + "response.id": "resp-current", + } { + if candidates[key] != expected { + t.Fatalf("expected conversation candidate %s=%q, got %q in %+v", key, expected, candidates[key], candidates) + } + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for audit event") + } +} + +func TestAuditMiddlewareReportsConversationMessageIDFromSSE(t *testing.T) { + gin.SetMode(gin.TestMode) + + events := make(chan auditservice.Event, 1) + auditServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var event auditservice.Event + if err := common.DecodeJson(r.Body, &event); err != nil { + t.Fatalf("failed to decode audit event: %v", err) + } + events <- event + w.WriteHeader(http.StatusAccepted) + })) + defer auditServer.Close() + + os.Setenv("AUDIT_ENDPOINT_URL", auditServer.URL) + os.Setenv("AUDIT_API_KEY", "test-key") + os.Setenv("AUDIT_TIMEOUT_SECONDS", "1") + os.Setenv("AUDIT_MAX_BODY_BYTES", "4096") + auditservice.InitFromEnv() + defer func() { + os.Unsetenv("AUDIT_ENDPOINT_URL") + os.Unsetenv("AUDIT_API_KEY") + os.Unsetenv("AUDIT_TIMEOUT_SECONDS") + os.Unsetenv("AUDIT_MAX_BODY_BYTES") + auditservice.InitFromEnv() + }() + + router := gin.New() + router.Use(Audit()) + router.POST("/v1/messages", func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + _, _ = c.Writer.WriteString("event: message_start\n") + _, _ = c.Writer.WriteString("data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg-stream\"}}\n\n") + _, _ = c.Writer.WriteString("data: [DONE]\n\n") + }) + + request := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"messages":[{"role":"user","content":"hello"}]}`)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + + router.ServeHTTP(recorder, request) + + select { + case event := <-events: + if event.Conversation.Candidates["response.message_id"] != "msg-stream" { + t.Fatalf("expected SSE response.message_id evidence, got %+v", event.Conversation.Candidates) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for audit event") + } +} + +func TestAuditMiddlewareReportsConversationResponseIDFromSSE(t *testing.T) { + gin.SetMode(gin.TestMode) + + events := make(chan auditservice.Event, 1) + auditServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var event auditservice.Event + if err := common.DecodeJson(r.Body, &event); err != nil { + t.Fatalf("failed to decode audit event: %v", err) + } + events <- event + w.WriteHeader(http.StatusAccepted) + })) + defer auditServer.Close() + + os.Setenv("AUDIT_ENDPOINT_URL", auditServer.URL) + os.Setenv("AUDIT_API_KEY", "test-key") + os.Setenv("AUDIT_TIMEOUT_SECONDS", "1") + os.Setenv("AUDIT_MAX_BODY_BYTES", "4096") + auditservice.InitFromEnv() + defer func() { + os.Unsetenv("AUDIT_ENDPOINT_URL") + os.Unsetenv("AUDIT_API_KEY") + os.Unsetenv("AUDIT_TIMEOUT_SECONDS") + os.Unsetenv("AUDIT_MAX_BODY_BYTES") + auditservice.InitFromEnv() + }() + + router := gin.New() + router.Use(Audit()) + router.POST("/v1/responses", func(c *gin.Context) { + c.Header("Content-Type", "text/event-stream") + _, _ = c.Writer.WriteString("data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-stream\"}}\n\n") + _, _ = c.Writer.WriteString("data: [DONE]\n\n") + }) + + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"input":"hello"}`)) + request.Header.Set("Content-Type", "application/json") + recorder := httptest.NewRecorder() + + router.ServeHTTP(recorder, request) + + select { + case event := <-events: + if event.Conversation.Candidates["response.id"] != "resp-stream" { + t.Fatalf("expected SSE response.id evidence, got %+v", event.Conversation.Candidates) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for audit event") + } +} diff --git a/service/audit/types.go b/service/audit/types.go index 9cf03eaabbd..4947164902e 100644 --- a/service/audit/types.go +++ b/service/audit/types.go @@ -20,21 +20,26 @@ func (c Config) Enabled() bool { } type Event struct { - Version string `json:"version"` - Event string `json:"event"` - RequestID string `json:"request_id,omitempty"` - Timestamp time.Time `json:"timestamp"` - Node string `json:"node,omitempty"` - Route string `json:"route,omitempty"` - User UserInfo `json:"user"` - Key KeyInfo `json:"key"` - Client ClientInfo `json:"client"` - Model ModelInfo `json:"model,omitempty"` - Billing BillingInfo `json:"billing,omitempty"` - Request Body `json:"request"` - Response Body `json:"response"` - DurationMS int64 `json:"duration_ms"` - Metadata map[string]string `json:"metadata,omitempty"` + Version string `json:"version"` + Event string `json:"event"` + RequestID string `json:"request_id,omitempty"` + Timestamp time.Time `json:"timestamp"` + Node string `json:"node,omitempty"` + Route string `json:"route,omitempty"` + User UserInfo `json:"user"` + Key KeyInfo `json:"key"` + Client ClientInfo `json:"client"` + Model ModelInfo `json:"model,omitempty"` + Billing BillingInfo `json:"billing,omitempty"` + Request Body `json:"request"` + Response Body `json:"response"` + DurationMS int64 `json:"duration_ms"` + Conversation ConversationInfo `json:"conversation,omitempty"` + Metadata map[string]string `json:"metadata,omitempty"` +} + +type ConversationInfo struct { + Candidates map[string]string `json:"candidates,omitempty"` } type UserInfo struct {