opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add internal subagent model override wiring (#24399)

Browse Source 
> Mux working on behalf of Mike.

## Summary
- add an enabled chat model config lookup by ID for internal callers
- keep `spawn_agent` unchanged while threading an internal model
override through child subagent chat creation
- extend chatd coverage for inherited bindings, plan mode, and internal
override behavior

## Validation
- `go test ./coderd/x/chatd ./coderd/database/dbauthz`
- `make lint`
This commit is contained in:
Michael Suchacz
2026-04-16 17:08:02 +02:00
committed by GitHub
parent eae9444dbe
commit 1092093e98
9 changed files with 256 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+7
View File
@@ -3026,6 +3026,13 @@ func (q *querier) GetEligibleProvisionerDaemonsByProvisionerJobIDs(ctx context.C
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetEligibleProvisionerDaemonsByProvisionerJobIDs)(ctx, provisionerJobIDs)
}
func (q *querier) GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return database.ChatModelConfig{}, err
}
return q.db.GetEnabledChatModelConfigByID(ctx, id)
}
func (q *querier) GetEnabledChatModelConfigs(ctx context.Context) ([]database.ChatModelConfig, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return nil, err
coderd/database/dbauthz/dbauthz_test.go
+5
View File
@@ -846,6 +846,11 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().GetChatWorkspaceTTL(gomock.Any()).Return("1h", nil).AnyTimes()
check.Args().Asserts()
}))
s.Run("GetEnabledChatModelConfigByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
config := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
dbm.EXPECT().GetEnabledChatModelConfigByID(gomock.Any(), config.ID).Return(config, nil).AnyTimes()
check.Args(config.ID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(config)
}))
s.Run("GetEnabledChatModelConfigs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
configA := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
configB := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
coderd/database/dbmetrics/querymetrics.go
+8
View File
@@ -1552,6 +1552,14 @@ func (m queryMetricsStore) GetEligibleProvisionerDaemonsByProvisionerJobIDs(ctx
return r0, r1
}
func (m queryMetricsStore) GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
start := time.Now()
r0, r1 := m.s.GetEnabledChatModelConfigByID(ctx, id)
m.queryLatencies.WithLabelValues("GetEnabledChatModelConfigByID").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetEnabledChatModelConfigByID").Inc()
return r0, r1
}
func (m queryMetricsStore) GetEnabledChatModelConfigs(ctx context.Context) ([]database.ChatModelConfig, error) {
start := time.Now()
r0, r1 := m.s.GetEnabledChatModelConfigs(ctx)
coderd/database/dbmock/dbmock.go
+15
View File
@@ -2867,6 +2867,21 @@ func (mr *MockStoreMockRecorder) GetEligibleProvisionerDaemonsByProvisionerJobID
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEligibleProvisionerDaemonsByProvisionerJobIDs", reflect.TypeOf((*MockStore)(nil).GetEligibleProvisionerDaemonsByProvisionerJobIDs), ctx, provisionerJobIds)
}
// GetEnabledChatModelConfigByID mocks base method.
func (m *MockStore) GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetEnabledChatModelConfigByID", ctx, id)
ret0, _ := ret[0].(database.ChatModelConfig)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetEnabledChatModelConfigByID indicates an expected call of GetEnabledChatModelConfigByID.
func (mr *MockStoreMockRecorder) GetEnabledChatModelConfigByID(ctx, id any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnabledChatModelConfigByID", reflect.TypeOf((*MockStore)(nil).GetEnabledChatModelConfigByID), ctx, id)
}
// GetEnabledChatModelConfigs mocks base method.
func (m *MockStore) GetEnabledChatModelConfigs(ctx context.Context) ([]database.ChatModelConfig, error) {
m.ctrl.T.Helper()
coderd/database/querier.go
+3
View File
@@ -352,6 +352,9 @@ type sqlcQuerier interface {
GetDeploymentWorkspaceAgentUsageStats(ctx context.Context, createdAt time.Time) (GetDeploymentWorkspaceAgentUsageStatsRow, error)
GetDeploymentWorkspaceStats(ctx context.Context) (GetDeploymentWorkspaceStatsRow, error)
GetEligibleProvisionerDaemonsByProvisionerJobIDs(ctx context.Context, provisionerJobIds []uuid.UUID) ([]GetEligibleProvisionerDaemonsByProvisionerJobIDsRow, error)
// Providers can be disabled independently of their model configs.
// Check both to ensure the selected config is actually usable.
GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error)
GetEnabledChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error)
GetEnabledChatProviders(ctx context.Context) ([]ChatProvider, error)
GetEnabledMCPServerConfigs(ctx context.Context) ([]MCPServerConfig, error)
coderd/database/queries.sql.go
+39
View File
@@ -4274,6 +4274,45 @@ func (q *sqlQuerier) GetDefaultChatModelConfig(ctx context.Context) (ChatModelCo
return i, err
}
const getEnabledChatModelConfigByID = `-- name: GetEnabledChatModelConfigByID :one
SELECT
cmc.id, cmc.provider, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options
FROM
chat_model_configs cmc
JOIN
chat_providers cp ON cp.provider = cmc.provider
WHERE
cmc.id = $1::uuid
AND cmc.deleted = FALSE
AND cmc.enabled = TRUE
AND cp.enabled = TRUE
`
// Providers can be disabled independently of their model configs.
// Check both to ensure the selected config is actually usable.
func (q *sqlQuerier) GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error) {
row := q.db.QueryRowContext(ctx, getEnabledChatModelConfigByID, id)
var i ChatModelConfig
err := row.Scan(
&i.ID,
&i.Provider,
&i.Model,
&i.DisplayName,
&i.CreatedBy,
&i.UpdatedBy,
&i.Enabled,
&i.IsDefault,
&i.Deleted,
&i.DeletedAt,
&i.CreatedAt,
&i.UpdatedAt,
&i.ContextLimit,
&i.CompressionThreshold,
&i.Options,
)
return i, err
}
const getEnabledChatModelConfigs = `-- name: GetEnabledChatModelConfigs :many
SELECT
cmc.id, cmc.provider, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options
coderd/database/queries/chatmodelconfigs.sql
+15
View File
@@ -46,6 +46,21 @@ ORDER BY
cmc.updated_at DESC,
cmc.id DESC;
-- name: GetEnabledChatModelConfigByID :one
SELECT
cmc.*
FROM
chat_model_configs cmc
-- Providers can be disabled independently of their model configs.
-- Check both to ensure the selected config is actually usable.
JOIN
chat_providers cp ON cp.provider = cmc.provider
WHERE
cmc.id = @id::uuid
AND cmc.deleted = FALSE
AND cmc.enabled = TRUE
AND cp.enabled = TRUE;
-- name: InsertChatModelConfig :one
INSERT INTO chat_model_configs (
provider,
coderd/x/chatd/subagent.go
+15 -9
View File
@@ -411,8 +411,9 @@ func parseSubagentToolChatID(raw string) (uuid.UUID, error) {
}
type childSubagentChatOptions struct {
chatMode database.NullChatMode
systemPrompt string
chatMode database.NullChatMode
systemPrompt string
modelConfigIDOverride *uuid.UUID
}
func (p *Server) createChildSubagentChat(
@@ -449,8 +450,13 @@ func (p *Server) createChildSubagentChatWithOptions(
if parent.RootChatID.Valid {
rootChatID = parent.RootChatID.UUID
}
if parent.LastModelConfigID == uuid.Nil {
return database.Chat{}, xerrors.New("parent chat model config id is required")
modelConfigID := parent.LastModelConfigID
if opts.modelConfigIDOverride != nil {
modelConfigID = *opts.modelConfigIDOverride
}
if modelConfigID == uuid.Nil {
return database.Chat{}, xerrors.New("model config is required")
}
mcpServerIDs := parent.MCPServerIDs
@@ -482,7 +488,7 @@ func (p *Server) createChildSubagentChatWithOptions(
AgentID: parent.AgentID,
ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true},
RootChatID: uuid.NullUUID{UUID: rootChatID, Valid: true},
LastModelConfigID: parent.LastModelConfigID,
LastModelConfigID: modelConfigID,
Title: title,
Mode: opts.chatMode,
PlanMode: parent.PlanMode,
@@ -528,7 +534,7 @@ func (p *Server) createChildSubagentChatWithOptions(
database.ChatMessageRoleSystem,
deploymentContent,
database.ChatMessageVisibilityModel,
parent.LastModelConfigID,
modelConfigID,
chatprompt.CurrentContentVersion,
))
}
@@ -543,7 +549,7 @@ func (p *Server) createChildSubagentChatWithOptions(
database.ChatMessageRoleSystem,
childSystemPromptContent,
database.ChatMessageVisibilityModel,
parent.LastModelConfigID,
modelConfigID,
chatprompt.CurrentContentVersion,
))
}
@@ -551,7 +557,7 @@ func (p *Server) createChildSubagentChatWithOptions(
database.ChatMessageRoleSystem,
workspaceAwarenessContent,
database.ChatMessageVisibilityModel,
parent.LastModelConfigID,
modelConfigID,
chatprompt.CurrentContentVersion,
))
if _, err := tx.InsertChatMessages(ctx, systemParams); err != nil {
@@ -578,7 +584,7 @@ func (p *Server) createChildSubagentChatWithOptions(
database.ChatMessageRoleUser,
userContent,
database.ChatMessageVisibilityBoth,
parent.LastModelConfigID,
modelConfigID,
chatprompt.CurrentContentVersion,
).withCreatedBy(parent.OwnerID))
if _, err := tx.InsertChatMessages(ctx, userParams); err != nil {
coderd/x/chatd/subagent_internal_test.go
+149
View File
@@ -156,6 +156,33 @@ func seedInternalChatDeps(
return user, org, model
}
func insertInternalChatModelConfig(
ctx context.Context,
t *testing.T,
db database.Store,
userID uuid.UUID,
model string,
enabled bool,
) database.ChatModelConfig {
t.Helper()
modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
Provider: "openai",
Model: model,
DisplayName: model,
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true},
Enabled: enabled,
IsDefault: false,
ContextLimit: 128000,
CompressionThreshold: 70,
Options: json.RawMessage(`{}`),
})
require.NoError(t, err)
return modelConfig
}
func seedWorkspaceBinding(
t *testing.T,
db database.Store,
@@ -256,6 +283,74 @@ func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) {
require.Equal(t, parentChat.AgentID, childChat.AgentID)
}
func createInternalParentChat(
ctx context.Context,
t *testing.T,
server *Server,
db database.Store,
orgID uuid.UUID,
userID uuid.UUID,
modelConfigID uuid.UUID,
title string,
) database.Chat {
t.Helper()
parent, err := server.CreateChat(ctx, CreateOptions{
OrganizationID: orgID,
OwnerID: userID,
Title: title,
ModelConfigID: modelConfigID,
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
})
require.NoError(t, err)
parentChat, err := db.GetChatByID(ctx, parent.ID)
require.NoError(t, err)
return parentChat
}
func runSpawnAgentTool(
ctx context.Context,
t *testing.T,
server *Server,
parentChat database.Chat,
args spawnAgentArgs,
) fantasy.ToolResponse {
t.Helper()
tools := server.subagentTools(ctx, func() database.Chat { return parentChat })
tool := findToolByName(tools, "spawn_agent")
require.NotNil(t, tool, "spawn_agent tool must be present")
input, err := json.Marshal(args)
require.NoError(t, err)
resp, err := tool.Run(ctx, fantasy.ToolCall{
ID: uuid.NewString(),
Name: "spawn_agent",
Input: string(input),
})
require.NoError(t, err)
return resp
}
func requireSpawnAgentChildChatID(t *testing.T, resp fantasy.ToolResponse) uuid.UUID {
t.Helper()
require.False(t, resp.IsError, "expected success but got: %s", resp.Content)
var result struct {
ChatID string `json:"chat_id"`
}
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
require.NotEmpty(t, result.ChatID, "response must contain chat_id")
childID, err := uuid.Parse(result.ChatID)
require.NoError(t, err)
return childID
}
func TestCreateChildSubagentChatCopiesPlanMode(t *testing.T) {
t.Parallel()
@@ -293,6 +388,60 @@ func TestCreateChildSubagentChatCopiesPlanMode(t *testing.T) {
require.Equal(t, planMode, childChat.PlanMode)
}
func TestSpawnAgent_InheritsParentModelWhenOmitted(t *testing.T) {
t.Parallel()
db, ps := dbtestutil.NewDB(t)
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
ctx := chatdTestContext(t)
user, org, model := seedInternalChatDeps(ctx, t, db)
parentChat := createInternalParentChat(
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-inherited-model",
)
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
Prompt: "delegate work",
})
childID := requireSpawnAgentChildChatID(t, resp)
childChat, err := db.GetChatByID(ctx, childID)
require.NoError(t, err)
require.Equal(t, parentChat.LastModelConfigID, childChat.LastModelConfigID)
}
func TestCreateChildSubagentChat_OverrideWorksWhenParentHasNoModel(t *testing.T) {
t.Parallel()
db, ps := dbtestutil.NewDB(t)
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
ctx := chatdTestContext(t)
user, org, model := seedInternalChatDeps(ctx, t, db)
overrideModel := insertInternalChatModelConfig(
ctx, t, db, user.ID, "override-no-parent-model-"+uuid.NewString(), true,
)
parentChat := createInternalParentChat(
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-no-model",
)
// The chats table enforces a foreign key for last_model_config_id, so
// use a synthetic parent value here to exercise the override path.
parentChat.LastModelConfigID = uuid.Nil
child, err := server.createChildSubagentChatWithOptions(
ctx,
parentChat,
"delegate work",
"",
childSubagentChatOptions{modelConfigIDOverride: &overrideModel.ID},
)
require.NoError(t, err)
childChat, err := db.GetChatByID(ctx, child.ID)
require.NoError(t, err)
require.Equal(t, overrideModel.ID, childChat.LastModelConfigID)
}
func TestSpawnComputerUseAgent_NoAnthropicProvider(t *testing.T) {
t.Parallel()