proxy

type ConfiguredHeader struct {
	// IsRequired, when true, causes the proxy to reject the request with a
	// bad-request error if the header cannot be resolved to a non-empty value.
	IsRequired bool

	// Name is the HTTP header name to send to the remote MCP server.
	Name string

	// StaticValue holds a fixed value (already decrypted if originally a
	// secret header). Leave empty when ValueFromRequestHeader is set.
	StaticValue string

	// ValueFromRequestHeader names a header on the inbound user request whose
	// value should be forwarded. Leave empty when StaticValue is set.
	ValueFromRequestHeader string
}

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

type Proxy struct {
	// GuardianPolicy is used to build a fresh, non-pooled HTTP client per
	// upstream request. Pooling is inappropriate here because each Proxy
	// instance handles a single upstream host and discards the connection
	// when the request finishes; a pooled transport would accumulate idle
	// connections across distinct Remote MCP Servers without ever reusing
	// them.
	GuardianPolicy *guardian.Policy
	Logger         *slog.Logger
	Tracer         trace.Tracer

	// NonStreamingTimeout bounds the connect+headers phase for every
	// upstream request, plus the body read for non-streaming responses.
	// For streaming (text/event-stream) responses this only bounds the
	// connect+headers phase; per-event idle bounds via StreamingTimeout
	// take over once headers are received. Callers must set an explicit
	// value; use [DefaultNonStreamingTimeout] for the package default.
	NonStreamingTimeout time.Duration

	// StreamingTimeout is the per-event idle bound applied to streaming
	// response bodies. The clock resets on every successful Read from
	// upstream, so an actively producing stream stays alive indefinitely
	// — only inactivity longer than this duration terminates the stream.
	// Callers must set an explicit value; use [DefaultStreamingTimeout]
	// for the package default.
	StreamingTimeout time.Duration

	// Metrics records per-request counters and histograms. Nil is safe and
	// disables metrics recording; tests that do not care about metrics pass
	// nil here.
	Metrics *Metrics
	// MaxBufferedBodyBytes bounds the size of any JSON body that is fully
	// read into memory before parsing — applied to both the inbound user
	// request and the upstream response. Streamed responses are not subject
	// to this cap; see [DefaultMaxBufferedBodyBytes] for the rationale.
	// Callers must set an explicit value; use [DefaultMaxBufferedBodyBytes]
	// for the package default.
	MaxBufferedBodyBytes int64

	// ServerID is the Remote MCP Server UUID. When set, it is attached to
	// every span emitted by the proxy so traces can be correlated across the
	// HTTP lifecycle and the proxy forward.
	ServerID string
	// RemoteURL is the upstream endpoint all requests are forwarded to.
	RemoteURL string
	// Headers are applied on top of any forwarded client headers when
	// constructing the upstream request.
	Headers []ConfiguredHeader

	// AuthorizationOverride is the Bearer token to set on the outgoing
	// Authorization header. The caller's incoming Authorization is
	// always dropped (Gram-issued credentials — API keys, OAuth tokens,
	// chat-session JWTs — are not meaningful upstream); when this field
	// is non-empty the proxy emits "Authorization: Bearer <override>"
	// instead. Use it for two flows:
	//
	//   - External OAuth: forward the caller's Bearer verbatim by
	//     setting this to the caller's own token (the upstream MCP
	//     server is the AS).
	//   - OAuth-proxy token swap: set this to a stored upstream
	//     credential resolved from the caller's Gram-issued OAuth
	//     token.
	//
	// Leave empty (default) to send no Authorization upstream.
	AuthorizationOverride string

	UserRequestInterceptors []UserRequestInterceptor

	// RemoteMessageInterceptors run for each JSON-RPC message arriving
	// from the remote MCP server: once per application/json POST response,
	// and once per parseable SSE event in a streamed response. Returning
	// a non-nil error blocks the message from being relayed to the user;
	// see [RemoteMessageInterceptor] for transport-specific rejection
	// semantics.
	RemoteMessageInterceptors []RemoteMessageInterceptor

	// ToolsCallRequestInterceptors run for inbound "tools/call" JSON-RPC
	// requests only, after the generic UserRequestInterceptors chain has
	// completed. Non-tools/call requests skip this loop entirely.
	ToolsCallRequestInterceptors []ToolsCallRequestInterceptor

	// ToolsCallResponseInterceptors run for "tools/call" JSON-RPC responses
	// only, after the generic RemoteMessageInterceptors chain has
	// completed. Dispatches from either the JSON response path or — for
	// SSE responses — when a terminal response event matching the
	// originating request ID is seen. Responses to non-tools/call requests
	// skip this loop entirely.
	ToolsCallResponseInterceptors []ToolsCallResponseInterceptor

	// ToolsListRequestInterceptors run for inbound "tools/list" JSON-RPC
	// requests only, after the generic UserRequestInterceptors chain has
	// completed. Non-tools/list requests skip this loop entirely.
	ToolsListRequestInterceptors []ToolsListRequestInterceptor

	// ToolsListResponseInterceptors run for "tools/list" JSON-RPC responses
	// only, after the generic RemoteMessageInterceptors chain has
	// completed. Dispatches from either the JSON response path or — for
	// SSE responses — when a terminal response event matching the
	// originating request ID is seen. Responses to non-tools/list requests
	// skip this loop entirely.
	ToolsListResponseInterceptors []ToolsListResponseInterceptor
}

type RejectError struct {
	// Code is the JSON-RPC error code carried back to the user. Use one of
	// the RejectCode* constants above when applicable.
	Code int

	// Message is the human-readable summary of the rejection. Surfaces in
	// the JSON-RPC error response's "message" field, so treat it as
	// user-facing. Interceptor authors should avoid putting internal
	// details (DB error strings, stack traces, secret-derived data) here;
	// log those via the interceptor's logger instead and pass a sanitized
	// summary as Message.
	Message string

	// Data is an optional structured payload. Surfaces in the JSON-RPC
	// error response's "data" field. Must be JSON-marshalable. Same
	// user-facing constraint as Message — sanitize before populating.
	Data any
}

type RemoteMessage struct {
	// UserHTTPRequest is the inbound user request that triggered this
	// message: the POST whose response carried the message, or the GET
	// that opened a server-initiated stream. Available so interceptors
	// can correlate messages back to their initiating request.
	UserHTTPRequest *http.Request

	// RemoteHTTPRequest is the outbound HTTP request the proxy built and
	// sent to the remote MCP server: the URL the proxy resolved to, the
	// method, and the headers the proxy applied (configured static and
	// secret headers, plus any forwarded user headers — minus the Gram
	// Authorization header, which is intentionally stripped). Available
	// so interceptors can inspect exactly what was sent on behalf of the
	// user.
	//
	// Prefer this over [http.Response.Request] on RemoteHTTPResponse:
	//
	//   - Redirects: the proxy uses the stdlib default redirect policy.
	//     If the remote responds with a 3xx, RemoteHTTPResponse.Request
	//     points to the LAST hop in the redirect chain, with potentially
	//     a different URL and a redirect-stripped header set.
	//     RemoteHTTPRequest is always the original outbound request the
	//     proxy built, regardless of redirects.
	//   - Transport mutation: [http.Transport] (and any wrapping
	//     transports such as otelhttp) may add or rewrite headers
	//     (Connection, Accept-Encoding, tracing headers) on the wire
	//     copy. RemoteHTTPRequest reflects the proxy's intent before
	//     transport-level mutation.
	//   - Body: per the stdlib contract, RemoteHTTPResponse.Request.Body
	//     is nil — the body has already been consumed.
	//
	// Interceptors must not read this request's body — it has already
	// been streamed upstream and the underlying reader is exhausted.
	RemoteHTTPRequest *http.Request

	// RemoteHTTPResponse is the upstream HTTP response. Available for
	// header inspection. Interceptors must not read its body — the proxy
	// has already consumed the byte stream and (for SSE) is mid-relay.
	RemoteHTTPResponse *http.Response

	// Message is the decoded JSON-RPC message — typically *jsonrpc.Request
	// for server-initiated requests, *jsonrpc.Response for responses (which
	// includes the terminal tools/call response in an SSE stream), or a
	// notification (a *jsonrpc.Request without an ID).
	Message jsonrpc.Message
}

type RemoteMessageInterceptor interface {
	// InterceptRemoteMessage is called once per JSON-RPC message from the
	// remote. Implementations may inspect msg and should return a non-nil
	// error to block the message from being relayed. The msg pointer is
	// freshly allocated per call and must not be retained past the call's
	// return.
	InterceptRemoteMessage(ctx context.Context, msg *RemoteMessage) error

	// Name returns a stable identifier for this interceptor, used for
	// tracing span attributes and log correlation.
	Name() string
}

type ToolsCallRequest struct {
	// Params is the decoded tools/call params. Arguments is retained as a
	// [json.RawMessage] so implementations can Unmarshal into tool-specific
	// argument schemas without a double-decode round-trip.
	Params *mcp.CallToolParamsRaw

	// UserRequest is the underlying request. Other interceptors in the generic
	// chain may already have observed it. Callers should prefer Params for
	// tools/call-specific data; UserRequest is exposed for RPC-level needs (JSON-RPC
	// ID, raw messages) and for forwarding control via the underlying HTTP
	// request.
	UserRequest *UserRequest
}

type ToolsCallRequestInterceptor interface {
	// InterceptToolsCallRequest is called with the parsed tools/call request.
	// Implementations may inspect call and should return a non-nil error to
	// reject the tool invocation; the interceptor's error is surfaced to the
	// user and the request is not forwarded to the remote server.
	InterceptToolsCallRequest(ctx context.Context, call *ToolsCallRequest) error

	// Name returns a stable identifier for this interceptor, used for tracing
	// span attributes and log correlation.
	Name() string
}

type ToolsCallResponse struct {
	// Error is the JSON-RPC protocol error when upstream returned an error
	// response (e.g. "tool not found", "method not found"). Mutually
	// exclusive with Result.
	Error *jsonrpc.Error

	// RemoteMessage is the underlying remote message. Other interceptors in the
	// generic chain may have observed it already.
	RemoteMessage *RemoteMessage

	// Request is the tools/call request this response is replying to.
	// Available so interceptors can correlate input and output without
	// re-parsing.
	Request *ToolsCallRequest

	// Result is the decoded tools/call result when upstream returned a
	// JSON-RPC success response. Check Result.IsError to distinguish
	// tool-level failures (the tool ran and reported an error) from tool-
	// level successes. Mutually exclusive with Error — exactly one of Result
	// and Error is non-nil.
	Result *mcp.CallToolResult
}

type ToolsCallResponseInterceptor interface {
	// InterceptToolsCallResponse is called with the parsed tools/call
	// response. Implementations may inspect call and should return a
	// non-nil error to reject the response; the interceptor's error is
	// surfaced to the user instead of the upstream payload.
	InterceptToolsCallResponse(ctx context.Context, call *ToolsCallResponse) error

	// Name returns a stable identifier for this interceptor, used for tracing
	// span attributes and log correlation.
	Name() string
}

type ToolsListRequest struct {
	// Params is the decoded tools/list params. Per the MCP spec, params may
	// be omitted entirely for tools/list; in that case Params is a
	// zero-valued [mcp.ListToolsParams].
	Params *mcp.ListToolsParams

	// UserRequest is the underlying request. Other interceptors in the generic
	// chain may already have observed it. Callers should prefer Params for
	// tools/list-specific data; UserRequest is exposed for RPC-level needs (JSON-RPC
	// ID, raw messages) and for forwarding control via the underlying HTTP
	// request.
	UserRequest *UserRequest
}

type ToolsListRequestInterceptor interface {
	// InterceptToolsListRequest is called with the parsed tools/list request.
	// Implementations may inspect list and should return a non-nil error to
	// reject the discovery request; the interceptor's error is surfaced to
	// the user and the request is not forwarded to the remote server.
	InterceptToolsListRequest(ctx context.Context, list *ToolsListRequest) error

	// Name returns a stable identifier for this interceptor, used for tracing
	// span attributes and log correlation.
	Name() string
}

type ToolsListResponse struct {
	// Error is the JSON-RPC protocol error when upstream returned an error
	// response (e.g. "method not found"). Mutually exclusive with Result.
	Error *jsonrpc.Error

	// RemoteMessage is the underlying remote message. Other interceptors in the
	// generic chain may have observed it already.
	RemoteMessage *RemoteMessage

	// Request is the tools/list request this response is replying to.
	// Available so interceptors can correlate input and output without
	// re-parsing.
	Request *ToolsListRequest

	// Result is the decoded tools/list result when upstream returned a
	// JSON-RPC success response. Mutually exclusive with Error — exactly one
	// of Result and Error is non-nil.
	Result *mcp.ListToolsResult
}

type ToolsListResponseInterceptor interface {
	// InterceptToolsListResponse is called with the parsed tools/list
	// response. Implementations may inspect list and should return a
	// non-nil error to reject the response; the interceptor's error is
	// surfaced to the user instead of the upstream payload.
	InterceptToolsListResponse(ctx context.Context, list *ToolsListResponse) error

	// Name returns a stable identifier for this interceptor, used for tracing
	// span attributes and log correlation.
	Name() string
}

type UserRequest struct {
	UserHTTPRequest *http.Request
	JSONRPCMessages []jsonrpc.Message
	// contains filtered or unexported fields
}

type UserRequestInterceptor interface {
	// InterceptUserRequest is called with the parsed inbound request.
	// Implementations may inspect req and should return a non-nil error to
	// reject the request; the interceptor's error is rendered as a
	// JSON-RPC error envelope back to the user and the request is not
	// forwarded to the remote server.
	InterceptUserRequest(ctx context.Context, req *UserRequest) error

	// Name returns a stable identifier for this interceptor, used for tracing
	// span attributes and log correlation.
	Name() string
}