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

Commitd9cabf3

Browse files
committed
fix
1 parent266d14b commitd9cabf3

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"
@@ -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