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

Commit820ebaa

Browse files
committed
fix
1 parentbf6c970 commit820ebaa

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"
@@ -1051,6 +1053,25 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
10511053
options.WorkspaceUsageTracker=tracker
10521054
defertracker.Close()
10531055

1056+
varaibridgeDaemon*aibridged.Server
1057+
// In-memory aibridge daemon.
1058+
ifvals.AI.BridgeConfig.Enabled {
1059+
ifexperiments.Enabled(codersdk.ExperimentAIBridge) {
1060+
aibridgeDaemon,err=newAIBridgeDaemon(coderAPI)
1061+
iferr!=nil {
1062+
returnxerrors.Errorf("create aibridged: %w",err)
1063+
}
1064+
1065+
coderAPI.RegisterInMemoryAIBridgedHTTPHandler(aibridgeDaemon)
1066+
}else {
1067+
logger.Warn(ctx,fmt.Sprintf("CODER_AIBRIDGE_ENABLED=true but experiment %q not enabled",codersdk.ExperimentAIBridge))
1068+
}
1069+
}else {
1070+
ifexperiments.Enabled(codersdk.ExperimentAIBridge) {
1071+
logger.Warn(ctx,"aibridge experiment enabled but CODER_AIBRIDGE_ENABLED=false")
1072+
}
1073+
}
1074+
10541075
// Wrap the server in middleware that redirects to the access URL if
10551076
// the request is not to a local IP.
10561077
varhandler http.Handler=coderAPI.RootHandler
@@ -1157,6 +1178,23 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
11571178
cliui.Errorf(inv.Stderr,"Notify systemd failed: %s",err)
11581179
}
11591180

1181+
// Stop accepting new connections to aibridged.
1182+
//
1183+
// When running as an in-memory daemon, the HTTP handler is wired into the
1184+
// coderd API and therefore is subject to its context. Calling shutdown on
1185+
// aibridged will NOT affect in-flight requests but those will be closed once
1186+
// the API server is shutdown below.
1187+
ifaibridgeDaemon!=nil {
1188+
cliui.Info(inv.Stdout,"Shutting down aibridge daemon...\n")
1189+
1190+
err=shutdownWithTimeout(aibridgeDaemon.Shutdown,5*time.Second)
1191+
iferr!=nil {
1192+
cliui.Errorf(inv.Stderr,"Graceful shutdown of aibridge daemon failed: %s\n",err)
1193+
}else {
1194+
cliui.Info(inv.Stdout,"Gracefully shut down aibridge daemon\n")
1195+
}
1196+
}
1197+
11601198
// Stop accepting new connections without interrupting
11611199
// in-flight requests, give in-flight requests 5 seconds to
11621200
// complete.
@@ -1519,6 +1557,40 @@ func newProvisionerDaemon(
15191557
}),nil
15201558
}
15211559

1560+
funcnewAIBridgeDaemon(coderAPI*coderd.API) (*aibridged.Server,error) {
1561+
ctx:=context.Background()
1562+
coderAPI.Logger.Debug(ctx,"starting in-memory aibridge daemon")
1563+
1564+
logger:=coderAPI.Logger.Named("aibridged")
1565+
1566+
// Setup supported providers.
1567+
providers:= []aibridge.Provider{
1568+
aibridge.NewOpenAIProvider(aibridge.ProviderConfig{
1569+
BaseURL:coderAPI.DeploymentValues.AI.BridgeConfig.OpenAI.BaseURL.String(),
1570+
Key:coderAPI.DeploymentValues.AI.BridgeConfig.OpenAI.Key.String(),
1571+
}),
1572+
aibridge.NewAnthropicProvider(aibridge.ProviderConfig{
1573+
BaseURL:coderAPI.DeploymentValues.AI.BridgeConfig.Anthropic.BaseURL.String(),
1574+
Key:coderAPI.DeploymentValues.AI.BridgeConfig.Anthropic.Key.String(),
1575+
}),
1576+
}
1577+
1578+
// Create pool for reusable stateful [aibridge.RequestBridge] instances (one per user).
1579+
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,logger.Named("pool"))// TODO: configurable.
1580+
iferr!=nil {
1581+
returnnil,xerrors.Errorf("create request pool: %w",err)
1582+
}
1583+
1584+
// Create daemon.
1585+
srv,err:=aibridged.New(ctx,pool,func(dialCtx context.Context) (aibridged.DRPCClient,error) {
1586+
returncoderAPI.CreateInMemoryAIBridgeServer(dialCtx)
1587+
},logger)
1588+
iferr!=nil {
1589+
returnnil,xerrors.Errorf("start in-memory aibridge daemon: %w",err)
1590+
}
1591+
returnsrv,nil
1592+
}
1593+
15221594
// nolint: revive
15231595
funcPrintLogo(inv*serpent.Invocation,daemonTitlestring) {
15241596
// 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"
@@ -1007,6 +1011,11 @@ func New(options *Options) *API {
10071011

10081012
// Experimental routes are not guaranteed to be stable and may change at any time.
10091013
r.Route("/api/experimental",func(r chi.Router) {
1014+
r.NotFound(func(rw http.ResponseWriter,_*http.Request) {httpapi.RouteNotFound(rw) })
1015+
1016+
// Only this group should be subject to apiKeyMiddleware; aibridged will mount its own
1017+
// router and handles key validation in a different fashion.
1018+
// See aibridged/http.go.
10101019
r.Group(func(r chi.Router) {
10111020
r.Use(apiKeyMiddleware)
10121021
r.Route("/aitasks",func(r chi.Router) {
@@ -1985,6 +1994,78 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
19851994
returnproto.NewDRPCProvisionerDaemonClient(clientSession),nil
19861995
}
19871996

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

0 commit comments

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