model-gateway/common/util/headers.go

package util

import (
	"context"

	"gitea.com/red-future/common/utils"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/util/gconv"
)

// AsyncCtx 固化异步上下文中的 token 和用户信息，避免请求结束后丢失
func AsyncCtx(ctx context.Context) context.Context {
	asyncCtx := context.WithoutCancel(ctx)

	if r := g.RequestFromCtx(ctx); r != nil {
		if token := r.Header.Get("Authorization"); token != "" {
			asyncCtx = context.WithValue(asyncCtx, "token", token)
		}
		if userInfo := r.Header.Get("X-User-Info"); userInfo != "" {
			asyncCtx = context.WithValue(asyncCtx, "xUserInfo", userInfo)
		}
	}

	if user, err := utils.GetUserInfo(ctx); err == nil && user != nil {
		asyncCtx = context.WithValue(asyncCtx, "user", user)
	}

	return asyncCtx
}

// ForwardHeaders 透传调用链路的头信息，优先使用 ctx 中的固化值
func ForwardHeaders(ctx context.Context) map[string]string {
	headers := make(map[string]string)
	SetHeaderFromContext(headers, ctx, "Authorization", "token")
	SetHeaderFromContext(headers, ctx, "X-User-Info", "xUserInfo")
	FallbackToRequestHeaders(headers, ctx)
	return headers
}

// SetHeaderFromContext 从上下文中设置 header
func SetHeaderFromContext(headers map[string]string, ctx context.Context, headerKey, ctxKey string) {
	if value, ok := ctx.Value(ctxKey).(string); ok && value != "" {
		headers[headerKey] = value
	}
}

// FallbackToRequestHeaders 从请求头中获取作为兜底
func FallbackToRequestHeaders(headers map[string]string, ctx context.Context) {
	r := g.RequestFromCtx(ctx)
	if r == nil {
		return
	}

	if headers["Authorization"] == "" {
		if token := r.Header.Get("Authorization"); token != "" {
			headers["Authorization"] = token
		}
	}

	if headers["X-User-Info"] == "" {
		if userInfo := r.Header.Get("X-User-Info"); userInfo != "" {
			headers["X-User-Info"] = userInfo
		}
	}
}

// SetTaskHeadersToCtx 把任务入库时保存的 header 信息注入 ctx，给 worker 调 OSS 用
func SetTaskHeadersToCtx(ctx context.Context, headers map[string]string) context.Context {
	if headers == nil {
		return ctx
	}
	if v := gconv.String(headers["Authorization"]); v != "" {
		ctx = context.WithValue(ctx, "token", v)
	}
	if v := gconv.String(headers["X-User-Info"]); v != "" {
		ctx = context.WithValue(ctx, "xUserInfo", v)
	}
	return ctx
}

// ParseStoredPayload 解析入库的 request_payload，拆出模型调用核心数据
func ParseStoredPayload(v map[string]any) map[string]any {
	if v == nil {
		return nil
	}
	if p, ok := v["payload"]; ok {
		return gconv.Map(p)
	}
	return v
}