Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit9c2ebdd

Browse files
committed
chore: reparent aibridge code under enterprise
1 parentaebdc28 commit9c2ebdd

File tree

9 files changed

+218
-201
lines changed

9 files changed

+218
-201
lines changed

‎cli/server.go‎

Lines changed: 0 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,7 @@ import (
5555

5656
"cdr.dev/slog"
5757
"cdr.dev/slog/sloggers/sloghuman"
58-
"github.com/coder/aibridge"
5958
"github.com/coder/coder/v2/coderd/pproflabel"
60-
"github.com/coder/coder/v2/enterprise/x/aibridged"
6159
"github.com/coder/pretty"
6260
"github.com/coder/quartz"
6361
"github.com/coder/retry"
@@ -1040,25 +1038,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
10401038
options.WorkspaceUsageTracker=tracker
10411039
defertracker.Close()
10421040

1043-
varaibridgeDaemon*aibridged.Server
1044-
// In-memory aibridge daemon.
1045-
ifvals.AI.BridgeConfig.Enabled {
1046-
ifexperiments.Enabled(codersdk.ExperimentAIBridge) {
1047-
aibridgeDaemon,err=newAIBridgeDaemon(coderAPI)
1048-
iferr!=nil {
1049-
returnxerrors.Errorf("create aibridged: %w",err)
1050-
}
1051-
1052-
coderAPI.RegisterInMemoryAIBridgedHTTPHandler(aibridgeDaemon)
1053-
}else {
1054-
logger.Warn(ctx,fmt.Sprintf("CODER_AIBRIDGE_ENABLED=true but experiment %q not enabled",codersdk.ExperimentAIBridge))
1055-
}
1056-
}else {
1057-
ifexperiments.Enabled(codersdk.ExperimentAIBridge) {
1058-
logger.Warn(ctx,"aibridge experiment enabled but CODER_AIBRIDGE_ENABLED=false")
1059-
}
1060-
}
1061-
10621041
// Wrap the server in middleware that redirects to the access URL if
10631042
// the request is not to a local IP.
10641043
varhandler http.Handler=coderAPI.RootHandler
@@ -1165,23 +1144,6 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
11651144
cliui.Errorf(inv.Stderr,"Notify systemd failed: %s",err)
11661145
}
11671146

1168-
// Stop accepting new connections to aibridged.
1169-
//
1170-
// When running as an in-memory daemon, the HTTP handler is wired into the
1171-
// coderd API and therefore is subject to its context. Calling shutdown on
1172-
// aibridged will NOT affect in-flight requests but those will be closed once
1173-
// the API server is shutdown below.
1174-
ifaibridgeDaemon!=nil {
1175-
cliui.Info(inv.Stdout,"Shutting down aibridge daemon...\n")
1176-
1177-
err=shutdownWithTimeout(aibridgeDaemon.Shutdown,5*time.Second)
1178-
iferr!=nil {
1179-
cliui.Errorf(inv.Stderr,"Graceful shutdown of aibridge daemon failed: %s\n",err)
1180-
}else {
1181-
cliui.Info(inv.Stdout,"Gracefully shut down aibridge daemon\n")
1182-
}
1183-
}
1184-
11851147
// Stop accepting new connections without interrupting
11861148
// in-flight requests, give in-flight requests 5 seconds to
11871149
// complete.
@@ -1544,40 +1506,6 @@ func newProvisionerDaemon(
15441506
}),nil
15451507
}
15461508

1547-
funcnewAIBridgeDaemon(coderAPI*coderd.API) (*aibridged.Server,error) {
1548-
ctx:=context.Background()
1549-
coderAPI.Logger.Debug(ctx,"starting in-memory aibridge daemon")
1550-
1551-
logger:=coderAPI.Logger.Named("aibridged")
1552-
1553-
// Setup supported providers.
1554-
providers:= []aibridge.Provider{
1555-
aibridge.NewOpenAIProvider(aibridge.ProviderConfig{
1556-
BaseURL:coderAPI.DeploymentValues.AI.BridgeConfig.OpenAI.BaseURL.String(),
1557-
Key:coderAPI.DeploymentValues.AI.BridgeConfig.OpenAI.Key.String(),
1558-
}),
1559-
aibridge.NewAnthropicProvider(aibridge.ProviderConfig{
1560-
BaseURL:coderAPI.DeploymentValues.AI.BridgeConfig.Anthropic.BaseURL.String(),
1561-
Key:coderAPI.DeploymentValues.AI.BridgeConfig.Anthropic.Key.String(),
1562-
}),
1563-
}
1564-
1565-
// Create pool for reusable stateful [aibridge.RequestBridge] instances (one per user).
1566-
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,logger.Named("pool"))// TODO: configurable.
1567-
iferr!=nil {
1568-
returnnil,xerrors.Errorf("create request pool: %w",err)
1569-
}
1570-
1571-
// Create daemon.
1572-
srv,err:=aibridged.New(ctx,pool,func(dialCtx context.Context) (aibridged.DRPCClient,error) {
1573-
returncoderAPI.CreateInMemoryAIBridgeServer(dialCtx)
1574-
},logger)
1575-
iferr!=nil {
1576-
returnnil,xerrors.Errorf("start in-memory aibridge daemon: %w",err)
1577-
}
1578-
returnsrv,nil
1579-
}
1580-
15811509
// nolint: revive
15821510
funcPrintLogo(inv*serpent.Invocation,daemonTitlestring) {
15831511
// Only print the logo in TTYs.

‎coderd/aibridged.go‎

Lines changed: 0 additions & 34 deletions
This file was deleted.

‎coderd/coderd.go‎

Lines changed: 1 addition & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,6 @@ import (
2525
"github.com/coder/coder/v2/coderd/prebuilds"
2626
"github.com/coder/coder/v2/coderd/usage"
2727
"github.com/coder/coder/v2/coderd/wsbuilder"
28-
"github.com/coder/coder/v2/enterprise/x/aibridged"
29-
aibridgedproto"github.com/coder/coder/v2/enterprise/x/aibridged/proto"
3028

3129
"github.com/andybalholm/brotli"
3230
"github.com/go-chi/chi/v5"
@@ -46,7 +44,6 @@ import (
4644
"tailscale.com/types/key"
4745
"tailscale.com/util/singleflight"
4846

49-
"github.com/coder/coder/v2/enterprise/x/aibridgedserver"
5047
"github.com/coder/coder/v2/provisionerd/proto"
5148

5249
"cdr.dev/slog"
@@ -1007,7 +1004,7 @@ func New(options *Options) *API {
10071004

10081005
// Only this group should be subject to apiKeyMiddleware; aibridged will mount its own
10091006
// router and handles key validation in a different fashion.
1010-
// See aibridged/http.go.
1007+
// Seeenterprise/x/aibridged/http.go.
10111008
r.Group(func(r chi.Router) {
10121009
r.Use(apiKeyMiddleware)
10131010
r.Route("/aitasks",func(r chi.Router) {
@@ -1987,78 +1984,6 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
19871984
returnproto.NewDRPCProvisionerDaemonClient(clientSession),nil
19881985
}
19891986

1990-
// CreateInMemoryAIBridgeServer creates a [aibridged.DRPCServer] and returns a
1991-
// [aibridged.DRPCClient] to it, connected over an in-memory transport.
1992-
// This server is responsible for all the Coder-specific functionality that aibridged
1993-
// requires such as persistence and retrieving configuration.
1994-
func (api*API)CreateInMemoryAIBridgeServer(dialCtx context.Context) (client aibridged.DRPCClient,errerror) {
1995-
// TODO(dannyk): implement options.
1996-
// TODO(dannyk): implement tracing.
1997-
// TODO(dannyk): implement API versioning.
1998-
1999-
clientSession,serverSession:=drpcsdk.MemTransportPipe()
2000-
deferfunc() {
2001-
iferr!=nil {
2002-
_=clientSession.Close()
2003-
_=serverSession.Close()
2004-
}
2005-
}()
2006-
2007-
mux:=drpcmux.New()
2008-
srv,err:=aibridgedserver.NewServer(api.ctx,api.Database,api.Logger.Named("aibridgedserver"),
2009-
api.AccessURL.String(),api.ExternalAuthConfigs,api.Experiments)
2010-
iferr!=nil {
2011-
returnnil,err
2012-
}
2013-
err=aibridgedproto.DRPCRegisterRecorder(mux,srv)
2014-
iferr!=nil {
2015-
returnnil,xerrors.Errorf("register recorder service: %w",err)
2016-
}
2017-
err=aibridgedproto.DRPCRegisterMCPConfigurator(mux,srv)
2018-
iferr!=nil {
2019-
returnnil,xerrors.Errorf("register MCP configurator service: %w",err)
2020-
}
2021-
err=aibridgedproto.DRPCRegisterAuthorizer(mux,srv)
2022-
iferr!=nil {
2023-
returnnil,xerrors.Errorf("register key validator service: %w",err)
2024-
}
2025-
server:=drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler:mux},
2026-
drpcserver.Options{
2027-
Manager:drpcsdk.DefaultDRPCOptions(nil),
2028-
Log:func(errerror) {
2029-
iferrors.Is(err,io.EOF) {
2030-
return
2031-
}
2032-
api.Logger.Debug(dialCtx,"aibridged drpc server error",slog.Error(err))
2033-
},
2034-
},
2035-
)
2036-
// in-mem pipes aren't technically "websockets" but they have the same properties as far as the
2037-
// API is concerned: they are long-lived connections that we need to close before completing
2038-
// shutdown of the API.
2039-
api.WebsocketWaitMutex.Lock()
2040-
api.WebsocketWaitGroup.Add(1)
2041-
api.WebsocketWaitMutex.Unlock()
2042-
gofunc() {
2043-
deferapi.WebsocketWaitGroup.Done()
2044-
// Here we pass the background context, since we want the server to keep serving until the
2045-
// client hangs up. The aibridged is local, in-mem, so there isn't a danger of losing contact with it and
2046-
// having a dead connection we don't know the status of.
2047-
err:=server.Serve(context.Background(),serverSession)
2048-
api.Logger.Info(dialCtx,"aibridge daemon disconnected",slog.Error(err))
2049-
// Close the sessions, so we don't leak goroutines serving them.
2050-
_=clientSession.Close()
2051-
_=serverSession.Close()
2052-
}()
2053-
2054-
return&aibridged.Client{
2055-
Conn:clientSession,
2056-
DRPCRecorderClient:aibridgedproto.NewDRPCRecorderClient(clientSession),
2057-
DRPCMCPConfiguratorClient:aibridgedproto.NewDRPCMCPConfiguratorClient(clientSession),
2058-
DRPCAuthorizerClient:aibridgedproto.NewDRPCAuthorizerClient(clientSession),
2059-
},nil
2060-
}
2061-
20621987
func (api*API)DERPMap()*tailcfg.DERPMap {
20631988
fn:=api.DERPMapper.Load()
20641989
iffn!=nil {

‎enterprise/cli/aibridged.go‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package cli
2+
3+
import (
4+
"context"
5+
6+
"golang.org/x/xerrors"
7+
8+
"github.com/coder/aibridge"
9+
"github.com/coder/coder/v2/enterprise/coderd"
10+
"github.com/coder/coder/v2/enterprise/x/aibridged"
11+
)
12+
13+
funcnewAIBridgeDaemon(coderAPI*coderd.API) (*aibridged.Server,error) {
14+
ctx:=context.Background()
15+
coderAPI.Logger.Debug(ctx,"starting in-memory aibridge daemon")
16+
17+
logger:=coderAPI.Logger.Named("aibridged")
18+
19+
// Setup supported providers.
20+
providers:= []aibridge.Provider{
21+
aibridge.NewOpenAIProvider(aibridge.ProviderConfig{
22+
BaseURL:coderAPI.DeploymentValues.AI.BridgeConfig.OpenAI.BaseURL.String(),
23+
Key:coderAPI.DeploymentValues.AI.BridgeConfig.OpenAI.Key.String(),
24+
}),
25+
aibridge.NewAnthropicProvider(aibridge.ProviderConfig{
26+
BaseURL:coderAPI.DeploymentValues.AI.BridgeConfig.Anthropic.BaseURL.String(),
27+
Key:coderAPI.DeploymentValues.AI.BridgeConfig.Anthropic.Key.String(),
28+
}),
29+
}
30+
31+
// Create pool for reusable stateful [aibridge.RequestBridge] instances (one per user).
32+
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,logger.Named("pool"))// TODO: configurable.
33+
iferr!=nil {
34+
returnnil,xerrors.Errorf("create request pool: %w",err)
35+
}
36+
37+
// Create daemon.
38+
srv,err:=aibridged.New(ctx,pool,func(dialCtx context.Context) (aibridged.DRPCClient,error) {
39+
returncoderAPI.CreateInMemoryAIBridgeServer(dialCtx)
40+
},logger)
41+
iferr!=nil {
42+
returnnil,xerrors.Errorf("start in-memory aibridge daemon: %w",err)
43+
}
44+
returnsrv,nil
45+
}

‎enterprise/cli/server.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"database/sql"
88
"encoding/base64"
99
"errors"
10+
"fmt"
1011
"io"
1112
"net/url"
1213

@@ -15,6 +16,7 @@ import (
1516
"tailscale.com/types/key"
1617

1718
"github.com/coder/coder/v2/coderd/database"
19+
"github.com/coder/coder/v2/codersdk"
1820
"github.com/coder/coder/v2/cryptorand"
1921
"github.com/coder/coder/v2/enterprise/audit"
2022
"github.com/coder/coder/v2/enterprise/audit/backends"
@@ -23,6 +25,7 @@ import (
2325
"github.com/coder/coder/v2/enterprise/coderd/usage"
2426
"github.com/coder/coder/v2/enterprise/dbcrypt"
2527
"github.com/coder/coder/v2/enterprise/trialer"
28+
"github.com/coder/coder/v2/enterprise/x/aibridged"
2629
"github.com/coder/coder/v2/tailnet"
2730
"github.com/coder/quartz"
2831
"github.com/coder/serpent"
@@ -143,6 +146,33 @@ func (r *RootCmd) Server(_ func()) *serpent.Command {
143146
}
144147
closers.Add(publisher)
145148

149+
experiments:=agplcoderd.ReadExperiments(options.Logger,options.DeploymentValues.Experiments.Value())
150+
151+
varaibridgeDaemon*aibridged.Server
152+
// In-memory aibridge daemon.
153+
ifoptions.DeploymentValues.AI.BridgeConfig.Enabled {
154+
ifexperiments.Enabled(codersdk.ExperimentAIBridge) {
155+
aibridgeDaemon,err=newAIBridgeDaemon(api)
156+
iferr!=nil {
157+
returnnil,nil,xerrors.Errorf("create aibridged: %w",err)
158+
}
159+
160+
api.RegisterInMemoryAIBridgedHTTPHandler(aibridgeDaemon)
161+
162+
// When running as an in-memory daemon, the HTTP handler is wired into the
163+
// coderd API and therefore is subject to its context. Calling Close() on
164+
// aibridged will NOT affect in-flight requests but those will be closed once
165+
// the API server is itself shutdown.
166+
closers.Add(aibridgeDaemon)
167+
}else {
168+
api.Logger.Warn(ctx,fmt.Sprintf("CODER_AIBRIDGE_ENABLED=true but experiment %q not enabled",codersdk.ExperimentAIBridge))
169+
}
170+
}else {
171+
ifexperiments.Enabled(codersdk.ExperimentAIBridge) {
172+
api.Logger.Warn(ctx,"aibridge experiment enabled but CODER_AIBRIDGE_ENABLED=false")
173+
}
174+
}
175+
146176
returnapi.AGPL,closers,nil
147177
})
148178

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp