aiproxy

package
v1.3.11 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Oct 12, 2026 License: Apache-2.0 Imports: 26 Imported by: 0

README

aiproxy

Authenticated reverse proxy for OpenAI-compatible APIs and MCP servers (streamable HTTP). Upstream URLs and keys stay server-side; clients only use the proxy with normal ResolveSpec auth.

Setup

p := aiproxy.New(aiproxy.Config{Prefix: "/ai"})
p.RegisterOpenAI(aiproxy.OpenAIUpstream{Upstream: aiproxy.Upstream{Name: "gpt", BaseURL: "https://api.openai.com/v1", APIKey: key}})
p.RegisterMCP(aiproxy.MCPUpstream{Upstream: aiproxy.Upstream{Name: "tools", BaseURL: "http://localhost:3000/mcp", APIKey: key}})
mux.Handle("/ai/", p.Handler(security.NewAuthMiddleware(securityList)))

Routes

Kind Client URL Upstream URL
OpenAI {prefix}/{name}/v1/* BaseURL + path after /v1 (BaseURL includes the version)
MCP {prefix}/{name} BaseURL as is

Client SDK base URL: https://host/ai/gpt/v1 (OpenAI), https://host/ai/tools (MCP). Query string, Mcp-Session-Id, SSE streaming all pass through.

Upstream fields

Field Notes
Name route segment, [A-Za-z0-9_.-], unique across kinds
BaseURL http/https
APIKey optional; never logged/serialized (redacted in String/JSON)
AuthHeader default Authorization
AuthFormat one %s; default Bearer %s for Authorization, else raw key
Headers static extra upstream headers
AllowedRoles any-of; empty = any authenticated user
RateLimit {PerSecond, Burst} per user per upstream; 429 + Retry-After
Timeout response-header timeout, default 120s (streams not cut)
AllowedModels (OpenAI) JSON body model must match; non-JSON writes rejected
AllowedTools (MCP) tools/call name must match (batches checked); tools/list not filtered

Security behavior

  • Auth required; guests rejected (401). Roles checked (403).
  • Client Authorization, Cookie, Proxy-Authorization, X-Api-Key (+ Config.StripHeaders) removed; key injected.
  • Response Set-Cookie, Www-Authenticate, redirect Location removed.
  • Upstream 401 becomes 502 (proxy credentials problem, not the client's); upstream errors hide URL.
  • JSON bodies inspected up to Config.MaxBodyBytes (10 MB), else 413.
  • Fails closed if Handler(nil).

Hooks (p.Hooks())

Hook When Context
BeforeProxy after auth/limit/inspection Request (headers editable), UserContext, Upstream, Kind, Model, Tools; abort via Abort/AbortMessage/AbortCode (default 403)
AfterProxy response finished/failed StatusCode, Duration, Usage, Error

Usage (OpenAI): from usage / response.usage; streams need stream_options.include_usage.

Audit

Config.Audit (AuditSink) gets one AuditRecord per handled request: user, upstream, method, client path (no query), model/tools, status, outcome (ok|upstream_error|denied|rate_limited), reason, duration, usage. No bodies, no keys. Default LogAuditSink (pkg/logger); NopAuditSink{} disables. Sinks run on the request path: keep them fast.

Metrics

Used automatically when the global metrics.Provider implements metrics.AIProxyRecorder (PrometheusProvider does): aiproxy_requests_total{upstream,kind,model,status,outcome}, aiproxy_request_duration_seconds{upstream,kind}, aiproxy_tokens_total{upstream,model,type}. model is client-supplied: bounded to 400 distinct pairs, then other.

Dynamic upstreams

Call Effect
Register(Definition) / Unregister(name) runtime add/remove
Reload(ctx, store) sync store-managed upstreams: add, replace changed, drop missing; unchanged keep rate-limit state
AutoReload(ctx, store, interval) Reload now, then periodically
NewProcStore(*sql.DB, proc) / NewProcStoreFromDatabase(common.Database, proc) calls stored procedure (PostgreSQL), default resolvespec_ai_proxies
  • Code-registered upstreams are never replaced/removed by Reload; a same-name stored row is skipped and reported.
  • Invalid rows are skipped (valid ones applied; last good definition kept); store failure changes nothing.
  • Procedure contract (like the security procedures): SELECT p_success, p_error, p_data FROM resolvespec_ai_proxies(); p_data = JSON array of {name, kind, base_url, api_key, auth_header, auth_format, headers, allowed_roles, allowed, rate_per_second, rate_burst, timeout_seconds, enabled}. Reference table + function: resolvespec_ai_proxies.sql (replace the body to source from anywhere).
  • Entries with enabled:false or undecodable JSON are skipped.
  • API keys come back in plaintext from the procedure: restrict execute rights on it.

Notes

  • Legacy MCP SSE transport is not supported.
  • GET /v1/models is not filtered by AllowedModels.

Documentation

Overview

Package aiproxy is an authenticated reverse proxy for OpenAI-compatible APIs and MCP servers (streamable HTTP). Upstream URLs and keys stay on the server; clients authenticate with the normal security layer and call the proxy instead.

p := aiproxy.New(aiproxy.Config{Prefix: "/ai"})
p.RegisterOpenAI(aiproxy.OpenAIUpstream{Upstream: aiproxy.Upstream{Name: "gpt", BaseURL: "https://api.openai.com/v1", APIKey: key}})
p.RegisterMCP(aiproxy.MCPUpstream{Upstream: aiproxy.Upstream{Name: "tools", BaseURL: "http://localhost:3000/mcp", APIKey: key}})
mux.Handle("/ai/", p.Handler(security.NewAuthMiddleware(securityList)))

Client routes: {prefix}/{name}/v1/... (OpenAI) and {prefix}/{name} (MCP).

Index

Constants

View Source
const DefaultProc = "resolvespec_ai_proxies"

DefaultProc is the stored procedure (function) ProcStore calls when no name is given.

Variables

This section is empty.

Functions

This section is empty.

Types

type AuditRecord

type AuditRecord struct {
	Time     time.Time
	UserID   int
	User     string
	RemoteID string
	Upstream string
	Kind     Kind
	Method   string
	Path     string // client path, no query string
	Model    string
	Tools    []string
	Status   int
	Outcome  Outcome
	Reason   string // why a request was refused
	Duration time.Duration
	Usage    Usage
	Error    string
}

AuditRecord is written once per handled request. It holds no bodies and no keys.

type AuditSink

type AuditSink interface {
	Record(AuditRecord)
}

AuditSink receives audit records. It runs on the request path, so keep it fast (queue slow work such as database writes).

type Config

type Config struct {
	// Prefix the handler is mounted under, e.g. "/ai". Empty when mounted with http.StripPrefix.
	Prefix string
	// MaxBodyBytes caps JSON request bodies that are inspected. Default 10 MB.
	MaxBodyBytes int64
	// StripHeaders are removed from client requests in addition to the defaults
	// (Authorization, Cookie, Proxy-Authorization, X-Api-Key).
	StripHeaders []string
	// Audit receives one record per request. Default: LogAuditSink. Use NopAuditSink to disable.
	Audit AuditSink
}

Config configures a Proxy.

type Definition

type Definition struct {
	Kind     Kind
	Upstream Upstream
	// Allowed holds AllowedModels (OpenAI) or AllowedTools (MCP).
	Allowed []string
}

Definition describes one upstream independent of how it was registered.

type HookContext

type HookContext struct {
	Context     context.Context
	Request     *http.Request // BeforeProxy may modify headers/query; the key is injected later
	UserContext *security.UserContext

	Upstream string
	Kind     Kind
	Model    string   // OpenAI: requested model
	Tools    []string // MCP: tools named by tools/call

	// AfterProxy only
	StatusCode int
	Duration   time.Duration
	Usage      Usage // OpenAI only; needs stream_options.include_usage for streams
	Error      error

	// BeforeProxy only
	Abort        bool
	AbortMessage string
	AbortCode    int // default 403
}

HookContext is passed to hooks.

type HookFunc

type HookFunc func(*HookContext) error

HookFunc is a hook. A returned error aborts a BeforeProxy request.

type HookRegistry

type HookRegistry struct {
	// contains filtered or unexported fields
}

HookRegistry holds hooks per type.

func NewHookRegistry

func NewHookRegistry() *HookRegistry

NewHookRegistry creates an empty registry.

func (*HookRegistry) Execute

func (r *HookRegistry) Execute(t HookType, ctx *HookContext) error

Execute runs the hooks in order and stops at the first error or abort.

func (*HookRegistry) Register

func (r *HookRegistry) Register(t HookType, h HookFunc)

Register adds a hook.

type HookType

type HookType string

HookType defines when a hook runs.

const (
	// BeforeProxy runs after auth, rate limit and request inspection, before the upstream call.
	BeforeProxy HookType = "before_proxy"
	// AfterProxy runs once the response finished (or failed). Errors are logged only.
	AfterProxy HookType = "after_proxy"
)

type Kind

type Kind string

Kind is the type of an upstream.

const (
	KindOpenAI Kind = "openai"
	KindMCP    Kind = "mcp"
)

type LogAuditSink

type LogAuditSink struct{}

LogAuditSink writes records through pkg/logger. It is the default sink.

func (LogAuditSink) Record

func (LogAuditSink) Record(r AuditRecord)

Record implements AuditSink.

type MCPUpstream

type MCPUpstream struct {
	Upstream
	// AllowedTools restricts tools/call to these tool names. Empty means any.
	AllowedTools []string
}

MCPUpstream is an MCP server speaking streamable HTTP.

type NopAuditSink

type NopAuditSink struct{}

NopAuditSink discards records (use it to switch the default logging off).

func (NopAuditSink) Record

func (NopAuditSink) Record(AuditRecord)

Record implements AuditSink.

type OpenAIUpstream

type OpenAIUpstream struct {
	Upstream
	// AllowedModels restricts the "model" a request may use. Empty means any.
	// When set, write requests must be JSON bodies carrying an allowed model.
	AllowedModels []string
}

OpenAIUpstream is an OpenAI-compatible API (OpenAI, Azure, Ollama, vLLM, ...).

type Outcome

type Outcome string

Outcome classifies a handled request.

const (
	OutcomeOK            Outcome = "ok"             // upstream answered < 400
	OutcomeUpstreamError Outcome = "upstream_error" // upstream answered >= 400 or failed
	OutcomeDenied        Outcome = "denied"         // refused by auth, roles, allowlist, size or a hook
	OutcomeRateLimited   Outcome = "rate_limited"
)

type ProcStore

type ProcStore struct {
	// contains filtered or unexported fields
}

ProcStore loads upstream definitions from a stored procedure (PostgreSQL), using the same contract as the security procedures:

SELECT p_success, p_error, p_data FROM resolvespec_ai_proxies()

p_data is a JSON array; each element has: name, kind ("openai"|"mcp"), base_url, and optionally api_key, auth_header, auth_format, headers (object), allowed_roles (array), allowed (array: models or tools), rate_per_second, rate_burst, timeout_seconds, enabled (default true). See resolvespec_ai_proxies.sql.

func NewProcStore

func NewProcStore(db *sql.DB, proc string) (*ProcStore, error)

NewProcStore calls proc through db. An empty proc means DefaultProc.

func NewProcStoreFromDatabase

func NewProcStoreFromDatabase(db common.Database, proc string) (*ProcStore, error)

NewProcStoreFromDatabase is NewProcStore for an application's common.Database. The connection is fetched on every call, so adapter reconnects are followed.

func (*ProcStore) List

func (s *ProcStore) List(ctx context.Context) ([]Definition, error)

List implements UpstreamStore. Disabled and undecodable entries are skipped.

type Proxy

type Proxy struct {
	// contains filtered or unexported fields
}

Proxy routes authenticated requests to registered upstreams.

func New

func New(cfg Config) *Proxy

New creates a Proxy.

func (*Proxy) AutoReload

func (p *Proxy) AutoReload(ctx context.Context, store UpstreamStore, interval time.Duration) error

AutoReload calls Reload now and then every interval until ctx ends. The first error is returned; later ones are logged.

func (*Proxy) Handler

func (p *Proxy) Handler(auth func(http.Handler) http.Handler) http.Handler

Handler returns the proxy handler. auth is the authentication middleware, normally security.NewAuthMiddleware(securityList). It is required: without it every request is refused.

func (*Proxy) Hooks

func (p *Proxy) Hooks() *HookRegistry

Hooks returns the hook registry.

func (*Proxy) Names

func (p *Proxy) Names() []string

Names lists registered upstream names (no URLs or keys).

func (*Proxy) Register

func (p *Proxy) Register(d Definition) error

Register registers an upstream from a Definition. The name must be unused.

func (*Proxy) RegisterMCP

func (p *Proxy) RegisterMCP(u MCPUpstream) error

RegisterMCP registers an MCP server (streamable HTTP).

func (*Proxy) RegisterOpenAI

func (p *Proxy) RegisterOpenAI(u OpenAIUpstream) error

RegisterOpenAI registers an OpenAI-compatible upstream.

func (*Proxy) Reload

func (p *Proxy) Reload(ctx context.Context, store UpstreamStore) error

Reload syncs the store-managed upstreams with the store.

  • New definitions are added, changed ones replaced, unchanged ones left alone (their rate-limit state and connections survive), missing ones removed.
  • Upstreams registered in code are never replaced or removed; a stored definition with the same name is skipped and reported.
  • Invalid definitions are skipped and reported in the returned error; valid ones are still applied. If the store itself fails, nothing changes.

func (*Proxy) Unregister

func (p *Proxy) Unregister(name string) bool

Unregister removes an upstream (code-registered or loaded). It reports whether it existed. In-flight requests finish normally.

type RateLimit

type RateLimit struct {
	PerSecond float64 // sustained requests per second
	Burst     int     // bucket size; defaults to ceil(PerSecond), min 1
}

RateLimit limits requests per user per upstream (token bucket).

type Upstream

type Upstream struct {
	// Name is the route segment clients use: {prefix}/{Name}/...
	Name string
	// BaseURL of the upstream. OpenAI: include the version path (https://api.openai.com/v1).
	// MCP: the full streamable HTTP endpoint (http://localhost:3000/mcp).
	BaseURL string
	// APIKey is sent to the upstream. Never exposed to clients or logs.
	APIKey string
	// AuthHeader carries the key. Default "Authorization". Applies to OpenAI and MCP upstreams.
	AuthHeader string
	// AuthFormat renders the header value; %s is the key (e.g. "Bearer %s", "Token %s", "%s").
	// Default: "Bearer %s" for the Authorization header, "%s" (raw key) for any other header.
	AuthFormat string
	// Headers are extra static headers sent to the upstream (applied before the key header).
	// Per-request values can be set in a BeforeProxy hook.
	Headers map[string]string
	// AllowedRoles: caller needs at least one. Empty means any authenticated user.
	AllowedRoles []string
	// RateLimit per user for this upstream. Nil means unlimited.
	RateLimit *RateLimit
	// Timeout waiting for the upstream response headers. Default 120s. Streams are not cut off.
	Timeout time.Duration
}

Upstream holds the settings shared by every upstream kind.

func (Upstream) GoString

func (u Upstream) GoString() string

GoString redacts the key (%#v).

func (Upstream) MarshalJSON

func (u Upstream) MarshalJSON() ([]byte, error)

MarshalJSON redacts the key.

func (Upstream) String

func (u Upstream) String() string

String redacts the key.

type UpstreamStore

type UpstreamStore interface {
	List(ctx context.Context) ([]Definition, error)
}

UpstreamStore supplies upstream definitions for Reload.

type Usage

type Usage struct {
	PromptTokens     int64
	CompletionTokens int64
	TotalTokens      int64
}

Usage is the token usage reported by an OpenAI upstream.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL