From 0f55c283f1f409bb3297136788f616c4d56f3a79 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Tue, 14 Jul 2026 13:50:25 +0100 Subject: [PATCH] fix: use backend-selected chat agent for desktop, git, terminal (#26959) --- coderd/database/dbauthz/dbauthz.go | 10 ++ coderd/database/dbauthz/dbauthz_test.go | 8 ++ coderd/database/dbmetrics/querymetrics.go | 8 ++ coderd/database/dbmock/dbmock.go | 15 +++ coderd/database/querier.go | 1 + coderd/database/queries.sql.go | 91 +++++++++++++++++++ coderd/database/queries/workspaceagents.sql | 26 ++++++ coderd/exp_chats.go | 80 ++++++++++++++-- coderd/exp_chats_internal_test.go | 63 +++++++++++++ coderd/workspaceagents_internal_test.go | 31 ++++--- .../{internal => }/agentselect/agentselect.go | 0 .../agentselect/agentselect_test.go | 2 +- coderd/x/chatd/chatd.go | 2 +- coderd/x/chatd/chattool/createworkspace.go | 2 +- coderd/x/chatd/chattool/startworkspace.go | 2 +- codersdk/chats.go | 4 +- site/src/api/typesGenerated.ts | 6 +- site/src/pages/AgentsPage/AgentChatPage.tsx | 72 +++++++++------ .../ChatConversation/chatHelpers.test.ts | 12 +-- .../ChatConversation/chatHelpers.ts | 13 +-- 20 files changed, 373 insertions(+), 75 deletions(-) rename coderd/x/chatd/{internal => }/agentselect/agentselect.go (100%) rename coderd/x/chatd/{internal => }/agentselect/agentselect_test.go (98%) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 4d23b86a09..183ae39d0e 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5436,6 +5436,16 @@ func (q *querier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Conte return q.db.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, workspace.ID) } +func (q *querier) GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx context.Context, workspaceIDs []uuid.UUID) ([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow, error) { + for _, workspaceID := range workspaceIDs { + if _, err := q.GetWorkspaceByID(ctx, workspaceID); err != nil { + return nil, err + } + } + + return q.db.GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx, workspaceIDs) +} + func (q *querier) GetWorkspaceAppByAgentIDAndSlug(ctx context.Context, arg database.GetWorkspaceAppByAgentIDAndSlugParams) (database.WorkspaceApp, error) { // If we can fetch the workspace, we can fetch the apps. Use the authorized call. if _, err := q.GetWorkspaceByAgentID(ctx, arg.AgentID); err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index d8fa56e6d7..dedf751e4d 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -3867,6 +3867,14 @@ func (s *MethodTestSuite) TestWorkspace() { dbm.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), ws.ID).Return([]database.WorkspaceAgent{agt}, nil).AnyTimes() check.Args(ws.ID).Asserts(ws, policy.ActionRead).Returns([]database.WorkspaceAgent{agt}) })) + s.Run("GetWorkspaceAgentsInLatestBuildByWorkspaceIDs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + ws := testutil.Fake(s.T(), faker, database.Workspace{}) + unauthorizedWorkspaceID := uuid.New() + ids := []uuid.UUID{ws.ID, unauthorizedWorkspaceID} + dbm.EXPECT().GetWorkspaceByID(gomock.Any(), ws.ID).Return(ws, nil).AnyTimes() + dbm.EXPECT().GetWorkspaceByID(gomock.Any(), unauthorizedWorkspaceID).Return(database.Workspace{}, sql.ErrNoRows).AnyTimes() + check.Args(ids).Asserts(ws, policy.ActionRead).Errors(xerrors.Errorf("fetch object: %w", sql.ErrNoRows)) + })) s.Run("GetWorkspaceByOwnerIDAndName", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { ws := testutil.Fake(s.T(), faker, database.Workspace{}) arg := database.GetWorkspaceByOwnerIDAndNameParams{ diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index b8b5397a7e..55f902ab76 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3577,6 +3577,14 @@ func (m queryMetricsStore) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx cont return r0, r1 } +func (m queryMetricsStore) GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx context.Context, workspaceIds []uuid.UUID) ([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow, error) { + start := time.Now() + r0, r1 := m.s.GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx, workspaceIds) + m.queryLatencies.WithLabelValues("GetWorkspaceAgentsInLatestBuildByWorkspaceIDs").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetWorkspaceAgentsInLatestBuildByWorkspaceIDs").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetWorkspaceAppByAgentIDAndSlug(ctx context.Context, arg database.GetWorkspaceAppByAgentIDAndSlugParams) (database.WorkspaceApp, error) { start := time.Now() r0, r1 := m.s.GetWorkspaceAppByAgentIDAndSlug(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index f56b257f3a..e684ba4426 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -6688,6 +6688,21 @@ func (mr *MockStoreMockRecorder) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ct return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentsInLatestBuildByWorkspaceID", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentsInLatestBuildByWorkspaceID), ctx, workspaceID) } +// GetWorkspaceAgentsInLatestBuildByWorkspaceIDs mocks base method. +func (m *MockStore) GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx context.Context, workspaceIds []uuid.UUID) ([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetWorkspaceAgentsInLatestBuildByWorkspaceIDs", ctx, workspaceIds) + ret0, _ := ret[0].([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetWorkspaceAgentsInLatestBuildByWorkspaceIDs indicates an expected call of GetWorkspaceAgentsInLatestBuildByWorkspaceIDs. +func (mr *MockStoreMockRecorder) GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx, workspaceIds any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWorkspaceAgentsInLatestBuildByWorkspaceIDs", reflect.TypeOf((*MockStore)(nil).GetWorkspaceAgentsInLatestBuildByWorkspaceIDs), ctx, workspaceIds) +} + // GetWorkspaceAppByAgentIDAndSlug mocks base method. func (m *MockStore) GetWorkspaceAppByAgentIDAndSlug(ctx context.Context, arg database.GetWorkspaceAppByAgentIDAndSlugParams) (database.WorkspaceApp, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 5dd4a96d91..5d09700112 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -930,6 +930,7 @@ type sqlcQuerier interface { GetWorkspaceAgentsCreatedAfter(ctx context.Context, createdAt time.Time) ([]WorkspaceAgent, error) GetWorkspaceAgentsForMetrics(ctx context.Context) ([]GetWorkspaceAgentsForMetricsRow, error) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) ([]WorkspaceAgent, error) + GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx context.Context, workspaceIds []uuid.UUID) ([]GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow, error) GetWorkspaceAppByAgentIDAndSlug(ctx context.Context, arg GetWorkspaceAppByAgentIDAndSlugParams) (WorkspaceApp, error) GetWorkspaceAppStatusesByAppIDs(ctx context.Context, ids []uuid.UUID) ([]WorkspaceAppStatus, error) GetWorkspaceAppsByAgentID(ctx context.Context, agentID uuid.UUID) ([]WorkspaceApp, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index e8e49861f8..f27075f7f5 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -33054,6 +33054,97 @@ func (q *sqlQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx context.Co return items, nil } +const getWorkspaceAgentsInLatestBuildByWorkspaceIDs = `-- name: GetWorkspaceAgentsInLatestBuildByWorkspaceIDs :many +SELECT + workspace_builds.workspace_id, + workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted +FROM + workspace_agents +JOIN + workspace_resources ON workspace_agents.resource_id = workspace_resources.id +JOIN + workspace_builds ON workspace_resources.job_id = workspace_builds.job_id +JOIN ( + SELECT + workspace_id, + MAX(build_number) AS build_number + FROM + workspace_builds + WHERE + workspace_id = ANY($1 :: uuid [ ]) + GROUP BY + workspace_id +) AS latest_builds ON + latest_builds.workspace_id = workspace_builds.workspace_id AND + latest_builds.build_number = workspace_builds.build_number +WHERE + workspace_agents.deleted = FALSE +` + +type GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow struct { + WorkspaceID uuid.UUID `db:"workspace_id" json:"workspace_id"` + WorkspaceAgent WorkspaceAgent `db:"workspace_agent" json:"workspace_agent"` +} + +func (q *sqlQuerier) GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx context.Context, workspaceIds []uuid.UUID) ([]GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow, error) { + rows, err := q.db.QueryContext(ctx, getWorkspaceAgentsInLatestBuildByWorkspaceIDs, pq.Array(workspaceIds)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow + for rows.Next() { + var i GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow + if err := rows.Scan( + &i.WorkspaceID, + &i.WorkspaceAgent.ID, + &i.WorkspaceAgent.CreatedAt, + &i.WorkspaceAgent.UpdatedAt, + &i.WorkspaceAgent.Name, + &i.WorkspaceAgent.FirstConnectedAt, + &i.WorkspaceAgent.LastConnectedAt, + &i.WorkspaceAgent.DisconnectedAt, + &i.WorkspaceAgent.ResourceID, + &i.WorkspaceAgent.AuthToken, + &i.WorkspaceAgent.AuthInstanceID, + &i.WorkspaceAgent.Architecture, + &i.WorkspaceAgent.EnvironmentVariables, + &i.WorkspaceAgent.OperatingSystem, + &i.WorkspaceAgent.InstanceMetadata, + &i.WorkspaceAgent.ResourceMetadata, + &i.WorkspaceAgent.Directory, + &i.WorkspaceAgent.Version, + &i.WorkspaceAgent.LastConnectedReplicaID, + &i.WorkspaceAgent.ConnectionTimeoutSeconds, + &i.WorkspaceAgent.TroubleshootingURL, + &i.WorkspaceAgent.MOTDFile, + &i.WorkspaceAgent.LifecycleState, + &i.WorkspaceAgent.ExpandedDirectory, + &i.WorkspaceAgent.LogsLength, + &i.WorkspaceAgent.LogsOverflowed, + &i.WorkspaceAgent.StartedAt, + &i.WorkspaceAgent.ReadyAt, + pq.Array(&i.WorkspaceAgent.Subsystems), + pq.Array(&i.WorkspaceAgent.DisplayApps), + &i.WorkspaceAgent.APIVersion, + &i.WorkspaceAgent.DisplayOrder, + &i.WorkspaceAgent.ParentID, + &i.WorkspaceAgent.APIKeyScope, + &i.WorkspaceAgent.Deleted, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getWorkspaceBuildAgentsByInstanceID = `-- name: GetWorkspaceBuildAgentsByInstanceID :many SELECT workspace_agents.id, workspace_agents.created_at, workspace_agents.updated_at, workspace_agents.name, workspace_agents.first_connected_at, workspace_agents.last_connected_at, workspace_agents.disconnected_at, workspace_agents.resource_id, workspace_agents.auth_token, workspace_agents.auth_instance_id, workspace_agents.architecture, workspace_agents.environment_variables, workspace_agents.operating_system, workspace_agents.instance_metadata, workspace_agents.resource_metadata, workspace_agents.directory, workspace_agents.version, workspace_agents.last_connected_replica_id, workspace_agents.connection_timeout_seconds, workspace_agents.troubleshooting_url, workspace_agents.motd_file, workspace_agents.lifecycle_state, workspace_agents.expanded_directory, workspace_agents.logs_length, workspace_agents.logs_overflowed, workspace_agents.started_at, workspace_agents.ready_at, workspace_agents.subsystems, workspace_agents.display_apps, workspace_agents.api_version, workspace_agents.display_order, workspace_agents.parent_id, workspace_agents.api_key_scope, workspace_agents.deleted, diff --git a/coderd/database/queries/workspaceagents.sql b/coderd/database/queries/workspaceagents.sql index 83534eb4e2..e5280252da 100644 --- a/coderd/database/queries/workspaceagents.sql +++ b/coderd/database/queries/workspaceagents.sql @@ -337,6 +337,32 @@ WHERE -- Filter out deleted sub agents. AND workspace_agents.deleted = FALSE; +-- name: GetWorkspaceAgentsInLatestBuildByWorkspaceIDs :many +SELECT + workspace_builds.workspace_id, + sqlc.embed(workspace_agents) +FROM + workspace_agents +JOIN + workspace_resources ON workspace_agents.resource_id = workspace_resources.id +JOIN + workspace_builds ON workspace_resources.job_id = workspace_builds.job_id +JOIN ( + SELECT + workspace_id, + MAX(build_number) AS build_number + FROM + workspace_builds + WHERE + workspace_id = ANY(@workspace_ids :: uuid [ ]) + GROUP BY + workspace_id +) AS latest_builds ON + latest_builds.workspace_id = workspace_builds.workspace_id AND + latest_builds.build_number = workspace_builds.build_number +WHERE + workspace_agents.deleted = FALSE; + -- name: GetWorkspaceAgentsByWorkspaceAndBuildNumber :many SELECT workspace_agents.* diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 8a3e1f6cfc..79ead7042a 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -1,6 +1,7 @@ package coderd import ( + "cmp" "context" "database/sql" "encoding/json" @@ -48,6 +49,7 @@ import ( "github.com/coder/coder/v2/coderd/workspaceapps" "github.com/coder/coder/v2/coderd/wsbuilder" "github.com/coder/coder/v2/coderd/x/chatd" + "github.com/coder/coder/v2/coderd/x/chatd/agentselect" "github.com/coder/coder/v2/coderd/x/chatd/chaterror" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" @@ -474,7 +476,59 @@ func (api *API) listChats(rw http.ResponseWriter, r *http.Request) { return } - httpapi.Write(ctx, rw, http.StatusOK, db2sdk.ChatRowsWithChildren(chatRows, childRows, diffStatusesByChatID)) + sdkChats := db2sdk.ChatRowsWithChildren(chatRows, childRows, diffStatusesByChatID) + api.enrichChatWithWorkspaceAgentIDs(ctx, sdkChats) + httpapi.Write(ctx, rw, http.StatusOK, sdkChats) +} + +// enrichChatWithWorkspaceAgentIDs fills missing AgentIDs for chats with a bound +// workspace, since chatd persists the binding lazily. Best-effort and +// response-only; on error the field stays null. +func (api *API) enrichChatWithWorkspaceAgentIDs(ctx context.Context, chats []codersdk.Chat) { + missingChats := make([]*codersdk.Chat, 0, len(chats)) + var workspaceIDs []uuid.UUID + addMissing := func(chat *codersdk.Chat) { + if chat.AgentID == nil && chat.WorkspaceID != nil { + missingChats = append(missingChats, chat) + workspaceIDs = append(workspaceIDs, *chat.WorkspaceID) + } + } + for i := range chats { + addMissing(&chats[i]) + for j := range chats[i].Children { + addMissing(&chats[i].Children[j]) + } + } + + slices.SortFunc(workspaceIDs, func(a, b uuid.UUID) int { + return cmp.Compare(a.String(), b.String()) + }) + ids := slices.Compact(workspaceIDs) + rows, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(ctx, ids) + if err != nil { + return + } + + agentsByWorkspace := make(map[uuid.UUID][]database.WorkspaceAgent) + for _, row := range rows { + agentsByWorkspace[row.WorkspaceID] = append(agentsByWorkspace[row.WorkspaceID], row.WorkspaceAgent) + } + agentIDs := make(map[uuid.UUID]uuid.UUID, len(agentsByWorkspace)) + for workspaceID, agents := range agentsByWorkspace { + agent, err := agentselect.FindChatAgent(agents) + if err != nil { + api.Logger.Debug(ctx, "failed to select chat agent for enrichment", slog.F("workspace_id", workspaceID), slog.Error(err)) + continue + } + agentIDs[workspaceID] = agent.ID + } + + for _, chat := range missingChats { + if agentID, ok := agentIDs[*chat.WorkspaceID]; ok { + id := agentID + chat.AgentID = &id + } + } } func (api *API) getChatDiffStatusesByChatID( @@ -2196,6 +2250,10 @@ func (api *API) getChat(rw http.ResponseWriter, r *http.Request) { sdkChat.Children = db2sdk.ChildChatRows(childRows, childDiffStatuses) } + enriched := []codersdk.Chat{sdkChat} + api.enrichChatWithWorkspaceAgentIDs(ctx, enriched) + sdkChat = enriched[0] + httpapi.Write(ctx, rw, http.StatusOK, sdkChat) } @@ -2459,9 +2517,11 @@ func (api *API) watchChatGit(rw http.ResponseWriter, r *http.Request) { }) return } - if len(agents) == 0 { + agent, err := agentselect.FindChatAgent(agents) + if err != nil { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: codersdk.ChatGitWatchWorkspaceNoAgentsMessage, + Message: codersdk.ChatGitWatchNoEligibleAgentMessage, + Detail: err.Error(), }) return } @@ -2469,7 +2529,7 @@ func (api *API) watchChatGit(rw http.ResponseWriter, r *http.Request) { apiAgent, err := db2sdk.WorkspaceAgent( api.DERPMap(), *api.TailnetCoordinator.Load(), - agents[0], + agent, nil, nil, nil, @@ -2493,7 +2553,7 @@ func (api *API) watchChatGit(rw http.ResponseWriter, r *http.Request) { dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second) defer dialCancel() - agentConn, release, err := api.agentProvider.AgentConn(dialCtx, agents[0].ID) + agentConn, release, err := api.agentProvider.AgentConn(dialCtx, agent.ID) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Internal error dialing workspace agent.", @@ -2614,9 +2674,11 @@ func (api *API) watchChatDesktop(rw http.ResponseWriter, r *http.Request) { }) return } - if len(agents) == 0 { + agent, err := agentselect.FindChatAgent(agents) + if err != nil { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Chat workspace has no agents.", + Message: codersdk.ChatGitWatchNoEligibleAgentMessage, + Detail: err.Error(), }) return } @@ -2624,7 +2686,7 @@ func (api *API) watchChatDesktop(rw http.ResponseWriter, r *http.Request) { apiAgent, err := db2sdk.WorkspaceAgent( api.DERPMap(), *api.TailnetCoordinator.Load(), - agents[0], + agent, nil, nil, nil, @@ -2648,7 +2710,7 @@ func (api *API) watchChatDesktop(rw http.ResponseWriter, r *http.Request) { dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second) defer dialCancel() - agentConn, release, err := api.agentProvider.AgentConn(dialCtx, agents[0].ID) + agentConn, release, err := api.agentProvider.AgentConn(dialCtx, agent.ID) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to dial workspace agent.", diff --git a/coderd/exp_chats_internal_test.go b/coderd/exp_chats_internal_test.go index 51a55b4918..1facd9ff97 100644 --- a/coderd/exp_chats_internal_test.go +++ b/coderd/exp_chats_internal_test.go @@ -3,12 +3,75 @@ package coderd import ( "testing" + "github.com/google/uuid" "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + "golang.org/x/xerrors" + "cdr.dev/slog/v3" + "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" ) +func TestEnrichMissingChatAgentIDs(t *testing.T) { + t.Parallel() + newAPI := func(t *testing.T) (*API, *dbmock.MockStore) { + t.Helper() + mDB := dbmock.NewMockStore(gomock.NewController(t)) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug) + return &API{Options: &Options{Database: mDB, Logger: logger}}, mDB + } + workspaceID, otherWorkspaceID := uuid.New(), uuid.New() + rootAgentID, otherAgentID := uuid.New(), uuid.New() + row := func(workspaceID, id uuid.UUID, parentID uuid.NullUUID, name string) database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow { + return database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{ + WorkspaceID: workspaceID, + WorkspaceAgent: database.WorkspaceAgent{ + ID: id, + ParentID: parentID, + Name: name, + }, + } + } + t.Run("batch selection and shared workspace", func(t *testing.T) { + t.Parallel() + api, mDB := newAPI(t) + mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), gomock.Any()).DoAndReturn(func(_ any, ids []uuid.UUID) ([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow, error) { + require.ElementsMatch(t, []uuid.UUID{workspaceID, otherWorkspaceID}, ids) + return []database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{ + row(workspaceID, uuid.New(), uuid.NullUUID{UUID: rootAgentID, Valid: true}, "sub"), row(workspaceID, rootAgentID, uuid.NullUUID{}, "root"), row(otherWorkspaceID, otherAgentID, uuid.NullUUID{}, "root"), + }, nil + }).Times(1) + chats := []codersdk.Chat{{WorkspaceID: &workspaceID, Children: []codersdk.Chat{{WorkspaceID: &workspaceID}}}, {WorkspaceID: &otherWorkspaceID}} + api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats) + require.Equal(t, rootAgentID, *chats[0].AgentID) + require.Equal(t, rootAgentID, *chats[0].Children[0].AgentID) + require.Equal(t, otherAgentID, *chats[1].AgentID) + }) + t.Run("query error", func(t *testing.T) { + t.Parallel() + api, mDB := newAPI(t) + mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), gomock.Any()).Return(nil, xerrors.New("boom")) + chats := []codersdk.Chat{{WorkspaceID: &workspaceID}, {WorkspaceID: &otherWorkspaceID}} + api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats) + require.Nil(t, chats[0].AgentID) + require.Nil(t, chats[1].AgentID) + }) + t.Run("selection error and skips bound or unbound", func(t *testing.T) { + t.Parallel() + api, mDB := newAPI(t) + mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), []uuid.UUID{workspaceID}).Return([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{row(workspaceID, uuid.New(), uuid.NullUUID{UUID: rootAgentID, Valid: true}, "sub")}, nil) + bound := otherAgentID + chats := []codersdk.Chat{{}, {WorkspaceID: &workspaceID}, {WorkspaceID: &workspaceID, AgentID: &bound}} + api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats) + require.Nil(t, chats[1].AgentID) + require.Equal(t, bound, *chats[2].AgentID) + }) +} + func TestValidateChatModelProviderOptions_AnthropicThinkingDisplay(t *testing.T) { t.Parallel() diff --git a/coderd/workspaceagents_internal_test.go b/coderd/workspaceagents_internal_test.go index f6b09c614b..c7697bbedf 100644 --- a/coderd/workspaceagents_internal_test.go +++ b/coderd/workspaceagents_internal_test.go @@ -304,8 +304,8 @@ func TestWatchChatGit(t *testing.T) { t.Run("DisconnectedAgentRejected", func(t *testing.T) { t.Parallel() - // This test ensures that a chat whose workspace agent is - // not connected returns a 400 error. + // The handler must reject a disconnected agent with a 400 + // and select the root agent when a sub-agent precedes it. var ( ctx = testutil.Context(t, testutil.WaitShort) @@ -317,6 +317,7 @@ func TestWatchChatGit(t *testing.T) { chatID = uuid.New() workspaceID = uuid.New() + subAgentID = uuid.New() agentID = uuid.New() resourceID = uuid.New() @@ -354,17 +355,25 @@ func TestWatchChatGit(t *testing.T) { ID: workspaceID, }, nil) - // And: Return an agent that is disconnected (no - // FirstConnectedAt). mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). - Return([]database.WorkspaceAgent{{ - ID: agentID, - ResourceID: resourceID, - LifecycleState: database.WorkspaceAgentLifecycleStateCreated, - }}, nil) + Return([]database.WorkspaceAgent{ + { + ID: subAgentID, + ParentID: uuid.NullUUID{UUID: agentID, Valid: true}, + Name: "dev-container", + ResourceID: resourceID, + LifecycleState: database.WorkspaceAgentLifecycleStateCreated, + }, + { + ID: agentID, + Name: "main", + ResourceID: resourceID, + LifecycleState: database.WorkspaceAgentLifecycleStateCreated, + }, + }, nil) - // And: Allow db2sdk.WorkspaceAgent to complete. - mCoordinator.EXPECT().Node(gomock.Any()).Return(nil) + // Node(agentID) proves the root agent was selected. + mCoordinator.EXPECT().Node(agentID).Return(nil) // And: We mount the HTTP handler. r.With(injectSystemActor, httpmw.ExtractChatParam(mDB)). diff --git a/coderd/x/chatd/internal/agentselect/agentselect.go b/coderd/x/chatd/agentselect/agentselect.go similarity index 100% rename from coderd/x/chatd/internal/agentselect/agentselect.go rename to coderd/x/chatd/agentselect/agentselect.go diff --git a/coderd/x/chatd/internal/agentselect/agentselect_test.go b/coderd/x/chatd/agentselect/agentselect_test.go similarity index 98% rename from coderd/x/chatd/internal/agentselect/agentselect_test.go rename to coderd/x/chatd/agentselect/agentselect_test.go index 84bbb5bee8..a19537625c 100644 --- a/coderd/x/chatd/internal/agentselect/agentselect_test.go +++ b/coderd/x/chatd/agentselect/agentselect_test.go @@ -8,7 +8,7 @@ import ( "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/x/chatd/internal/agentselect" + "github.com/coder/coder/v2/coderd/x/chatd/agentselect" ) func TestFindChatAgent(t *testing.T) { diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 6c866765f9..8390441acb 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -39,6 +39,7 @@ import ( "github.com/coder/coder/v2/coderd/util/xjson" "github.com/coder/coder/v2/coderd/webpush" "github.com/coder/coder/v2/coderd/workspacestats" + "github.com/coder/coder/v2/coderd/x/chatd/agentselect" "github.com/coder/coder/v2/coderd/x/chatd/chatadvisor" "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" "github.com/coder/coder/v2/coderd/x/chatd/chaterror" @@ -48,7 +49,6 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/coderd/x/chatd/chattool" - "github.com/coder/coder/v2/coderd/x/chatd/internal/agentselect" "github.com/coder/coder/v2/coderd/x/chatd/mcpclient" "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" skillspkg "github.com/coder/coder/v2/coderd/x/skills" diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index a20db20ac5..9b740fe378 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -19,7 +19,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/httpapi/httperror" "github.com/coder/coder/v2/coderd/util/namesgenerator" - "github.com/coder/coder/v2/coderd/x/chatd/internal/agentselect" + "github.com/coder/coder/v2/coderd/x/chatd/agentselect" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" ) diff --git a/coderd/x/chatd/chattool/startworkspace.go b/coderd/x/chatd/chattool/startworkspace.go index 24b55348e6..650513d4be 100644 --- a/coderd/x/chatd/chattool/startworkspace.go +++ b/coderd/x/chatd/chattool/startworkspace.go @@ -11,7 +11,7 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/httpapi/httperror" - "github.com/coder/coder/v2/coderd/x/chatd/internal/agentselect" + "github.com/coder/coder/v2/coderd/x/chatd/agentselect" "github.com/coder/coder/v2/codersdk" ) diff --git a/codersdk/chats.go b/codersdk/chats.go index c35af0c4ce..b416fbaf43 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -1629,7 +1629,7 @@ type ChatDiffContents struct { const ( ChatGitWatchNoWorkspaceMessage = "Chat has no workspace to watch." ChatGitWatchWorkspaceNotFoundMessage = "Chat workspace not found." - ChatGitWatchWorkspaceNoAgentsMessage = "Chat workspace has no agents." + ChatGitWatchNoEligibleAgentMessage = "No eligible agent found for chat workspace." // ChatGitWatchAgentStatePrefix is the common prefix of the // message produced by ChatGitWatchAgentStateMessage. The CLI // uses it as a mechanical fingerprint for the "agent not yet @@ -1654,7 +1654,7 @@ func IsChatGitWatchFallbackMessage(msg string) bool { switch trimmed { case ChatGitWatchNoWorkspaceMessage, ChatGitWatchWorkspaceNotFoundMessage, - ChatGitWatchWorkspaceNoAgentsMessage: + ChatGitWatchNoEligibleAgentMessage: return true } return strings.HasPrefix(trimmed, ChatGitWatchAgentStatePrefix) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 901e65f6ab..edf614927b 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2300,7 +2300,8 @@ export const ChatGitWatchAgentStatePrefix = "Agent state is "; * IsChatGitWatchFallbackMessage instead of coupling to exact wording. * Keep these in sync with coderd/exp_chats.go. */ -export const ChatGitWatchNoWorkspaceMessage = "Chat has no workspace to watch."; +export const ChatGitWatchNoEligibleAgentMessage = + "No eligible agent found for chat workspace."; // From codersdk/chats.go /** @@ -2312,8 +2313,7 @@ export const ChatGitWatchNoWorkspaceMessage = "Chat has no workspace to watch."; * IsChatGitWatchFallbackMessage instead of coupling to exact wording. * Keep these in sync with coderd/exp_chats.go. */ -export const ChatGitWatchWorkspaceNoAgentsMessage = - "Chat workspace has no agents."; +export const ChatGitWatchNoWorkspaceMessage = "Chat has no workspace to watch."; // From codersdk/chats.go /** diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 2bcd8b5293..90364d68f3 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -1,4 +1,11 @@ -import { type FC, useEffect, useLayoutEffect, useRef, useState } from "react"; +import { + type FC, + useEffect, + useEffectEvent, + useLayoutEffect, + useRef, + useState, +} from "react"; import { useInfiniteQuery, @@ -763,6 +770,7 @@ const AgentChatPage: FC = () => { enabled: Boolean(parentChatID), }); const workspaceId = chatQuery.data?.workspace_id; + const chatAgentId = chatQuery.data?.agent_id; const workspaceQuery = useQuery({ ...workspaceById(workspaceId ?? ""), enabled: Boolean(workspaceId), @@ -835,6 +843,36 @@ const AgentChatPage: FC = () => { // Subscribe to live workspace updates so that agent status changes // (e.g. connected/disconnected) are reflected without a page refresh. + const applyWatchedWorkspaceUpdate = useEffectEvent( + (watchedWorkspaceId: string, next: TypesGen.Workspace) => { + queryClient.setQueryData( + workspaceByIdKey(watchedWorkspaceId), + (prev) => { + // Return the same reference when nothing the UI + // reads has changed. This prevents react-query + // from notifying subscribers and avoids a full + // AgentChatPage re-render on every heartbeat. + const prevAgent = getWorkspaceAgent(prev, chatAgentId); + const nextAgent = getWorkspaceAgent(next, chatAgentId); + if ( + prev && + prev.latest_build.status === next.latest_build.status && + prev.health.healthy === next.health.healthy && + prev.name === next.name && + prev.owner_name === next.owner_name && + prevAgent?.id === nextAgent?.id && + prevAgent?.status === nextAgent?.status && + prevAgent?.name === nextAgent?.name && + prevAgent?.expanded_directory === nextAgent?.expanded_directory && + prevAgent?.lifecycle_state === nextAgent?.lifecycle_state + ) { + return prev; + } + return next; + }, + ); + }, + ); useEffect(() => { if (!workspaceId) { return; @@ -847,33 +885,9 @@ const AgentChatPage: FC = () => { return; } if (event.parsedMessage.type === "data") { - const next = event.parsedMessage.data as TypesGen.Workspace; - queryClient.setQueryData( - workspaceByIdKey(workspaceId), - (prev) => { - // Return the same reference when nothing the UI - // reads has changed. This prevents react-query - // from notifying subscribers and avoids a full - // AgentChatPage re-render on every heartbeat. - const prevAgent = getWorkspaceAgent(prev, undefined); - const nextAgent = getWorkspaceAgent(next, undefined); - if ( - prev && - prev.latest_build.status === next.latest_build.status && - prev.health.healthy === next.health.healthy && - prev.name === next.name && - prev.owner_name === next.owner_name && - prevAgent?.id === nextAgent?.id && - prevAgent?.status === nextAgent?.status && - prevAgent?.name === nextAgent?.name && - prevAgent?.expanded_directory === - nextAgent?.expanded_directory && - prevAgent?.lifecycle_state === nextAgent?.lifecycle_state - ) { - return prev; - } - return next; - }, + applyWatchedWorkspaceUpdate( + workspaceId, + event.parsedMessage.data as TypesGen.Workspace, ); } }); @@ -891,7 +905,7 @@ const AgentChatPage: FC = () => { }); }, [workspaceId, queryClient]); const sshConfigQuery = useQuery(deploymentSSHConfig()); - const workspaceAgent = getWorkspaceAgent(workspace, undefined); + const workspaceAgent = getWorkspaceAgent(workspace, chatAgentId); const { proxy } = useProxy(); const chatRecord = chatQuery.data; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.test.ts index 3d3f2c544c..d6a88b30d4 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.test.ts @@ -234,18 +234,14 @@ describe("getWorkspaceAgent", () => { ); }); - it("returns the first agent when workspaceAgentId does not match", () => { + it("returns undefined when workspaceAgentId does not match", () => { const ws = buildWorkspace([buildAgent("a1"), buildAgent("a2")]); - expect(getWorkspaceAgent(ws, "no-match")).toEqual( - expect.objectContaining({ id: "a1" }), - ); + expect(getWorkspaceAgent(ws, "no-match")).toBeUndefined(); }); - it("returns the first agent when workspaceAgentId is undefined", () => { + it("returns undefined when workspaceAgentId is undefined", () => { const ws = buildWorkspace([buildAgent("a1")]); - expect(getWorkspaceAgent(ws, undefined)).toEqual( - expect.objectContaining({ id: "a1" }), - ); + expect(getWorkspaceAgent(ws, undefined)).toBeUndefined(); }); it("collects agents from multiple resources", () => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts index 083c4d8210..a03af4ac31 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatHelpers.ts @@ -1,5 +1,5 @@ import type * as TypesGen from "#/api/typesGenerated"; -import { getWorkspaceAgents } from "#/utils/workspace"; +import { findWorkspaceAgent } from "#/utils/workspace"; import type { AgentContextUsage } from "../AgentChatInput"; import type { ModelSelectorOption } from "../ChatElements"; import { asString } from "../ChatElements/runtimeTypeUtils"; @@ -96,12 +96,7 @@ export const getWorkspaceAgent = ( workspace: TypesGen.Workspace | undefined, workspaceAgentId: string | undefined, ): TypesGen.WorkspaceAgent | undefined => { - if (!workspace) { - return undefined; - } - const agents = getWorkspaceAgents(workspace); - if (agents.length === 0) { - return undefined; - } - return agents.find((agent) => agent.id === workspaceAgentId) ?? agents[0]; + return workspace && workspaceAgentId + ? findWorkspaceAgent(workspace, workspaceAgentId) + : undefined; };