57 lines
1.7 KiB
Go
57 lines
1.7 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import "context"
|
||
|
|
|
||
|
|
// HTTPUpstreamProfile marks HTTP upstream requests that need provider-specific
|
||
|
|
// transport policy.
|
||
|
|
type HTTPUpstreamProfile string
|
||
|
|
|
||
|
|
const (
|
||
|
|
HTTPUpstreamProfileDefault HTTPUpstreamProfile = ""
|
||
|
|
HTTPUpstreamProfileOpenAI HTTPUpstreamProfile = "openai"
|
||
|
|
)
|
||
|
|
|
||
|
|
type httpUpstreamProfileContextKey struct{}
|
||
|
|
type httpUpstreamDisableRedirectsContextKey struct{}
|
||
|
|
|
||
|
|
// WithHTTPUpstreamProfile injects an upstream transport profile into ctx.
|
||
|
|
func WithHTTPUpstreamProfile(ctx context.Context, profile HTTPUpstreamProfile) context.Context {
|
||
|
|
if ctx == nil {
|
||
|
|
ctx = context.Background()
|
||
|
|
}
|
||
|
|
if profile == HTTPUpstreamProfileDefault {
|
||
|
|
return ctx
|
||
|
|
}
|
||
|
|
return context.WithValue(ctx, httpUpstreamProfileContextKey{}, profile)
|
||
|
|
}
|
||
|
|
|
||
|
|
// HTTPUpstreamProfileFromContext resolves the upstream transport profile from ctx.
|
||
|
|
func HTTPUpstreamProfileFromContext(ctx context.Context) HTTPUpstreamProfile {
|
||
|
|
if ctx == nil {
|
||
|
|
return HTTPUpstreamProfileDefault
|
||
|
|
}
|
||
|
|
profile, ok := ctx.Value(httpUpstreamProfileContextKey{}).(HTTPUpstreamProfile)
|
||
|
|
if !ok {
|
||
|
|
return HTTPUpstreamProfileDefault
|
||
|
|
}
|
||
|
|
switch profile {
|
||
|
|
case HTTPUpstreamProfileOpenAI:
|
||
|
|
return profile
|
||
|
|
default:
|
||
|
|
return HTTPUpstreamProfileDefault
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithHTTPUpstreamRedirectsDisabled prevents credential-bearing probes from
|
||
|
|
// following redirects through the shared upstream client.
|
||
|
|
func WithHTTPUpstreamRedirectsDisabled(ctx context.Context) context.Context {
|
||
|
|
if ctx == nil {
|
||
|
|
ctx = context.Background()
|
||
|
|
}
|
||
|
|
return context.WithValue(ctx, httpUpstreamDisableRedirectsContextKey{}, true)
|
||
|
|
}
|
||
|
|
|
||
|
|
func HTTPUpstreamRedirectsDisabled(ctx context.Context) bool {
|
||
|
|
return ctx != nil && ctx.Value(httpUpstreamDisableRedirectsContextKey{}) == true
|
||
|
|
}
|