diff --git a/aibridged/aibridged.go b/aibridged/aibridged.go new file mode 100644 index 0000000000000..1100e713a1430 --- /dev/null +++ b/aibridged/aibridged.go @@ -0,0 +1,306 @@ +package aibridged + +import ( + "context" + "errors" + "io" + "net/http" + "sync" + "sync/atomic" + "time" + + "github.com/hashicorp/yamux" + "github.com/valyala/fasthttp/fasthttputil" + "golang.org/x/xerrors" + + "cdr.dev/slog" + "github.com/coder/aibridge" + "github.com/coder/retry" + + "github.com/coder/coder/v2/aibridged/proto" + "github.com/coder/coder/v2/codersdk" +) + +type Dialer func(ctx context.Context) (proto.DRPCRecorderClient, error) + +// Server is the implementation which fulfills the proto.DRPCRecorderServer interface. +// It is responsible for communication with the +type Server struct { + clientDialer Dialer + clientCh chan proto.DRPCRecorderClient + + requestBridgePool pooler + + logger slog.Logger + wg sync.WaitGroup + + // initConnectionCh will receive when the daemon connects to coderd for the + // first time. + initConnectionCh chan struct{} + initConnectionOnce sync.Once + + // closeContext is canceled when we start closing. + closeContext context.Context + closeCancel context.CancelFunc + closeOnce sync.Once + // closeError stores the error when closing to return to subsequent callers + closeError error + // closingB is set to true when we start closing + closing atomic.Bool + shutdownOnce sync.Once + // shuttingDownCh will receive when we start graceful shutdown + shuttingDownCh chan struct{} +} + +var _ proto.DRPCRecorderServer = &Server{} + +func New(rpcDialer Dialer, requestBridgePool pooler, logger slog.Logger) (*Server, error) { + if rpcDialer == nil { + return nil, xerrors.Errorf("nil rpcDialer given") + } + + ctx, cancel := context.WithCancel(context.Background()) + daemon := &Server{ + logger: logger, + clientDialer: rpcDialer, + requestBridgePool: requestBridgePool, + clientCh: make(chan proto.DRPCRecorderClient), + closeContext: ctx, + closeCancel: cancel, + initConnectionCh: make(chan struct{}), + shuttingDownCh: make(chan struct{}), + } + + daemon.wg.Add(1) + go daemon.connect() + + return daemon, nil +} + +// Connect establishes a connection to coderd. +func (s *Server) connect() { + defer s.logger.Debug(s.closeContext, "connect loop exited") + defer s.wg.Done() + logConnect := s.logger.With(slog.F("context", "aibridged.server")).Debug + // An exponential back-off occurs when the connection is failing to dial. + // This is to prevent server spam in case of a coderd outage. +connectLoop: + for retrier := retry.New(50*time.Millisecond, 10*time.Second); retrier.Wait(s.closeContext); { + // It's possible for the aibridge daemon to be shut down + // before the wait is complete! + if s.isClosed() { + return + } + s.logger.Debug(s.closeContext, "dialing coderd") + client, err := s.clientDialer(s.closeContext) + if err != nil { + if errors.Is(err, context.Canceled) { + return + } + var sdkErr *codersdk.Error + // If something is wrong with our auth, stop trying to connect. + if errors.As(err, &sdkErr) && sdkErr.StatusCode() == http.StatusForbidden { + s.logger.Error(s.closeContext, "not authorized to dial coderd", slog.Error(err)) + return + } + if s.isClosed() { + return + } + s.logger.Warn(s.closeContext, "coderd client failed to dial", slog.Error(err)) + continue + } + + // TODO: log this with INFO level when we implement external aibridge daemons. + logConnect(s.closeContext, "successfully connected to coderd") + retrier.Reset() + s.initConnectionOnce.Do(func() { + close(s.initConnectionCh) + }) + + // serve the client until we are closed or it disconnects + for { + select { + case <-s.closeContext.Done(): + client.DRPCConn().Close() + return + case <-client.DRPCConn().Closed(): + logConnect(s.closeContext, "connection to coderd closed") + continue connectLoop + case s.clientCh <- client: + continue + } + } + } +} + +func (s *Server) Client() (proto.DRPCRecorderClient, error) { + select { + case <-s.closeContext.Done(): + return nil, xerrors.New("context closed") + case <-s.shuttingDownCh: + // Shutting down should return a nil client and unblock + return nil, xerrors.New("shutting down") + case client := <-s.clientCh: + return client, nil + } +} + +// GetRequestHandler retrieves a (possibly reused) *aibridge.RequestBridge from the pool, for the given user. +func (s *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handler, error) { + if s.requestBridgePool == nil { + return nil, xerrors.New("nil requestBridgePool") + } + + recorder := aibridge.NewRecorder(s.logger.Named("recorder"), func() (aibridge.Recorder, error) { + client, err := s.Client() + if err != nil { + return nil, xerrors.Errorf("acquire client: %w", err) + } + + return &translator{client: client}, nil + }) + + reqBridge, err := s.requestBridgePool.Acquire(ctx, req, recorder) + if err != nil { + return nil, xerrors.Errorf("acquire request bridge: %w", err) + } + + return reqBridge, nil +} + +func (s *Server) RecordSession(ctx context.Context, in *proto.RecordSessionRequest) (*proto.RecordSessionResponse, error) { + out, err := clientDoWithRetries(ctx, s.Client, func(ctx context.Context, client proto.DRPCRecorderClient) (*proto.RecordSessionResponse, error) { + return client.RecordSession(ctx, in) + }) + if err != nil { + return nil, err + } + return out, nil +} + +func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse, error) { + out, err := clientDoWithRetries(ctx, s.Client, func(ctx context.Context, client proto.DRPCRecorderClient) (*proto.RecordTokenUsageResponse, error) { + return client.RecordTokenUsage(ctx, in) + }) + if err != nil { + return nil, err + } + return out, nil +} + +func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) { + out, err := clientDoWithRetries(ctx, s.Client, func(ctx context.Context, client proto.DRPCRecorderClient) (*proto.RecordPromptUsageResponse, error) { + return client.RecordPromptUsage(ctx, in) + }) + if err != nil { + return nil, err + } + return out, nil +} + +func (s *Server) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageRequest) (*proto.RecordToolUsageResponse, error) { + out, err := clientDoWithRetries(ctx, s.Client, func(ctx context.Context, client proto.DRPCRecorderClient) (*proto.RecordToolUsageResponse, error) { + return client.RecordToolUsage(ctx, in) + }) + if err != nil { + return nil, err + } + return out, nil +} + +// NOTE: mostly copypasta from provisionerd; might be work abstracting. +func retryable(err error) bool { + return xerrors.Is(err, yamux.ErrSessionShutdown) || xerrors.Is(err, io.EOF) || xerrors.Is(err, fasthttputil.ErrInmemoryListenerClosed) || + // annoyingly, dRPC sometimes returns context.Canceled if the transport was closed, even if the context for + // the RPC *is not canceled*. Retrying is fine if the RPC context is not canceled. + xerrors.Is(err, context.Canceled) +} + +// clientDoWithRetries runs the function f with a client, and retries with +// backoff until either the error returned is not retryable() or the context +// expires. +// NOTE: mostly copypasta from provisionerd; might be work abstracting. +func clientDoWithRetries[T any](ctx context.Context, + getClient func() (proto.DRPCRecorderClient, error), + f func(context.Context, proto.DRPCRecorderClient) (T, error), +) (ret T, _ error) { + for retrier := retry.New(25*time.Millisecond, 5*time.Second); retrier.Wait(ctx); { + var empty T + client, err := getClient() + if err != nil { + if retryable(err) { + continue + } + return empty, err + } + resp, err := f(ctx, client) + if retryable(err) { + continue + } + return resp, err + } + return ret, ctx.Err() +} + +// isClosed returns whether the API is closed or not. +func (s *Server) isClosed() bool { + select { + case <-s.closeContext.Done(): + return true + default: + return false + } +} + +// closeWithError closes aibridged once; subsequent calls will return the error err. +func (s *Server) closeWithError(err error) error { + s.closing.Store(true) + s.closeOnce.Do(func() { + s.closeCancel() + s.logger.Debug(context.Background(), "waiting for goroutines to exit") + s.wg.Wait() + s.logger.Debug(context.Background(), "closing server with error", slog.Error(err)) + s.closeError = err + }) + + return s.closeError +} + +// Close ends the aibridge daemon. +func (s *Server) Close() error { + if s == nil { + return nil + } + + s.logger.Info(s.closeContext, "closing aibridged") + return s.closeWithError(nil) +} + +// Shutdown waits for all exiting in-flight requests to complete, or the context to expire, whichever comes first. +func (s *Server) Shutdown(ctx context.Context) error { + if s == nil { + return nil + } + + var err error + s.shutdownOnce.Do(func() { + close(s.shuttingDownCh) + + select { + case <-ctx.Done(): + s.logger.Warn(ctx, "graceful shutdown failed", slog.Error(ctx.Err())) + err = ctx.Err() + return + default: + } + + s.logger.Info(ctx, "shutting down aibridged pool") + if err = s.requestBridgePool.Shutdown(ctx); err != nil && errors.Is(err, http.ErrServerClosed) { + s.logger.Error(ctx, "shutdown failed with error", slog.Error(err)) + return + } + + s.logger.Info(ctx, "gracefully shutdown") + }) + return err +} diff --git a/aibridged/context.go b/aibridged/context.go new file mode 100644 index 0000000000000..b20e758b61b78 --- /dev/null +++ b/aibridged/context.go @@ -0,0 +1,5 @@ +package aibridged + +type ( + ContextKeyBridgeAPIKey struct{} +) diff --git a/aibridged/mcp.go b/aibridged/mcp.go new file mode 100644 index 0000000000000..520e5bf51291f --- /dev/null +++ b/aibridged/mcp.go @@ -0,0 +1,121 @@ +package aibridged + +import ( + "context" + "net/url" + + "github.com/google/uuid" + "github.com/hashicorp/go-multierror" + "golang.org/x/xerrors" + + "cdr.dev/slog" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/externalauth" +) + +type MCPServerConfig struct { + Name, URL, AccessToken string + ValidateFn func(ctx context.Context) (bool, error) + RefreshFn func(ctx context.Context) (bool, error) +} + +type store interface { + GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error) +} + +type MCPConfigurator interface { + GetMCPConfigs(ctx context.Context, sessionKey string, userID uuid.UUID) ([]*MCPServerConfig, error) +} + +type StoreMCPConfigurator struct { + accessURL string + store store + logger slog.Logger + externalAuthConfigs []*externalauth.Config +} + +func NewStoreMCPConfigurator(accessURL string, store store, externalAuthConfigs []*externalauth.Config, logger slog.Logger) *StoreMCPConfigurator { + return &StoreMCPConfigurator{accessURL: accessURL, store: store, logger: logger, externalAuthConfigs: externalAuthConfigs} +} + +func (m *StoreMCPConfigurator) GetMCPConfigs(ctx context.Context, sessionKey string, userID uuid.UUID) ([]*MCPServerConfig, error) { + var out []*MCPServerConfig + var merr multierror.Error + + coder, err := m.getCoderMCPServerConfig(sessionKey) + if err != nil { + merr.Errors = append(merr.Errors, xerrors.Errorf("get coder MCP server config: %w", err)) + } else { + out = append(out, coder) + } + + others, err := m.getExternalAuthMCPServerConfigs(ctx, m.logger, m.externalAuthConfigs, userID) + if err != nil { + merr.Errors = append(merr.Errors, xerrors.Errorf("get external auth MCP server config: %w", err)) + } else { + out = append(out, others...) + } + + return out, merr.ErrorOrNil() +} + +func (m *StoreMCPConfigurator) getCoderMCPServerConfig(sessionKey string) (*MCPServerConfig, error) { + mcpURL, err := url.JoinPath(m.accessURL, "/api/experimental/mcp/http") + if err != nil { + return nil, xerrors.Errorf("build MCP URL: %w", err) + } + + return &MCPServerConfig{ + Name: "coder", + URL: mcpURL, + AccessToken: sessionKey, + ValidateFn: func(_ context.Context) (bool, error) { + // No-op since request would not proceed if session key was invalid. + return true, nil + }, + }, nil +} + +func (m *StoreMCPConfigurator) getExternalAuthMCPServerConfigs(ctx context.Context, logger slog.Logger, externalAuthConfigs []*externalauth.Config, userID uuid.UUID) ([]*MCPServerConfig, error) { + externalAuthLinks, err := m.store.GetExternalAuthLinksByUserID(ctx, userID) + if err != nil { + return nil, xerrors.Errorf("load external auth links: %w", err) + } + + if len(externalAuthLinks) == 0 { + return nil, nil + } + + cfgs := make([]*MCPServerConfig, 0, len(externalAuthLinks)) + + for _, link := range externalAuthLinks { + var externalAuthConfig *externalauth.Config + for _, eac := range externalAuthConfigs { + if eac.ID == link.ProviderID { + externalAuthConfig = eac + break + } + } + + if externalAuthConfig == nil { + logger.Warn(ctx, "failed to find external auth config matching known external auth link", slog.F("id", link.ProviderID)) + continue + } + + cfgs = append(cfgs, &MCPServerConfig{ + Name: link.ProviderID, + URL: externalAuthConfig.MCPURL, + AccessToken: link.OAuthAccessToken, + ValidateFn: func(ctx context.Context) (bool, error) { + valid, _, err := externalAuthConfig.ValidateToken(ctx, link.OAuthToken()) + if err != nil { + return false, xerrors.Errorf("validate token for %q MCP init: %w", link.ProviderID, err) + } + return valid, nil + }, + // TODO: implement RefreshFn. + }) + } + + return cfgs, nil +} diff --git a/aibridged/middleware.go b/aibridged/middleware.go new file mode 100644 index 0000000000000..64a1b14aa79d9 --- /dev/null +++ b/aibridged/middleware.go @@ -0,0 +1,83 @@ +package aibridged + +import ( + "bytes" + "context" + "crypto/subtle" + "net/http" + + "cdr.dev/slog" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/httpmw" + "github.com/coder/coder/v2/coderd/rbac" +) + +// AuthMiddleware extracts and validates authorization tokens for AI bridge endpoints. +// It supports both Bearer tokens in Authorization headers and Coder session tokens +// from cookies/headers following the same patterns as existing Coder authentication. +func AuthMiddleware(db database.Store, logger slog.Logger) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + token := extractAuthToken(r) + if token == "" { + http.Error(rw, "Authorization token required", http.StatusUnauthorized) + return + } + + // Validate token using httpmw.APIKeyFromRequest + key, _, ok := httpmw.APIKeyFromRequest(ctx, db, func(*http.Request) string { + return token + }, &http.Request{}) + + if !ok { + http.Error(rw, "Invalid or expired session token", http.StatusUnauthorized) + return + } + + // Inject the initiator's RBAC subject into the scope so all actions occur on their behalf. + actor, _, err := httpmw.UserRBACSubject(ctx, db, key.UserID, rbac.ScopeAll) + if err != nil { + logger.Error(ctx, "failed to setup user RBAC context", slog.Error(err), slog.F("user_id", key.UserID), slog.F("key_id", key.ID)) + http.Error(rw, "internal server error", http.StatusInternalServerError) // Don't leak reason as this might have security implications. + return + } + ctx = dbauthz.As(ctx, actor) + + // TODO: I'd prefer if we didn't have to do this, or at least in this fashion. + // Inject the API key into the context to later be used to authenticate against the Coder MCP server. + ctx = context.WithValue(ctx, ContextKeyBridgeAPIKey{}, token) + + // Pass request with modify context including the request token. + next.ServeHTTP(rw, r.WithContext(ctx)) + }) + } +} + +// extractAuthToken extracts authorization token from HTTP request using multiple sources. +// These sources represent the different ways clients authenticate against AI providers. +// It checks Authorization header (Bearer token), X-Api-Key header, and Coder session headers and cookies. +func extractAuthToken(r *http.Request) string { + // 1. Check Authorization header for Bearer token. + authHeader := r.Header.Get("Authorization") + if authHeader != "" { + bearer := []byte("bearer ") + hdr := []byte(authHeader) + + // Use case-insensitive comparison for Bearer token. + if len(hdr) >= len(bearer) && subtle.ConstantTimeCompare(bytes.ToLower(hdr[:len(bearer)]), bearer) == 1 { + return string(hdr[len(bearer):]) + } + } + + // 2. Check X-Api-Key header. + apiKeyHeader := r.Header.Get("X-Api-Key") + if apiKeyHeader != "" { + return apiKeyHeader + } + + // 3. Fall back to Coder's standard token extraction. + return httpmw.APITokenFromRequest(r) +} diff --git a/aibridged/pool.go b/aibridged/pool.go new file mode 100644 index 0000000000000..5f0772bfaaf7f --- /dev/null +++ b/aibridged/pool.go @@ -0,0 +1,185 @@ +package aibridged + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/dgraph-io/ristretto/v2" + "github.com/google/uuid" + "golang.org/x/sync/errgroup" + "golang.org/x/xerrors" + "tailscale.com/util/singleflight" + + "cdr.dev/slog" + "github.com/coder/coder/v2/codersdk" + + "github.com/coder/aibridge" +) + +const ( + bridgeCacheTTL = time.Hour // TODO: configurable. + cacheCost = 1 // We can't know the actual size in bytes of the value (it'll change over time). +) + +// pooler describes a pool of *aibridge.RequestBridge instances from which instances can be retrieved. +// One *aibridge.RequestBridge instance is created per given key. +type pooler interface { + Acquire(ctx context.Context, req Request, recorder aibridge.Recorder) (*aibridge.RequestBridge, error) + Shutdown(ctx context.Context) error +} + +type CachedBridgePool struct { + cache *ristretto.Cache[string, *aibridge.RequestBridge] + providers []aibridge.Provider + mcpCfg MCPConfigurator + logger slog.Logger + + singleflight *singleflight.Group[string, *aibridge.RequestBridge] + + shutDownOnce sync.Once + shuttingDownCh chan struct{} +} + +func NewCachedBridgePool(cfg codersdk.AIBridgeConfig, instances int64, mcpCfg MCPConfigurator, logger slog.Logger) (*CachedBridgePool, error) { + cache, err := ristretto.NewCache(&ristretto.Config[string, *aibridge.RequestBridge]{ + // TODO: the cost seems to actually take into account the size of the object in bytes...? Stop at breakpoint and see. + NumCounters: instances * 10, // Docs suggest setting this 10x number of keys. + MaxCost: instances * cacheCost, // Up to n instances. + BufferItems: 64, // Sticking with recommendation from docs. + }) + if err != nil { + return nil, xerrors.Errorf("create cache: %w", err) + } + + return &CachedBridgePool{ + cache: cache, + providers: []aibridge.Provider{ + aibridge.NewOpenAIProvider(cfg.OpenAI.BaseURL.String(), cfg.OpenAI.Key.String()), + aibridge.NewAnthropicProvider(cfg.Anthropic.BaseURL.String(), cfg.Anthropic.Key.String()), + }, + mcpCfg: mcpCfg, + logger: logger, + + singleflight: &singleflight.Group[string, *aibridge.RequestBridge]{}, + + shuttingDownCh: make(chan struct{}), + }, nil +} + +// Acquire retrieves or creates a Bridge instance per given key. +// +// Each returned Bridge is safe for concurrent use. +// Each Bridge is stateful because it has MCP clients which maintain sessions to the configured MCP server. +func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, recorder aibridge.Recorder) (*aibridge.RequestBridge, error) { + if err := ctx.Err(); err != nil { + return nil, xerrors.Errorf("acquire: %w", err) + } + + select { + case <-p.shuttingDownCh: + return nil, xerrors.New("pool shutting down") + default: + } + + // Fast path. + bridge, ok := p.cache.Get(req.InitiatorID.String()) + if ok && bridge != nil { + // TODO: remove. + p.logger.Debug(ctx, "reusing existing bridge", slog.F("ptr", fmt.Sprintf("%p", bridge))) + + // Set key again to refresh its TTL. + // + // It's possible that two calls can race here, but since they'll both be setting the same value and + // approximately the same TTL, we don't really care. We could debounce this to prevent unnecessary writes + // but it'll likely never be an issue. + p.cache.SetWithTTL(req.InitiatorID.String(), bridge, cacheCost, bridgeCacheTTL) + + return bridge, nil + } + + // Slow path. + // Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value. + // TODO: track startup time since it adds latency to first request (histogram count will also help us see how often this occurs). + instance, err, _ := p.singleflight.Do(req.InitiatorID.String(), func() (*aibridge.RequestBridge, error) { + tools, err := p.setupMCPServerProxies(ctx, req.SessionKey, req.InitiatorID) + if err != nil { + p.logger.Warn(ctx, "failed to load tools", slog.Error(err)) + } + + bridge, err = aibridge.NewRequestBridge(ctx, p.providers, p.logger, recorder, aibridge.NewInjectedToolManager(tools)) + if err != nil { + return nil, xerrors.Errorf("create new request bridge: %w", err) + } + // TODO: remove. + p.logger.Debug(ctx, "created new request bridge", slog.F("ptr", fmt.Sprintf("%p", bridge))) + + p.cache.SetWithTTL(req.InitiatorID.String(), bridge, cacheCost, bridgeCacheTTL) + + return bridge, nil + }) + + return instance, err +} + +func (p *CachedBridgePool) setupMCPServerProxies(ctx context.Context, key string, userID uuid.UUID) ([]*aibridge.MCPServerProxy, error) { + var ( + proxies []*aibridge.MCPServerProxy + eg errgroup.Group + ) + + cfgs, err := p.mcpCfg.GetMCPConfigs(ctx, key, userID) + if err != nil { + return nil, xerrors.Errorf("get mcp configs: %w", err) + } + + for _, c := range cfgs { + eg.Go(func() error { + valid, err := c.ValidateFn(ctx) + if !valid { + if c.RefreshFn != nil { + // TODO: refresh token. + return xerrors.Errorf("%q token is not valid and cannot be refreshed currently: %w", c.Name, err) + } + return xerrors.Errorf("%q external auth token invalid: %w", c.Name, err) + } + + linkBridge, err := aibridge.NewMCPServerProxy(c.Name, c.URL, map[string]string{ + "Authorization": fmt.Sprintf("Bearer %s", c.AccessToken), + }, p.logger.Named(fmt.Sprintf("mcp-bridge-%s", c.Name))) + if err != nil { + return xerrors.Errorf("%s MCP bridge setup: %w", c.Name, err) + } + proxies = append(proxies, linkBridge) + + ctx, cancel := context.WithTimeout(ctx, time.Second*30) + defer cancel() + + err = linkBridge.Init(ctx) + if err == nil { + return nil + } + return xerrors.Errorf("%s MCP init: %w", c.Name, err) + }) + } + + if err := eg.Wait(); err != nil { + // Still return proxies even if there's an error; some is better than none. + return proxies, xerrors.Errorf("MCP proxy init: %w", err) + } + + return proxies, nil +} + +// Shutdown will close the cache which will trigger eviction of all the Bridge entries. +func (p *CachedBridgePool) Shutdown(_ context.Context) error { + p.shutDownOnce.Do(func() { + // Prevent new requests from being served. + close(p.shuttingDownCh) + + p.cache.Close() + }) + + return nil +} diff --git a/aibridged/proto/aibridge.pb.go b/aibridged/proto/aibridge.pb.go new file mode 100644 index 0000000000000..e2d718d98533b --- /dev/null +++ b/aibridged/proto/aibridge.pb.go @@ -0,0 +1,758 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.30.0 +// protoc v4.23.4 +// source: proto/aibridge.proto + +package proto + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + anypb "google.golang.org/protobuf/types/known/anypb" + reflect "reflect" + sync "sync" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type RecordSessionRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + InitiatorId string `protobuf:"bytes,2,opt,name=initiator_id,json=initiatorId,proto3" json:"initiator_id,omitempty"` + Provider string `protobuf:"bytes,3,opt,name=provider,proto3" json:"provider,omitempty"` + Model string `protobuf:"bytes,4,opt,name=model,proto3" json:"model,omitempty"` +} + +func (x *RecordSessionRequest) Reset() { + *x = RecordSessionRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordSessionRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordSessionRequest) ProtoMessage() {} + +func (x *RecordSessionRequest) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[0] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordSessionRequest.ProtoReflect.Descriptor instead. +func (*RecordSessionRequest) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{0} +} + +func (x *RecordSessionRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *RecordSessionRequest) GetInitiatorId() string { + if x != nil { + return x.InitiatorId + } + return "" +} + +func (x *RecordSessionRequest) GetProvider() string { + if x != nil { + return x.Provider + } + return "" +} + +func (x *RecordSessionRequest) GetModel() string { + if x != nil { + return x.Model + } + return "" +} + +type RecordSessionResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *RecordSessionResponse) Reset() { + *x = RecordSessionResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordSessionResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordSessionResponse) ProtoMessage() {} + +func (x *RecordSessionResponse) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[1] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordSessionResponse.ProtoReflect.Descriptor instead. +func (*RecordSessionResponse) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{1} +} + +type RecordTokenUsageRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + MsgId string `protobuf:"bytes,2,opt,name=msg_id,json=msgId,proto3" json:"msg_id,omitempty"` + InputTokens int64 `protobuf:"varint,3,opt,name=input_tokens,json=inputTokens,proto3" json:"input_tokens,omitempty"` + OutputTokens int64 `protobuf:"varint,4,opt,name=output_tokens,json=outputTokens,proto3" json:"output_tokens,omitempty"` + Metadata map[string]*anypb.Any `protobuf:"bytes,5,rep,name=metadata,proto3" json:"metadata,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` +} + +func (x *RecordTokenUsageRequest) Reset() { + *x = RecordTokenUsageRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordTokenUsageRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordTokenUsageRequest) ProtoMessage() {} + +func (x *RecordTokenUsageRequest) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[2] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordTokenUsageRequest.ProtoReflect.Descriptor instead. +func (*RecordTokenUsageRequest) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{2} +} + +func (x *RecordTokenUsageRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *RecordTokenUsageRequest) GetMsgId() string { + if x != nil { + return x.MsgId + } + return "" +} + +func (x *RecordTokenUsageRequest) GetInputTokens() int64 { + if x != nil { + return x.InputTokens + } + return 0 +} + +func (x *RecordTokenUsageRequest) GetOutputTokens() int64 { + if x != nil { + return x.OutputTokens + } + return 0 +} + +func (x *RecordTokenUsageRequest) GetMetadata() map[string]*anypb.Any { + if x != nil { + return x.Metadata + } + return nil +} + +type RecordTokenUsageResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *RecordTokenUsageResponse) Reset() { + *x = RecordTokenUsageResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordTokenUsageResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordTokenUsageResponse) ProtoMessage() {} + +func (x *RecordTokenUsageResponse) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[3] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordTokenUsageResponse.ProtoReflect.Descriptor instead. +func (*RecordTokenUsageResponse) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{3} +} + +type RecordPromptUsageRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + MsgId string `protobuf:"bytes,2,opt,name=msg_id,json=msgId,proto3" json:"msg_id,omitempty"` + Prompt string `protobuf:"bytes,3,opt,name=prompt,proto3" json:"prompt,omitempty"` + Metadata map[string]*anypb.Any `protobuf:"bytes,4,rep,name=metadata,proto3" json:"metadata,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` +} + +func (x *RecordPromptUsageRequest) Reset() { + *x = RecordPromptUsageRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordPromptUsageRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordPromptUsageRequest) ProtoMessage() {} + +func (x *RecordPromptUsageRequest) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[4] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordPromptUsageRequest.ProtoReflect.Descriptor instead. +func (*RecordPromptUsageRequest) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{4} +} + +func (x *RecordPromptUsageRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *RecordPromptUsageRequest) GetMsgId() string { + if x != nil { + return x.MsgId + } + return "" +} + +func (x *RecordPromptUsageRequest) GetPrompt() string { + if x != nil { + return x.Prompt + } + return "" +} + +func (x *RecordPromptUsageRequest) GetMetadata() map[string]*anypb.Any { + if x != nil { + return x.Metadata + } + return nil +} + +type RecordPromptUsageResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *RecordPromptUsageResponse) Reset() { + *x = RecordPromptUsageResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordPromptUsageResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordPromptUsageResponse) ProtoMessage() {} + +func (x *RecordPromptUsageResponse) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[5] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordPromptUsageResponse.ProtoReflect.Descriptor instead. +func (*RecordPromptUsageResponse) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{5} +} + +type RecordToolUsageRequest struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields + + SessionId string `protobuf:"bytes,1,opt,name=session_id,json=sessionId,proto3" json:"session_id,omitempty"` + MsgId string `protobuf:"bytes,2,opt,name=msg_id,json=msgId,proto3" json:"msg_id,omitempty"` + Tool string `protobuf:"bytes,3,opt,name=tool,proto3" json:"tool,omitempty"` + Input string `protobuf:"bytes,4,opt,name=input,proto3" json:"input,omitempty"` + Injected bool `protobuf:"varint,5,opt,name=injected,proto3" json:"injected,omitempty"` + Metadata map[string]*anypb.Any `protobuf:"bytes,6,rep,name=metadata,proto3" json:"metadata,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"` +} + +func (x *RecordToolUsageRequest) Reset() { + *x = RecordToolUsageRequest{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordToolUsageRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordToolUsageRequest) ProtoMessage() {} + +func (x *RecordToolUsageRequest) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[6] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordToolUsageRequest.ProtoReflect.Descriptor instead. +func (*RecordToolUsageRequest) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{6} +} + +func (x *RecordToolUsageRequest) GetSessionId() string { + if x != nil { + return x.SessionId + } + return "" +} + +func (x *RecordToolUsageRequest) GetMsgId() string { + if x != nil { + return x.MsgId + } + return "" +} + +func (x *RecordToolUsageRequest) GetTool() string { + if x != nil { + return x.Tool + } + return "" +} + +func (x *RecordToolUsageRequest) GetInput() string { + if x != nil { + return x.Input + } + return "" +} + +func (x *RecordToolUsageRequest) GetInjected() bool { + if x != nil { + return x.Injected + } + return false +} + +func (x *RecordToolUsageRequest) GetMetadata() map[string]*anypb.Any { + if x != nil { + return x.Metadata + } + return nil +} + +type RecordToolUsageResponse struct { + state protoimpl.MessageState + sizeCache protoimpl.SizeCache + unknownFields protoimpl.UnknownFields +} + +func (x *RecordToolUsageResponse) Reset() { + *x = RecordToolUsageResponse{} + if protoimpl.UnsafeEnabled { + mi := &file_proto_aibridge_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) + } +} + +func (x *RecordToolUsageResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*RecordToolUsageResponse) ProtoMessage() {} + +func (x *RecordToolUsageResponse) ProtoReflect() protoreflect.Message { + mi := &file_proto_aibridge_proto_msgTypes[7] + if protoimpl.UnsafeEnabled && x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use RecordToolUsageResponse.ProtoReflect.Descriptor instead. +func (*RecordToolUsageResponse) Descriptor() ([]byte, []int) { + return file_proto_aibridge_proto_rawDescGZIP(), []int{7} +} + +var File_proto_aibridge_proto protoreflect.FileDescriptor + +var file_proto_aibridge_proto_rawDesc = []byte{ + 0x0a, 0x14, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x61, 0x69, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x05, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x1a, 0x19, 0x67, + 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x61, + 0x6e, 0x79, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x8a, 0x01, 0x0a, 0x14, 0x52, 0x65, 0x63, + 0x6f, 0x72, 0x64, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, + 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, 0x64, + 0x12, 0x21, 0x0a, 0x0c, 0x69, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x74, 0x6f, 0x72, 0x5f, 0x69, 0x64, + 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0b, 0x69, 0x6e, 0x69, 0x74, 0x69, 0x61, 0x74, 0x6f, + 0x72, 0x49, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x18, + 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x70, 0x72, 0x6f, 0x76, 0x69, 0x64, 0x65, 0x72, 0x12, + 0x14, 0x0a, 0x05, 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, + 0x6d, 0x6f, 0x64, 0x65, 0x6c, 0x22, 0x17, 0x0a, 0x15, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x53, + 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0xb4, + 0x02, 0x0a, 0x17, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, + 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x65, + 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, + 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x15, 0x0a, 0x06, 0x6d, 0x73, 0x67, + 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x6d, 0x73, 0x67, 0x49, 0x64, + 0x12, 0x21, 0x0a, 0x0c, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x5f, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x73, + 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x0b, 0x69, 0x6e, 0x70, 0x75, 0x74, 0x54, 0x6f, 0x6b, + 0x65, 0x6e, 0x73, 0x12, 0x23, 0x0a, 0x0d, 0x6f, 0x75, 0x74, 0x70, 0x75, 0x74, 0x5f, 0x74, 0x6f, + 0x6b, 0x65, 0x6e, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x03, 0x52, 0x0c, 0x6f, 0x75, 0x74, 0x70, + 0x75, 0x74, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x73, 0x12, 0x48, 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, + 0x64, 0x61, 0x74, 0x61, 0x18, 0x05, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x2c, 0x2e, 0x70, 0x72, 0x6f, + 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, + 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x2e, 0x4d, 0x65, 0x74, 0x61, 0x64, + 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, + 0x74, 0x61, 0x1a, 0x51, 0x0a, 0x0d, 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, + 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, + 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x2a, 0x0a, 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x18, 0x02, + 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, + 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x41, 0x6e, 0x79, 0x52, 0x05, 0x76, 0x61, 0x6c, 0x75, + 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0x1a, 0x0a, 0x18, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, + 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, + 0x65, 0x22, 0x86, 0x02, 0x0a, 0x18, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, + 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1d, + 0x0a, 0x0a, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, + 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x15, 0x0a, + 0x06, 0x6d, 0x73, 0x67, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x6d, + 0x73, 0x67, 0x49, 0x64, 0x12, 0x16, 0x0a, 0x06, 0x70, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x18, 0x03, + 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x70, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x12, 0x49, 0x0a, 0x08, + 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x04, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x2d, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, + 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x2e, + 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x08, 0x6d, + 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x1a, 0x51, 0x0a, 0x0d, 0x4d, 0x65, 0x74, 0x61, 0x64, + 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x2a, 0x0a, 0x05, 0x76, 0x61, + 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x67, 0x6f, 0x6f, 0x67, + 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x41, 0x6e, 0x79, 0x52, + 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0x1b, 0x0a, 0x19, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, + 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0xb0, 0x02, 0x0a, 0x16, 0x52, 0x65, 0x63, 0x6f, + 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, + 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x73, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x49, + 0x64, 0x12, 0x15, 0x0a, 0x06, 0x6d, 0x73, 0x67, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, + 0x09, 0x52, 0x05, 0x6d, 0x73, 0x67, 0x49, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x6f, 0x6f, 0x6c, + 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x74, 0x6f, 0x6f, 0x6c, 0x12, 0x14, 0x0a, 0x05, + 0x69, 0x6e, 0x70, 0x75, 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x69, 0x6e, 0x70, + 0x75, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x69, 0x6e, 0x6a, 0x65, 0x63, 0x74, 0x65, 0x64, 0x18, 0x05, + 0x20, 0x01, 0x28, 0x08, 0x52, 0x08, 0x69, 0x6e, 0x6a, 0x65, 0x63, 0x74, 0x65, 0x64, 0x12, 0x47, + 0x0a, 0x08, 0x6d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x18, 0x06, 0x20, 0x03, 0x28, 0x0b, + 0x32, 0x2b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, + 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x2e, + 0x4d, 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x52, 0x08, 0x6d, + 0x65, 0x74, 0x61, 0x64, 0x61, 0x74, 0x61, 0x1a, 0x51, 0x0a, 0x0d, 0x4d, 0x65, 0x74, 0x61, 0x64, + 0x61, 0x74, 0x61, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x10, 0x0a, 0x03, 0x6b, 0x65, 0x79, 0x18, + 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6b, 0x65, 0x79, 0x12, 0x2a, 0x0a, 0x05, 0x76, 0x61, + 0x6c, 0x75, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x67, 0x6f, 0x6f, 0x67, + 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x41, 0x6e, 0x79, 0x52, + 0x05, 0x76, 0x61, 0x6c, 0x75, 0x65, 0x3a, 0x02, 0x38, 0x01, 0x22, 0x19, 0x0a, 0x17, 0x52, 0x65, + 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, + 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x32, 0xd5, 0x02, 0x0a, 0x08, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, + 0x65, 0x72, 0x12, 0x4a, 0x0a, 0x0d, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x53, 0x65, 0x73, 0x73, + 0x69, 0x6f, 0x6e, 0x12, 0x1b, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, + 0x72, 0x64, 0x53, 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, + 0x1a, 0x1c, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x53, + 0x65, 0x73, 0x73, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x53, + 0x0a, 0x10, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, + 0x67, 0x65, 0x12, 0x1e, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, + 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, + 0x73, 0x74, 0x1a, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, + 0x64, 0x54, 0x6f, 0x6b, 0x65, 0x6e, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, + 0x6e, 0x73, 0x65, 0x12, 0x56, 0x0a, 0x11, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, + 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1f, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, 0x61, + 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x20, 0x2e, 0x70, 0x72, 0x6f, 0x74, + 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x50, 0x72, 0x6f, 0x6d, 0x70, 0x74, 0x55, 0x73, + 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x50, 0x0a, 0x0f, 0x52, + 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x12, 0x1d, + 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, + 0x6c, 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, + 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2e, 0x52, 0x65, 0x63, 0x6f, 0x72, 0x64, 0x54, 0x6f, 0x6f, 0x6c, + 0x55, 0x73, 0x61, 0x67, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x42, 0x21, 0x5a, + 0x1f, 0x67, 0x69, 0x74, 0x68, 0x75, 0x62, 0x2e, 0x63, 0x6f, 0x6d, 0x2f, 0x63, 0x6f, 0x64, 0x65, + 0x72, 0x2f, 0x61, 0x69, 0x62, 0x72, 0x69, 0x64, 0x67, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, + 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, +} + +var ( + file_proto_aibridge_proto_rawDescOnce sync.Once + file_proto_aibridge_proto_rawDescData = file_proto_aibridge_proto_rawDesc +) + +func file_proto_aibridge_proto_rawDescGZIP() []byte { + file_proto_aibridge_proto_rawDescOnce.Do(func() { + file_proto_aibridge_proto_rawDescData = protoimpl.X.CompressGZIP(file_proto_aibridge_proto_rawDescData) + }) + return file_proto_aibridge_proto_rawDescData +} + +var file_proto_aibridge_proto_msgTypes = make([]protoimpl.MessageInfo, 11) +var file_proto_aibridge_proto_goTypes = []interface{}{ + (*RecordSessionRequest)(nil), // 0: proto.RecordSessionRequest + (*RecordSessionResponse)(nil), // 1: proto.RecordSessionResponse + (*RecordTokenUsageRequest)(nil), // 2: proto.RecordTokenUsageRequest + (*RecordTokenUsageResponse)(nil), // 3: proto.RecordTokenUsageResponse + (*RecordPromptUsageRequest)(nil), // 4: proto.RecordPromptUsageRequest + (*RecordPromptUsageResponse)(nil), // 5: proto.RecordPromptUsageResponse + (*RecordToolUsageRequest)(nil), // 6: proto.RecordToolUsageRequest + (*RecordToolUsageResponse)(nil), // 7: proto.RecordToolUsageResponse + nil, // 8: proto.RecordTokenUsageRequest.MetadataEntry + nil, // 9: proto.RecordPromptUsageRequest.MetadataEntry + nil, // 10: proto.RecordToolUsageRequest.MetadataEntry + (*anypb.Any)(nil), // 11: google.protobuf.Any +} +var file_proto_aibridge_proto_depIdxs = []int32{ + 8, // 0: proto.RecordTokenUsageRequest.metadata:type_name -> proto.RecordTokenUsageRequest.MetadataEntry + 9, // 1: proto.RecordPromptUsageRequest.metadata:type_name -> proto.RecordPromptUsageRequest.MetadataEntry + 10, // 2: proto.RecordToolUsageRequest.metadata:type_name -> proto.RecordToolUsageRequest.MetadataEntry + 11, // 3: proto.RecordTokenUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 11, // 4: proto.RecordPromptUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 11, // 5: proto.RecordToolUsageRequest.MetadataEntry.value:type_name -> google.protobuf.Any + 0, // 6: proto.Recorder.RecordSession:input_type -> proto.RecordSessionRequest + 2, // 7: proto.Recorder.RecordTokenUsage:input_type -> proto.RecordTokenUsageRequest + 4, // 8: proto.Recorder.RecordPromptUsage:input_type -> proto.RecordPromptUsageRequest + 6, // 9: proto.Recorder.RecordToolUsage:input_type -> proto.RecordToolUsageRequest + 1, // 10: proto.Recorder.RecordSession:output_type -> proto.RecordSessionResponse + 3, // 11: proto.Recorder.RecordTokenUsage:output_type -> proto.RecordTokenUsageResponse + 5, // 12: proto.Recorder.RecordPromptUsage:output_type -> proto.RecordPromptUsageResponse + 7, // 13: proto.Recorder.RecordToolUsage:output_type -> proto.RecordToolUsageResponse + 10, // [10:14] is the sub-list for method output_type + 6, // [6:10] is the sub-list for method input_type + 6, // [6:6] is the sub-list for extension type_name + 6, // [6:6] is the sub-list for extension extendee + 0, // [0:6] is the sub-list for field type_name +} + +func init() { file_proto_aibridge_proto_init() } +func file_proto_aibridge_proto_init() { + if File_proto_aibridge_proto != nil { + return + } + if !protoimpl.UnsafeEnabled { + file_proto_aibridge_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordSessionRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_proto_aibridge_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordSessionResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_proto_aibridge_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordTokenUsageRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_proto_aibridge_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordTokenUsageResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_proto_aibridge_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordPromptUsageRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_proto_aibridge_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordPromptUsageResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_proto_aibridge_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordToolUsageRequest); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + file_proto_aibridge_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} { + switch v := v.(*RecordToolUsageResponse); i { + case 0: + return &v.state + case 1: + return &v.sizeCache + case 2: + return &v.unknownFields + default: + return nil + } + } + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + GoPackagePath: reflect.TypeOf(x{}).PkgPath(), + RawDescriptor: file_proto_aibridge_proto_rawDesc, + NumEnums: 0, + NumMessages: 11, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_proto_aibridge_proto_goTypes, + DependencyIndexes: file_proto_aibridge_proto_depIdxs, + MessageInfos: file_proto_aibridge_proto_msgTypes, + }.Build() + File_proto_aibridge_proto = out.File + file_proto_aibridge_proto_rawDesc = nil + file_proto_aibridge_proto_goTypes = nil + file_proto_aibridge_proto_depIdxs = nil +} diff --git a/aibridged/proto/aibridge.proto b/aibridged/proto/aibridge.proto new file mode 100644 index 0000000000000..deaf1bd1f42d8 --- /dev/null +++ b/aibridged/proto/aibridge.proto @@ -0,0 +1,58 @@ +syntax = "proto3"; +option go_package = "github.com/coder/coder/v2/aibridged/proto"; + +package proto; + +import "google/protobuf/any.proto"; + +message RecordSessionRequest { + string session_id = 1; + string initiator_id = 2; + string provider = 3; + string model = 4; +} + +message RecordSessionResponse {} + +message RecordTokenUsageRequest { + string session_id = 1; + string msg_id = 2; + int64 input_tokens = 3; + int64 output_tokens = 4; + map metadata = 5; +} +message RecordTokenUsageResponse {} + +message RecordPromptUsageRequest { + string session_id = 1; + string msg_id = 2; + string prompt = 3; + map metadata = 4; +} +message RecordPromptUsageResponse {} + +message RecordToolUsageRequest { + string session_id = 1; + string msg_id = 2; + string tool = 3; + string input = 4; + bool injected = 5; + map metadata = 6; +} +message RecordToolUsageResponse {} + +service Recorder { + rpc RecordSession(RecordSessionRequest) returns (RecordSessionResponse); + rpc RecordTokenUsage(RecordTokenUsageRequest) returns (RecordTokenUsageResponse); + rpc RecordPromptUsage(RecordPromptUsageRequest) returns (RecordPromptUsageResponse); + rpc RecordToolUsage(RecordToolUsageRequest) returns (RecordToolUsageResponse); +} + +// Just leaving some ideas here. +// The point is that this will be a *separate* service from Recorder, not to be accepted by RequestBridge but by some other component. +// +// service Stats { +// rpc AggregateTokenUsageByUser(TokenAggregationRequest) returns (TokenAggregationResponse); +// rpc AggregateToolUsageByUser(ToolAggregationRequest) returns (ToolAggregationResponse); +// rpc CategorizePromptsByUser(PromptCategorizationRequest) returns (PromptCategorizationResponse); +// } diff --git a/aibridged/proto/aibridge_drpc.pb.go b/aibridged/proto/aibridge_drpc.pb.go new file mode 100644 index 0000000000000..266ac493f3666 --- /dev/null +++ b/aibridged/proto/aibridge_drpc.pb.go @@ -0,0 +1,231 @@ +// Code generated by protoc-gen-go-drpc. DO NOT EDIT. +// protoc-gen-go-drpc version: v0.0.34 +// source: proto/aibridge.proto + +package proto + +import ( + context "context" + errors "errors" + protojson "google.golang.org/protobuf/encoding/protojson" + proto "google.golang.org/protobuf/proto" + drpc "storj.io/drpc" + drpcerr "storj.io/drpc/drpcerr" +) + +type drpcEncoding_File_proto_aibridge_proto struct{} + +func (drpcEncoding_File_proto_aibridge_proto) Marshal(msg drpc.Message) ([]byte, error) { + return proto.Marshal(msg.(proto.Message)) +} + +func (drpcEncoding_File_proto_aibridge_proto) MarshalAppend(buf []byte, msg drpc.Message) ([]byte, error) { + return proto.MarshalOptions{}.MarshalAppend(buf, msg.(proto.Message)) +} + +func (drpcEncoding_File_proto_aibridge_proto) Unmarshal(buf []byte, msg drpc.Message) error { + return proto.Unmarshal(buf, msg.(proto.Message)) +} + +func (drpcEncoding_File_proto_aibridge_proto) JSONMarshal(msg drpc.Message) ([]byte, error) { + return protojson.Marshal(msg.(proto.Message)) +} + +func (drpcEncoding_File_proto_aibridge_proto) JSONUnmarshal(buf []byte, msg drpc.Message) error { + return protojson.Unmarshal(buf, msg.(proto.Message)) +} + +type DRPCRecorderClient interface { + DRPCConn() drpc.Conn + + RecordSession(ctx context.Context, in *RecordSessionRequest) (*RecordSessionResponse, error) + RecordTokenUsage(ctx context.Context, in *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) + RecordPromptUsage(ctx context.Context, in *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error) + RecordToolUsage(ctx context.Context, in *RecordToolUsageRequest) (*RecordToolUsageResponse, error) +} + +type drpcRecorderClient struct { + cc drpc.Conn +} + +func NewDRPCRecorderClient(cc drpc.Conn) DRPCRecorderClient { + return &drpcRecorderClient{cc} +} + +func (c *drpcRecorderClient) DRPCConn() drpc.Conn { return c.cc } + +func (c *drpcRecorderClient) RecordSession(ctx context.Context, in *RecordSessionRequest) (*RecordSessionResponse, error) { + out := new(RecordSessionResponse) + err := c.cc.Invoke(ctx, "/proto.Recorder/RecordSession", drpcEncoding_File_proto_aibridge_proto{}, in, out) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *drpcRecorderClient) RecordTokenUsage(ctx context.Context, in *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) { + out := new(RecordTokenUsageResponse) + err := c.cc.Invoke(ctx, "/proto.Recorder/RecordTokenUsage", drpcEncoding_File_proto_aibridge_proto{}, in, out) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *drpcRecorderClient) RecordPromptUsage(ctx context.Context, in *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error) { + out := new(RecordPromptUsageResponse) + err := c.cc.Invoke(ctx, "/proto.Recorder/RecordPromptUsage", drpcEncoding_File_proto_aibridge_proto{}, in, out) + if err != nil { + return nil, err + } + return out, nil +} + +func (c *drpcRecorderClient) RecordToolUsage(ctx context.Context, in *RecordToolUsageRequest) (*RecordToolUsageResponse, error) { + out := new(RecordToolUsageResponse) + err := c.cc.Invoke(ctx, "/proto.Recorder/RecordToolUsage", drpcEncoding_File_proto_aibridge_proto{}, in, out) + if err != nil { + return nil, err + } + return out, nil +} + +type DRPCRecorderServer interface { + RecordSession(context.Context, *RecordSessionRequest) (*RecordSessionResponse, error) + RecordTokenUsage(context.Context, *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) + RecordPromptUsage(context.Context, *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error) + RecordToolUsage(context.Context, *RecordToolUsageRequest) (*RecordToolUsageResponse, error) +} + +type DRPCRecorderUnimplementedServer struct{} + +func (s *DRPCRecorderUnimplementedServer) RecordSession(context.Context, *RecordSessionRequest) (*RecordSessionResponse, error) { + return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) +} + +func (s *DRPCRecorderUnimplementedServer) RecordTokenUsage(context.Context, *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) { + return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) +} + +func (s *DRPCRecorderUnimplementedServer) RecordPromptUsage(context.Context, *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error) { + return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) +} + +func (s *DRPCRecorderUnimplementedServer) RecordToolUsage(context.Context, *RecordToolUsageRequest) (*RecordToolUsageResponse, error) { + return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented) +} + +type DRPCRecorderDescription struct{} + +func (DRPCRecorderDescription) NumMethods() int { return 4 } + +func (DRPCRecorderDescription) Method(n int) (string, drpc.Encoding, drpc.Receiver, interface{}, bool) { + switch n { + case 0: + return "/proto.Recorder/RecordSession", drpcEncoding_File_proto_aibridge_proto{}, + func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { + return srv.(DRPCRecorderServer). + RecordSession( + ctx, + in1.(*RecordSessionRequest), + ) + }, DRPCRecorderServer.RecordSession, true + case 1: + return "/proto.Recorder/RecordTokenUsage", drpcEncoding_File_proto_aibridge_proto{}, + func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { + return srv.(DRPCRecorderServer). + RecordTokenUsage( + ctx, + in1.(*RecordTokenUsageRequest), + ) + }, DRPCRecorderServer.RecordTokenUsage, true + case 2: + return "/proto.Recorder/RecordPromptUsage", drpcEncoding_File_proto_aibridge_proto{}, + func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { + return srv.(DRPCRecorderServer). + RecordPromptUsage( + ctx, + in1.(*RecordPromptUsageRequest), + ) + }, DRPCRecorderServer.RecordPromptUsage, true + case 3: + return "/proto.Recorder/RecordToolUsage", drpcEncoding_File_proto_aibridge_proto{}, + func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) { + return srv.(DRPCRecorderServer). + RecordToolUsage( + ctx, + in1.(*RecordToolUsageRequest), + ) + }, DRPCRecorderServer.RecordToolUsage, true + default: + return "", nil, nil, nil, false + } +} + +func DRPCRegisterRecorder(mux drpc.Mux, impl DRPCRecorderServer) error { + return mux.Register(impl, DRPCRecorderDescription{}) +} + +type DRPCRecorder_RecordSessionStream interface { + drpc.Stream + SendAndClose(*RecordSessionResponse) error +} + +type drpcRecorder_RecordSessionStream struct { + drpc.Stream +} + +func (x *drpcRecorder_RecordSessionStream) SendAndClose(m *RecordSessionResponse) error { + if err := x.MsgSend(m, drpcEncoding_File_proto_aibridge_proto{}); err != nil { + return err + } + return x.CloseSend() +} + +type DRPCRecorder_RecordTokenUsageStream interface { + drpc.Stream + SendAndClose(*RecordTokenUsageResponse) error +} + +type drpcRecorder_RecordTokenUsageStream struct { + drpc.Stream +} + +func (x *drpcRecorder_RecordTokenUsageStream) SendAndClose(m *RecordTokenUsageResponse) error { + if err := x.MsgSend(m, drpcEncoding_File_proto_aibridge_proto{}); err != nil { + return err + } + return x.CloseSend() +} + +type DRPCRecorder_RecordPromptUsageStream interface { + drpc.Stream + SendAndClose(*RecordPromptUsageResponse) error +} + +type drpcRecorder_RecordPromptUsageStream struct { + drpc.Stream +} + +func (x *drpcRecorder_RecordPromptUsageStream) SendAndClose(m *RecordPromptUsageResponse) error { + if err := x.MsgSend(m, drpcEncoding_File_proto_aibridge_proto{}); err != nil { + return err + } + return x.CloseSend() +} + +type DRPCRecorder_RecordToolUsageStream interface { + drpc.Stream + SendAndClose(*RecordToolUsageResponse) error +} + +type drpcRecorder_RecordToolUsageStream struct { + drpc.Stream +} + +func (x *drpcRecorder_RecordToolUsageStream) SendAndClose(m *RecordToolUsageResponse) error { + if err := x.MsgSend(m, drpcEncoding_File_proto_aibridge_proto{}); err != nil { + return err + } + return x.CloseSend() +} diff --git a/aibridged/request.go b/aibridged/request.go new file mode 100644 index 0000000000000..83d21c7ed596a --- /dev/null +++ b/aibridged/request.go @@ -0,0 +1,8 @@ +package aibridged + +import "github.com/google/uuid" + +type Request struct { + SessionKey string + InitiatorID, RequestID uuid.UUID +} diff --git a/aibridged/translator.go b/aibridged/translator.go new file mode 100644 index 0000000000000..9d4e2ac7fc0e2 --- /dev/null +++ b/aibridged/translator.go @@ -0,0 +1,105 @@ +package aibridged + +import ( + "context" + "encoding/json" + "fmt" + + "golang.org/x/xerrors" + "google.golang.org/protobuf/types/known/anypb" + "google.golang.org/protobuf/types/known/structpb" + + "github.com/coder/coder/v2/aibridged/proto" + + "github.com/coder/aibridge" +) + +var _ aibridge.Recorder = &translator{} + +// translator satisfies the aibridge.Recorder interface and translates calls into dRPC calls to aibridgedserver. +type translator struct { + client proto.DRPCRecorderClient +} + +func (t *translator) RecordSession(ctx context.Context, req *aibridge.SessionRequest) error { + _, err := t.client.RecordSession(ctx, &proto.RecordSessionRequest{ + SessionId: req.SessionID, + InitiatorId: req.InitiatorID, + Provider: req.Provider, + Model: req.Model, + }) + return err +} + +func (t *translator) RecordPromptUsage(ctx context.Context, req *aibridge.PromptUsageRequest) error { + _, err := t.client.RecordPromptUsage(ctx, &proto.RecordPromptUsageRequest{ + SessionId: req.SessionID, + MsgId: req.MsgID, + Prompt: req.Prompt, + Metadata: marshalForProto(req.Metadata), + }) + return err +} + +func (t *translator) RecordTokenUsage(ctx context.Context, req *aibridge.TokenUsageRequest) error { + _, err := t.client.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ + SessionId: req.SessionID, + MsgId: req.MsgID, + InputTokens: req.Input, + OutputTokens: req.Output, + Metadata: marshalForProto(req.Metadata), + }) + return err +} + +func (t *translator) RecordToolUsage(ctx context.Context, req *aibridge.ToolUsageRequest) error { + serialized, err := json.Marshal(req.Args) + if err != nil { + return xerrors.Errorf("serialize tool %q args: %w", req.Name, err) + } + + _, err = t.client.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{ + SessionId: req.SessionID, + MsgId: req.MsgID, + Tool: req.Name, + Input: string(serialized), + Injected: req.Injected, + Metadata: marshalForProto(req.Metadata), + }) + return err +} + +// marshalForProto will attempt to convert from aibridge.Metadata into a proto-friendly map[string]*anypb.Any. +// If any marshaling fails, rather return a map with the error details since we don't want to fail Record* funcs if metadata can't encode, +// since it's, well, metadata. +func marshalForProto(in aibridge.Metadata) map[string]*anypb.Any { + out := make(map[string]*anypb.Any, len(in)) + if len(in) == 0 { + return out + } + + // Instead of returning error, just encode error into metadata. + encodeErr := func(err error) map[string]*anypb.Any { + errVal, _ := anypb.New(structpb.NewStringValue(err.Error())) + mdVal, _ := anypb.New(structpb.NewStringValue(fmt.Sprintf("%+v", in))) + return map[string]*anypb.Any{ + "error": errVal, + "metadata": mdVal, + } + } + + for k, v := range in { + sv, err := structpb.NewValue(v) + if err != nil { + return encodeErr(err) + } + + av, err := anypb.New(sv) + if err != nil { + return encodeErr(err) + } + + out[k] = av + } + return out +} diff --git a/cli/server.go b/cli/server.go index f9e744761b22e..65ceaff1a22fb 100644 --- a/cli/server.go +++ b/cli/server.go @@ -62,6 +62,8 @@ import ( "github.com/coder/serpent" "github.com/coder/wgtunnel/tunnelsdk" + "github.com/coder/coder/v2/aibridged" + aibridgedproto "github.com/coder/coder/v2/aibridged/proto" "github.com/coder/coder/v2/coderd/entitlements" "github.com/coder/coder/v2/coderd/notifications/reports" "github.com/coder/coder/v2/coderd/runtimeconfig" @@ -1093,6 +1095,17 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd. } }() + // Built in aibridge daemons. + if vals.AI.BridgeConfig.Enabled { + // NOTE: this shares the same context as the HTTP API server because it gets wired into it. + // See coderd/aibridge.go. + srv, err := newAIBridgeServer(shutdownConnsCtx, coderAPI) + if err != nil { + return xerrors.Errorf("create aibridged: %w", err) + } + coderAPI.AIBridgeServer = srv + } + // Updates the systemd status from activating to activated. _, err = daemon.SdNotify(false, daemon.SdNotifyReady) if err != nil { @@ -1146,11 +1159,28 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd. cliui.Errorf(inv.Stderr, "Notify systemd failed: %s", err) } + // Stop accepting new connections to AI Bridge without interrupting in-flight requests. + // + // We have to shut down AI Bridge server before the API server, because AI Bridge server is wired into + // API server and therefore share a context. + if coderAPI.AIBridgeServer != nil { + r.Verbosef(inv, "Shutting down aibridge daemon...") + + // Give in-flight requests 5 seconds to complete before shutting down. + err = shutdownWithTimeout(coderAPI.AIBridgeServer.Shutdown, 5*time.Second) + if err != nil { + cliui.Errorf(inv.Stderr, "Shutdown aibridge daemon failed: %s\n", err) + } else { + _ = coderAPI.AIBridgeServer.Close() + r.Verbosef(inv, "Gracefully shut down aibridge daemon") + } + } + // Stop accepting new connections without interrupting // in-flight requests, give in-flight requests 5 seconds to // complete. cliui.Info(inv.Stdout, "Shutting down API server..."+"\n") - err = shutdownWithTimeout(httpServer.Shutdown, 3*time.Second) + err = shutdownWithTimeout(httpServer.Shutdown, 3*time.Second) // TODO: why do we call shutdownWithTimeout on httpServer.Shutdown *twice*? if err != nil { cliui.Errorf(inv.Stderr, "API server shutdown took longer than 3s: %s\n", err) } else { @@ -1313,6 +1343,24 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd. return serverCmd } +func newAIBridgeServer(ctx context.Context, coderAPI *coderd.API) (*aibridged.Server, error) { + mcpCfg := aibridged.NewStoreMCPConfigurator( + coderAPI.DeploymentValues.AccessURL.String(), coderAPI.Database, + coderAPI.Options.ExternalAuthConfigs, coderAPI.Logger.Named("aibridged.mcp"), + ) + pool, err := aibridged.NewCachedBridgePool(coderAPI.DeploymentValues.AI.BridgeConfig, 100, mcpCfg, coderAPI.Logger.Named("aibridge-pool")) // TODO: configurable size. + if err != nil { + return nil, xerrors.Errorf("create aibridge pool: %w", err) + } + daemon, err := aibridged.New(func(dialCtx context.Context) (aibridgedproto.DRPCRecorderClient, error) { + return coderAPI.CreateInMemoryAIBridgeDaemon(dialCtx) + }, pool, coderAPI.Logger.Named("aibridged")) + if err != nil { + return nil, xerrors.Errorf("create aibridge daemon: %w", err) + } + return daemon, nil +} + // templateHelpers builds a set of functions which can be called in templates. // We build them here to avoid an import cycle by using coderd.Options in notifications.Manager. // We can later use this to inject whitelabel fields when app name / logo URL are overridden. @@ -2711,6 +2759,8 @@ func parseExternalAuthProvidersFromEnv(prefix string, environ []string) ([]coder provider.DisplayName = v.Value case "DISPLAY_ICON": provider.DisplayIcon = v.Value + case "MCP_URL": + provider.MCPURL = v.Value } providers[providerNum] = provider } diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden index 9c949532398ac..8602b25760cf9 100644 --- a/cli/testdata/coder_server_--help.golden +++ b/cli/testdata/coder_server_--help.golden @@ -76,6 +76,11 @@ OPTIONS: Periodically check for new releases of Coder and inform the owner. The check is performed once per day. +AI BRIDGE OPTIONS: + --ai-bridge-enabled bool, $CODER_AI_BRIDGE_ENABLED (default: true) + Whether to start an in-memory aibridged instance ('ai-bridge' + experiment must be enabled, too). + CLIENT OPTIONS: These options change the behavior of how clients interact with the Coder. Clients include the Coder CLI, Coder Desktop, IDE extensions, and the web UI. diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index e23274e442078..d03e589af157a 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -709,3 +709,20 @@ workspace_prebuilds: # limit; disabled when set to zero. # (default: 3, type: int) failure_hard_limit: 3 +ai_bridge: + # Whether to start an in-memory aibridged instance ('ai-bridge' experiment must be + # enabled, too). + # (default: true, type: bool) + enabled: true + # TODO. + # (default: https://api.openai.com/v1/, type: string) + openai_base_url: https://api.openai.com/v1/ + # TODO. + # (default: , type: string) + openai_key: "" + # TODO. + # (default: https://api.anthropic.com/, type: string) + base_url: https://api.anthropic.com/ + # TODO. + # (default: , type: string) + key: "" diff --git a/coderd/aibridge.go b/coderd/aibridge.go new file mode 100644 index 0000000000000..cd9a7ea210326 --- /dev/null +++ b/coderd/aibridge.go @@ -0,0 +1,63 @@ +package coderd + +import ( + "net/http" + + "cdr.dev/slog" + + "github.com/google/uuid" + + "github.com/coder/aibridge" + "github.com/coder/coder/v2/aibridged" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/httpmw" +) + +// bridgeAIRequest handles requests destined for an upstream AI provider; aibridged intercepts these requests +// and applies a governance layer. +// +// See also: aibridged/middleware.go. +func (api *API) bridgeAIRequest(rw http.ResponseWriter, r *http.Request) { + srv := api.AIBridgeServer + if srv == nil { + http.Error(rw, "no AI bridge daemon running", http.StatusBadGateway) + return + } + + ctx := r.Context() + + actor, set := dbauthz.ActorFromContext(ctx) + if !set { + api.Logger.Error(ctx, "missing dbauthz actor in context") + http.Error(rw, "unauthorized", http.StatusUnauthorized) + return + } + + userID, err := uuid.Parse(actor.ID) + if err != nil { + api.Logger.Error(ctx, "actor ID is not a uuid", slog.Error(err), slog.F("user_id", actor.ID)) + http.Error(rw, "internal server error", http.StatusInternalServerError) + return + } + + sessionKey, ok := ctx.Value(aibridged.ContextKeyBridgeAPIKey{}).(string) + if sessionKey == "" || !ok { + http.Error(rw, "unable to retrieve request session key", http.StatusBadRequest) + return + } + + // Rewire request context to include actor. + r = r.WithContext(aibridge.AsActor(ctx, actor.ID, aibridge.Metadata{"email": actor.Email})) + + handler, err := srv.GetRequestHandler(ctx, aibridged.Request{ + SessionKey: sessionKey, + InitiatorID: userID, + RequestID: httpmw.RequestID(r), + }) + if err != nil { + api.Logger.Error(ctx, "failed to handle request", slog.Error(err)) + http.Error(rw, "failed to handle request", http.StatusInternalServerError) + return + } + http.StripPrefix("/api/v2/aibridge", handler).ServeHTTP(rw, r) +} diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go new file mode 100644 index 0000000000000..97368cf99296b --- /dev/null +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -0,0 +1,136 @@ +package aibridgedserver + +import ( + "context" + "encoding/json" + + "github.com/google/uuid" + "golang.org/x/xerrors" + "google.golang.org/protobuf/types/known/anypb" + "google.golang.org/protobuf/types/known/structpb" + + "cdr.dev/slog" + + "github.com/coder/coder/v2/aibridged/proto" + "github.com/coder/coder/v2/coderd/database" +) + +var _ proto.DRPCRecorderServer = &Server{} + +type Server struct { + // lifecycleCtx must be tied to the API server's lifecycle + // as when the API server shuts down, we want to cancel any + // long-running operations. + lifecycleCtx context.Context + store database.Store + logger slog.Logger +} + +func NewServer(lifecycleCtx context.Context, store database.Store, logger slog.Logger) (*Server, error) { + return &Server{ + lifecycleCtx: lifecycleCtx, + store: store, + logger: logger.Named("aibridgedserver"), + }, nil +} + +func (s *Server) RecordSession(ctx context.Context, in *proto.RecordSessionRequest) (*proto.RecordSessionResponse, error) { + sessID, err := uuid.Parse(in.GetSessionId()) + if err != nil { + return nil, xerrors.Errorf("invalid session ID %q: %w", in.GetSessionId(), err) + } + initID, err := uuid.Parse(in.GetInitiatorId()) + if err != nil { + return nil, xerrors.Errorf("invalid initiator ID %q: %w", in.GetInitiatorId(), err) + } + + err = s.store.InsertAIBridgeSession(ctx, database.InsertAIBridgeSessionParams{ + ID: sessID, + InitiatorID: initID, + Provider: in.Provider, + Model: in.Model, + }) + if err != nil { + return nil, xerrors.Errorf("start session: %w", err) + } + + return &proto.RecordSessionResponse{}, nil +} + +func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse, error) { + sessID, err := uuid.Parse(in.GetSessionId()) + if err != nil { + return nil, xerrors.Errorf("failed to parse session_id %q: %w", in.GetSessionId(), err) + } + + err = s.store.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{ + ID: uuid.New(), + SessionID: sessID, + ProviderID: in.GetMsgId(), + InputTokens: in.GetInputTokens(), + OutputTokens: in.GetOutputTokens(), + Metadata: s.marshalMetadata(in.GetMetadata()), + }) + if err != nil { + return nil, xerrors.Errorf("insert token usage: %w", err) + } + return &proto.RecordTokenUsageResponse{}, nil +} + +func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) { + sessID, err := uuid.Parse(in.GetSessionId()) + if err != nil { + return nil, xerrors.Errorf("failed to parse session_id %q: %w", in.GetSessionId(), err) + } + + err = s.store.InsertAIBridgeUserPrompt(ctx, database.InsertAIBridgeUserPromptParams{ + ID: uuid.New(), + SessionID: sessID, + ProviderID: in.GetMsgId(), + Prompt: in.GetPrompt(), + Metadata: s.marshalMetadata(in.GetMetadata()), + }) + if err != nil { + return nil, xerrors.Errorf("insert user prompt: %w", err) + } + return &proto.RecordPromptUsageResponse{}, nil +} + +func (s *Server) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageRequest) (*proto.RecordToolUsageResponse, error) { + sessID, err := uuid.Parse(in.GetSessionId()) + if err != nil { + return nil, xerrors.Errorf("failed to parse session_id %q: %w", in.GetSessionId(), err) + } + + err = s.store.InsertAIBridgeToolUsage(ctx, database.InsertAIBridgeToolUsageParams{ + ID: uuid.New(), + SessionID: sessID, + ProviderID: in.GetMsgId(), + Tool: in.GetTool(), + Input: in.GetInput(), + Injected: in.GetInjected(), + Metadata: s.marshalMetadata(in.GetMetadata()), + }) + if err != nil { + return nil, xerrors.Errorf("insert tool usage: %w", err) + } + return &proto.RecordToolUsageResponse{}, nil +} + +func (s *Server) marshalMetadata(in map[string]*anypb.Any) []byte { + mdMap := make(map[string]any, len(in)) + for k, v := range in { + if v == nil { + continue + } + var sv structpb.Value + if err := v.UnmarshalTo(&sv); err == nil { + mdMap[k] = sv.AsInterface() + } + } + out, err := json.Marshal(mdMap) + if err != nil { + s.logger.Warn(s.lifecycleCtx, "failed to marshal metadata", slog.Error(err)) + } + return out +} diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 00478e029e084..f3626fad0359d 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -11054,6 +11054,50 @@ const docTemplate = `{ } } }, + "codersdk.AIBridgeAnthropicConfig": { + "type": "object", + "properties": { + "base_url": { + "type": "string" + }, + "key": { + "type": "string" + } + } + }, + "codersdk.AIBridgeConfig": { + "type": "object", + "properties": { + "anthropic": { + "$ref": "#/definitions/codersdk.AIBridgeAnthropicConfig" + }, + "enabled": { + "type": "boolean" + }, + "openai": { + "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig" + } + } + }, + "codersdk.AIBridgeOpenAIConfig": { + "type": "object", + "properties": { + "base_url": { + "type": "string" + }, + "key": { + "type": "string" + } + } + }, + "codersdk.AIConfig": { + "type": "object", + "properties": { + "bridge": { + "$ref": "#/definitions/codersdk.AIBridgeConfig" + } + } + }, "codersdk.APIKey": { "type": "object", "required": [ @@ -12671,6 +12715,9 @@ const docTemplate = `{ "agent_stat_refresh_interval": { "type": "integer" }, + "ai": { + "$ref": "#/definitions/codersdk.AIConfig" + }, "allow_workspace_renames": { "type": "boolean" }, @@ -12998,9 +13045,11 @@ const docTemplate = `{ "web-push", "oauth2", "mcp-server-http", - "workspace-sharing" + "workspace-sharing", + "ai-bridge" ], "x-enum-comments": { + "ExperimentAIBridge": "Enables AI Bridge functionality.", "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentExample": "This isn't used for anything.", "ExperimentMCPServerHTTP": "Enables the MCP HTTP server functionality.", @@ -13018,7 +13067,8 @@ const docTemplate = `{ "ExperimentWebPush", "ExperimentOAuth2", "ExperimentMCPServerHTTP", - "ExperimentWorkspaceSharing" + "ExperimentWorkspaceSharing", + "ExperimentAIBridge" ] }, "codersdk.ExternalAgentCredentials": { @@ -13116,6 +13166,9 @@ const docTemplate = `{ "description": "ID is a unique identifier for the auth config.\nIt defaults to ` + "`" + `type` + "`" + ` when not provided.", "type": "string" }, + "mcp_url": { + "type": "string" + }, "no_refresh": { "type": "boolean" }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 3dfa9fdf9792d..af282e44147c5 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -9800,6 +9800,50 @@ } } }, + "codersdk.AIBridgeAnthropicConfig": { + "type": "object", + "properties": { + "base_url": { + "type": "string" + }, + "key": { + "type": "string" + } + } + }, + "codersdk.AIBridgeConfig": { + "type": "object", + "properties": { + "anthropic": { + "$ref": "#/definitions/codersdk.AIBridgeAnthropicConfig" + }, + "enabled": { + "type": "boolean" + }, + "openai": { + "$ref": "#/definitions/codersdk.AIBridgeOpenAIConfig" + } + } + }, + "codersdk.AIBridgeOpenAIConfig": { + "type": "object", + "properties": { + "base_url": { + "type": "string" + }, + "key": { + "type": "string" + } + } + }, + "codersdk.AIConfig": { + "type": "object", + "properties": { + "bridge": { + "$ref": "#/definitions/codersdk.AIBridgeConfig" + } + } + }, "codersdk.APIKey": { "type": "object", "required": [ @@ -11326,6 +11370,9 @@ "agent_stat_refresh_interval": { "type": "integer" }, + "ai": { + "$ref": "#/definitions/codersdk.AIConfig" + }, "allow_workspace_renames": { "type": "boolean" }, @@ -11646,9 +11693,11 @@ "web-push", "oauth2", "mcp-server-http", - "workspace-sharing" + "workspace-sharing", + "ai-bridge" ], "x-enum-comments": { + "ExperimentAIBridge": "Enables AI Bridge functionality.", "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentExample": "This isn't used for anything.", "ExperimentMCPServerHTTP": "Enables the MCP HTTP server functionality.", @@ -11666,7 +11715,8 @@ "ExperimentWebPush", "ExperimentOAuth2", "ExperimentMCPServerHTTP", - "ExperimentWorkspaceSharing" + "ExperimentWorkspaceSharing", + "ExperimentAIBridge" ] }, "codersdk.ExternalAgentCredentials": { @@ -11764,6 +11814,9 @@ "description": "ID is a unique identifier for the auth config.\nIt defaults to `type` when not provided.", "type": "string" }, + "mcp_url": { + "type": "string" + }, "no_refresh": { "type": "boolean" }, diff --git a/coderd/coderd.go b/coderd/coderd.go index 724952bde7bb9..aee9baa940028 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -20,6 +20,8 @@ import ( "sync/atomic" "time" + "github.com/coder/coder/v2/aibridged" + aibridgedproto "github.com/coder/coder/v2/aibridged/proto" "github.com/coder/coder/v2/coderd/oauth2provider" "github.com/coder/coder/v2/coderd/pproflabel" "github.com/coder/coder/v2/coderd/prebuilds" @@ -44,6 +46,9 @@ import ( "tailscale.com/types/key" "tailscale.com/util/singleflight" + "github.com/coder/coder/v2/coderd/aibridgedserver" + "github.com/coder/coder/v2/provisionerd/proto" + "cdr.dev/slog" "github.com/coder/quartz" "github.com/coder/serpent" @@ -95,7 +100,6 @@ import ( "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/healthsdk" - "github.com/coder/coder/v2/provisionerd/proto" "github.com/coder/coder/v2/provisionersdk" "github.com/coder/coder/v2/site" "github.com/coder/coder/v2/tailnet" @@ -1582,6 +1586,14 @@ func New(options *Options) *API { r.Route("/init-script", func(r chi.Router) { r.Get("/{os}/{arch}", api.initScript) }) + r.Route("/aibridge", func(r chi.Router) { + r.Use( + httpmw.RequireExperiment(api.Experiments, codersdk.ExperimentAIBridge), + aibridged.AuthMiddleware(api.Database, api.Logger), + ) + r.HandleFunc("/openai/*", api.bridgeAIRequest) + r.HandleFunc("/anthropic/*", api.bridgeAIRequest) + }) }) if options.SwaggerEndpoint { @@ -1742,6 +1754,8 @@ type API struct { // dbRolluper rolls up template usage stats from raw agent and app // stats. This is used to provide insights in the WebUI. dbRolluper *dbrollup.Rolluper + + AIBridgeServer *aibridged.Server } // Close waits for all WebSocket connections to drain before returning. @@ -1972,6 +1986,63 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n return proto.NewDRPCProvisionerDaemonClient(clientSession), nil } +func (api *API) CreateInMemoryAIBridgeDaemon(dialCtx context.Context) (client aibridgedproto.DRPCRecorderClient, err error) { + // TODO(dannyk): implement options. + // TODO(dannyk): implement tracing. + + clientSession, serverSession := drpcsdk.MemTransportPipe() + defer func() { + if err != nil { + _ = clientSession.Close() + _ = serverSession.Close() + } + }() + + // TODO(dannyk): implement API versioning. + + mux := drpcmux.New() + api.Logger.Debug(dialCtx, "starting in-memory AI bridge daemon") + logger := api.Logger.Named("inmem-aibridged") + srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Logger) + if err != nil { + return nil, err + } + err = aibridgedproto.DRPCRegisterRecorder(mux, srv) + if err != nil { + return nil, err + } + server := drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler: mux}, + drpcserver.Options{ + Manager: drpcsdk.DefaultDRPCOptions(nil), + Log: func(err error) { + if xerrors.Is(err, io.EOF) { + return + } + logger.Debug(dialCtx, "drpc server error", slog.Error(err)) + }, + }, + ) + // in-mem pipes aren't technically "websockets" but they have the same properties as far as the + // API is concerned: they are long-lived connections that we need to close before completing + // shutdown of the API. + api.WebsocketWaitMutex.Lock() + api.WebsocketWaitGroup.Add(1) + api.WebsocketWaitMutex.Unlock() + go func() { + defer api.WebsocketWaitGroup.Done() + // Here we pass the background context, since we want the server to keep serving until the + // client hangs up. The aibridged is local, in-mem, so there isn't a danger of losing contact with it and + // having a dead connection we don't know the status of. + err := server.Serve(context.Background(), serverSession) + logger.Info(dialCtx, "AI bridge daemon disconnected", slog.Error(err)) + // close the sessions, so we don't leak goroutines serving them. + _ = clientSession.Close() + _ = serverSession.Close() + }() + + return aibridgedproto.NewDRPCRecorderClient(clientSession), nil +} + func (api *API) DERPMap() *tailcfg.DERPMap { fn := api.DERPMapper.Load() if fn != nil { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 46cdac5e7b71b..29e98e63b6cd5 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -3710,6 +3710,26 @@ func (q *querier) GetWorkspacesEligibleForTransition(ctx context.Context, now ti return q.db.GetWorkspacesEligibleForTransition(ctx, now) } +func (q *querier) InsertAIBridgeSession(ctx context.Context, arg database.InsertAIBridgeSessionParams) error { + // TODO: authz. + return q.db.InsertAIBridgeSession(ctx, arg) +} + +func (q *querier) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error { + // TODO: authz. + return q.db.InsertAIBridgeTokenUsage(ctx, arg) +} + +func (q *querier) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error { + // TODO: authz. + return q.db.InsertAIBridgeToolUsage(ctx, arg) +} + +func (q *querier) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error { + // TODO: authz. + return q.db.InsertAIBridgeUserPrompt(ctx, arg) +} + func (q *querier) InsertAPIKey(ctx context.Context, arg database.InsertAPIKeyParams) (database.APIKey, error) { return insert(q.log, q.auth, rbac.ResourceApiKey.WithOwner(arg.UserID.String()), diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 4b5e953d771dd..2751c03b6efa4 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2140,6 +2140,34 @@ func (m queryMetricsStore) GetWorkspacesEligibleForTransition(ctx context.Contex return workspaces, err } +func (m queryMetricsStore) InsertAIBridgeSession(ctx context.Context, arg database.InsertAIBridgeSessionParams) error { + start := time.Now() + r0 := m.s.InsertAIBridgeSession(ctx, arg) + m.queryLatencies.WithLabelValues("InsertAIBridgeSession").Observe(time.Since(start).Seconds()) + return r0 +} + +func (m queryMetricsStore) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error { + start := time.Now() + r0 := m.s.InsertAIBridgeTokenUsage(ctx, arg) + m.queryLatencies.WithLabelValues("InsertAIBridgeTokenUsage").Observe(time.Since(start).Seconds()) + return r0 +} + +func (m queryMetricsStore) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error { + start := time.Now() + r0 := m.s.InsertAIBridgeToolUsage(ctx, arg) + m.queryLatencies.WithLabelValues("InsertAIBridgeToolUsage").Observe(time.Since(start).Seconds()) + return r0 +} + +func (m queryMetricsStore) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error { + start := time.Now() + r0 := m.s.InsertAIBridgeUserPrompt(ctx, arg) + m.queryLatencies.WithLabelValues("InsertAIBridgeUserPrompt").Observe(time.Since(start).Seconds()) + return r0 +} + func (m queryMetricsStore) InsertAPIKey(ctx context.Context, arg database.InsertAPIKeyParams) (database.APIKey, error) { start := time.Now() key, err := m.s.InsertAPIKey(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 02415d6cb8ea4..24e0719ac0270 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -4575,6 +4575,62 @@ func (mr *MockStoreMockRecorder) InTx(arg0, arg1 any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InTx", reflect.TypeOf((*MockStore)(nil).InTx), arg0, arg1) } +// InsertAIBridgeSession mocks base method. +func (m *MockStore) InsertAIBridgeSession(ctx context.Context, arg database.InsertAIBridgeSessionParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InsertAIBridgeSession", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// InsertAIBridgeSession indicates an expected call of InsertAIBridgeSession. +func (mr *MockStoreMockRecorder) InsertAIBridgeSession(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertAIBridgeSession", reflect.TypeOf((*MockStore)(nil).InsertAIBridgeSession), ctx, arg) +} + +// InsertAIBridgeTokenUsage mocks base method. +func (m *MockStore) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InsertAIBridgeTokenUsage", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// InsertAIBridgeTokenUsage indicates an expected call of InsertAIBridgeTokenUsage. +func (mr *MockStoreMockRecorder) InsertAIBridgeTokenUsage(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertAIBridgeTokenUsage", reflect.TypeOf((*MockStore)(nil).InsertAIBridgeTokenUsage), ctx, arg) +} + +// InsertAIBridgeToolUsage mocks base method. +func (m *MockStore) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InsertAIBridgeToolUsage", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// InsertAIBridgeToolUsage indicates an expected call of InsertAIBridgeToolUsage. +func (mr *MockStoreMockRecorder) InsertAIBridgeToolUsage(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertAIBridgeToolUsage", reflect.TypeOf((*MockStore)(nil).InsertAIBridgeToolUsage), ctx, arg) +} + +// InsertAIBridgeUserPrompt mocks base method. +func (m *MockStore) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InsertAIBridgeUserPrompt", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// InsertAIBridgeUserPrompt indicates an expected call of InsertAIBridgeUserPrompt. +func (mr *MockStoreMockRecorder) InsertAIBridgeUserPrompt(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertAIBridgeUserPrompt", reflect.TypeOf((*MockStore)(nil).InsertAIBridgeUserPrompt), ctx, arg) +} + // InsertAPIKey mocks base method. func (m *MockStore) InsertAPIKey(ctx context.Context, arg database.InsertAPIKeyParams) (database.APIKey, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 066fe0b1b8847..edf951c7321c0 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -814,6 +814,44 @@ BEGIN END; $$; +CREATE TABLE aibridge_sessions ( + id uuid NOT NULL, + initiator_id uuid NOT NULL, + provider text NOT NULL, + model text NOT NULL, + created_at timestamp with time zone DEFAULT now() +); + +CREATE TABLE aibridge_token_usages ( + id uuid NOT NULL, + session_id uuid NOT NULL, + provider_id text NOT NULL, + input_tokens bigint NOT NULL, + output_tokens bigint NOT NULL, + metadata jsonb, + created_at timestamp with time zone DEFAULT now() +); + +CREATE TABLE aibridge_tool_usages ( + id uuid NOT NULL, + session_id uuid NOT NULL, + provider_id text NOT NULL, + tool text NOT NULL, + input text NOT NULL, + injected boolean DEFAULT false NOT NULL, + metadata jsonb, + created_at timestamp with time zone DEFAULT now() +); + +CREATE TABLE aibridge_user_prompts ( + id uuid NOT NULL, + session_id uuid NOT NULL, + provider_id text NOT NULL, + prompt text NOT NULL, + metadata jsonb, + created_at timestamp with time zone DEFAULT now() +); + CREATE TABLE api_keys ( id text NOT NULL, hashed_secret bytea NOT NULL, @@ -2531,6 +2569,18 @@ ALTER TABLE ONLY workspace_resource_metadata ALTER COLUMN id SET DEFAULT nextval ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); +ALTER TABLE ONLY aibridge_sessions + ADD CONSTRAINT aibridge_sessions_pkey PRIMARY KEY (id); + +ALTER TABLE ONLY aibridge_token_usages + ADD CONSTRAINT aibridge_token_usages_pkey PRIMARY KEY (id); + +ALTER TABLE ONLY aibridge_tool_usages + ADD CONSTRAINT aibridge_tool_usages_pkey PRIMARY KEY (id); + +ALTER TABLE ONLY aibridge_user_prompts + ADD CONSTRAINT aibridge_user_prompts_pkey PRIMARY KEY (id); + ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_pkey PRIMARY KEY (id); @@ -2820,6 +2870,24 @@ CREATE INDEX idx_agent_stats_created_at ON workspace_agent_stats USING btree (cr CREATE INDEX idx_agent_stats_user_id ON workspace_agent_stats USING btree (user_id); +CREATE INDEX idx_aibridge_sessions_model ON aibridge_sessions USING btree (model); + +CREATE INDEX idx_aibridge_sessions_provider ON aibridge_sessions USING btree (provider); + +CREATE INDEX idx_aibridge_token_usages_session_id ON aibridge_token_usages USING btree (session_id); + +CREATE INDEX idx_aibridge_token_usages_session_provider_id ON aibridge_token_usages USING btree (session_id, provider_id); + +CREATE INDEX idx_aibridge_tool_usages_session_id ON aibridge_tool_usages USING btree (session_id); + +CREATE INDEX idx_aibridge_tool_usages_session_provider_id ON aibridge_tool_usages USING btree (session_id, provider_id); + +CREATE INDEX idx_aibridge_tool_usages_tool ON aibridge_tool_usages USING btree (tool); + +CREATE INDEX idx_aibridge_user_prompts_session_id ON aibridge_user_prompts USING btree (session_id); + +CREATE INDEX idx_aibridge_user_prompts_session_provider_id ON aibridge_user_prompts USING btree (session_id, provider_id); + CREATE UNIQUE INDEX idx_api_key_name ON api_keys USING btree (user_id, token_name) WHERE (login_type = 'token'::login_type); CREATE INDEX idx_api_keys_user ON api_keys USING btree (user_id); @@ -3056,6 +3124,18 @@ COMMENT ON TRIGGER workspace_agent_name_unique_trigger ON workspace_agents IS 'U the uniqueness requirement. A trigger allows us to enforce uniqueness going forward without requiring a migration to clean up historical data.'; +ALTER TABLE ONLY aibridge_sessions + ADD CONSTRAINT aibridge_sessions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id) ON DELETE CASCADE; + +ALTER TABLE ONLY aibridge_token_usages + ADD CONSTRAINT aibridge_token_usages_session_id_fkey FOREIGN KEY (session_id) REFERENCES aibridge_sessions(id) ON DELETE CASCADE; + +ALTER TABLE ONLY aibridge_tool_usages + ADD CONSTRAINT aibridge_tool_usages_session_id_fkey FOREIGN KEY (session_id) REFERENCES aibridge_sessions(id) ON DELETE CASCADE; + +ALTER TABLE ONLY aibridge_user_prompts + ADD CONSTRAINT aibridge_user_prompts_session_id_fkey FOREIGN KEY (session_id) REFERENCES aibridge_sessions(id) ON DELETE CASCADE; + ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_user_id_uuid_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index 33aa8edd69032..d2bf59c1c356d 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -6,6 +6,10 @@ type ForeignKeyConstraint string // ForeignKeyConstraint enums. const ( + ForeignKeyAibridgeSessionsInitiatorID ForeignKeyConstraint = "aibridge_sessions_initiator_id_fkey" // ALTER TABLE ONLY aibridge_sessions ADD CONSTRAINT aibridge_sessions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id) ON DELETE CASCADE; + ForeignKeyAibridgeTokenUsagesSessionID ForeignKeyConstraint = "aibridge_token_usages_session_id_fkey" // ALTER TABLE ONLY aibridge_token_usages ADD CONSTRAINT aibridge_token_usages_session_id_fkey FOREIGN KEY (session_id) REFERENCES aibridge_sessions(id) ON DELETE CASCADE; + ForeignKeyAibridgeToolUsagesSessionID ForeignKeyConstraint = "aibridge_tool_usages_session_id_fkey" // ALTER TABLE ONLY aibridge_tool_usages ADD CONSTRAINT aibridge_tool_usages_session_id_fkey FOREIGN KEY (session_id) REFERENCES aibridge_sessions(id) ON DELETE CASCADE; + ForeignKeyAibridgeUserPromptsSessionID ForeignKeyConstraint = "aibridge_user_prompts_session_id_fkey" // ALTER TABLE ONLY aibridge_user_prompts ADD CONSTRAINT aibridge_user_prompts_session_id_fkey FOREIGN KEY (session_id) REFERENCES aibridge_sessions(id) ON DELETE CASCADE; ForeignKeyAPIKeysUserIDUUID ForeignKeyConstraint = "api_keys_user_id_uuid_fkey" // ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_user_id_uuid_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyConnectionLogsOrganizationID ForeignKeyConstraint = "connection_logs_organization_id_fkey" // ALTER TABLE ONLY connection_logs ADD CONSTRAINT connection_logs_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; ForeignKeyConnectionLogsWorkspaceID ForeignKeyConstraint = "connection_logs_workspace_id_fkey" // ALTER TABLE ONLY connection_logs ADD CONSTRAINT connection_logs_workspace_id_fkey FOREIGN KEY (workspace_id) REFERENCES workspaces(id) ON DELETE CASCADE; diff --git a/coderd/database/migrations/000362_aibridge.down.sql b/coderd/database/migrations/000362_aibridge.down.sql new file mode 100644 index 0000000000000..ef4b81bd08b0c --- /dev/null +++ b/coderd/database/migrations/000362_aibridge.down.sql @@ -0,0 +1,4 @@ +DROP TABLE IF EXISTS aibridge_tool_usages CASCADE; +DROP TABLE IF EXISTS aibridge_user_prompts CASCADE; +DROP TABLE IF EXISTS aibridge_token_usages CASCADE; +DROP TABLE IF EXISTS aibridge_sessions CASCADE; diff --git a/coderd/database/migrations/000362_aibridge.up.sql b/coderd/database/migrations/000362_aibridge.up.sql new file mode 100644 index 0000000000000..71362d7d79dfd --- /dev/null +++ b/coderd/database/migrations/000362_aibridge.up.sql @@ -0,0 +1,50 @@ +CREATE TABLE IF NOT EXISTS aibridge_sessions ( + id UUID PRIMARY KEY, + initiator_id uuid NOT NULL REFERENCES users(id) ON DELETE CASCADE, + provider TEXT NOT NULL, + model TEXT NOT NULL, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +CREATE INDEX idx_aibridge_sessions_provider ON aibridge_sessions(provider); +CREATE INDEX idx_aibridge_sessions_model ON aibridge_sessions(model); + +CREATE TABLE IF NOT EXISTS aibridge_token_usages ( + id UUID PRIMARY KEY, + session_id UUID NOT NULL REFERENCES aibridge_sessions(id) ON DELETE CASCADE, + provider_id TEXT NOT NULL, + input_tokens BIGINT NOT NULL, + output_tokens BIGINT NOT NULL, + metadata JSONB DEFAULT NULL, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +CREATE INDEX idx_aibridge_token_usages_session_id ON aibridge_token_usages(session_id); +CREATE INDEX idx_aibridge_token_usages_session_provider_id ON aibridge_token_usages(session_id, provider_id); + +CREATE TABLE IF NOT EXISTS aibridge_user_prompts ( + id UUID PRIMARY KEY, + session_id UUID NOT NULL REFERENCES aibridge_sessions(id) ON DELETE CASCADE, + provider_id TEXT NOT NULL, + prompt TEXT NOT NULL, + metadata JSONB DEFAULT NULL, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +CREATE INDEX idx_aibridge_user_prompts_session_id ON aibridge_user_prompts(session_id); +CREATE INDEX idx_aibridge_user_prompts_session_provider_id ON aibridge_user_prompts(session_id, provider_id); + +CREATE TABLE IF NOT EXISTS aibridge_tool_usages ( + id UUID PRIMARY KEY, + session_id UUID NOT NULL REFERENCES aibridge_sessions(id) ON DELETE CASCADE, + provider_id TEXT NOT NULL, + tool TEXT NOT NULL, + input TEXT NOT NULL, + injected BOOLEAN NOT NULL DEFAULT FALSE, + metadata JSONB DEFAULT NULL, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +CREATE INDEX idx_aibridge_tool_usages_session_id ON aibridge_tool_usages(session_id); +CREATE INDEX idx_aibridge_tool_usages_tool ON aibridge_tool_usages(tool); +CREATE INDEX idx_aibridge_tool_usages_session_provider_id ON aibridge_tool_usages(session_id, provider_id); diff --git a/coderd/database/models.go b/coderd/database/models.go index effd436f4d18d..b30473d3b4124 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -2968,6 +2968,44 @@ type APIKey struct { TokenName string `db:"token_name" json:"token_name"` } +type AibridgeSession struct { + ID uuid.UUID `db:"id" json:"id"` + InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` + Provider string `db:"provider" json:"provider"` + Model string `db:"model" json:"model"` + CreatedAt sql.NullTime `db:"created_at" json:"created_at"` +} + +type AibridgeTokenUsage struct { + ID uuid.UUID `db:"id" json:"id"` + SessionID uuid.UUID `db:"session_id" json:"session_id"` + ProviderID string `db:"provider_id" json:"provider_id"` + InputTokens int64 `db:"input_tokens" json:"input_tokens"` + OutputTokens int64 `db:"output_tokens" json:"output_tokens"` + Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"` + CreatedAt sql.NullTime `db:"created_at" json:"created_at"` +} + +type AibridgeToolUsage struct { + ID uuid.UUID `db:"id" json:"id"` + SessionID uuid.UUID `db:"session_id" json:"session_id"` + ProviderID string `db:"provider_id" json:"provider_id"` + Tool string `db:"tool" json:"tool"` + Input string `db:"input" json:"input"` + Injected bool `db:"injected" json:"injected"` + Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"` + CreatedAt sql.NullTime `db:"created_at" json:"created_at"` +} + +type AibridgeUserPrompt struct { + ID uuid.UUID `db:"id" json:"id"` + SessionID uuid.UUID `db:"session_id" json:"session_id"` + ProviderID string `db:"provider_id" json:"provider_id"` + Prompt string `db:"prompt" json:"prompt"` + Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"` + CreatedAt sql.NullTime `db:"created_at" json:"created_at"` +} + type AuditLog struct { ID uuid.UUID `db:"id" json:"id"` Time time.Time `db:"time" json:"time"` diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 28ed7609c53d6..fb12bf17a6c1c 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -483,6 +483,10 @@ type sqlcQuerier interface { GetWorkspacesAndAgentsByOwnerID(ctx context.Context, ownerID uuid.UUID) ([]GetWorkspacesAndAgentsByOwnerIDRow, error) GetWorkspacesByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceTable, error) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]GetWorkspacesEligibleForTransitionRow, error) + InsertAIBridgeSession(ctx context.Context, arg InsertAIBridgeSessionParams) error + InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) error + InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBridgeToolUsageParams) error + InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIBridgeUserPromptParams) error InsertAPIKey(ctx context.Context, arg InsertAPIKeyParams) (APIKey, error) // We use the organization_id as the id // for simplicity since all users is diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 2f56b422f350b..3c3fb60e2c3fe 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -111,6 +111,115 @@ func (q *sqlQuerier) ActivityBumpWorkspace(ctx context.Context, arg ActivityBump return err } +const insertAIBridgeSession = `-- name: InsertAIBridgeSession :exec +INSERT INTO aibridge_sessions (id, initiator_id, provider, model) +VALUES ($1::uuid, $2::uuid, $3, $4) +` + +type InsertAIBridgeSessionParams struct { + ID uuid.UUID `db:"id" json:"id"` + InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` + Provider string `db:"provider" json:"provider"` + Model string `db:"model" json:"model"` +} + +func (q *sqlQuerier) InsertAIBridgeSession(ctx context.Context, arg InsertAIBridgeSessionParams) error { + _, err := q.db.ExecContext(ctx, insertAIBridgeSession, + arg.ID, + arg.InitiatorID, + arg.Provider, + arg.Model, + ) + return err +} + +const insertAIBridgeTokenUsage = `-- name: InsertAIBridgeTokenUsage :exec +INSERT INTO aibridge_token_usages ( + id, session_id, provider_id, input_tokens, output_tokens, metadata +) VALUES ( + $1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb) +) +` + +type InsertAIBridgeTokenUsageParams struct { + ID uuid.UUID `db:"id" json:"id"` + SessionID uuid.UUID `db:"session_id" json:"session_id"` + ProviderID string `db:"provider_id" json:"provider_id"` + InputTokens int64 `db:"input_tokens" json:"input_tokens"` + OutputTokens int64 `db:"output_tokens" json:"output_tokens"` + Metadata json.RawMessage `db:"metadata" json:"metadata"` +} + +func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) error { + _, err := q.db.ExecContext(ctx, insertAIBridgeTokenUsage, + arg.ID, + arg.SessionID, + arg.ProviderID, + arg.InputTokens, + arg.OutputTokens, + arg.Metadata, + ) + return err +} + +const insertAIBridgeToolUsage = `-- name: InsertAIBridgeToolUsage :exec +INSERT INTO aibridge_tool_usages ( + id, session_id, provider_id, tool, input, injected, metadata +) VALUES ( + $1, $2, $3, $4, $5, $6, COALESCE($7::jsonb, '{}'::jsonb) +) +` + +type InsertAIBridgeToolUsageParams struct { + ID uuid.UUID `db:"id" json:"id"` + SessionID uuid.UUID `db:"session_id" json:"session_id"` + ProviderID string `db:"provider_id" json:"provider_id"` + Tool string `db:"tool" json:"tool"` + Input string `db:"input" json:"input"` + Injected bool `db:"injected" json:"injected"` + Metadata json.RawMessage `db:"metadata" json:"metadata"` +} + +func (q *sqlQuerier) InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBridgeToolUsageParams) error { + _, err := q.db.ExecContext(ctx, insertAIBridgeToolUsage, + arg.ID, + arg.SessionID, + arg.ProviderID, + arg.Tool, + arg.Input, + arg.Injected, + arg.Metadata, + ) + return err +} + +const insertAIBridgeUserPrompt = `-- name: InsertAIBridgeUserPrompt :exec +INSERT INTO aibridge_user_prompts ( + id, session_id, provider_id, prompt, metadata +) VALUES ( + $1, $2, $3, $4, COALESCE($5::jsonb, '{}'::jsonb) +) +` + +type InsertAIBridgeUserPromptParams struct { + ID uuid.UUID `db:"id" json:"id"` + SessionID uuid.UUID `db:"session_id" json:"session_id"` + ProviderID string `db:"provider_id" json:"provider_id"` + Prompt string `db:"prompt" json:"prompt"` + Metadata json.RawMessage `db:"metadata" json:"metadata"` +} + +func (q *sqlQuerier) InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIBridgeUserPromptParams) error { + _, err := q.db.ExecContext(ctx, insertAIBridgeUserPrompt, + arg.ID, + arg.SessionID, + arg.ProviderID, + arg.Prompt, + arg.Metadata, + ) + return err +} + const deleteAPIKeyByID = `-- name: DeleteAPIKeyByID :exec DELETE FROM api_keys diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql new file mode 100644 index 0000000000000..4089b5c9dec32 --- /dev/null +++ b/coderd/database/queries/aibridge.sql @@ -0,0 +1,24 @@ +-- name: InsertAIBridgeSession :exec +INSERT INTO aibridge_sessions (id, initiator_id, provider, model) +VALUES (@id::uuid, @initiator_id::uuid, @provider, @model); + +-- name: InsertAIBridgeTokenUsage :exec +INSERT INTO aibridge_token_usages ( + id, session_id, provider_id, input_tokens, output_tokens, metadata +) VALUES ( + @id, @session_id, @provider_id, @input_tokens, @output_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb) +); + +-- name: InsertAIBridgeUserPrompt :exec +INSERT INTO aibridge_user_prompts ( + id, session_id, provider_id, prompt, metadata +) VALUES ( + @id, @session_id, @provider_id, @prompt, COALESCE(@metadata::jsonb, '{}'::jsonb) +); + +-- name: InsertAIBridgeToolUsage :exec +INSERT INTO aibridge_tool_usages ( + id, session_id, provider_id, tool, input, injected, metadata +) VALUES ( + @id, @session_id, @provider_id, @tool, @input, @injected, COALESCE(@metadata::jsonb, '{}'::jsonb) +); diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 1b0b13ea2ba5a..52fe5f9adc975 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -7,6 +7,10 @@ type UniqueConstraint string // UniqueConstraint enums. const ( UniqueAgentStatsPkey UniqueConstraint = "agent_stats_pkey" // ALTER TABLE ONLY workspace_agent_stats ADD CONSTRAINT agent_stats_pkey PRIMARY KEY (id); + UniqueAibridgeSessionsPkey UniqueConstraint = "aibridge_sessions_pkey" // ALTER TABLE ONLY aibridge_sessions ADD CONSTRAINT aibridge_sessions_pkey PRIMARY KEY (id); + UniqueAibridgeTokenUsagesPkey UniqueConstraint = "aibridge_token_usages_pkey" // ALTER TABLE ONLY aibridge_token_usages ADD CONSTRAINT aibridge_token_usages_pkey PRIMARY KEY (id); + UniqueAibridgeToolUsagesPkey UniqueConstraint = "aibridge_tool_usages_pkey" // ALTER TABLE ONLY aibridge_tool_usages ADD CONSTRAINT aibridge_tool_usages_pkey PRIMARY KEY (id); + UniqueAibridgeUserPromptsPkey UniqueConstraint = "aibridge_user_prompts_pkey" // ALTER TABLE ONLY aibridge_user_prompts ADD CONSTRAINT aibridge_user_prompts_pkey PRIMARY KEY (id); UniqueAPIKeysPkey UniqueConstraint = "api_keys_pkey" // ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_pkey PRIMARY KEY (id); UniqueAuditLogsPkey UniqueConstraint = "audit_logs_pkey" // ALTER TABLE ONLY audit_logs ADD CONSTRAINT audit_logs_pkey PRIMARY KEY (id); UniqueConnectionLogsPkey UniqueConstraint = "connection_logs_pkey" // ALTER TABLE ONLY connection_logs ADD CONSTRAINT connection_logs_pkey PRIMARY KEY (id); diff --git a/coderd/externalauth/externalauth.go b/coderd/externalauth/externalauth.go index 24ebe13d03074..9c2e5d1a78a58 100644 --- a/coderd/externalauth/externalauth.go +++ b/coderd/externalauth/externalauth.go @@ -81,6 +81,9 @@ type Config struct { // AppInstallationsURL is an API endpoint that returns a list of // installations for the user. This is used for GitHub Apps. AppInstallationsURL string + // MCPURL is the endpoint that clients must use to communicate with the associated + // MCP server. + MCPURL string } // GenerateTokenExtra generates the extra token data to store in the database. @@ -620,6 +623,7 @@ func ConvertConfig(instrument *promoauth.Factory, entries []codersdk.ExternalAut DisplayName: entry.DisplayName, DisplayIcon: entry.DisplayIcon, ExtraTokenKeys: entry.ExtraTokenKeys, + MCPURL: entry.MCPURL, } if entry.DeviceFlow { diff --git a/coderd/httpmw/cors.go b/coderd/httpmw/cors.go index 218aab6609f60..2d59eb55fabc5 100644 --- a/coderd/httpmw/cors.go +++ b/coderd/httpmw/cors.go @@ -50,19 +50,29 @@ func Cors(allowAll bool, origins ...string) func(next http.Handler) http.Handler permissiveCors := cors.Handler(cors.Options{ AllowedOrigins: []string{"*"}, AllowedMethods: []string{ + http.MethodHead, http.MethodGet, http.MethodPost, http.MethodDelete, http.MethodOptions, }, AllowedHeaders: []string{ - "Content-Type", "Accept", + "Content-Type", "Authorization", - "x-api-key", + "X-Api-Key", + "X-Requested-With", + "Last-Event-ID", + // MCP headers "Mcp-Session-Id", "MCP-Protocol-Version", - "Last-Event-ID", + // Provider-specific headers + "OpenAI-Organization", + "OpenAI-Beta", + "Anthropic-Version", + "Anthropic-Beta", + "anthropic-version", + "anthropic-beta", }, ExposedHeaders: []string{ "Content-Type", @@ -77,10 +87,11 @@ func Cors(allowAll bool, origins ...string) func(next http.Handler) http.Handler return func(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Use permissive CORS for OAuth2, MCP, and well-known endpoints + // Use permissive CORS for OAuth2, MCP, well-known, and AI bridge endpoints if strings.HasPrefix(r.URL.Path, "/oauth2/") || strings.HasPrefix(r.URL.Path, "/api/experimental/mcp/") || - strings.HasPrefix(r.URL.Path, "/.well-known/oauth-") { + strings.HasPrefix(r.URL.Path, "/.well-known/oauth-") || + strings.HasPrefix(r.URL.Path, "/api/v2/aibridge") { permissiveCors(next).ServeHTTP(w, r) return } diff --git a/codersdk/deployment.go b/codersdk/deployment.go index a70a6b55500d2..bc1a08d80d25b 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -500,6 +500,7 @@ type DeploymentValues struct { WorkspaceHostnameSuffix serpent.String `json:"workspace_hostname_suffix,omitempty" typescript:",notnull"` Prebuilds PrebuildsConfig `json:"workspace_prebuilds,omitempty" typescript:",notnull"` HideAITasks serpent.Bool `json:"hide_ai_tasks,omitempty" typescript:",notnull"` + AI AIConfig `json:"ai,omitempty"` Config serpent.YAMLConfigPath `json:"config,omitempty" typescript:",notnull"` WriteConfig serpent.Bool `json:"write_config,omitempty" typescript:",notnull"` @@ -736,6 +737,7 @@ type ExternalAuthConfig struct { ExtraTokenKeys []string `json:"-" yaml:"extra_token_keys"` DeviceFlow bool `json:"device_flow" yaml:"device_flow"` DeviceCodeURL string `json:"device_code_url" yaml:"device_code_url"` + MCPURL string `json:"mcp_url" yaml:"mcp_url"` // Regex allows API requesters to match an auth config by // a string (e.g. coder.com) instead of by it's type. // @@ -1158,6 +1160,10 @@ func (c *DeploymentValues) Options() serpent.OptionSet { Parent: &deploymentGroupNotifications, YAML: "inbox", } + deploymentGroupAIBridge = serpent.Group{ + Name: "AI Bridge", + YAML: "ai_bridge", + } ) httpAddress := serpent.Option{ @@ -3208,11 +3214,88 @@ Write out the current server config as YAML to stdout.`, Group: &deploymentGroupClient, YAML: "hideAITasks", }, + + // AI Bridge Options + { + Name: "AI Bridge Enabled", + Description: "Whether to start an in-memory aibridged instance ('ai-bridge' experiment must be enabled, too).", + Flag: "ai-bridge-enabled", + Env: "CODER_AI_BRIDGE_ENABLED", + Value: &c.AI.BridgeConfig.Enabled, + Default: "true", + Group: &deploymentGroupAIBridge, + YAML: "enabled", + Hidden: false, + }, + { + Name: "AI Bridge OpenAI Base URL", + Description: "TODO.", + Flag: "ai-bridge-openai-base-url", + Env: "CODER_AI_BRIDGE_OPENAI_BASE_URL", + Value: &c.AI.BridgeConfig.OpenAI.BaseURL, + Default: "https://api.openai.com/v1/", + Group: &deploymentGroupAIBridge, + YAML: "openai_base_url", + Hidden: true, + }, + { + Name: "AI Bridge OpenAI Key", + Description: "TODO.", + Flag: "ai-bridge-openai-key", + Env: "CODER_AI_BRIDGE_OPENAI_KEY", + Value: &c.AI.BridgeConfig.OpenAI.Key, + Default: "", + Group: &deploymentGroupAIBridge, + YAML: "openai_key", + Hidden: true, + }, + { + Name: "AI Bridge Anthropic Base URL", + Description: "TODO.", + Flag: "ai-bridge-anthropic-base-url", + Env: "CODER_AI_BRIDGE_ANTHROPIC_BASE_URL", + Value: &c.AI.BridgeConfig.Anthropic.BaseURL, + Default: "https://api.anthropic.com/", + Group: &deploymentGroupAIBridge, + YAML: "base_url", + Hidden: true, + }, + { + Name: "AI Bridge Anthropic KEY", + Description: "TODO.", + Flag: "ai-bridge-anthropic-key", + Env: "CODER_AI_BRIDGE_ANTHROPIC_KEY", + Value: &c.AI.BridgeConfig.Anthropic.Key, + Default: "", + Group: &deploymentGroupAIBridge, + YAML: "key", + Hidden: true, + }, } return opts } +type AIBridgeConfig struct { + Enabled serpent.Bool `json:"enabled" typescript:",notnull"` + OpenAI AIBridgeOpenAIConfig `json:"openai" typescript:",notnull"` + Anthropic AIBridgeAnthropicConfig `json:"anthropic" typescript:",notnull"` +} + +type AIBridgeOpenAIConfig struct { + BaseURL serpent.String `json:"base_url" typescript:",notnull"` + Key serpent.String `json:"key" typescript:",notnull"` +} + +type AIBridgeAnthropicConfig struct { + BaseURL serpent.String `json:"base_url" typescript:",notnull"` + Key serpent.String `json:"key" typescript:",notnull"` +} + +type AIConfig struct { + BridgeConfig AIBridgeConfig `json:"bridge,omitempty"` +} + type SupportConfig struct { Links serpent.Struct[[]LinkConfig] `json:"links" typescript:",notnull"` } @@ -3436,6 +3519,7 @@ const ( ExperimentOAuth2 Experiment = "oauth2" // Enables OAuth2 provider functionality. ExperimentMCPServerHTTP Experiment = "mcp-server-http" // Enables the MCP HTTP server functionality. ExperimentWorkspaceSharing Experiment = "workspace-sharing" // Enables updating workspace ACLs for sharing with users and groups. + ExperimentAIBridge Experiment = "ai-bridge" // Enables AI Bridge functionality. ) func (e Experiment) DisplayName() string { @@ -3456,6 +3540,8 @@ func (e Experiment) DisplayName() string { return "MCP HTTP Server Functionality" case ExperimentWorkspaceSharing: return "Workspace Sharing" + case ExperimentAIBridge: + return "AI Bridge" default: // Split on hyphen and convert to title case // e.g. "web-push" -> "Web Push", "mcp-server-http" -> "Mcp Server Http" @@ -3474,6 +3560,7 @@ var ExperimentsKnown = Experiments{ ExperimentOAuth2, ExperimentMCPServerHTTP, ExperimentWorkspaceSharing, + ExperimentAIBridge, } // ExperimentsSafe should include all experiments that are safe for diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index 72543f6774dfd..3923b8ca8c88f 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -161,6 +161,19 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ "user": {} }, "agent_stat_refresh_interval": 0, + "ai": { + "bridge": { + "anthropic": { + "base_url": "string", + "key": "string" + }, + "enabled": true, + "openai": { + "base_url": "string", + "key": "string" + } + } + }, "allow_workspace_renames": true, "autobuild_poll_interval": 0, "browser_only": true, @@ -241,6 +254,7 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ "display_icon": "string", "display_name": "string", "id": "string", + "mcp_url": "string", "no_refresh": true, "regex": "string", "scopes": [ diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 99e852b3fe4b9..296c76c4508df 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -335,6 +335,86 @@ | `groups` | array of [codersdk.Group](#codersdkgroup) | false | | | | `users` | array of [codersdk.ReducedUser](#codersdkreduceduser) | false | | | +## codersdk.AIBridgeAnthropicConfig + +```json +{ + "base_url": "string", + "key": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|------------|--------|----------|--------------|-------------| +| `base_url` | string | false | | | +| `key` | string | false | | | + +## codersdk.AIBridgeConfig + +```json +{ + "anthropic": { + "base_url": "string", + "key": "string" + }, + "enabled": true, + "openai": { + "base_url": "string", + "key": "string" + } +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|-------------|----------------------------------------------------------------------|----------|--------------|-------------| +| `anthropic` | [codersdk.AIBridgeAnthropicConfig](#codersdkaibridgeanthropicconfig) | false | | | +| `enabled` | boolean | false | | | +| `openai` | [codersdk.AIBridgeOpenAIConfig](#codersdkaibridgeopenaiconfig) | false | | | + +## codersdk.AIBridgeOpenAIConfig + +```json +{ + "base_url": "string", + "key": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|------------|--------|----------|--------------|-------------| +| `base_url` | string | false | | | +| `key` | string | false | | | + +## codersdk.AIConfig + +```json +{ + "bridge": { + "anthropic": { + "base_url": "string", + "key": "string" + }, + "enabled": true, + "openai": { + "base_url": "string", + "key": "string" + } + } +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------|----------------------------------------------------|----------|--------------|-------------| +| `bridge` | [codersdk.AIBridgeConfig](#codersdkaibridgeconfig) | false | | | + ## codersdk.APIKey ```json @@ -2152,6 +2232,19 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "user": {} }, "agent_stat_refresh_interval": 0, + "ai": { + "bridge": { + "anthropic": { + "base_url": "string", + "key": "string" + }, + "enabled": true, + "openai": { + "base_url": "string", + "key": "string" + } + } + }, "allow_workspace_renames": true, "autobuild_poll_interval": 0, "browser_only": true, @@ -2232,6 +2325,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "display_icon": "string", "display_name": "string", "id": "string", + "mcp_url": "string", "no_refresh": true, "regex": "string", "scopes": [ @@ -2639,6 +2733,19 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "user": {} }, "agent_stat_refresh_interval": 0, + "ai": { + "bridge": { + "anthropic": { + "base_url": "string", + "key": "string" + }, + "enabled": true, + "openai": { + "base_url": "string", + "key": "string" + } + } + }, "allow_workspace_renames": true, "autobuild_poll_interval": 0, "browser_only": true, @@ -2719,6 +2826,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "display_icon": "string", "display_name": "string", "id": "string", + "mcp_url": "string", "no_refresh": true, "regex": "string", "scopes": [ @@ -3017,6 +3125,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o | `address` | [serpent.HostPort](#serpenthostport) | false | | Deprecated: Use HTTPAddress or TLS.Address instead. | | `agent_fallback_troubleshooting_url` | [serpent.URL](#serpenturl) | false | | | | `agent_stat_refresh_interval` | integer | false | | | +| `ai` | [codersdk.AIConfig](#codersdkaiconfig) | false | | | | `allow_workspace_renames` | boolean | false | | | | `autobuild_poll_interval` | integer | false | | | | `browser_only` | boolean | false | | | @@ -3321,6 +3430,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o | `oauth2` | | `mcp-server-http` | | `workspace-sharing` | +| `ai-bridge` | ## codersdk.ExternalAgentCredentials @@ -3419,6 +3529,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "display_icon": "string", "display_name": "string", "id": "string", + "mcp_url": "string", "no_refresh": true, "regex": "string", "scopes": [ @@ -3443,6 +3554,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o | `display_icon` | string | false | | Display icon is a URL to an icon to display in the UI. | | `display_name` | string | false | | Display name is shown in the UI to identify the auth config. | | `id` | string | false | | ID is a unique identifier for the auth config. It defaults to `type` when not provided. | +| `mcp_url` | string | false | | | | `no_refresh` | boolean | false | | | |`regex`|string|false||Regex allows API requesters to match an auth config by a string (e.g. coder.com) instead of by it's type. Git clone makes use of this by parsing the URL from: 'Username for "https://github.com":' And sending it to the Coder server to match against the Regex.| @@ -12698,6 +12810,7 @@ None "display_icon": "string", "display_name": "string", "id": "string", + "mcp_url": "string", "no_refresh": true, "regex": "string", "scopes": [ diff --git a/docs/reference/cli/server.md b/docs/reference/cli/server.md index 8d601cace5d1d..e3b760555f1ca 100644 --- a/docs/reference/cli/server.md +++ b/docs/reference/cli/server.md @@ -1636,3 +1636,14 @@ How often to reconcile workspace prebuilds state. | Default | false | Hide AI tasks from the dashboard. + +### --ai-bridge-enabled + +| | | +|-------------|---------------------------------------| +| Type | bool | +| Environment | $CODER_AI_BRIDGE_ENABLED | +| YAML | ai_bridge.enabled | +| Default | true | + +Whether to start an in-memory aibridged instance ('ai-bridge' experiment must be enabled, too). diff --git a/enterprise/cli/testdata/coder_server_--help.golden b/enterprise/cli/testdata/coder_server_--help.golden index e86cb52692ec3..c11876dda8628 100644 --- a/enterprise/cli/testdata/coder_server_--help.golden +++ b/enterprise/cli/testdata/coder_server_--help.golden @@ -77,6 +77,11 @@ OPTIONS: Periodically check for new releases of Coder and inform the owner. The check is performed once per day. +AI BRIDGE OPTIONS: + --ai-bridge-enabled bool, $CODER_AI_BRIDGE_ENABLED (default: true) + Whether to start an in-memory aibridged instance ('ai-bridge' + experiment must be enabled, too). + CLIENT OPTIONS: These options change the behavior of how clients interact with the Coder. Clients include the Coder CLI, Coder Desktop, IDE extensions, and the web UI. diff --git a/go.mod b/go.mod index 3f9d92aa54c0e..258dbe8dc2b51 100644 --- a/go.mod +++ b/go.mod @@ -196,7 +196,7 @@ require ( go.uber.org/mock v0.5.0 go4.org/netipx v0.0.0-20230728180743-ad4cb58a6516 golang.org/x/crypto v0.41.0 - golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 + golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b golang.org/x/mod v0.27.0 golang.org/x/net v0.43.0 golang.org/x/oauth2 v0.30.0 @@ -208,14 +208,14 @@ require ( golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da google.golang.org/api v0.246.0 google.golang.org/grpc v1.74.2 - google.golang.org/protobuf v1.36.6 + google.golang.org/protobuf v1.36.7 gopkg.in/DataDog/dd-trace-go.v1 v1.74.0 gopkg.in/natefinch/lumberjack.v2 v2.2.1 gopkg.in/yaml.v3 v3.0.1 gvisor.dev/gvisor v0.0.0-20240509041132-65b30f7869dc kernel.org/pub/linux/libs/security/libcap/cap v1.2.73 - storj.io/drpc v0.0.33 - tailscale.com v1.80.3 + storj.io/drpc v0.0.34 + tailscale.com v1.86.5 ) require ( @@ -466,7 +466,11 @@ require ( require github.com/coder/clistat v1.0.0 -require github.com/SherClockHolmes/webpush-go v1.4.0 +require ( + github.com/SherClockHolmes/webpush-go v1.4.0 + github.com/coder/aibridge v0.0.0 + github.com/dgraph-io/ristretto/v2 v2.3.0 +) require ( github.com/charmbracelet/colorprofile v0.2.3-0.20250311203215-f60798e515dc // indirect @@ -484,9 +488,13 @@ require ( github.com/coder/preview v1.0.3 github.com/fsnotify/fsnotify v1.9.0 github.com/go-git/go-git/v5 v5.16.2 - github.com/mark3labs/mcp-go v0.37.0 + github.com/mark3labs/mcp-go v0.38.0 + github.com/tidwall/sjson v1.2.5 // indirect ) +// aibridge-related deps and directives. // TODO: remove. +replace github.com/coder/aibridge v0.0.0 => /home/coder/aibridge + require ( cel.dev/expr v0.24.0 // indirect cloud.google.com/go v0.120.0 // indirect @@ -509,7 +517,6 @@ require ( github.com/buger/jsonparser v1.1.1 // indirect github.com/charmbracelet/x/exp/slice v0.0.0-20250327172914-2fdc97757edf // indirect github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 // indirect - github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da // indirect github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect github.com/esiqveland/notify v0.13.3 // indirect @@ -524,15 +531,14 @@ require ( github.com/klauspost/cpuid/v2 v2.2.10 // indirect github.com/moby/sys/user v0.4.0 // indirect github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646 // indirect - github.com/openai/openai-go v1.7.0 // indirect + github.com/openai/openai-go v1.12.0 // indirect github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect github.com/puzpuzpuz/xsync/v3 v3.5.1 // indirect github.com/samber/lo v1.50.0 // indirect github.com/sergeymakinen/go-bmp v1.0.0 // indirect github.com/sergeymakinen/go-ico v1.0.0-beta.0 // indirect github.com/spiffe/go-spiffe/v2 v2.5.0 // indirect - github.com/tidwall/sjson v1.2.5 // indirect - github.com/tmaxmax/go-sse v0.10.0 // indirect + github.com/tmaxmax/go-sse v0.11.0 // indirect github.com/ulikunitz/xz v0.5.12 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect diff --git a/go.sum b/go.sum index 4bc0e0336ab06..8da3d9e9d3ac8 100644 --- a/go.sum +++ b/go.sum @@ -977,8 +977,8 @@ github.com/dblohm7/wingoes v0.0.0-20240820181039-f2b84150679e h1:L+XrFvD0vBIBm+W github.com/dblohm7/wingoes v0.0.0-20240820181039-f2b84150679e/go.mod h1:SUxUaAK/0UG5lYyZR1L1nC4AaYYvSSYTWQSH3FPcxKU= github.com/dgraph-io/badger/v4 v4.7.0 h1:Q+J8HApYAY7UMpL8d9owqiB+odzEc0zn/aqOD9jhc6Y= github.com/dgraph-io/badger/v4 v4.7.0/go.mod h1:He7TzG3YBy3j4f5baj5B7Zl2XyfNe5bl4Udl0aPemVA= -github.com/dgraph-io/ristretto/v2 v2.2.0 h1:bkY3XzJcXoMuELV8F+vS8kzNgicwQFAaGINAEJdWGOM= -github.com/dgraph-io/ristretto/v2 v2.2.0/go.mod h1:RZrm63UmcBAaYWC1DotLYBmTvgkrs0+XhBd7Npn7/zI= +github.com/dgraph-io/ristretto/v2 v2.3.0 h1:qTQ38m7oIyd4GAed/QkUZyPFNMnvVWyazGXRwvOt5zk= +github.com/dgraph-io/ristretto/v2 v2.3.0/go.mod h1:gpoRV3VzrEY1a9dWAYV6T1U7YzfgttXdd/ZzL1s9OZM= github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da h1:aIftn67I1fkbMa512G+w+Pxci9hJPB8oMnkcP3iZF38= github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw= @@ -1511,8 +1511,8 @@ github.com/makeworld-the-better-one/dither/v2 v2.4.0 h1:Az/dYXiTcwcRSe59Hzw4RI1r github.com/makeworld-the-better-one/dither/v2 v2.4.0/go.mod h1:VBtN8DXO7SNtyGmLiGA7IsFeKrBkQPze1/iAeM95arc= github.com/marekm4/color-extractor v1.2.1 h1:3Zb2tQsn6bITZ8MBVhc33Qn1k5/SEuZ18mrXGUqIwn0= github.com/marekm4/color-extractor v1.2.1/go.mod h1:90VjmiHI6M8ez9eYUaXLdcKnS+BAOp7w+NpwBdkJmpA= -github.com/mark3labs/mcp-go v0.37.0 h1:BywvZLPRT6Zx6mMG/MJfxLSZQkTGIcJSEGKsvr4DsoQ= -github.com/mark3labs/mcp-go v0.37.0/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g= +github.com/mark3labs/mcp-go v0.38.0 h1:E5tmJiIXkhwlV0pLAwAT0O5ZjUZSISE/2Jxg+6vpq4I= +github.com/mark3labs/mcp-go v0.38.0/go.mod h1:T7tUa2jO6MavG+3P25Oy/jR7iCeJPHImCZHRymCn39g= github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU= github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE= github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc= @@ -1625,8 +1625,8 @@ github.com/open-telemetry/opentelemetry-collector-contrib/pkg/sampling v0.120.1 github.com/open-telemetry/opentelemetry-collector-contrib/pkg/sampling v0.120.1/go.mod h1:01TvyaK8x640crO2iFwW/6CFCZgNsOvOGH3B5J239m0= github.com/open-telemetry/opentelemetry-collector-contrib/processor/probabilisticsamplerprocessor v0.120.1 h1:TCyOus9tym82PD1VYtthLKMVMlVyRwtDI4ck4SR2+Ok= github.com/open-telemetry/opentelemetry-collector-contrib/processor/probabilisticsamplerprocessor v0.120.1/go.mod h1:Z/S1brD5gU2Ntht/bHxBVnGxXKTvZDr0dNv/riUzPmY= -github.com/openai/openai-go v1.7.0 h1:M1JfDjQgo3d3PsLyZgpGUG0wUAaUAitqJPM4Rl56dCA= -github.com/openai/openai-go v1.7.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y= +github.com/openai/openai-go v1.12.0 h1:NBQCnXzqOTv5wsgNC36PrFEiskGfO5wccfCWDo9S1U0= +github.com/openai/openai-go v1.12.0/go.mod h1:g461MYGXEXBVdV5SaR/5tNzNbSfwTBBefwc+LlDCK0Y= github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= @@ -1825,8 +1825,8 @@ github.com/tklauser/go-sysconf v0.3.15 h1:VE89k0criAymJ/Os65CSn1IXaol+1wrsFHEB8O github.com/tklauser/go-sysconf v0.3.15/go.mod h1:Dmjwr6tYFIseJw7a3dRLJfsHAMXZ3nEnL/aZY+0IuI4= github.com/tklauser/numcpus v0.10.0 h1:18njr6LDBk1zuna922MgdjQuJFjrdppsZG60sHGfjso= github.com/tklauser/numcpus v0.10.0/go.mod h1:BiTKazU708GQTYF4mB+cmlpT2Is1gLk7XVuEeem8LsQ= -github.com/tmaxmax/go-sse v0.10.0 h1:j9F93WB4Hxt8wUf6oGffMm4dutALvUPoDDxfuDQOSqA= -github.com/tmaxmax/go-sse v0.10.0/go.mod h1:u/2kZQR1tyngo1lKaNCj1mJmhXGZWS1Zs5yiSOD+Eg8= +github.com/tmaxmax/go-sse v0.11.0 h1:nogmJM6rJUoOLoAwEKeQe5XlVpt9l7N82SS1jI7lWFg= +github.com/tmaxmax/go-sse v0.11.0/go.mod h1:u/2kZQR1tyngo1lKaNCj1mJmhXGZWS1Zs5yiSOD+Eg8= github.com/u-root/gobusybox/src v0.0.0-20240225013946-a274a8d5d83a h1:eg5FkNoQp76ZsswyGZ+TjYqA/rhKefxK8BW7XOlQsxo= github.com/u-root/gobusybox/src v0.0.0-20240225013946-a274a8d5d83a/go.mod h1:e/8TmrdreH0sZOw2DFKBaUV7bvDWRq6SeM9PzkuVM68= github.com/u-root/u-root v0.14.0 h1:Ka4T10EEML7dQ5XDvO9c3MBN8z4nuSnGjcd1jmU2ivg= @@ -2031,8 +2031,8 @@ golang.org/x/exp v0.0.0-20200119233911-0405dc783f0a/go.mod h1:2RIsYlXP63K8oxa1u0 golang.org/x/exp v0.0.0-20200207192155-f17229e696bd/go.mod h1:J/WKrq2StrnmMY6+EHIKF9dgMWnmCNThgcyBT1FY9mM= golang.org/x/exp v0.0.0-20200224162631-6cc2880d07d6/go.mod h1:3jZMyOhIsHpP37uCMkUooju7aAi5cS1Q23tOzKc+0MU= golang.org/x/exp v0.0.0-20220827204233-334a2380cb91/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE= -golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0 h1:R84qjqJb5nVJMxqWYb3np9L5ZsaDtB+a39EqjV0JSUM= -golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0/go.mod h1:S9Xr4PYopiDyqSyp5NjCrhFrqg6A5zA2E/iPHPhqnS8= +golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b h1:DXr+pvt3nC887026GRP39Ej11UATqWDmWuS99x26cD0= +golang.org/x/exp v0.0.0-20250819193227-8b4c13bb791b/go.mod h1:4QTo5u+SEIbbKW1RacMZq1YEfOBqeXa19JeshGi+zc4= golang.org/x/image v0.0.0-20180708004352-c73c2afc3b81/go.mod h1:ux5Hcp/YLpHSI86hEcLt0YII63i6oz57MZXIpbrjZUs= golang.org/x/image v0.0.0-20190227222117-0694c2d4d067/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js= golang.org/x/image v0.0.0-20190802002840-cff245a6509b/go.mod h1:FeLwcggjj3mMvU+oOTbSwawSJRM1uh48EjtB4UJZlP0= @@ -2707,8 +2707,8 @@ google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqw google.golang.org/protobuf v1.29.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= -google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY= -google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY= +google.golang.org/protobuf v1.36.7 h1:IgrO7UwFQGJdRNXH/sQux4R1Dj1WAKcLElzeeRaXV2A= +google.golang.org/protobuf v1.36.7/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY= gopkg.in/DataDog/dd-trace-go.v1 v1.74.0 h1:wScziU1ff6Bnyr8MEyxATPSLJdnLxKz3p6RsA8FUaek= gopkg.in/DataDog/dd-trace-go.v1 v1.74.0/go.mod h1:ReNBsNfnsjVC7GsCe80zRcykL/n+nxvsNrg3NbjuleM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= @@ -2800,5 +2800,5 @@ sigs.k8s.io/yaml v1.4.0 h1:Mk1wCc2gy/F0THH0TAp1QYyJNzRm2KCLy3o5ASXVI5E= sigs.k8s.io/yaml v1.4.0/go.mod h1:Ejl7/uTz7PSA4eKMyQCUTnhZYNmLIl+5c2lQPGR2BPY= software.sslmate.com/src/go-pkcs12 v0.2.0 h1:nlFkj7bTysH6VkC4fGphtjXRbezREPgrHuJG20hBGPE= software.sslmate.com/src/go-pkcs12 v0.2.0/go.mod h1:23rNcYsMabIc1otwLpTkCCPwUq6kQsTyowttG/as0kQ= -storj.io/drpc v0.0.33 h1:yCGZ26r66ZdMP0IcTYsj7WDAUIIjzXk6DJhbhvt9FHI= -storj.io/drpc v0.0.33/go.mod h1:vR804UNzhBa49NOJ6HeLjd2H3MakC1j5Gv8bsOQT6N4= +storj.io/drpc v0.0.34 h1:q9zlQKfJ5A7x8NQNFk8x7eKUF78FMhmAbZLnFK+og7I= +storj.io/drpc v0.0.34/go.mod h1:Y9LZaa8esL1PW2IDMqJE7CFSNq7d5bQ3RI7mGPtmKMg= diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index f35dfdb1235c8..232e71e2952e3 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -6,6 +6,30 @@ export interface ACLAvailable { readonly groups: readonly Group[]; } +// From codersdk/deployment.go +export interface AIBridgeAnthropicConfig { + readonly base_url: string; + readonly key: string; +} + +// From codersdk/deployment.go +export interface AIBridgeConfig { + readonly enabled: boolean; + readonly openai: AIBridgeOpenAIConfig; + readonly anthropic: AIBridgeAnthropicConfig; +} + +// From codersdk/deployment.go +export interface AIBridgeOpenAIConfig { + readonly base_url: string; + readonly key: string; +} + +// From codersdk/deployment.go +export interface AIConfig { + readonly bridge?: AIBridgeConfig; +} + // From codersdk/aitasks.go export const AITaskPromptParameterName = "AI Prompt"; @@ -827,6 +851,7 @@ export interface DeploymentValues { readonly workspace_hostname_suffix?: string; readonly workspace_prebuilds?: PrebuildsConfig; readonly hide_ai_tasks?: boolean; + readonly ai?: AIConfig; readonly config?: string; readonly write_config?: boolean; readonly address?: string; @@ -918,6 +943,7 @@ export const EntitlementsWarningHeader = "X-Coder-Entitlements-Warning"; // From codersdk/deployment.go export type Experiment = + | "ai-bridge" | "auto-fill-parameters" | "example" | "mcp-server-http" @@ -928,6 +954,7 @@ export type Experiment = | "workspace-usage"; export const Experiments: Experiment[] = [ + "ai-bridge", "auto-fill-parameters", "example", "mcp-server-http", @@ -976,6 +1003,7 @@ export interface ExternalAuthConfig { readonly scopes: readonly string[]; readonly device_flow: boolean; readonly device_code_url: string; + readonly mcp_url: string; readonly regex: string; readonly display_name: string; readonly display_icon: string; diff --git a/site/src/pages/DeploymentSettingsPage/ExternalAuthSettingsPage/ExternalAuthSettingsPageView.stories.tsx b/site/src/pages/DeploymentSettingsPage/ExternalAuthSettingsPage/ExternalAuthSettingsPageView.stories.tsx index 5184219b38cca..4fbfa60abcc63 100644 --- a/site/src/pages/DeploymentSettingsPage/ExternalAuthSettingsPage/ExternalAuthSettingsPageView.stories.tsx +++ b/site/src/pages/DeploymentSettingsPage/ExternalAuthSettingsPage/ExternalAuthSettingsPageView.stories.tsx @@ -23,6 +23,7 @@ const meta: Meta = { device_code_url: "", display_icon: "", display_name: "GitHub", + mcp_url: "", }, ], },