67c173cbcb
Muse now forwards business context alongside each per-user token request, so new-api needs to ingest those tags for relay selection, billing audit, and log correlation without breaking its own internal request-id semantics. This adds a dedicated Muse request-context middleware, carries Muse tags into RelayInfo, and appends them into consume/error log payloads. Muse request ids are stored as separate business metadata instead of overwriting the internal request id used for pre-consume idempotency. Constraint: The relay layer must keep working for both /v1 and /v1beta request paths Constraint: Client-supplied Muse request ids cannot replace the internal request id used by new-api billing/idempotency logic Rejected: Overwriting common.RequestIdKey with X-Request-Id | mixes external business ids with internal idempotency ids Rejected: Adding new log table columns in phase 1 | this task only needs relay/log propagation, so other JSON is the minimal compatible path Confidence: medium Scope-risk: moderate Reversibility: clean Directive: Keep Muse business identifiers in dedicated context keys and log other fields; do not reuse internal request-id slots for external correlation ids Tested: go test ./middleware ./relay/common ./controller -run MuseRequestContext -count=1 Tested: go test ./middleware ./relay/common ./controller -count=1 Not-tested: Live relay call from Muse into a running new-api instance
114 lines
2.7 KiB
Go
114 lines
2.7 KiB
Go
package middleware
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
const (
|
|
museHeaderUserID = "X-Muse-User-Id"
|
|
museHeaderWorkID = "X-Muse-Work-Id"
|
|
museHeaderRequest = "X-Request-Id"
|
|
museHeaderTraceID = "X-Trace-Id"
|
|
museHeaderScene = "X-Muse-Scene"
|
|
jsonContentType = "application/json"
|
|
)
|
|
|
|
type museMetadataEnvelope struct {
|
|
Metadata map[string]any `json:"metadata"`
|
|
}
|
|
|
|
func MuseRequestContext() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
metadata := extractMuseMetadata(c)
|
|
|
|
museUserID := firstNonBlank(
|
|
c.GetHeader(museHeaderUserID),
|
|
stringifyMetadata(metadata["muse_user_id"]),
|
|
)
|
|
museWorkID := firstNonBlank(
|
|
c.GetHeader(museHeaderWorkID),
|
|
stringifyMetadata(metadata["work_id"]),
|
|
stringifyMetadata(metadata["muse_work_id"]),
|
|
)
|
|
museScene := firstNonBlank(
|
|
c.GetHeader(museHeaderScene),
|
|
stringifyMetadata(metadata["scene"]),
|
|
)
|
|
requestID := firstNonBlank(
|
|
c.GetHeader(museHeaderRequest),
|
|
stringifyMetadata(metadata["request_id"]),
|
|
)
|
|
traceID := firstNonBlank(
|
|
c.GetHeader(museHeaderTraceID),
|
|
stringifyMetadata(metadata["trace_id"]),
|
|
)
|
|
|
|
if museUserID != "" {
|
|
common.SetContextKey(c, constant.ContextKeyMuseUserID, museUserID)
|
|
}
|
|
if museWorkID != "" {
|
|
common.SetContextKey(c, constant.ContextKeyMuseWorkID, museWorkID)
|
|
}
|
|
if requestID != "" {
|
|
common.SetContextKey(c, constant.ContextKeyMuseRequestID, requestID)
|
|
}
|
|
if museScene != "" {
|
|
common.SetContextKey(c, constant.ContextKeyMuseScene, museScene)
|
|
}
|
|
if traceID != "" {
|
|
common.SetContextKey(c, constant.ContextKeyMuseTraceID, traceID)
|
|
}
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
func extractMuseMetadata(c *gin.Context) map[string]any {
|
|
if c == nil || c.Request == nil {
|
|
return nil
|
|
}
|
|
if !strings.Contains(strings.ToLower(c.GetHeader("Content-Type")), jsonContentType) {
|
|
return nil
|
|
}
|
|
if !needsMuseMetadataRead(c) {
|
|
return nil
|
|
}
|
|
var envelope museMetadataEnvelope
|
|
if err := common.UnmarshalBodyReusable(c, &envelope); err != nil {
|
|
return nil
|
|
}
|
|
return envelope.Metadata
|
|
}
|
|
|
|
func needsMuseMetadataRead(c *gin.Context) bool {
|
|
return c.GetHeader(museHeaderUserID) == "" ||
|
|
c.GetHeader(museHeaderWorkID) == "" ||
|
|
c.GetHeader(museHeaderRequest) == "" ||
|
|
c.GetHeader(museHeaderTraceID) == "" ||
|
|
c.GetHeader(museHeaderScene) == ""
|
|
}
|
|
|
|
func firstNonBlank(values ...string) string {
|
|
for _, value := range values {
|
|
if strings.TrimSpace(value) != "" {
|
|
return strings.TrimSpace(value)
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func stringifyMetadata(value any) string {
|
|
switch v := value.(type) {
|
|
case nil:
|
|
return ""
|
|
case string:
|
|
return strings.TrimSpace(v)
|
|
default:
|
|
return strings.TrimSpace(fmt.Sprintf("%v", v))
|
|
}
|
|
}
|