aibridged

package
v2.34.1 Latest Latest
Warning

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

 Go to latest Published: Jun 8, 2026 License: AGPL-3.0 Imports: 38 Imported by: 0

Documentation

Index

Constants

View Source 
const (
	InternalMCPServerID = "coder"
)

Variables

View Source 
var (
	ErrNoAuthKey             = xerrors.New("no authentication key provided")
	ErrConnect               = xerrors.New("could not connect to coderd")
	ErrUnauthorized          = xerrors.New("unauthorized")
	ErrAcquireRequestHandler = xerrors.New("failed to acquire request handler")
)
View Source 
var (
	ErrEmptyConfig  = xerrors.New("empty config given")
	ErrCompileRegex = xerrors.New("compile tool regex")
)
View Source 
var DefaultPoolOptions = PoolOptions{MaxItems: 5000, TTL: time.Minute * 15}

Functions

func NewTransportFactory 

func NewTransportFactory(handler http.Handler) aibridge.TransportFactory

NewTransportFactory returns an [aibridge.TransportFactory] whose RoundTripper dispatches requests to handler in-process, streaming the response body through an io.Pipe so SSE/NDJSON/chunked responses propagate token-by-token just as they would over the wire.

handler is typically the aibridged HTTP entrypoint registered via [API.RegisterInMemoryAIBridgedHTTPHandler].

func SubscribeProviderReload 

func SubscribeProviderReload(
	ctx context.Context,
	ps dbpubsub.Pubsub,
	reloader ProviderReloader,
	logger slog.Logger,
) (func(), error)

SubscribeProviderReload refreshes once, then on AI provider changes.

func WriteProviderInfoSnapshot 

func WriteProviderInfoSnapshot(info *prometheus.GaugeVec, outcomes []ProviderOutcome)

WriteProviderInfoSnapshot Resets info and writes one series per outcome. Both aibridged and aibridgeproxyd use this so the provider_info recording contract stays in one place.

Types

type CachedBridgePool 

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

func NewCachedBridgePool 

func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, logger slog.Logger, metrics *aibridge.Metrics, tracer trace.Tracer) (*CachedBridgePool, error)

func (*CachedBridgePool) Acquire 

func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn ClientFunc, mcpProxyFactory MCPProxyBuilder) (_ http.Handler, outErr error)

Acquire retrieves or creates a [*aibridge.RequestBridge] instance per given key.

Each returned [*aibridge.RequestBridge] is safe for concurrent use. Each [*aibridge.RequestBridge] is stateful because it has MCP clients which maintain sessions to the configured MCP server.

func (*CachedBridgePool) CacheMetrics 

func (p *CachedBridgePool) CacheMetrics() PoolMetrics

func (*CachedBridgePool) ReplaceProviders 

func (p *CachedBridgePool) ReplaceProviders(providers []aibridge.Provider)

ReplaceProviders swaps the provider snapshot used by future Acquires. It is safe to call concurrently with Acquire and is a no-op after Shutdown.

func (*CachedBridgePool) Shutdown 

func (p *CachedBridgePool) Shutdown(_ context.Context) error

Shutdown will close the cache which will trigger eviction of all the Bridge entries.

type Client 

type Client struct {
	proto.DRPCRecorderClient
	proto.DRPCMCPConfiguratorClient
	proto.DRPCAuthorizerClient

	Conn drpc.Conn
}

func (*Client) DRPCConn 

func (c *Client) DRPCConn() drpc.Conn

type ClientFunc 

type ClientFunc func() (DRPCClient, error)

type DRPCClient 

type DRPCClient interface {
	proto.DRPCRecorderClient
	proto.DRPCMCPConfiguratorClient
	proto.DRPCAuthorizerClient
}

DRPCClient is the union of various service interfaces the client must support.

type DRPCServer 

type DRPCServer interface {
	proto.DRPCRecorderServer
	proto.DRPCMCPConfiguratorServer
	proto.DRPCAuthorizerServer
}

type Dialer 

type Dialer func(ctx context.Context) (DRPCClient, error)

type MCPProxyBuilder deprecated 

type MCPProxyBuilder interface {
	// Build creates a [mcp.ServerProxier] for the given request initiator.
	// At minimum, the Coder MCP server will be proxied.
	// The SessionKey from [Request] is used to authenticate against the Coder MCP server.
	//
	// NOTE: the [mcp.ServerProxier] instance may be proxying one or more MCP servers.
	Build(ctx context.Context, req Request, tracer trace.Tracer) (mcp.ServerProxier, error)
}

Deprecated: Injected MCP in AI Bridge is deprecated and will be removed in a future release.

type MCPProxyFactory deprecated 

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

Deprecated: Injected MCP in AI Bridge is deprecated and will be removed in a future release.

func NewMCPProxyFactory 

func NewMCPProxyFactory(logger slog.Logger, tracer trace.Tracer, clientFn ClientFunc) *MCPProxyFactory

func (*MCPProxyFactory) Build 

func (m *MCPProxyFactory) Build(ctx context.Context, req Request, tracer trace.Tracer) (mcp.ServerProxier, error)

type Metrics 

type Metrics struct {

	// ProviderInfo is one series per configured provider; value is
	// always 1 and the status label carries the alertable signal.
	// Labels: provider_name, provider_type, status.
	ProviderInfo *prometheus.GaugeVec

	// ProvidersLastReloadTimestampSeconds is the unix timestamp of the
	// last reload attempt, success or failure.
	ProvidersLastReloadTimestampSeconds prometheus.Gauge

	// ProvidersLastReloadSuccessTimestampSeconds is the unix timestamp
	// of the last reload that successfully refreshed the pool. A gap
	// against ProvidersLastReloadTimestampSeconds means the loop is
	// firing but the refresh function is failing.
	ProvidersLastReloadSuccessTimestampSeconds prometheus.Gauge
	// contains filtered or unexported fields
}

Metrics is the prometheus surface for aibridged provider reloads.

func NewMetrics 

func NewMetrics(reg prometheus.Registerer) *Metrics

NewMetrics registers the provider metrics against reg.

func (*Metrics) RecordReloadAttempt 

func (m *Metrics) RecordReloadAttempt()

RecordReloadAttempt stamps the attempt-time gauge at the start of a reload. A reload that hangs mid-flight is detected by watching the gap between this gauge and ProvidersLastReloadSuccessTimestampSeconds.

func (*Metrics) RecordReloadSuccess 

func (m *Metrics) RecordReloadSuccess(outcomes []ProviderOutcome)

RecordReloadSuccess rewrites the ProviderInfo GaugeVec from the outcomes and stamps the success-time gauge. Reset clears series for providers that have left the configuration so they don't linger as stale.

func (*Metrics) Unregister 

func (m *Metrics) Unregister()

Unregister removes the provider metrics from the registerer.

type PoolMetrics 

type PoolMetrics interface {
	Hits() uint64
	Misses() uint64
	KeysAdded() uint64
	KeysEvicted() uint64
}

type PoolOptions 

type PoolOptions struct {
	MaxItems int64
	TTL      time.Duration
}

type Pooler 

type Pooler interface {
	Acquire(ctx context.Context, req Request, clientFn ClientFunc, mcpBootstrapper MCPProxyBuilder) (http.Handler, error)
	// ReplaceProviders swaps the providers used to construct future
	// RequestBridge instances and clears the cache. Disabled providers
	// must be included; the bridge serves a 503 sentinel on their
	// routes.
	ReplaceProviders(providers []aibridge.Provider)
	Shutdown(ctx context.Context) error
}

Pooler describes a pool of [*aibridge.RequestBridge] instances from which instances can be retrieved. One [*aibridge.RequestBridge] instance is created per given key.

type ProviderOutcome 

type ProviderOutcome struct {
	Name   string
	Type   string
	Status ProviderStatus
	Err    error
}

ProviderOutcome classifies one ai_providers row, including disabled rows (which the pool keeps as 503 stubs) and errored rows (which the pool excludes). Err is populated only when Status == ProviderStatusError; the build error is already logged at the call site.

type ProviderReloader 

type ProviderReloader interface {
	Reload(ctx context.Context) error
}

ProviderReloader refreshes a component's provider snapshot.

type ProviderStatus 

type ProviderStatus string

ProviderStatus is the lifecycle state of a configured AI provider. 

const (
	// ProviderStatusEnabled indicates the provider is configured and
	// valid, and is included in the active pool snapshot.
	ProviderStatusEnabled ProviderStatus = "enabled"
	// ProviderStatusDisabled indicates the provider is configured but
	// intentionally turned off by an operator.
	ProviderStatusDisabled ProviderStatus = "disabled"
	// ProviderStatusError indicates the provider is configured but
	// cannot be constructed (missing keys, unsupported type, malformed
	// settings).
	ProviderStatusError ProviderStatus = "error"
)

type Request 

type Request struct {
	SessionKey  string
	APIKeyID    string
	InitiatorID uuid.UUID
}

type Server 

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

Server provides the AI Bridge functionality. It is responsible for:

  • receiving requests on /api/v2/aibridged/*
  • manipulating the requests
  • relaying requests to upstream AI services and relaying responses to caller

It requires a Dialer to provide a DRPCClient implementation to communicate with a DRPCServer implementation, to persist state and perform other functions.

func New 

func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger, tracer trace.Tracer) (*Server, error)

func (*Server) Client 

func (s *Server) Client() (DRPCClient, error)

func (*Server) Close 

func (s *Server) Close() error

Close shuts down the server with a timeout of 5s.

func (*Server) GetRequestHandler 

func (s *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handler, error)

GetRequestHandler retrieves a (possibly reused) [*aibridge.RequestBridge] from the pool, for the given user.

func (*Server) ServeHTTP 

func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request)

ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge. This function will validate that the given API key may be used to perform the request.

An [aibridge.RequestBridge] instance is acquired from a pool based on the API key's owner (referred to as the "initiator"); this instance is responsible for the AI Bridge-specific handling of the request.

A DRPCClient is provided to the [aibridge.RequestBridge] instance so that data can be passed up to a DRPCServer for persistence.

func (*Server) Shutdown 

func (s *Server) Shutdown(ctx context.Context) error

Shutdown waits for all exiting in-flight requests to complete, or the context to expire, whichever comes first.

Source Files

View all Source files

Directories

Path Synopsis
Package aibridgedmock is a generated GoMock package.
 Package aibridgedmock is a generated GoMock package.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.