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

Commite7f3628

Browse files
committed
feat: aibridged init & API mount
1 parentd5ea768 commite7f3628

File tree

3 files changed

+145
-1
lines changed

3 files changed

+145
-1
lines changed

‎cli/server.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1152,6 +1152,24 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
11521152
cliui.Errorf(inv.Stderr,"Notify systemd failed: %s",err)
11531153
}
11541154

1155+
// Stop accepting new connections to aibridged.
1156+
//
1157+
// When running as an in-memory daemon, the HTTP handler is wired into the
1158+
// coderd API and therefore is subject to its context. Calling shutdown on
1159+
// aibridged will NOT affect in-flight requests but those will be closed once
1160+
// the API server is shutdown below.
1161+
ifcurrent:=coderAPI.AIBridgeDaemon.Load();current!=nil {
1162+
cliui.Info(inv.Stdout,"Shutting down aibridge daemon...\n")
1163+
1164+
err=shutdownWithTimeout((*current).Shutdown,5*time.Second)
1165+
iferr!=nil {
1166+
cliui.Errorf(inv.Stderr,"Graceful shutdown of aibridge daemon failed: %s\n",err)
1167+
}else {
1168+
_= (*current).Close()
1169+
cliui.Info(inv.Stdout,"Gracefully shut down aibridge daemon\n")
1170+
}
1171+
}
1172+
11551173
// Stop accepting new connections without interrupting
11561174
// in-flight requests, give in-flight requests 5 seconds to
11571175
// complete.

‎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"
@@ -632,6 +636,7 @@ func New(options *Options) *API {
632636
api.PortSharer.Store(&portsharing.DefaultPortSharer)
633637
api.PrebuildsClaimer.Store(&prebuilds.DefaultClaimer)
634638
api.PrebuildsReconciler.Store(&prebuilds.DefaultReconciler)
639+
api.AIBridgeDaemon.Store(&aibridged.DefaultServer)
635640
buildInfo:= codersdk.BuildInfoResponse{
636641
ExternalURL:buildinfo.ExternalURL(),
637642
Version:buildinfo.Version(),
@@ -1766,6 +1771,8 @@ type API struct {
17661771
// dbRolluper rolls up template usage stats from raw agent and app
17671772
// stats. This is used to provide insights in the WebUI.
17681773
dbRolluper*dbrollup.Rolluper
1774+
1775+
AIBridgeDaemon atomic.Pointer[aibridged.Server]
17691776
}
17701777

17711778
// Close waits for all WebSocket connections to drain before returning.
@@ -1824,6 +1831,10 @@ func (api *API) Close() error {
18241831
(*current).Stop(ctx,nil)
18251832
}
18261833

1834+
ifcurrent:=api.AIBridgeDaemon.Load();current!=nil {
1835+
_= (*current).Close()
1836+
}
1837+
18271838
returnnil
18281839
}
18291840

@@ -1997,6 +2008,76 @@ func (api *API) CreateInMemoryTaggedProvisionerDaemon(dialCtx context.Context, n
19972008
returnproto.NewDRPCProvisionerDaemonClient(clientSession),nil
19982009
}
19992010

2011+
func (api*API)CreateInMemoryAIBridgeDaemon(dialCtx context.Context) (client aibridged.DRPCClient,errerror) {
2012+
// TODO(dannyk): implement options.
2013+
// TODO(dannyk): implement tracing.
2014+
// TODO(dannyk): implement API versioning.
2015+
2016+
clientSession,serverSession:=drpcsdk.MemTransportPipe()
2017+
deferfunc() {
2018+
iferr!=nil {
2019+
_=clientSession.Close()
2020+
_=serverSession.Close()
2021+
}
2022+
}()
2023+
2024+
mux:=drpcmux.New()
2025+
api.Logger.Debug(dialCtx,"starting in-memory aibridge daemon")
2026+
logger:=api.Logger.Named("inmem-aibridged")
2027+
srv,err:=aibridgedserver.NewServer(api.ctx,api.Database,logger,
2028+
api.DeploymentValues.AccessURL.String(),api.ExternalAuthConfigs)
2029+
iferr!=nil {
2030+
returnnil,err
2031+
}
2032+
err=aibridgedproto.DRPCRegisterRecorder(mux,srv)
2033+
iferr!=nil {
2034+
returnnil,xerrors.Errorf("register recorder service: %w",err)
2035+
}
2036+
err=aibridgedproto.DRPCRegisterMCPConfigurator(mux,srv)
2037+
iferr!=nil {
2038+
returnnil,xerrors.Errorf("register MCP configurator service: %w",err)
2039+
}
2040+
err=aibridgedproto.DRPCRegisterAuthenticator(mux,srv)
2041+
iferr!=nil {
2042+
returnnil,xerrors.Errorf("register authenticator service: %w",err)
2043+
}
2044+
server:=drpcserver.NewWithOptions(&tracing.DRPCHandler{Handler:mux},
2045+
drpcserver.Options{
2046+
Manager:drpcsdk.DefaultDRPCOptions(nil),
2047+
Log:func(errerror) {
2048+
ifxerrors.Is(err,io.EOF) {
2049+
return
2050+
}
2051+
logger.Debug(dialCtx,"drpc server error",slog.Error(err))
2052+
},
2053+
},
2054+
)
2055+
// in-mem pipes aren't technically "websockets" but they have the same properties as far as the
2056+
// API is concerned: they are long-lived connections that we need to close before completing
2057+
// shutdown of the API.
2058+
api.WebsocketWaitMutex.Lock()
2059+
api.WebsocketWaitGroup.Add(1)
2060+
api.WebsocketWaitMutex.Unlock()
2061+
gofunc() {
2062+
deferapi.WebsocketWaitGroup.Done()
2063+
// Here we pass the background context, since we want the server to keep serving until the
2064+
// client hangs up. The aibridged is local, in-mem, so there isn't a danger of losing contact with it and
2065+
// having a dead connection we don't know the status of.
2066+
err:=server.Serve(context.Background(),serverSession)
2067+
logger.Info(dialCtx,"aibridge daemon disconnected",slog.Error(err))
2068+
// Close the sessions, so we don't leak goroutines serving them.
2069+
_=clientSession.Close()
2070+
_=serverSession.Close()
2071+
}()
2072+
2073+
return&aibridged.Client{
2074+
Conn:clientSession,
2075+
DRPCRecorderClient:aibridgedproto.NewDRPCRecorderClient(clientSession),
2076+
DRPCMCPConfiguratorClient:aibridgedproto.NewDRPCMCPConfiguratorClient(clientSession),
2077+
DRPCAuthenticatorClient:aibridgedproto.NewDRPCAuthenticatorClient(clientSession),
2078+
},nil
2079+
}
2080+
20002081
func (api*API)DERPMap()*tailcfg.DERPMap {
20012082
fn:=api.DERPMapper.Load()
20022083
iffn!=nil {

‎enterprise/coderd/coderd.go‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"sync/atomic"
1414
"time"
1515

16+
"github.com/coder/coder/v2/aibridged"
1617
"github.com/coder/coder/v2/buildinfo"
1718
"github.com/coder/coder/v2/coderd/appearance"
1819
"github.com/coder/coder/v2/coderd/database"
@@ -61,6 +62,8 @@ import (
6162
"github.com/coder/coder/v2/enterprise/tailnet"
6263
"github.com/coder/coder/v2/provisionerd/proto"
6364
agpltailnet"github.com/coder/coder/v2/tailnet"
65+
66+
"github.com/coder/aibridge"
6467
)
6568

6669
// New constructs an Enterprise coderd API instance.
@@ -617,6 +620,20 @@ func New(ctx context.Context, options *Options) (_ *API, err error) {
617620
returnnil,xerrors.Errorf("unable to register license metrics collector")
618621
}
619622

623+
// In-memory aibridge daemons.
624+
// TODO: license entitlement.
625+
ifapi.DeploymentValues.AI.BridgeConfig.Enabled {
626+
ifapi.AGPL.Experiments.Enabled(codersdk.ExperimentAIBridge) {
627+
srv,err:=newAIBridgeServer(api.AGPL)
628+
iferr!=nil {
629+
returnnil,xerrors.Errorf("create aibridged: %w",err)
630+
}
631+
api.AGPL.AIBridgeDaemon.Store(&srv)
632+
}else {
633+
api.Logger.Error(ctx,fmt.Sprintf("aibridge enabled but experiment %q not enabled",codersdk.ExperimentAIBridge))
634+
}
635+
}
636+
620637
err=api.updateEntitlements(ctx)
621638
iferr!=nil {
622639
returnnil,xerrors.Errorf("update entitlements: %w",err)
@@ -1275,3 +1292,31 @@ func (api *API) setupPrebuilds(featureEnabled bool) (agplprebuilds.Reconciliatio
12751292
api.Logger.Named("prebuilds"),quartz.NewReal(),api.PrometheusRegistry,api.NotificationsEnqueuer,api.AGPL.BuildUsageChecker)
12761293
returnreconciler,prebuilds.NewEnterpriseClaimer(api.Database)
12771294
}
1295+
1296+
funcnewAIBridgeServer(coderAPI*coderd.API) (aibridged.Server,error) {
1297+
srv,err:=aibridged.New(
1298+
func(dialCtx context.Context) (aibridged.DRPCClient,error) {
1299+
returncoderAPI.CreateInMemoryAIBridgeDaemon(dialCtx)
1300+
},
1301+
convertAIBridgeDeploymentValues(coderAPI.DeploymentValues.AI.BridgeConfig),
1302+
coderAPI.Logger.Named("aibridged"),
1303+
)
1304+
iferr!=nil {
1305+
returnnil,xerrors.Errorf("create aibridge daemon: %w",err)
1306+
}
1307+
returnsrv,nil
1308+
}
1309+
1310+
funcconvertAIBridgeDeploymentValues(vals codersdk.AIBridgeConfig) aibridge.Config {
1311+
return aibridge.Config{
1312+
OpenAI: aibridge.ProviderConfig{
1313+
BaseURL:vals.OpenAI.BaseURL.String(),
1314+
Key:vals.OpenAI.Key.String(),
1315+
},
1316+
Anthropic: aibridge.ProviderConfig{
1317+
BaseURL:vals.Anthropic.BaseURL.String(),
1318+
Key:vals.Anthropic.Key.String(),
1319+
},
1320+
CacheSize:100,// TODO: configurable.
1321+
}
1322+
}

0 commit comments

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