Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 7d7fe06

Browse files
dannykoppingdannykopping
committed
feat: aibridged init & API mount
1 parent 8632adb commit 7d7fe06

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+
if current := 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+
if err != 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+
if current := api.AIBridgeDaemon.Load(); current != nil {
1835+
_ = (*current).Close()
1836+
}
1837+
18271838
return nil
18281839
}
18291840

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

2011+
func (api *API) CreateInMemoryAIBridgeDaemon(dialCtx context.Context) (client aibridged.DRPCClient, err error) {
2012+
// TODO(dannyk): implement options.
2013+
// TODO(dannyk): implement tracing.
2014+
// TODO(dannyk): implement API versioning.
2015+
2016+
clientSession, serverSession := drpcsdk.MemTransportPipe()
2017+
defer func() {
2018+
if err != 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+
if err != nil {
2030+
return nil, err
2031+
}
2032+
err = aibridgedproto.DRPCRegisterRecorder(mux, srv)
2033+
if err != nil {
2034+
return nil, xerrors.Errorf("register recorder service: %w", err)
2035+
}
2036+
err = aibridgedproto.DRPCRegisterMCPConfigurator(mux, srv)
2037+
if err != nil {
2038+
return nil, xerrors.Errorf("register MCP configurator service: %w", err)
2039+
}
2040+
err = aibridgedproto.DRPCRegisterAuthenticator(mux, srv)
2041+
if err != nil {
2042+
return nil, 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(err error) {
2048+
if xerrors.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+
go func() {
2062+
defer api.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
if fn != 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
return nil, xerrors.Errorf("unable to register license metrics collector")
618621
}
619622

623+
// In-memory aibridge daemons.
624+
// TODO: license entitlement.
625+
if api.DeploymentValues.AI.BridgeConfig.Enabled {
626+
if api.AGPL.Experiments.Enabled(codersdk.ExperimentAIBridge) {
627+
srv, err := newAIBridgeServer(api.AGPL)
628+
if err != nil {
629+
return nil, 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
if err != nil {
622639
return nil, 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
return reconciler, prebuilds.NewEnterpriseClaimer(api.Database)
12771294
}
1295+
1296+
func newAIBridgeServer(coderAPI *coderd.API) (aibridged.Server, error) {
1297+
srv, err := aibridged.New(
1298+
func(dialCtx context.Context) (aibridged.DRPCClient, error) {
1299+
return coderAPI.CreateInMemoryAIBridgeDaemon(dialCtx)
1300+
},
1301+
convertAIBridgeDeploymentValues(coderAPI.DeploymentValues.AI.BridgeConfig),
1302+
coderAPI.Logger.Named("aibridged"),
1303+
)
1304+
if err != nil {
1305+
return nil, xerrors.Errorf("create aibridge daemon: %w", err)
1306+
}
1307+
return srv, nil
1308+
}
1309+
1310+
func convertAIBridgeDeploymentValues(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)