From fc9bff7107acb7b44e499e93dc89c9a99d0a5b44 Mon Sep 17 00:00:00 2001 From: Danny Kopping Date: Thu, 25 Sep 2025 15:40:25 +0200 Subject: [PATCH] feat: add aibridged package (#19797) Addresses https://github.com/coder/internal/issues/987 --- Makefile | 12 +- coderd/coderd.go | 42 +-- coderd/database/dbauthz/dbauthz.go | 1 + enterprise/x/aibridged/aibridged.go | 185 ++++++++++ enterprise/x/aibridged/aibridged_test.go | 339 ++++++++++++++++++ .../x/aibridged/aibridgedmock/clientmock.go | 162 +++++++++ enterprise/x/aibridged/aibridgedmock/doc.go | 4 + .../x/aibridged/aibridgedmock/poolmock.go | 72 ++++ enterprise/x/aibridged/client.go | 34 ++ enterprise/x/aibridged/http.go | 97 +++++ enterprise/x/aibridged/pool.go | 175 +++++++++ enterprise/x/aibridged/pool_test.go | 78 ++++ enterprise/x/aibridged/request.go | 8 + enterprise/x/aibridged/server.go | 9 + enterprise/x/aibridged/translator.go | 119 ++++++ enterprise/x/aibridged/utils_test.go | 23 ++ .../x/aibridgedserver/aibridgedserver.go | 7 +- go.mod | 18 +- go.sum | 35 +- 19 files changed, 1377 insertions(+), 43 deletions(-) create mode 100644 enterprise/x/aibridged/aibridged.go create mode 100644 enterprise/x/aibridged/aibridged_test.go create mode 100644 enterprise/x/aibridged/aibridgedmock/clientmock.go create mode 100644 enterprise/x/aibridged/aibridgedmock/doc.go create mode 100644 enterprise/x/aibridged/aibridgedmock/poolmock.go create mode 100644 enterprise/x/aibridged/client.go create mode 100644 enterprise/x/aibridged/http.go create mode 100644 enterprise/x/aibridged/pool.go create mode 100644 enterprise/x/aibridged/pool_test.go create mode 100644 enterprise/x/aibridged/request.go create mode 100644 enterprise/x/aibridged/server.go create mode 100644 enterprise/x/aibridged/translator.go create mode 100644 enterprise/x/aibridged/utils_test.go diff --git a/Makefile b/Makefile index 526ddaf5b8..ed746b0f71 100644 --- a/Makefile +++ b/Makefile @@ -635,6 +635,10 @@ TAILNETTEST_MOCKS := \ tailnet/tailnettest/workspaceupdatesprovidermock.go \ tailnet/tailnettest/subscriptionmock.go +AIBRIDGED_MOCKS := \ + enterprise/x/aibridged/aibridgedmock/clientmock.go \ + enterprise/x/aibridged/aibridgedmock/poolmock.go + GEN_FILES := \ tailnet/proto/tailnet.pb.go \ agent/proto/agent.pb.go \ @@ -660,7 +664,8 @@ GEN_FILES := \ agent/agentcontainers/acmock/acmock.go \ agent/agentcontainers/dcspec/dcspec_gen.go \ coderd/httpmw/loggermw/loggermock/loggermock.go \ - codersdk/workspacesdk/agentconnmock/agentconnmock.go + codersdk/workspacesdk/agentconnmock/agentconnmock.go \ + $(AIBRIDGED_MOCKS) # all gen targets should be added here and to gen/mark-fresh gen: gen/db gen/golden-files $(GEN_FILES) @@ -713,6 +718,7 @@ gen/mark-fresh: agent/agentcontainers/dcspec/dcspec_gen.go \ coderd/httpmw/loggermw/loggermock/loggermock.go \ codersdk/workspacesdk/agentconnmock/agentconnmock.go \ + $(AIBRIDGED_MOCKS) \ " for file in $$files; do @@ -760,6 +766,10 @@ codersdk/workspacesdk/agentconnmock/agentconnmock.go: codersdk/workspacesdk/agen go generate ./codersdk/workspacesdk/agentconnmock/ touch "$@" +$(AIBRIDGED_MOCKS): enterprise/x/aibridged/client.go enterprise/x/aibridged/pool.go + go generate ./enterprise/x/aibridged/aibridgedmock/ + touch "$@" + agent/agentcontainers/dcspec/dcspec_gen.go: \ node_modules/.installed \ agent/agentcontainers/dcspec/devContainer.base.schema.json \ diff --git a/coderd/coderd.go b/coderd/coderd.go index fe7e8bc0f4..d8a52c6f6e 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -999,29 +999,31 @@ func New(options *Options) *API { // Experimental routes are not guaranteed to be stable and may change at any time. r.Route("/api/experimental", func(r chi.Router) { - r.Use(apiKeyMiddleware) - r.Route("/aitasks", func(r chi.Router) { - r.Get("/prompts", api.aiTasksPrompts) - }) - r.Route("/tasks", func(r chi.Router) { - r.Use(apiRateLimiter) + r.Group(func(r chi.Router) { + r.Use(apiKeyMiddleware) + r.Route("/aitasks", func(r chi.Router) { + r.Get("/prompts", api.aiTasksPrompts) + }) + r.Route("/tasks", func(r chi.Router) { + r.Use(apiRateLimiter) - r.Get("/", api.tasksList) + r.Get("/", api.tasksList) - r.Route("/{user}", func(r chi.Router) { - r.Use(httpmw.ExtractOrganizationMembersParam(options.Database, api.HTTPAuth.Authorize)) - r.Get("/{id}", api.taskGet) - r.Delete("/{id}", api.taskDelete) - r.Post("/{id}/send", api.taskSend) - r.Post("/", api.tasksCreate) + r.Route("/{user}", func(r chi.Router) { + r.Use(httpmw.ExtractOrganizationMembersParam(options.Database, api.HTTPAuth.Authorize)) + r.Get("/{id}", api.taskGet) + r.Delete("/{id}", api.taskDelete) + r.Post("/{id}/send", api.taskSend) + r.Post("/", api.tasksCreate) + }) + }) + r.Route("/mcp", func(r chi.Router) { + r.Use( + httpmw.RequireExperimentWithDevBypass(api.Experiments, codersdk.ExperimentOAuth2, codersdk.ExperimentMCPServerHTTP), + ) + // MCP HTTP transport endpoint with mandatory authentication + r.Mount("/http", api.mcpHTTPHandler()) }) - }) - r.Route("/mcp", func(r chi.Router) { - r.Use( - httpmw.RequireExperimentWithDevBypass(api.Experiments, codersdk.ExperimentOAuth2, codersdk.ExperimentMCPServerHTTP), - ) - // MCP HTTP transport endpoint with mandatory authentication - r.Mount("/http", api.mcpHTTPHandler()) }) }) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 477a0b83c1..e0da9c5ac1 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -570,6 +570,7 @@ var ( DisplayName: "AIBridge Daemon", Site: rbac.Permissions(map[string][]policy.Action{ rbac.ResourceUser.Type: { + policy.ActionRead, // Required to validate API key owner is active. policy.ActionReadPersonal, // Required to read users' external auth links. // TODO: this is too broad; reduce scope to just external_auth_links by creating separate resource. }, rbac.ResourceApiKey.Type: {policy.ActionRead}, // Validate API keys. diff --git a/enterprise/x/aibridged/aibridged.go b/enterprise/x/aibridged/aibridged.go new file mode 100644 index 0000000000..ee1df6084f --- /dev/null +++ b/enterprise/x/aibridged/aibridged.go @@ -0,0 +1,185 @@ +package aibridged + +import ( + "context" + "errors" + "net/http" + "sync" + "time" + + "golang.org/x/xerrors" + + "cdr.dev/slog" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/retry" +) + +// Server provides the AI Bridge functionality. +// It is responsible for: +// - receiving requests on /api/experimental/aibridged/* // TODO: update endpoint once out of experimental +// - manipulating the requests +// - relaying requests to upstream AI services and relaying responses to caller +// +// It requires a [Dialer] to provide a [DRPCClient] implementation to +// communicate with a [DRPCServer] implementation, to persist state and perform other functions. +type Server struct { + clientDialer Dialer + clientCh chan DRPCClient + + // A pool of [aibridge.RequestBridge] instances, which service incoming requests. + 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 + + // lifecycleCtx is canceled when we start closing. + lifecycleCtx context.Context + // cancelFn closes the lifecycleCtx. + cancelFn func() + + shutdownOnce sync.Once +} + +func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger) (*Server, error) { + if rpcDialer == nil { + return nil, xerrors.Errorf("nil rpcDialer given") + } + + ctx, cancel := context.WithCancel(ctx) + daemon := &Server{ + logger: logger, + clientDialer: rpcDialer, + requestBridgePool: pool, + clientCh: make(chan DRPCClient), + lifecycleCtx: ctx, + cancelFn: cancel, + initConnectionCh: 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.lifecycleCtx, "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.lifecycleCtx); { + // It's possible for the aibridge daemon to be shut down + // before the wait is complete! + if s.isShutdown() { + return + } + s.logger.Debug(s.lifecycleCtx, "dialing coderd") + client, err := s.clientDialer(s.lifecycleCtx) + 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.lifecycleCtx, "not authorized to dial coderd", slog.Error(err)) + return + } + if s.isShutdown() { + return + } + s.logger.Warn(s.lifecycleCtx, "coderd client failed to dial", slog.Error(err)) + continue + } + + // TODO: log this with INFO level when we implement external aibridge daemons. + logConnect(s.lifecycleCtx, "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.lifecycleCtx.Done(): + client.DRPCConn().Close() + return + case <-client.DRPCConn().Closed(): + logConnect(s.lifecycleCtx, "connection to coderd closed") + continue connectLoop + case s.clientCh <- client: + continue + } + } + } +} + +func (s *Server) Client() (DRPCClient, error) { + select { + case <-s.lifecycleCtx.Done(): + return nil, xerrors.New("context closed") + 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") + } + + reqBridge, err := s.requestBridgePool.Acquire(ctx, req, s.Client) + if err != nil { + return nil, xerrors.Errorf("acquire request bridge: %w", err) + } + + return reqBridge, nil +} + +// isShutdown returns whether the Server is shutdown or not. +func (s *Server) isShutdown() bool { + select { + case <-s.lifecycleCtx.Done(): + return true + default: + return false + } +} + +// 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 { + var err error + s.shutdownOnce.Do(func() { + s.cancelFn() + + // Wait for any outstanding connections to terminate. + s.wg.Wait() + + 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 request pool") + if err = s.requestBridgePool.Shutdown(ctx); err != nil { + s.logger.Error(ctx, "request pool shutdown failed with error", slog.Error(err)) + } + + s.logger.Info(ctx, "gracefully shutdown") + }) + return err +} diff --git a/enterprise/x/aibridged/aibridged_test.go b/enterprise/x/aibridged/aibridged_test.go new file mode 100644 index 0000000000..0d6541260c --- /dev/null +++ b/enterprise/x/aibridged/aibridged_test.go @@ -0,0 +1,339 @@ +package aibridged_test + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + "golang.org/x/xerrors" + "storj.io/drpc" + + "cdr.dev/slog/sloggers/slogtest" + "github.com/coder/aibridge" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/x/aibridged" + mock "github.com/coder/coder/v2/enterprise/x/aibridged/aibridgedmock" + "github.com/coder/coder/v2/enterprise/x/aibridged/proto" + "github.com/coder/coder/v2/testutil" +) + +func newTestServer(t *testing.T) (*aibridged.Server, *mock.MockDRPCClient, *mock.MockPooler) { + t.Helper() + + logger := slogtest.Make(t, nil) + ctrl := gomock.NewController(t) + client := mock.NewMockDRPCClient(ctrl) + pool := mock.NewMockPooler(ctrl) + + conn := &mockDRPCConn{} + client.EXPECT().DRPCConn().AnyTimes().Return(conn) + pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil) + + srv, err := aibridged.New( + t.Context(), + pool, + func(ctx context.Context) (aibridged.DRPCClient, error) { + return client, nil + }, + logger) + require.NoError(t, err, "create new aibridged") + t.Cleanup(func() { + srv.Shutdown(context.Background()) + }) + + return srv, client, pool +} + +// mockDRPCConn is a mock implementation of drpc.Conn +type mockDRPCConn struct{} + +func (*mockDRPCConn) Close() error { return nil } +func (*mockDRPCConn) Closed() <-chan struct{} { ch := make(chan struct{}); return ch } +func (*mockDRPCConn) Transport() drpc.Transport { return nil } +func (*mockDRPCConn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, out drpc.Message) error { + return nil +} + +func (*mockDRPCConn) NewStream(ctx context.Context, rpc string, enc drpc.Encoding) (drpc.Stream, error) { + // nolint:nilnil // Chillchill. + return nil, nil +} + +func TestServeHTTP_FailureModes(t *testing.T) { + t.Parallel() + + defaultHeaders := map[string]string{"Authorization": "Bearer key"} + + cases := []struct { + name string + reqHeaders map[string]string + applyMocksFn func(client *mock.MockDRPCClient, pool *mock.MockPooler) + dialerFn aibridged.Dialer + contextFn func() context.Context + expectedErr error + expectedStatus int + }{ + // Authnz-related failures. + { + name: "no auth key", + reqHeaders: make(map[string]string), + expectedErr: aibridged.ErrNoAuthKey, + expectedStatus: http.StatusBadRequest, + }, + { + name: "unrecognized header", + reqHeaders: map[string]string{ + codersdk.SessionTokenHeader: "key", // Coder-Session-Token is not supported; requests originate with AI clients, not coder CLI. + }, + applyMocksFn: func(client *mock.MockDRPCClient, _ *mock.MockPooler) {}, + expectedErr: aibridged.ErrNoAuthKey, + expectedStatus: http.StatusBadRequest, + }, + { + name: "unauthorized", + applyMocksFn: func(client *mock.MockDRPCClient, _ *mock.MockPooler) { + client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("not authorized")) + }, + expectedErr: aibridged.ErrUnauthorized, + expectedStatus: http.StatusForbidden, + }, + { + name: "invalid key owner ID", + applyMocksFn: func(client *mock.MockDRPCClient, _ *mock.MockPooler) { + client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: "oops"}, nil) + }, + expectedErr: aibridged.ErrUnauthorized, + expectedStatus: http.StatusForbidden, + }, + + // TODO: coderd connection-related failures. + + // Pool-related failures. + { + name: "pool instance", + applyMocksFn: func(client *mock.MockDRPCClient, pool *mock.MockPooler) { + // Should pass authorization. + client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + // But fail when acquiring a pool instance. + pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().Return(nil, xerrors.New("oops")) + }, + expectedErr: aibridged.ErrAcquireRequestHandler, + expectedStatus: http.StatusInternalServerError, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + srv, client, pool := newTestServer(t) + conn := &mockDRPCConn{} + client.EXPECT().DRPCConn().AnyTimes().Return(conn) + + if tc.applyMocksFn != nil { + tc.applyMocksFn(client, pool) + } + + httpSrv := httptest.NewServer(srv) + + ctx := testutil.Context(t, testutil.WaitShort) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, httpSrv.URL+"/openai/v1/chat/completions", nil) + require.NoError(t, err, "make request to test server") + + headers := defaultHeaders + if tc.reqHeaders != nil { + headers = tc.reqHeaders + } + for k, v := range headers { + req.Header.Set(k, v) + } + + resp, err := http.DefaultClient.Do(req) + t.Cleanup(func() { + if resp == nil || resp.Body == nil { + return + } + resp.Body.Close() + }) + require.NoError(t, err) + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err, "read response body") + require.Contains(t, string(body), tc.expectedErr.Error()) + require.Equal(t, tc.expectedStatus, resp.StatusCode) + }) + } +} + +func TestExtractAuthToken(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + headers map[string]string + expectedKey string + }{ + { + name: "none", + }, + { + name: "authorization/invalid", + headers: map[string]string{"authorization": "invalid"}, + }, + { + name: "authorization/bearer empty", + headers: map[string]string{"authorization": "bearer"}, + }, + { + name: "authorization/bearer ok", + headers: map[string]string{"authorization": "bearer key"}, + expectedKey: "key", + }, + { + name: "authorization/case", + headers: map[string]string{"AUTHORIZATION": "BEARer key"}, + expectedKey: "key", + }, + { + name: "x-api-key/empty", + headers: map[string]string{"X-Api-Key": ""}, + }, + { + name: "x-api-key/ok", + headers: map[string]string{"X-Api-Key": "key"}, + expectedKey: "key", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + headers := make(http.Header, len(tc.headers)) + for k, v := range tc.headers { + headers.Add(k, v) + } + key := aibridged.ExtractAuthToken(headers) + require.Equal(t, tc.expectedKey, key) + }) + } +} + +var _ http.Handler = &mockHandler{} + +type mockHandler struct{} + +func (*mockHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) { + rw.WriteHeader(http.StatusOK) + _, _ = rw.Write([]byte(r.URL.Path)) +} + +// TestRouting validates that a request which originates with aibridged will be handled +// by coder/aibridge's handling logic in a provider-specific manner. +// We must validate that logic that pertains to coder/coder is exercised. +// aibridge will only handle certain routes; we don't need to test these exhaustively +// (that's coder/aibridge's responsibility), but we do need to validate that it handles +// requests correctly. +func TestRouting(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + path string + expectedStatus int + expectedHits int // Expected hits to the upstream server. + }{ + { + name: "unsupported", + path: "/this-route-does-not-exist", + expectedStatus: http.StatusNotFound, + expectedHits: 0, + }, + { + name: "openai chat completions", + path: "/openai/v1/chat/completions", + expectedStatus: http.StatusTeapot, // Nonsense status to indicate server was hit. + expectedHits: 1, + }, + { + name: "anthropic messages", + path: "/anthropic/v1/messages", + expectedStatus: http.StatusTeapot, // Nonsense status to indicate server was hit. + expectedHits: 1, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Setup mock upstream AI server. + upstreamSrv := &mockAIUpstreamServer{} + openaiSrv := httptest.NewServer(upstreamSrv) + antSrv := httptest.NewServer(upstreamSrv) + t.Cleanup(openaiSrv.Close) + t.Cleanup(antSrv.Close) + + // Setup. + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + ctrl := gomock.NewController(t) + client := mock.NewMockDRPCClient(ctrl) + + providers := []aibridge.Provider{ + aibridge.NewOpenAIProvider(aibridge.ProviderConfig{BaseURL: openaiSrv.URL}), + aibridge.NewAnthropicProvider(aibridge.ProviderConfig{BaseURL: antSrv.URL}), + } + pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger) + require.NoError(t, err) + conn := &mockDRPCConn{} + client.EXPECT().DRPCConn().AnyTimes().Return(conn) + + client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId: uuid.NewString()}, nil) + client.EXPECT().GetMCPServerConfigs(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.GetMCPServerConfigsResponse{}, nil) + // This is the only recording we really care about in this test. This is called before the provider-specific logic processes + // the incoming request, and anything beyond that is the responsibility of coder/aibridge to test. + var interceptionID string + client.EXPECT().RecordInterception(gomock.Any(), gomock.Any()).Times(tc.expectedHits).DoAndReturn(func(ctx context.Context, in *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) { + interceptionID = in.GetId() + return &proto.RecordInterceptionResponse{}, nil + }) + + // Given: aibridged is started. + srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) { + return client, nil + }, logger) + require.NoError(t, err, "create new aibridged") + t.Cleanup(func() { + _ = srv.Shutdown(testutil.Context(t, testutil.WaitShort)) + }) + + // When: a request is made to aibridged. + ctx := testutil.Context(t, testutil.WaitShort) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, tc.path, bytes.NewBufferString(`{}`)) + require.NoError(t, err, "make request to test server") + req.Header.Add("Authorization", "Bearer key") + req.Header.Add("Accept", "application/json") + + // When: aibridged handles the request. + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + // Then: the upstream server will have received a number of hits. + // NOTE: we *expect* the interceptions to fail because [mockAIUpstreamServer] returns a nonsense status code. + // We only need to test that the request was routed, NOT processed. + require.Equal(t, tc.expectedStatus, rec.Code) + assert.EqualValues(t, tc.expectedHits, upstreamSrv.Hits()) + if tc.expectedHits > 0 { + _, err = uuid.Parse(interceptionID) + require.NoError(t, err, "parse interception ID") + } + }) + } +} diff --git a/enterprise/x/aibridged/aibridgedmock/clientmock.go b/enterprise/x/aibridged/aibridgedmock/clientmock.go new file mode 100644 index 0000000000..2210051689 --- /dev/null +++ b/enterprise/x/aibridged/aibridgedmock/clientmock.go @@ -0,0 +1,162 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/coder/coder/v2/enterprise/x/aibridged (interfaces: DRPCClient) +// +// Generated by this command: +// +// mockgen -destination ./clientmock.go -package aibridgedmock github.com/coder/coder/v2/enterprise/x/aibridged DRPCClient +// + +// Package aibridgedmock is a generated GoMock package. +package aibridgedmock + +import ( + context "context" + reflect "reflect" + + proto "github.com/coder/coder/v2/enterprise/x/aibridged/proto" + gomock "go.uber.org/mock/gomock" + drpc "storj.io/drpc" +) + +// MockDRPCClient is a mock of DRPCClient interface. +type MockDRPCClient struct { + ctrl *gomock.Controller + recorder *MockDRPCClientMockRecorder + isgomock struct{} +} + +// MockDRPCClientMockRecorder is the mock recorder for MockDRPCClient. +type MockDRPCClientMockRecorder struct { + mock *MockDRPCClient +} + +// NewMockDRPCClient creates a new mock instance. +func NewMockDRPCClient(ctrl *gomock.Controller) *MockDRPCClient { + mock := &MockDRPCClient{ctrl: ctrl} + mock.recorder = &MockDRPCClientMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockDRPCClient) EXPECT() *MockDRPCClientMockRecorder { + return m.recorder +} + +// DRPCConn mocks base method. +func (m *MockDRPCClient) DRPCConn() drpc.Conn { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DRPCConn") + ret0, _ := ret[0].(drpc.Conn) + return ret0 +} + +// DRPCConn indicates an expected call of DRPCConn. +func (mr *MockDRPCClientMockRecorder) DRPCConn() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DRPCConn", reflect.TypeOf((*MockDRPCClient)(nil).DRPCConn)) +} + +// GetMCPServerAccessTokensBatch mocks base method. +func (m *MockDRPCClient) GetMCPServerAccessTokensBatch(ctx context.Context, in *proto.GetMCPServerAccessTokensBatchRequest) (*proto.GetMCPServerAccessTokensBatchResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetMCPServerAccessTokensBatch", ctx, in) + ret0, _ := ret[0].(*proto.GetMCPServerAccessTokensBatchResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetMCPServerAccessTokensBatch indicates an expected call of GetMCPServerAccessTokensBatch. +func (mr *MockDRPCClientMockRecorder) GetMCPServerAccessTokensBatch(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMCPServerAccessTokensBatch", reflect.TypeOf((*MockDRPCClient)(nil).GetMCPServerAccessTokensBatch), ctx, in) +} + +// GetMCPServerConfigs mocks base method. +func (m *MockDRPCClient) GetMCPServerConfigs(ctx context.Context, in *proto.GetMCPServerConfigsRequest) (*proto.GetMCPServerConfigsResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetMCPServerConfigs", ctx, in) + ret0, _ := ret[0].(*proto.GetMCPServerConfigsResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetMCPServerConfigs indicates an expected call of GetMCPServerConfigs. +func (mr *MockDRPCClientMockRecorder) GetMCPServerConfigs(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetMCPServerConfigs", reflect.TypeOf((*MockDRPCClient)(nil).GetMCPServerConfigs), ctx, in) +} + +// IsAuthorized mocks base method. +func (m *MockDRPCClient) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest) (*proto.IsAuthorizedResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IsAuthorized", ctx, in) + ret0, _ := ret[0].(*proto.IsAuthorizedResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// IsAuthorized indicates an expected call of IsAuthorized. +func (mr *MockDRPCClientMockRecorder) IsAuthorized(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsAuthorized", reflect.TypeOf((*MockDRPCClient)(nil).IsAuthorized), ctx, in) +} + +// RecordInterception mocks base method. +func (m *MockDRPCClient) RecordInterception(ctx context.Context, in *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RecordInterception", ctx, in) + ret0, _ := ret[0].(*proto.RecordInterceptionResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RecordInterception indicates an expected call of RecordInterception. +func (mr *MockDRPCClientMockRecorder) RecordInterception(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordInterception", reflect.TypeOf((*MockDRPCClient)(nil).RecordInterception), ctx, in) +} + +// RecordPromptUsage mocks base method. +func (m *MockDRPCClient) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RecordPromptUsage", ctx, in) + ret0, _ := ret[0].(*proto.RecordPromptUsageResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RecordPromptUsage indicates an expected call of RecordPromptUsage. +func (mr *MockDRPCClientMockRecorder) RecordPromptUsage(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordPromptUsage", reflect.TypeOf((*MockDRPCClient)(nil).RecordPromptUsage), ctx, in) +} + +// RecordTokenUsage mocks base method. +func (m *MockDRPCClient) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RecordTokenUsage", ctx, in) + ret0, _ := ret[0].(*proto.RecordTokenUsageResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RecordTokenUsage indicates an expected call of RecordTokenUsage. +func (mr *MockDRPCClientMockRecorder) RecordTokenUsage(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordTokenUsage", reflect.TypeOf((*MockDRPCClient)(nil).RecordTokenUsage), ctx, in) +} + +// RecordToolUsage mocks base method. +func (m *MockDRPCClient) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageRequest) (*proto.RecordToolUsageResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "RecordToolUsage", ctx, in) + ret0, _ := ret[0].(*proto.RecordToolUsageResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// RecordToolUsage indicates an expected call of RecordToolUsage. +func (mr *MockDRPCClientMockRecorder) RecordToolUsage(ctx, in any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordToolUsage", reflect.TypeOf((*MockDRPCClient)(nil).RecordToolUsage), ctx, in) +} diff --git a/enterprise/x/aibridged/aibridgedmock/doc.go b/enterprise/x/aibridged/aibridgedmock/doc.go new file mode 100644 index 0000000000..3d3f56c055 --- /dev/null +++ b/enterprise/x/aibridged/aibridgedmock/doc.go @@ -0,0 +1,4 @@ +package aibridgedmock + +//go:generate mockgen -destination ./clientmock.go -package aibridgedmock github.com/coder/coder/v2/enterprise/x/aibridged DRPCClient +//go:generate mockgen -destination ./poolmock.go -package aibridgedmock github.com/coder/coder/v2/enterprise/x/aibridged Pooler diff --git a/enterprise/x/aibridged/aibridgedmock/poolmock.go b/enterprise/x/aibridged/aibridgedmock/poolmock.go new file mode 100644 index 0000000000..80ce4a35b9 --- /dev/null +++ b/enterprise/x/aibridged/aibridgedmock/poolmock.go @@ -0,0 +1,72 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/coder/coder/v2/enterprise/x/aibridged (interfaces: Pooler) +// +// Generated by this command: +// +// mockgen -destination ./poolmock.go -package aibridgedmock github.com/coder/coder/v2/enterprise/x/aibridged Pooler +// + +// Package aibridgedmock is a generated GoMock package. +package aibridgedmock + +import ( + context "context" + http "net/http" + reflect "reflect" + + aibridged "github.com/coder/coder/v2/enterprise/x/aibridged" + gomock "go.uber.org/mock/gomock" +) + +// MockPooler is a mock of Pooler interface. +type MockPooler struct { + ctrl *gomock.Controller + recorder *MockPoolerMockRecorder + isgomock struct{} +} + +// MockPoolerMockRecorder is the mock recorder for MockPooler. +type MockPoolerMockRecorder struct { + mock *MockPooler +} + +// NewMockPooler creates a new mock instance. +func NewMockPooler(ctrl *gomock.Controller) *MockPooler { + mock := &MockPooler{ctrl: ctrl} + mock.recorder = &MockPoolerMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockPooler) EXPECT() *MockPoolerMockRecorder { + return m.recorder +} + +// Acquire mocks base method. +func (m *MockPooler) Acquire(ctx context.Context, req aibridged.Request, clientFn aibridged.ClientFunc) (http.Handler, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Acquire", ctx, req, clientFn) + ret0, _ := ret[0].(http.Handler) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Acquire indicates an expected call of Acquire. +func (mr *MockPoolerMockRecorder) Acquire(ctx, req, clientFn any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Acquire", reflect.TypeOf((*MockPooler)(nil).Acquire), ctx, req, clientFn) +} + +// Shutdown mocks base method. +func (m *MockPooler) Shutdown(ctx context.Context) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Shutdown", ctx) + ret0, _ := ret[0].(error) + return ret0 +} + +// Shutdown indicates an expected call of Shutdown. +func (mr *MockPoolerMockRecorder) Shutdown(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Shutdown", reflect.TypeOf((*MockPooler)(nil).Shutdown), ctx) +} diff --git a/enterprise/x/aibridged/client.go b/enterprise/x/aibridged/client.go new file mode 100644 index 0000000000..3004a84df9 --- /dev/null +++ b/enterprise/x/aibridged/client.go @@ -0,0 +1,34 @@ +package aibridged + +import ( + "context" + + "storj.io/drpc" + + "github.com/coder/coder/v2/enterprise/x/aibridged/proto" +) + +type Dialer func(ctx context.Context) (DRPCClient, error) + +type ClientFunc func() (DRPCClient, error) + +// DRPCClient is the union of various service interfaces the client must support. +type DRPCClient interface { + proto.DRPCRecorderClient + proto.DRPCMCPConfiguratorClient + proto.DRPCAuthorizerClient +} + +var _ DRPCClient = &Client{} + +type Client struct { + proto.DRPCRecorderClient + proto.DRPCMCPConfiguratorClient + proto.DRPCAuthorizerClient + + Conn drpc.Conn +} + +func (c *Client) DRPCConn() drpc.Conn { + return c.Conn +} diff --git a/enterprise/x/aibridged/http.go b/enterprise/x/aibridged/http.go new file mode 100644 index 0000000000..43f4ba7670 --- /dev/null +++ b/enterprise/x/aibridged/http.go @@ -0,0 +1,97 @@ +package aibridged + +import ( + "net/http" + "strings" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog" + "github.com/coder/aibridge" + "github.com/coder/coder/v2/enterprise/x/aibridged/proto" +) + +var _ http.Handler = &Server{} + +var ( + ErrNoAuthKey = xerrors.New("no authentication key provided") + ErrConnect = xerrors.New("could not connect to coderd") + ErrUnauthorized = xerrors.New("unauthorized") + ErrAcquireRequestHandler = xerrors.New("failed to acquire request handler") +) + +// ServeHTTP is the entrypoint for requests which will be intercepted by AI Bridge. +// This function will validate that the given API key may be used to perform the request. +// +// An [aibridge.RequestBridge] instance is acquired from a pool based on the API key's +// owner (referred to as the "initiator"); this instance is responsible for the +// AI Bridge-specific handling of the request. +// +// A [DRPCClient] is provided to the [aibridge.RequestBridge] instance so that data can +// be passed up to a [DRPCServer] for persistence. +func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + logger := s.logger.With(slog.F("path", r.URL.Path)) + + key := strings.TrimSpace(ExtractAuthToken(r.Header)) + if key == "" { + logger.Warn(ctx, "no auth key provided") + http.Error(rw, ErrNoAuthKey.Error(), http.StatusBadRequest) + return + } + + client, err := s.Client() + if err != nil { + logger.Warn(ctx, "failed to connect to coderd", slog.Error(err)) + http.Error(rw, ErrConnect.Error(), http.StatusServiceUnavailable) + return + } + + resp, err := client.IsAuthorized(ctx, &proto.IsAuthorizedRequest{Key: key}) + if err != nil { + logger.Warn(ctx, "key authorization check failed", slog.Error(err)) + http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden) + return + } + + // Rewire request context to include actor. + r = r.WithContext(aibridge.AsActor(ctx, resp.GetOwnerId(), nil)) + + id, err := uuid.Parse(resp.GetOwnerId()) + if err != nil { + logger.Warn(ctx, "failed to parse user ID", slog.Error(err), slog.F("id", resp.GetOwnerId())) + http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden) + return + } + + handler, err := s.GetRequestHandler(ctx, Request{ + SessionKey: key, + InitiatorID: id, + }) + if err != nil { + logger.Warn(ctx, "failed to acquire request handler", slog.Error(err)) + http.Error(rw, ErrAcquireRequestHandler.Error(), http.StatusInternalServerError) + return + } + + handler.ServeHTTP(rw, r) +} + +// ExtractAuthToken extracts authorization token from HTTP request using multiple sources. +// These sources represent the different ways clients authenticate against AI providers. +// It checks the Authorization header (Bearer token) and X-Api-Key header. +// If neither are present, an empty string is returned. +func ExtractAuthToken(header http.Header) string { + if auth := strings.TrimSpace(header.Get("Authorization")); auth != "" { + fields := strings.Fields(auth) + if len(fields) == 2 && strings.EqualFold(fields[0], "Bearer") { + return fields[1] + } + } + if apiKey := strings.TrimSpace(header.Get("X-Api-Key")); apiKey != "" { + return apiKey + } + return "" +} diff --git a/enterprise/x/aibridged/pool.go b/enterprise/x/aibridged/pool.go new file mode 100644 index 0000000000..01de7059a3 --- /dev/null +++ b/enterprise/x/aibridged/pool.go @@ -0,0 +1,175 @@ +package aibridged + +import ( + "context" + "errors" + "net/http" + "sync" + "time" + + "github.com/dgraph-io/ristretto/v2" + "golang.org/x/xerrors" + "tailscale.com/util/singleflight" + + "cdr.dev/slog" + + "github.com/coder/aibridge" +) + +const ( + 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, clientFn ClientFunc) (http.Handler, error) + Shutdown(ctx context.Context) error +} + +type PoolMetrics interface { + Hits() uint64 + Misses() uint64 + KeysAdded() uint64 + KeysEvicted() uint64 +} + +type PoolOptions struct { + MaxItems int64 + TTL time.Duration +} + +var DefaultPoolOptions = PoolOptions{MaxItems: 100, TTL: time.Minute * 15} + +var _ Pooler = &CachedBridgePool{} + +type CachedBridgePool struct { + cache *ristretto.Cache[string, *aibridge.RequestBridge] + providers []aibridge.Provider + logger slog.Logger + options PoolOptions + + singleflight *singleflight.Group[string, *aibridge.RequestBridge] + + shutDownOnce sync.Once + shuttingDownCh chan struct{} +} + +func NewCachedBridgePool(options PoolOptions, providers []aibridge.Provider, logger slog.Logger) (*CachedBridgePool, error) { + cache, err := ristretto.NewCache(&ristretto.Config[string, *aibridge.RequestBridge]{ + NumCounters: options.MaxItems * 10, // Docs suggest setting this 10x number of keys. + MaxCost: options.MaxItems * cacheCost, // Up to n instances. + IgnoreInternalCost: true, // Don't try estimate cost using bytes (ristretto does this naïvely anyway, just using the size of the value struct not the REAL memory usage). + BufferItems: 64, // Sticking with recommendation from docs. + Metrics: true, // Collect metrics (only used in tests, for now). + OnEvict: func(item *ristretto.Item[*aibridge.RequestBridge]) { + if item == nil || item.Value == nil { + return + } + + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), time.Second*5) + defer shutdownCancel() + + // Run the eviction in the background since ristretto blocks sets until a free slot is available. + go func() { + if err := item.Value.Shutdown(shutdownCtx); err != nil { + if errors.Is(err, context.DeadlineExceeded) { + logger.Debug(shutdownCtx, "bridge shutdown timed out") + } else { + logger.Debug(shutdownCtx, "bridge shutdown failed", slog.Error(err)) + } + } + }() + }, + }) + if err != nil { + return nil, xerrors.Errorf("create cache: %w", err) + } + + return &CachedBridgePool{ + cache: cache, + providers: providers, + logger: logger, + options: options, + + singleflight: &singleflight.Group[string, *aibridge.RequestBridge]{}, + + shuttingDownCh: make(chan struct{}), + }, nil +} + +// Acquire retrieves or creates a [*aibridge.RequestBridge] instance per given key. +// +// Each returned [*aibridge.RequestBridge] is safe for concurrent use. +// Each [*aibridge.RequestBridge] is stateful because it has MCP clients which maintain sessions to the configured MCP server. +func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn ClientFunc) (http.Handler, 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: + } + + // Wait for all buffered writes to be applied, otherwise multiple calls in quick succession + // may visit the slow path unnecessarily. + defer p.cache.Wait() + + recorder := aibridge.NewRecorder(p.logger.Named("recorder"), func() (aibridge.Recorder, error) { + client, err := clientFn() + if err != nil { + return nil, xerrors.Errorf("acquire client: %w", err) + } + + return &recorderTranslation{client: client}, nil + }) + + // Fast path. + bridge, ok := p.cache.Get(req.InitiatorID.String()) + if ok && bridge != nil { + // TODO: future improvement: + // Once we can detect token expiry against an MCP server, we no longer need to let these instances + // expire after the original TTL; we can extend the TTL on each Acquire() call. + // For now, we need to let the instance expiry to keep the MCP connections fresh. + + 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) { + bridge, err := aibridge.NewRequestBridge(ctx, p.providers, p.logger, recorder, nil) + if err != nil { + return nil, xerrors.Errorf("create new request bridge: %w", err) + } + + p.cache.SetWithTTL(req.InitiatorID.String(), bridge, cacheCost, p.options.TTL) + + return bridge, nil + }) + + return instance, err +} + +func (p *CachedBridgePool) Metrics() PoolMetrics { + if p.cache == nil { + return nil + } + + return p.cache.Metrics +} + +// 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/enterprise/x/aibridged/pool_test.go b/enterprise/x/aibridged/pool_test.go new file mode 100644 index 0000000000..c3442a80de --- /dev/null +++ b/enterprise/x/aibridged/pool_test.go @@ -0,0 +1,78 @@ +package aibridged_test + +import ( + "context" + _ "embed" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "cdr.dev/slog/sloggers/slogtest" + "github.com/coder/coder/v2/enterprise/x/aibridged" + mock "github.com/coder/coder/v2/enterprise/x/aibridged/aibridgedmock" +) + +// TestPool validates the published behavior of [aibridged.CachedBridgePool]. +// It is not meant to be an exhaustive test of the internal cache's functionality, +// since that is already covered by its library. +func TestPool(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, nil) + + ctrl := gomock.NewController(t) + client := mock.NewMockDRPCClient(ctrl) + + opts := aibridged.PoolOptions{MaxItems: 1, TTL: time.Second} + pool, err := aibridged.NewCachedBridgePool(opts, nil, logger) + require.NoError(t, err) + t.Cleanup(func() { pool.Shutdown(context.Background()) }) + + id, id2 := uuid.New(), uuid.New() + clientFn := func() (aibridged.DRPCClient, error) { + return client, nil + } + + // Acquiring a pool instance will create one the first time it sees an + // initiator ID... + inst, err := pool.Acquire(t.Context(), aibridged.Request{ + SessionKey: "key", + InitiatorID: id, + }, clientFn) + require.NoError(t, err, "acquire pool instance") + + // ...and it will return it when acquired again. + instB, err := pool.Acquire(t.Context(), aibridged.Request{ + SessionKey: "key", + InitiatorID: id, + }, clientFn) + require.NoError(t, err, "acquire pool instance") + require.Same(t, inst, instB) + + metrics := pool.Metrics() + require.EqualValues(t, 1, metrics.KeysAdded()) + require.EqualValues(t, 0, metrics.KeysEvicted()) + require.EqualValues(t, 1, metrics.Hits()) + require.EqualValues(t, 1, metrics.Misses()) + + // But that key will be evicted when a new initiator is seen (maxItems=1): + inst2, err := pool.Acquire(t.Context(), aibridged.Request{ + SessionKey: "key", + InitiatorID: id2, + }, clientFn) + require.NoError(t, err, "acquire pool instance") + require.NotSame(t, inst, inst2) + + metrics = pool.Metrics() + require.EqualValues(t, 2, metrics.KeysAdded()) + require.EqualValues(t, 1, metrics.KeysEvicted()) + require.EqualValues(t, 1, metrics.Hits()) + require.EqualValues(t, 2, metrics.Misses()) + + // TODO: add test for expiry. + // This requires Go 1.25's [synctest](https://pkg.go.dev/testing/synctest) since the + // internal cache lib cannot be tested using coder/quartz. +} diff --git a/enterprise/x/aibridged/request.go b/enterprise/x/aibridged/request.go new file mode 100644 index 0000000000..29196adb88 --- /dev/null +++ b/enterprise/x/aibridged/request.go @@ -0,0 +1,8 @@ +package aibridged + +import "github.com/google/uuid" + +type Request struct { + SessionKey string + InitiatorID uuid.UUID +} diff --git a/enterprise/x/aibridged/server.go b/enterprise/x/aibridged/server.go new file mode 100644 index 0000000000..713ea2a0cd --- /dev/null +++ b/enterprise/x/aibridged/server.go @@ -0,0 +1,9 @@ +package aibridged + +import "github.com/coder/coder/v2/enterprise/x/aibridged/proto" + +type DRPCServer interface { + proto.DRPCRecorderServer + proto.DRPCMCPConfiguratorServer + proto.DRPCAuthorizerServer +} diff --git a/enterprise/x/aibridged/translator.go b/enterprise/x/aibridged/translator.go new file mode 100644 index 0000000000..57d574f1bf --- /dev/null +++ b/enterprise/x/aibridged/translator.go @@ -0,0 +1,119 @@ +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" + "google.golang.org/protobuf/types/known/timestamppb" + + "github.com/coder/coder/v2/coderd/util/ptr" + "github.com/coder/coder/v2/enterprise/x/aibridged/proto" + + "github.com/coder/aibridge" +) + +var _ aibridge.Recorder = &recorderTranslation{} + +// recorderTranslation satisfies the aibridge.Recorder interface and translates calls into dRPC calls to aibridgedserver. +type recorderTranslation struct { + client proto.DRPCRecorderClient +} + +func (t *recorderTranslation) RecordInterception(ctx context.Context, req *aibridge.InterceptionRecord) error { + _, err := t.client.RecordInterception(ctx, &proto.RecordInterceptionRequest{ + Id: req.ID, + InitiatorId: req.InitiatorID, + Provider: req.Provider, + Model: req.Model, + Metadata: marshalForProto(req.Metadata), + StartedAt: timestamppb.New(req.StartedAt), + }) + return err +} + +func (t *recorderTranslation) RecordPromptUsage(ctx context.Context, req *aibridge.PromptUsageRecord) error { + _, err := t.client.RecordPromptUsage(ctx, &proto.RecordPromptUsageRequest{ + InterceptionId: req.InterceptionID, + MsgId: req.MsgID, + Prompt: req.Prompt, + Metadata: marshalForProto(req.Metadata), + CreatedAt: timestamppb.New(req.CreatedAt), + }) + return err +} + +func (t *recorderTranslation) RecordTokenUsage(ctx context.Context, req *aibridge.TokenUsageRecord) error { + _, err := t.client.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{ + InterceptionId: req.InterceptionID, + MsgId: req.MsgID, + InputTokens: req.Input, + OutputTokens: req.Output, + Metadata: marshalForProto(req.Metadata), + CreatedAt: timestamppb.New(req.CreatedAt), + }) + return err +} + +func (t *recorderTranslation) RecordToolUsage(ctx context.Context, req *aibridge.ToolUsageRecord) error { + serialized, err := json.Marshal(req.Args) + if err != nil { + return xerrors.Errorf("serialize tool %q args: %w", req.Tool, err) + } + + var invErr *string + if req.InvocationError != nil { + invErr = ptr.Ref(req.InvocationError.Error()) + } + + _, err = t.client.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{ + InterceptionId: req.InterceptionID, + MsgId: req.MsgID, + ServerUrl: req.ServerURL, + Tool: req.Tool, + Input: string(serialized), + Injected: req.Injected, + InvocationError: invErr, + Metadata: marshalForProto(req.Metadata), + CreatedAt: timestamppb.New(req.CreatedAt), + }) + 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/enterprise/x/aibridged/utils_test.go b/enterprise/x/aibridged/utils_test.go new file mode 100644 index 0000000000..2989f7b661 --- /dev/null +++ b/enterprise/x/aibridged/utils_test.go @@ -0,0 +1,23 @@ +package aibridged_test + +import ( + "net/http" + "sync/atomic" +) + +var _ http.Handler = &mockAIUpstreamServer{} + +type mockAIUpstreamServer struct { + hitCounter atomic.Int32 +} + +func (m *mockAIUpstreamServer) ServeHTTP(rw http.ResponseWriter, _ *http.Request) { + m.hitCounter.Add(1) + + rw.WriteHeader(http.StatusTeapot) + _, _ = rw.Write([]byte(`i am a teapot`)) +} + +func (m *mockAIUpstreamServer) Hits() int32 { + return m.hitCounter.Load() +} diff --git a/enterprise/x/aibridgedserver/aibridgedserver.go b/enterprise/x/aibridgedserver/aibridgedserver.go index 51483c435f..c3b0149989 100644 --- a/enterprise/x/aibridgedserver/aibridgedserver.go +++ b/enterprise/x/aibridgedserver/aibridgedserver.go @@ -25,6 +25,7 @@ import ( "github.com/coder/coder/v2/coderd/httpmw" codermcp "github.com/coder/coder/v2/coderd/mcp" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/enterprise/x/aibridged" "github.com/coder/coder/v2/enterprise/x/aibridged/proto" ) @@ -47,11 +48,7 @@ var ( ErrNoExternalAuthLinkFound = xerrors.New("no external auth link found") ) -var ( - _ proto.DRPCAuthorizerServer = &Server{} - _ proto.DRPCMCPConfiguratorServer = &Server{} - _ proto.DRPCRecorderServer = &Server{} -) +var _ aibridged.DRPCServer = &Server{} type store interface { // Recorder-related queries. diff --git a/go.mod b/go.mod index cca6cf79a9..1ad97ecb58 100644 --- a/go.mod +++ b/go.mod @@ -413,7 +413,7 @@ require ( github.com/tchap/go-patricia/v2 v2.3.2 // indirect github.com/tcnksm/go-httpstat v0.2.0 // indirect github.com/tdewolff/parse/v2 v2.8.3 // indirect - github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/match v1.2.0 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tinylib/msgp v1.2.5 // indirect github.com/tklauser/go-sysconf v0.3.15 // indirect @@ -475,15 +475,17 @@ require ( ) require ( - github.com/anthropics/anthropic-sdk-go v1.11.0 + github.com/anthropics/anthropic-sdk-go v1.12.0 github.com/brianvoe/gofakeit/v7 v7.7.1 github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225 + github.com/coder/aibridge v0.1.0 github.com/coder/aisdk-go v0.0.9 github.com/coder/preview v1.0.4 + github.com/dgraph-io/ristretto/v2 v2.3.0 github.com/fsnotify/fsnotify v1.9.0 github.com/go-git/go-git/v5 v5.16.2 github.com/icholy/replace v0.6.0 - github.com/mark3labs/mcp-go v0.32.0 + github.com/mark3labs/mcp-go v0.38.0 ) require ( @@ -507,11 +509,12 @@ require ( github.com/aquasecurity/trivy v0.61.1-0.20250407075540-f1329c7ea1aa // indirect github.com/aquasecurity/trivy-checks v1.11.3-0.20250604022615-9a7efa7c9169 // indirect github.com/aws/aws-sdk-go v1.55.7 // indirect + github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/bgentry/go-netrc v0.0.0-20140422174119-9fd32a8b3d3d // indirect + github.com/buger/jsonparser v1.1.1 // indirect github.com/cenkalti/backoff/v5 v5.0.2 // 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 @@ -522,12 +525,14 @@ require ( github.com/gorilla/websocket v1.5.4-0.20250319132907-e064f32e3674 // indirect github.com/hashicorp/go-getter v1.7.9 // indirect github.com/hashicorp/go-safetemp v1.0.0 // indirect + github.com/invopop/jsonschema v0.13.0 // indirect github.com/jackmordaunt/icns/v3 v3.0.1 // indirect github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect 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/openai/openai-go/v2 v2.7.0 // indirect github.com/package-url/packageurl-go v0.1.3 // indirect github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect github.com/puzpuzpuz/xsync/v3 v3.5.1 // indirect @@ -536,9 +541,10 @@ require ( 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.15 // indirect github.com/vektah/gqlparser/v2 v2.5.28 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect github.com/zeebo/xxh3 v1.0.2 // indirect go.opentelemetry.io/contrib/detectors/gcp v1.37.0 // indirect diff --git a/go.sum b/go.sum index 32a763a7c7..1b581287f3 100644 --- a/go.sum +++ b/go.sum @@ -722,8 +722,8 @@ github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwTo github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4= -github.com/anthropics/anthropic-sdk-go v1.11.0 h1:Ic72+3QvCwZA/nHgkVNXkfU7ktEEv9CZUnBech9HuK4= -github.com/anthropics/anthropic-sdk-go v1.11.0/go.mod h1:WTz31rIUHUHqai2UslPpw5CwXrQP3geYBioRV4WOLvE= +github.com/anthropics/anthropic-sdk-go v1.12.0 h1:xPqlGnq7rWrTiHazIvCiumA0u7mGQnwDQtvA1M82h9U= +github.com/anthropics/anthropic-sdk-go v1.12.0/go.mod h1:WTz31rIUHUHqai2UslPpw5CwXrQP3geYBioRV4WOLvE= github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY= github.com/apache/arrow/go/v10 v10.0.1/go.mod h1:YvhnlEePVnBS4+0z3fhPfUy7W1Ikj0Ih0vcRo/gZ1M0= github.com/apache/arrow/go/v11 v11.0.0/go.mod h1:Eg5OsL5H+e299f7u5ssuXsuHQVEGC4xei5aX110hRiI= @@ -794,6 +794,8 @@ github.com/aymanbagabas/go-udiff v0.2.0 h1:TK0fH4MteXUDspT88n8CKzvK0X9O2xu9yQjWp github.com/aymanbagabas/go-udiff v0.2.0/go.mod h1:RE4Ex0qsGkTAJoQdQQCA0uG+nAzJO/pI/QwceO5fgrA= github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuPk= github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/bep/clocks v0.5.0 h1:hhvKVGLPQWRVsBP/UB7ErrHYIO42gINVbvqxvYTPVps= @@ -836,6 +838,8 @@ github.com/bramvdbogaerde/go-scp v1.5.0 h1:a9BinAjTfQh273eh7vd3qUgmBC+bx+3TRDtkZ github.com/bramvdbogaerde/go-scp v1.5.0/go.mod h1:on2aH5AxaFb2G0N5Vsdy6B0Ml7k9HuHSwfo1y0QzAbQ= github.com/brianvoe/gofakeit/v7 v7.7.1 h1:Z74GFLZz57rAUHjpNbaKOr8c7nXdUohsiwF/jhkqE0k= github.com/brianvoe/gofakeit/v7 v7.7.1/go.mod h1:QXuPeBw164PJCzCUZVmgpgHJ3Llj49jSLVkKPMtxtxA= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/bytecodealliance/wasmtime-go/v3 v3.0.2 h1:3uZCA/BLTIu+DqCfguByNMJa2HVHpXvjfy0Dy7g6fuA= github.com/bytecodealliance/wasmtime-go/v3 v3.0.2/go.mod h1:RnUjnIXxEJcL6BgCvNyzCCRzZcxCgsZCi+RNlvYor5Q= github.com/cakturk/go-netstat v0.0.0-20200220111822-e5b49efee7a5 h1:BjkPE3785EwPhhyuFkbINB+2a1xATwk8SNDWnJiD41g= @@ -907,6 +911,8 @@ github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 h1:aQ3y1lwWyqYPiWZThqv github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8= github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225 h1:tRIViZ5JRmzdOEo5wUWngaGEFBG8OaE1o2GIHN5ujJ8= github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225/go.mod h1:rNLVpYgEVeu1Zk29K64z6Od8RBP9DwqCu9OfCzh8MR4= +github.com/coder/aibridge v0.1.0 h1:UKkO1V7v3YYWLXuB3jQlCre5v7N3VOPI4KYgKZeM/bA= +github.com/coder/aibridge v0.1.0/go.mod h1:GWc0Owtlzz5iMHosDm6FhbO+SoG5W+VeOKyP9p9g9ZM= github.com/coder/aisdk-go v0.0.9 h1:Vzo/k2qwVGLTR10ESDeP2Ecek1SdPfZlEjtTfMveiVo= github.com/coder/aisdk-go v0.0.9/go.mod h1:KF6/Vkono0FJJOtWtveh5j7yfNrSctVTpwgweYWSp5M= github.com/coder/bubbletea v1.2.2-0.20241212190825-007a1cdb2c41 h1:SBN/DA63+ZHwuWwPHPYoCZ/KLAjHv5g4h2MS4f2/MTI= @@ -985,8 +991,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= @@ -1431,6 +1437,8 @@ github.com/illarion/gonotify v1.0.1 h1:F1d+0Fgbq/sDWjj/r66ekjDG+IDeecQKUFH4wNwso github.com/illarion/gonotify v1.0.1/go.mod h1:zt5pmDofZpU1f8aqlK0+95eQhoEAn/d4G4B/FjVW4jE= github.com/insomniacslk/dhcp v0.0.0-20231206064809-8c70d406f6d2 h1:9K06NfxkBh25x56yVhWWlKFE8YpicaSfHwoV8SFbueA= github.com/insomniacslk/dhcp v0.0.0-20231206064809-8c70d406f6d2/go.mod h1:3A9PQ1cunSDF/1rbTq99Ts4pVnycWg+vlPkfeD2NLFI= +github.com/invopop/jsonschema v0.13.0 h1:KvpoAJWEjR3uD9Kbm2HWJmqsEaHt8lBUpd0qHcIi21E= +github.com/invopop/jsonschema v0.13.0/go.mod h1:ffZ5Km5SWWRAIN6wbDXItl95euhFz2uON45H2qjYt+0= github.com/jackmordaunt/icns/v3 v3.0.1 h1:xxot6aNuGrU+lNgxz5I5H0qSeCjNKp8uTXB1j8D4S3o= github.com/jackmordaunt/icns/v3 v3.0.1/go.mod h1:5sHL59nqTd2ynTnowxB/MDQFhKNqkK8X687uKNygaSQ= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOlocH6Fxy8MmwDt+yVQYULKfN0RoTN8A= @@ -1521,8 +1529,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.32.0 h1:fgwmbfL2gbd67obg57OfV2Dnrhs1HtSdlY/i5fn7MU8= -github.com/mark3labs/mcp-go v0.32.0/go.mod h1:rXqOudj/djTORU/ThxYx8fqEVj/5pvTuuebQ2RC7uk4= +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= @@ -1635,8 +1643,10 @@ 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/openai/openai-go/v2 v2.7.0 h1:/8MSFCXcasin7AyuWQ2au6FraXL71gzAs+VfbMv+J3k= +github.com/openai/openai-go/v2 v2.7.0/go.mod h1:jrJs23apqJKKbT+pqtFgNKpRju/KP9zpUTZhz3GElQE= 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 +1835,9 @@ github.com/tetratelabs/wazero v1.9.0/go.mod h1:TSbcXCfFP0L2FGkRPxHphadXPjo1T6W+C github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= -github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM= +github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= @@ -1838,8 +1849,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= @@ -1875,6 +1886,8 @@ github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAh github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds= github.com/wagslane/go-password-validator v0.3.0 h1:vfxOPzGHkz5S146HDpavl0cw1DSVP061Ry2PX0/ON6I= github.com/wagslane/go-password-validator v0.3.0/go.mod h1:TI1XJ6T5fRdRnHqHt14pvy1tNVnrwe7m3/f1f2fDphQ= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/wlynxg/anet v0.0.3/go.mod h1:eay5PRQr7fIVAMbTbchTnO9gG65Hg/uYGdc7mguHxoA= github.com/woodsbury/decimal128 v1.3.0 h1:8pffMNWIlC0O5vbyHWFZAt5yWvWcrHA+3ovIIjVWss0= github.com/woodsbury/decimal128 v1.3.0/go.mod h1:C5UTmyTjW3JftjUFzOVhC20BEQa2a4ZKOB5I6Zjb+ds=