Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 11 additions & 11 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -53,16 +53,16 @@ endif
tailnet/tailnettest/coordinateemock.go \
tailnet/tailnettest/workspaceupdatesprovidermock.go \
tailnet/tailnettest/subscriptionmock.go \
enterprise/aibridged/aibridgedmock/clientmock.go \
enterprise/aibridged/aibridgedmock/poolmock.go \
coderd/aibridged/aibridgedmock/clientmock.go \
coderd/aibridged/aibridgedmock/poolmock.go \
tailnet/proto/tailnet.pb.go \
agent/proto/agent.pb.go \
agent/agentsocket/proto/agentsocket.pb.go \
agent/boundarylogproxy/codec/boundary.pb.go \
provisionersdk/proto/provisioner.pb.go \
provisionerd/proto/provisionerd.pb.go \
vpn/vpn.pb.go \
enterprise/aibridged/proto/aibridged.pb.go \
coderd/aibridged/proto/aibridged.pb.go \
site/src/api/typesGenerated.ts \
site/e2e/provisionerGenerated.ts \
site/src/api/chatModelOptionsGenerated.json \
Expand Down Expand Up @@ -956,8 +956,8 @@ TAILNETTEST_MOCKS := \
tailnet/tailnettest/subscriptionmock.go

AIBRIDGED_MOCKS := \
enterprise/aibridged/aibridgedmock/clientmock.go \
enterprise/aibridged/aibridgedmock/poolmock.go
coderd/aibridged/aibridgedmock/clientmock.go \
coderd/aibridged/aibridgedmock/poolmock.go

GEN_FILES := \
tailnet/proto/tailnet.pb.go \
Expand All @@ -967,7 +967,7 @@ GEN_FILES := \
provisionersdk/proto/provisioner.pb.go \
provisionerd/proto/provisionerd.pb.go \
vpn/vpn.pb.go \
enterprise/aibridged/proto/aibridged.pb.go \
coderd/aibridged/proto/aibridged.pb.go \
$(DB_GEN_FILES) \
$(SITE_GEN_FILES) \
coderd/rbac/object_gen.go \
Expand Down Expand Up @@ -1032,7 +1032,7 @@ gen/mark-fresh:
agent/agentsocket/proto/agentsocket.pb.go \
agent/boundarylogproxy/codec/boundary.pb.go \
vpn/vpn.pb.go \
enterprise/aibridged/proto/aibridged.pb.go \
coderd/aibridged/proto/aibridged.pb.go \
coderd/database/dump.sql \
coderd/database/querier.go \
coderd/database/unique_constraint.go \
Expand Down Expand Up @@ -1121,8 +1121,8 @@ codersdk/workspacesdk/agentconnmock/agentconnmock.go: codersdk/workspacesdk/agen
./scripts/format_go_file.sh "$@"
touch "$@"

$(AIBRIDGED_MOCKS): enterprise/aibridged/client.go enterprise/aibridged/pool.go
go generate ./enterprise/aibridged/aibridgedmock/
$(AIBRIDGED_MOCKS): coderd/aibridged/client.go coderd/aibridged/pool.go
go generate ./coderd/aibridged/aibridgedmock/
touch "$@"

agent/agentcontainers/dcspec/dcspec_gen.go: \
Expand Down Expand Up @@ -1189,13 +1189,13 @@ agent/boundarylogproxy/codec/boundary.pb.go: agent/boundarylogproxy/codec/bounda
--go_opt=paths=source_relative \
./agent/boundarylogproxy/codec/boundary.proto

enterprise/aibridged/proto/aibridged.pb.go: enterprise/aibridged/proto/aibridged.proto
coderd/aibridged/proto/aibridged.pb.go: coderd/aibridged/proto/aibridged.proto
./scripts/atomic_protoc.sh \
--go_out=. \
--go_opt=paths=source_relative \
--go-drpc_out=. \
--go-drpc_opt=paths=source_relative \
./enterprise/aibridged/proto/aibridged.proto
./coderd/aibridged/proto/aibridged.proto

site/src/api/typesGenerated.ts: site/node_modules/.installed $(wildcard scripts/apitypings/*) \
$(shell find ./codersdk $(FIND_EXCLUSIONS) -type f -name '*.go') \
Expand Down
8 changes: 4 additions & 4 deletions enterprise/cli/aibridged.go → cli/aibridged.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,10 @@ import (
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/config"
"github.com/coder/coder/v2/aibridge/keypool"
"github.com/coder/coder/v2/coderd"
"github.com/coder/coder/v2/coderd/aibridged"
"github.com/coder/coder/v2/coderd/tracing"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/enterprise/aibridged"
"github.com/coder/coder/v2/enterprise/coderd"
"github.com/coder/quartz"
)

Expand Down Expand Up @@ -44,13 +44,13 @@ func newAIBridgeDaemon(coderAPI *coderd.API, providers []aibridge.Provider) (*ai
return srv, nil
}

// buildProviders constructs the list of AI providers from config.
// BuildProviders constructs the list of AI providers from config.
// It merges legacy single-provider env vars and indexed provider configs:
// 1. Legacy providers (from CODER_AI_GATEWAY_OPENAI_KEY, etc.) are added first.
// If a legacy name conflicts with an indexed provider, startup fails with
// a clear error asking the admin to remove one or the other.
// 2. Indexed providers (from CODER_AI_GATEWAY_PROVIDER_<N>_*) are added next.
func buildProviders(cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) {
func BuildProviders(cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) {
var cbConfig *config.CircuitBreaker
if cfg.CircuitBreakerEnabled.Value() {
cbConfig = &config.CircuitBreaker{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ func TestBuildProviders(t *testing.T) {

t.Run("EmptyConfig", func(t *testing.T) {
t.Parallel()
providers, err := buildProviders(codersdk.AIBridgeConfig{})
providers, err := BuildProviders(codersdk.AIBridgeConfig{})
require.NoError(t, err)
assert.Empty(t, providers)
})
Expand All @@ -30,7 +30,7 @@ func TestBuildProviders(t *testing.T) {
cfg.LegacyOpenAI.Key = serpent.String("sk-openai")
cfg.LegacyAnthropic.Key = serpent.String("sk-anthropic")

providers, err := buildProviders(cfg)
providers, err := BuildProviders(cfg)
require.NoError(t, err)

names := providerNames(providers)
Expand Down Expand Up @@ -59,7 +59,7 @@ func TestBuildProviders(t *testing.T) {
},
}

providers, err := buildProviders(cfg)
providers, err := BuildProviders(cfg)
require.NoError(t, err)

names := providerNames(providers)
Expand All @@ -77,7 +77,7 @@ func TestBuildProviders(t *testing.T) {
}
cfg.LegacyOpenAI.Key = serpent.String("sk-legacy")

_, err := buildProviders(cfg)
_, err := BuildProviders(cfg)
require.Error(t, err)
assert.Contains(t, err.Error(), "conflicts with indexed provider")
})
Expand All @@ -91,7 +91,7 @@ func TestBuildProviders(t *testing.T) {
}
cfg.LegacyAnthropic.Key = serpent.String("sk-legacy")

_, err := buildProviders(cfg)
_, err := BuildProviders(cfg)
require.Error(t, err)
assert.Contains(t, err.Error(), "conflicts with indexed provider")
})
Expand All @@ -106,7 +106,7 @@ func TestBuildProviders(t *testing.T) {
cfg.LegacyOpenAI.Key = serpent.String("sk-openai")
cfg.LegacyAnthropic.Key = serpent.String("sk-anthropic")

providers, err := buildProviders(cfg)
providers, err := BuildProviders(cfg)
require.NoError(t, err)

names := providerNames(providers)
Expand All @@ -123,7 +123,7 @@ func TestBuildProviders(t *testing.T) {
cfg.LegacyBedrock.AccessKey = serpent.String("AKID")
cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret")

providers, err := buildProviders(cfg)
providers, err := BuildProviders(cfg)
require.NoError(t, err)

names := providerNames(providers)
Expand All @@ -139,7 +139,7 @@ func TestBuildProviders(t *testing.T) {
cfg.LegacyBedrock.AccessKey = serpent.String("AKID")
cfg.LegacyBedrock.AccessKeySecret = serpent.String("secret")

providers, err := buildProviders(cfg)
providers, err := BuildProviders(cfg)
require.NoError(t, err)
require.Len(t, providers, 1)

Expand All @@ -156,7 +156,7 @@ func TestBuildProviders(t *testing.T) {
},
}

_, err := buildProviders(cfg)
_, err := BuildProviders(cfg)
require.Error(t, err)
assert.Contains(t, err.Error(), "unknown provider type")
})
Expand All @@ -173,7 +173,7 @@ func TestBuildProviders(t *testing.T) {
},
}

providers, err := buildProviders(cfg)
providers, err := BuildProviders(cfg)
require.NoError(t, err)
require.Len(t, providers, 3)

Expand All @@ -195,7 +195,7 @@ func TestBuildProviders(t *testing.T) {
},
}

providers, err := buildProviders(cfg)
providers, err := BuildProviders(cfg)
require.NoError(t, err)
require.Len(t, providers, 1)

Expand All @@ -211,73 +211,3 @@ func providerNames(providers []aibridge.Provider) []string {
}
return names
}

func TestDomainsFromProviders(t *testing.T) {
t.Parallel()

t.Run("ExtractsHostnames", func(t *testing.T) {
t.Parallel()

providers, err := buildProviders(codersdk.AIBridgeConfig{
Providers: []codersdk.AIProviderConfig{
{Type: aibridge.ProviderOpenAI, Name: "openai", Keys: []string{"k"}},
{Type: aibridge.ProviderAnthropic, Name: "anthropic", Keys: []string{"k"}},
{Type: aibridge.ProviderOpenAI, Name: "custom", Keys: []string{"k"}, BaseURL: "https://custom-llm.example.com:8443/api"},
},
})
require.NoError(t, err)

domains, mapping := domainsFromProviders(providers)

assert.Contains(t, domains, "api.openai.com")
assert.Contains(t, domains, "api.anthropic.com")
assert.Contains(t, domains, "custom-llm.example.com")

assert.Equal(t, "openai", mapping("api.openai.com"))
assert.Equal(t, "anthropic", mapping("api.anthropic.com"))
assert.Equal(t, "custom", mapping("custom-llm.example.com"))
assert.Empty(t, mapping("unknown.com"))
})

t.Run("DeduplicatesSameHost", func(t *testing.T) {
t.Parallel()

providers, err := buildProviders(codersdk.AIBridgeConfig{
Providers: []codersdk.AIProviderConfig{
{Type: aibridge.ProviderOpenAI, Name: "first", Keys: []string{"k"}, BaseURL: "https://api.example.com/v1"},
{Type: aibridge.ProviderOpenAI, Name: "second", Keys: []string{"k"}, BaseURL: "https://api.example.com/v2"},
},
})
require.NoError(t, err)

domains, mapping := domainsFromProviders(providers)

// Count occurrences of api.example.com.
count := 0
for _, d := range domains {
if d == "api.example.com" {
count++
}
}
assert.Equal(t, 1, count)
// First provider wins.
assert.Equal(t, "first", mapping("api.example.com"))
})

t.Run("CaseInsensitive", func(t *testing.T) {
t.Parallel()

providers, err := buildProviders(codersdk.AIBridgeConfig{
Providers: []codersdk.AIProviderConfig{
{Type: aibridge.ProviderOpenAI, Name: "provider", Keys: []string{"k"}, BaseURL: "https://API.Example.COM/v1"},
},
})
require.NoError(t, err)

domains, mapping := domainsFromProviders(providers)

assert.Contains(t, domains, "api.example.com")
assert.Equal(t, "provider", mapping("API.Example.COM"))
assert.Equal(t, "provider", mapping("api.example.com"))
})
}
23 changes: 23 additions & 0 deletions cli/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -1026,6 +1026,29 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
return xerrors.Errorf("seed ai providers from env: %w", err)
}

// In-memory aibridge daemon. Registered on coderd so chatd can
// dispatch LLM requests via the in-process transport without
// crossing the gated /api/v2/aibridge HTTP route. The HTTP route
// itself is registered (and license-gated) only by enterprise/coderd;
// in AGPL builds it does not exist at all. The daemon starts here
// unconditionally when the bridge feature is enabled by config so
// chatd can use it regardless of license entitlement.
if vals.AI.BridgeConfig.Enabled.Value() {
providers, err := BuildProviders(vals.AI.BridgeConfig)
if err != nil {
return xerrors.Errorf("build AI providers: %w", err)
}
aibridgeDaemon, err := newAIBridgeDaemon(coderAPI, providers)
if err != nil {
return xerrors.Errorf("create aibridged: %w", err)
}
coderAPI.RegisterInMemoryAIBridgedHTTPHandler(aibridgeDaemon)
// The handler is bound to coderAPI's lifecycle; Close() on the
// daemon does not affect in-flight requests but is needed to
// release pool/recorder resources at shutdown.
defer aibridgeDaemon.Close()
}

if vals.Prometheus.Enable {
// Agent metrics require reference to the tailnet coordinator, so must be initiated after Coder API.
closeAgentsFunc, err := prometheusmetrics.Agents(ctx, logger, options.PrometheusRegistry, coderAPI.Database, &coderAPI.TailnetCoordinator, coderAPI.DERPMap, coderAPI.Options.AgentInactiveDisconnectTimeout, 0)
Expand Down
25 changes: 17 additions & 8 deletions enterprise/coderd/aibridged.go → coderd/aibridged.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,22 @@ import (
"storj.io/drpc/drpcserver"

"cdr.dev/slog/v3"
"github.com/coder/coder/v2/coderd/aibridged"
aibridgedproto "github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/coderd/aibridgedserver"
"github.com/coder/coder/v2/coderd/tracing"
"github.com/coder/coder/v2/codersdk/drpcsdk"
"github.com/coder/coder/v2/enterprise/aibridged"
aibridgedproto "github.com/coder/coder/v2/enterprise/aibridged/proto"
"github.com/coder/coder/v2/enterprise/aibridgedserver"
)

// GetAIBridgedHandler returns the in-memory aibridge HTTP handler set by
// [API.RegisterInMemoryAIBridgedHTTPHandler], or nil if the daemon has not
// been wired in. Used by the enterprise /api/v2/aibridge route (license-gated)
// to forward requests into the same in-memory handler that chatd dispatches
// to in-process.
func (api *API) GetAIBridgedHandler() http.Handler {
return api.aibridgedHandler
}

// RegisterInMemoryAIBridgedHTTPHandler mounts [aibridged.Server]'s HTTP router onto
// [API]'s router, so that requests to aibridged will be relayed from Coder's API server
// to the in-memory aibridged.
Expand Down Expand Up @@ -48,7 +57,7 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai

mux := drpcmux.New()
srv, err := aibridgedserver.NewServer(api.ctx, api.Database, api.Logger.Named("aibridgedserver"),
api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.AGPL.Experiments, api.aiSeatTracker)
api.AccessURL.String(), api.DeploymentValues.AI.BridgeConfig, api.ExternalAuthConfigs, api.Experiments, api.AISeatTracker)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -78,11 +87,11 @@ func (api *API) CreateInMemoryAIBridgeServer(dialCtx context.Context) (client ai
// in-mem pipes aren't technically "websockets" but they have the same properties as far as the
// API is concerned: they are long-lived connections that we need to close before completing
// shutdown of the API.
api.AGPL.WebsocketWaitMutex.Lock()
api.AGPL.WebsocketWaitGroup.Add(1)
api.AGPL.WebsocketWaitMutex.Unlock()
api.WebsocketWaitMutex.Lock()
api.WebsocketWaitGroup.Add(1)
api.WebsocketWaitMutex.Unlock()
go func() {
defer api.AGPL.WebsocketWaitGroup.Done()
defer api.WebsocketWaitGroup.Done()
// Here we pass the background context, since we want the server to keep serving until the
// client hangs up. The aibridged is local, in-mem, so there isn't a danger of losing contact with it and
// having a dead connection we don't know the status of.
Expand Down
File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,10 @@ import (
"github.com/coder/coder/v2/aibridge"
"github.com/coder/coder/v2/aibridge/intercept"
agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
"github.com/coder/coder/v2/coderd/aibridged"
mock "github.com/coder/coder/v2/coderd/aibridged/aibridgedmock"
"github.com/coder/coder/v2/coderd/aibridged/proto"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/enterprise/aibridged"
mock "github.com/coder/coder/v2/enterprise/aibridged/aibridgedmock"
"github.com/coder/coder/v2/enterprise/aibridged/proto"
"github.com/coder/coder/v2/testutil"
)

Expand Down

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading