From 1092093e98912b798f5c5115a46c77216f4ffc22 Mon Sep 17 00:00:00 2001 From: Michael Suchacz <203725896+ibetitsmike@users.noreply.github.com> Date: Thu, 16 Apr 2026 17:08:02 +0200 Subject: [PATCH] feat: add internal subagent model override wiring (#24399) > 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` --- coderd/database/dbauthz/dbauthz.go | 7 + coderd/database/dbauthz/dbauthz_test.go | 5 + coderd/database/dbmetrics/querymetrics.go | 8 + coderd/database/dbmock/dbmock.go | 15 ++ coderd/database/querier.go | 3 + coderd/database/queries.sql.go | 39 +++++ coderd/database/queries/chatmodelconfigs.sql | 15 ++ coderd/x/chatd/subagent.go | 24 +-- coderd/x/chatd/subagent_internal_test.go | 149 +++++++++++++++++++ 9 files changed, 256 insertions(+), 9 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 8d9a95d909..31ab74c85e 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -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 diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 1ea0ff5b94..032ca903ce 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -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{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 2a1ffac7ad..9e9be6c736 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -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) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index ace93212e0..54c877387b 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -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() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 7c508dce71..a50c13d68b 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -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) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 89c180c957..30dfd05714 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -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 diff --git a/coderd/database/queries/chatmodelconfigs.sql b/coderd/database/queries/chatmodelconfigs.sql index ec719760ad..1463b96cf3 100644 --- a/coderd/database/queries/chatmodelconfigs.sql +++ b/coderd/database/queries/chatmodelconfigs.sql @@ -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, diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index 8aa541b5a7..73ecec9bfd 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -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 { diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 06ddce0226..17f13eb236 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -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()