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

Commit11e7d49

Browse files
committed
fix
1 parentb2d6770 commit11e7d49

File tree

3 files changed

+188
-1
lines changed

3 files changed

+188
-1
lines changed

‎cli/server.go‎

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

5656
"cdr.dev/slog"
5757
"cdr.dev/slog/sloggers/sloghuman"
58+
"github.com/coder/aibridge"
59+
"github.com/coder/coder/v2/aibridged"
5860
"github.com/coder/coder/v2/coderd/pproflabel"
5961
"github.com/coder/pretty"
6062
"github.com/coder/quartz"
@@ -1038,6 +1040,25 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
10381040
options.WorkspaceUsageTracker=tracker
10391041
defertracker.Close()
10401042

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+
10411062
// Wrap the server in middleware that redirects to the access URL if
10421063
// the request is not to a local IP.
10431064
varhandler http.Handler=coderAPI.RootHandler
@@ -1144,6 +1165,23 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
11441165
cliui.Errorf(inv.Stderr,"Notify systemd failed: %s",err)
11451166
}
11461167

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+
11471185
// Stop accepting new connections without interrupting
11481186
// in-flight requests, give in-flight requests 5 seconds to
11491187
// complete.
@@ -1506,6 +1544,40 @@ func newProvisionerDaemon(
15061544
}),nil
15071545
}
15081546

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+
15091581
// nolint: revive
15101582
funcPrintLogo(inv*serpent.Invocation,daemonTitlestring) {
15111583
// Only print the logo in TTYs.

‎coderd/aibridged.go‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
package coderd
2+
3+
import (
4+
"net/http"
5+
6+
"github.com/go-chi/chi/v5"
7+
8+
"github.com/coder/coder/v2/aibridged"
9+
"github.com/coder/coder/v2/coderd/httpmw"
10+
"github.com/coder/coder/v2/codersdk"
11+
)
12+
13+
// RegisterInMemoryAIBridgedHTTPHandler mounts [aibridged.Server]'s HTTP router onto
14+
// [API]'s router, so that requests to aibridged will be relayed from Coder's API server
15+
// to the in-memory aibridged.
16+
func (api*API)RegisterInMemoryAIBridgedHTTPHandler(srv*aibridged.Server) {
17+
ifsrv==nil {
18+
panic("aibridged cannot be nil")
19+
}
20+
21+
ifapi.RootHandler==nil {
22+
panic("api.RootHandler cannot be nil")
23+
}
24+
25+
aibridgeEndpoint:="/api/experimental/aibridge"
26+
27+
r:=chi.NewRouter()
28+
r.Group(func(r chi.Router) {
29+
r.Use(httpmw.RequireExperiment(api.Experiments,codersdk.ExperimentAIBridge))
30+
r.HandleFunc("/*",http.StripPrefix(aibridgeEndpoint,srv).ServeHTTP)
31+
})
32+
33+
api.RootHandler.Mount(aibridgeEndpoint,r)
34+
}

‎coderd/coderd.go‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ import (
2020
"sync/atomic"
2121
"time"
2222

23+
"github.com/coder/coder/v2/aibridged"
24+
aibridgedproto"github.com/coder/coder/v2/aibridged/proto"
2325
"github.com/coder/coder/v2/coderd/oauth2provider"
2426
"github.com/coder/coder/v2/coderd/pproflabel"
2527
"github.com/coder/coder/v2/coderd/prebuilds"
@@ -44,6 +46,9 @@ import (
4446
"tailscale.com/types/key"
4547
"tailscale.com/util/singleflight"
4648

49+
"github.com/coder/coder/v2/coderd/aibridgedserver"
50+
"github.com/coder/coder/v2/provisionerd/proto"
51+
4752
"cdr.dev/slog"
4853
"github.com/coder/quartz"
4954
"github.com/coder/serpent"
@@ -95,7 +100,6 @@ import (
95100
"github.com/coder/coder/v2/coderd/workspacestats"
96101
"github.com/coder/coder/v2/codersdk"
97102
"github.com/coder/coder/v2/codersdk/healthsdk"
98-
"github.com/coder/coder/v2/provisionerd/proto"
99103
"github.com/coder/coder/v2/provisionersdk"
100104
"github.com/coder/coder/v2/site"
101105
"github.com/coder/coder/v2/tailnet"
@@ -999,6 +1003,11 @@ func New(options *Options) *API {
9991003

10001004
// Experimental routes are not guaranteed to be stable and may change at any time.
10011005
r.Route("/api/experimental",func(r chi.Router) {
1006+
r.NotFound(func(rw http.ResponseWriter,_*http.Request) {httpapi.RouteNotFound(rw) })
1007+
1008+
// Only this group should be subject to apiKeyMiddleware; aibridged will mount its own
1009+
// router and handles key validation in a different fashion.
1010+
// See aibridged/http.go.
10021011
r.Group(func(r chi.Router) {
10031012
r.Use(apiKeyMiddleware)
10041013
r.Route("/aitasks",func(r chi.Router) {
@@ -1977,6 +1986,78 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
19771986
returnproto.NewDRPCProvisionerDaemonClient(clientSession),nil
19781987
}
19791988

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

0 commit comments

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