From a74015fc85cb121d828336126e1e4af7368973c4 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Wed, 6 May 2026 11:05:16 +0100 Subject: [PATCH] refactor: make store and chatID explicit parameter arguments in chattools (#24850) Fixes CODAGT-175 Addresses a review finding in https://github.com/coder/coder/pull/23827 that the nil-guards for both `database.Store` and `chatID` are both dead code in practice in the `chattool` package. - Modifies the return signatures require passing both `database.Store` and `chatID` explicitly as positional arguments instead of just parameter struct keys. - Drops the nil-guards for `database.Store` and `chatID`. --- coderd/x/chatd/chatd.go | 11 +- coderd/x/chatd/chattool/createworkspace.go | 135 ++++++++---------- .../x/chatd/chattool/createworkspace_test.go | 114 ++++++++------- coderd/x/chatd/chattool/listtemplates.go | 7 +- coderd/x/chatd/chattool/listtemplates_test.go | 43 ++++-- coderd/x/chatd/chattool/readtemplate.go | 7 +- coderd/x/chatd/chattool/readtemplate_test.go | 4 +- coderd/x/chatd/chattool/startworkspace.go | 45 +++--- .../x/chatd/chattool/startworkspace_test.go | 56 ++------ 9 files changed, 199 insertions(+), 223 deletions(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index bd502f9fcf..ff551bb76f 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -5950,17 +5950,16 @@ func (p *Server) appendRootChatTools( } tools = append(tools, - chattool.ListTemplates(opts.chat.OrganizationID, p.db, chattool.ListTemplatesOptions{ + chattool.ListTemplates(p.db, opts.chat.OrganizationID, chattool.ListTemplatesOptions{ OwnerID: opts.chat.OwnerID, AllowedTemplateIDs: p.chatTemplateAllowlist, }), - chattool.ReadTemplate(opts.chat.OrganizationID, p.db, chattool.ReadTemplateOptions{ + chattool.ReadTemplate(p.db, opts.chat.OrganizationID, chattool.ReadTemplateOptions{ OwnerID: opts.chat.OwnerID, AllowedTemplateIDs: p.chatTemplateAllowlist, }), - chattool.CreateWorkspace(opts.chat.OrganizationID, p.db, chattool.CreateWorkspaceOptions{ + chattool.CreateWorkspace(p.db, opts.chat.OrganizationID, opts.chat.ID, chattool.CreateWorkspaceOptions{ OwnerID: opts.chat.OwnerID, - ChatID: opts.chat.ID, CreateFn: p.createWorkspaceFn, AgentConnFn: chattool.AgentConnFunc(p.agentConnFn), AgentInactiveDisconnectTimeout: p.agentInactiveDisconnectTimeout, @@ -5969,10 +5968,8 @@ func (p *Server) appendRootChatTools( Logger: p.logger, AllowedTemplateIDs: p.chatTemplateAllowlist, }), - chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: p.db, + chattool.StartWorkspace(p.db, opts.chat.ID, chattool.StartWorkspaceOptions{ OwnerID: opts.chat.OwnerID, - ChatID: opts.chat.ID, StartFn: p.startWorkspaceFn, AgentConnFn: chattool.AgentConnFunc(p.agentConnFn), WorkspaceMu: opts.workspaceMu, diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index cbd83352f8..4a95a6e639 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -63,7 +63,6 @@ type AgentConnFunc func( // CreateWorkspaceOptions configures the create_workspace tool. type CreateWorkspaceOptions struct { OwnerID uuid.UUID - ChatID uuid.UUID CreateFn CreateWorkspaceFn AgentConnFn AgentConnFunc AgentInactiveDisconnectTimeout time.Duration @@ -85,7 +84,8 @@ type createWorkspaceArgs struct { // workspace that is building or running, it returns the existing // workspace instead of creating a new one. A mutex prevents parallel // calls from creating duplicate workspaces. -func CreateWorkspace(organizationID uuid.UUID, db database.Store, options CreateWorkspaceOptions) fantasy.AgentTool { +// db must not be nil and chatID must not be uuid.Nil. +func CreateWorkspace(db database.Store, organizationID, chatID uuid.UUID, options CreateWorkspaceOptions) fantasy.AgentTool { return fantasy.NewAgentTool( "create_workspace", "Create a new workspace from a template. Requires a "+ @@ -99,9 +99,6 @@ func CreateWorkspace(organizationID uuid.UUID, db database.Store, options Create "workspace that is building or running, the existing "+ "workspace is returned.", func(ctx context.Context, args createWorkspaceArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - if db == nil { - return fantasy.NewTextErrorResponse("database is not configured"), nil - } if options.CreateFn == nil { return fantasy.NewTextErrorResponse("workspace creator is not configured"), nil } @@ -129,7 +126,7 @@ func CreateWorkspace(organizationID uuid.UUID, db database.Store, options Create } // Check for an existing workspace on the chat. - check := options.checkExistingWorkspace(ctx, db) + check := options.checkExistingWorkspace(ctx, db, chatID) if check.Err != nil { if check.FailedBuildID != uuid.Nil { return buildToolResponse(newBuildError(check.Err.Error(), check.FailedBuildID)), nil @@ -232,31 +229,29 @@ func CreateWorkspace(organizationID uuid.UUID, db database.Store, options Create // later fails. The checkExistingWorkspace recovery // path handles failed workspaces by allowing // re-creation. - if options.ChatID != uuid.Nil { - updatedChat, err := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ - ID: options.ChatID, - WorkspaceID: uuid.NullUUID{ - UUID: workspace.ID, - Valid: true, - }, - BuildID: uuid.NullUUID{ - UUID: workspace.LatestBuild.ID, - Valid: workspace.LatestBuild.ID != uuid.Nil, - }, - // AgentID is left null because the build hasn't - // completed yet. The chatd runtime binds it once - // the agent comes online. - AgentID: uuid.NullUUID{}, - }) - if err != nil { - options.Logger.Error(ctx, "failed to persist chat workspace association", - slog.F("chat_id", options.ChatID), - slog.F("workspace_id", workspace.ID), - slog.Error(err), - ) - } else if options.OnChatUpdated != nil { - options.OnChatUpdated(updatedChat) - } + updatedChat, err := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ + ID: chatID, + WorkspaceID: uuid.NullUUID{ + UUID: workspace.ID, + Valid: true, + }, + BuildID: uuid.NullUUID{ + UUID: workspace.LatestBuild.ID, + Valid: workspace.LatestBuild.ID != uuid.Nil, + }, + // AgentID is left null because the build hasn't + // completed yet. The chatd runtime binds it once + // the agent comes online. + AgentID: uuid.NullUUID{}, + }) + if err != nil { + options.Logger.Error(ctx, "failed to persist chat workspace association", + slog.F("chat_id", chatID), + slog.F("workspace_id", workspace.ID), + slog.Error(err), + ) + } else if options.OnChatUpdated != nil { + options.OnChatUpdated(updatedChat) } // Wait for the build to complete and the agent to @@ -312,7 +307,7 @@ func CreateWorkspace(organizationID uuid.UUID, db database.Store, options Create // and the connection usually times out before the // agent is reachable. if options.OnChatUpdated != nil { - if latest, err := db.GetChatByID(ctx, options.ChatID); err == nil { + if latest, err := db.GetChatByID(ctx, chatID); err == nil { options.OnChatUpdated(latest) } } @@ -335,7 +330,7 @@ type existingWorkspaceResult struct { Err error } -// checkExistingWorkspace checks whether the configured chat +// checkExistingWorkspace checks whether the given chat // already has a usable workspace. Returns an // existingWorkspaceResult with Done set when the caller should // return early (workspace exists and is alive or building). @@ -344,12 +339,8 @@ type existingWorkspaceResult struct { func (o CreateWorkspaceOptions) checkExistingWorkspace( ctx context.Context, db database.Store, + chatID uuid.UUID, ) existingWorkspaceResult { - if o.ChatID == uuid.Nil { - return existingWorkspaceResult{} - } - - chatID := o.ChatID agentConnFn := o.AgentConnFn agentInactiveDisconnectTimeout := o.AgentInactiveDisconnectTimeout @@ -388,7 +379,7 @@ func (o CreateWorkspaceOptions) checkExistingWorkspace( // Build is in progress. Publish the build ID so the // frontend can start streaming logs, then wait. updatedChat, bindErr := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ - ID: o.ChatID, + ID: chatID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, BuildID: uuid.NullUUID{ UUID: build.ID, @@ -398,7 +389,7 @@ func (o CreateWorkspaceOptions) checkExistingWorkspace( }) if bindErr != nil { o.Logger.Error(ctx, "failed to persist build ID on chat binding", - slog.F("chat_id", o.ChatID), + slog.F("chat_id", chatID), slog.F("build_id", build.ID), slog.Error(bindErr), ) @@ -590,46 +581,42 @@ func waitForAgentReady( } // Phase 2: poll lifecycle until startup scripts finish. - if db != nil { - scriptCtx, scriptCancel := context.WithTimeout(ctx, startupScriptTimeout) - defer scriptCancel() + scriptCtx, scriptCancel := context.WithTimeout(ctx, startupScriptTimeout) + defer scriptCancel() - ticker := time.NewTicker(startupScriptPollInterval) - defer ticker.Stop() + ticker := time.NewTicker(startupScriptPollInterval) + defer ticker.Stop() - var lastState database.WorkspaceAgentLifecycleState - for { - row, err := db.GetWorkspaceAgentLifecycleStateByID(scriptCtx, agentID) - if err == nil { - lastState = row.LifecycleState - switch lastState { - case database.WorkspaceAgentLifecycleStateCreated, - database.WorkspaceAgentLifecycleStateStarting: - // Still in progress, keep polling. - case database.WorkspaceAgentLifecycleStateReady: - return result - default: - // Terminal non-ready state. - result["startup_scripts"] = "startup_scripts_failed" - result["lifecycle_state"] = string(lastState) - return result - } - } - - select { - case <-scriptCtx.Done(): - if errors.Is(scriptCtx.Err(), context.DeadlineExceeded) { - result["startup_scripts"] = "startup_scripts_timeout" - } else { - result["startup_scripts"] = "startup_scripts_unknown" - } + var lastState database.WorkspaceAgentLifecycleState + for { + row, err := db.GetWorkspaceAgentLifecycleStateByID(scriptCtx, agentID) + if err == nil { + lastState = row.LifecycleState + switch lastState { + case database.WorkspaceAgentLifecycleStateCreated, + database.WorkspaceAgentLifecycleStateStarting: + // Still in progress, keep polling. + case database.WorkspaceAgentLifecycleStateReady: + return result + default: + // Terminal non-ready state. + result["startup_scripts"] = "startup_scripts_failed" + result["lifecycle_state"] = string(lastState) return result - case <-ticker.C: } } - } - return result + select { + case <-scriptCtx.Done(): + if errors.Is(scriptCtx.Err(), context.DeadlineExceeded) { + result["startup_scripts"] = "startup_scripts_timeout" + } else { + result["startup_scripts"] = "startup_scripts_unknown" + } + return result + case <-ticker.C: + } + } } func generatedWorkspaceName(seed string) string { diff --git a/coderd/x/chatd/chattool/createworkspace_test.go b/coderd/x/chatd/chattool/createworkspace_test.go index cacf908526..f7f9e7e0d3 100644 --- a/coderd/x/chatd/chattool/createworkspace_test.go +++ b/coderd/x/chatd/chattool/createworkspace_test.go @@ -107,17 +107,6 @@ func TestWaitForAgentReady(t *testing.T) { result := waitForAgentReady(context.Background(), db, agentID, nil) require.Empty(t, result) }) - - t.Run("NilDB", func(t *testing.T) { - t.Parallel() - - connFn := func(ctx context.Context, id uuid.UUID) (workspacesdk.AgentConn, func(), error) { - return nil, func() {}, nil - } - - result := waitForAgentReady(context.Background(), nil, uuid.New(), connFn) - require.Empty(t, result) - }) } func TestCreateWorkspace_PrefersChatSuffixAgent(t *testing.T) { @@ -128,6 +117,7 @@ func TestCreateWorkspace_PrefersChatSuffixAgent(t *testing.T) { ownerID := uuid.New() orgID := uuid.New() + chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() @@ -135,6 +125,14 @@ func TestCreateWorkspace_PrefersChatSuffixAgent(t *testing.T) { fallbackAgentID := uuid.New() chatAgentID := uuid.New() + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ID: chatID}, nil) + + db.EXPECT(). + UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). + Return(database.Chat{ID: chatID}, nil) + db.EXPECT(). GetAuthorizationUserRoles(gomock.Any(), ownerID). Return(database.GetAuthorizationUserRolesRow{ @@ -196,7 +194,7 @@ func TestCreateWorkspace_PrefersChatSuffixAgent(t *testing.T) { return nil, func() {}, nil } - tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: createFn, @@ -286,10 +284,9 @@ func TestCreateWorkspace_ReturnsSelectionErrorImmediately(t *testing.T) { {ID: uuid.New(), Name: "beta-coderd-chat", DisplayOrder: 1}, }, nil) - tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, - ChatID: chatID, CreateFn: func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{ ID: workspaceID, @@ -333,11 +330,20 @@ func TestCreateWorkspace_PostCreationBuildFailure(t *testing.T) { ownerID := uuid.New() orgID := uuid.New() + chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ID: chatID}, nil) + + db.EXPECT(). + UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). + Return(database.Chat{ID: chatID}, nil) + db.EXPECT(). GetAuthorizationUserRoles(gomock.Any(), ownerID). Return(database.GetAuthorizationUserRolesRow{ @@ -387,10 +393,9 @@ func TestCreateWorkspace_PostCreationBuildFailure(t *testing.T) { }, nil } - tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, - ChatID: uuid.Nil, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), @@ -420,8 +425,13 @@ func TestCreateWorkspace_ResponderErrorPreservesStructuredFields(t *testing.T) { ownerID := uuid.New() orgID := uuid.New() + chatID := uuid.New() templateID := uuid.New() + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ID: chatID}, nil) + db.EXPECT(). GetAuthorizationUserRoles(gomock.Any(), ownerID). Return(database.GetAuthorizationUserRolesRow{ @@ -442,7 +452,7 @@ func TestCreateWorkspace_ResponderErrorPreservesStructuredFields(t *testing.T) { GetChatWorkspaceTTL(gomock.Any()). Return("0s", nil) - tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { return codersdk.Workspace{}, httperror.NewResponseError(400, codersdk.Response{ @@ -521,11 +531,20 @@ func TestCreateWorkspace_GlobalTTL(t *testing.T) { ownerID := uuid.New() orgID := uuid.New() + chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ID: chatID}, nil) + + db.EXPECT(). + UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). + Return(database.Chat{ID: chatID}, nil) + db.EXPECT(). GetAuthorizationUserRoles(gomock.Any(), ownerID). Return(database.GetAuthorizationUserRolesRow{ @@ -577,10 +596,9 @@ func TestCreateWorkspace_GlobalTTL(t *testing.T) { }, nil } - tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, - ChatID: uuid.Nil, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), @@ -648,10 +666,9 @@ func TestCreateWorkspace_RejectsCrossOrgTemplate(t *testing.T) { }, nil) createCalled := false - tool := CreateWorkspace(chatOrgID, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, chatOrgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, - ChatID: chatID, CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { createCalled = true return codersdk.Workspace{}, nil @@ -711,8 +728,8 @@ func TestCheckExistingWorkspace_ConnectedAgent(t *testing.T) { return nil, nil, xerrors.New("unexpected agent dial") } - options := testCheckExistingWorkspaceOptions(chatID, connFn) - check := options.checkExistingWorkspace(context.Background(), db) + options := testCheckExistingWorkspaceOptions(connFn) + check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.True(t, check.Done) @@ -804,8 +821,8 @@ func TestCheckExistingWorkspace_InProgressBuildReturnsBuildID(t *testing.T) { GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{}, nil) - options := testCheckExistingWorkspaceOptions(chatID, nil) - check := options.checkExistingWorkspace(context.Background(), db) + options := testCheckExistingWorkspaceOptions(nil) + check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.True(t, check.Done) @@ -887,8 +904,8 @@ func TestCheckExistingWorkspace_InProgressBuildFailureReturnsBuildID(t *testing. WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }, nil) - options := testCheckExistingWorkspaceOptions(chatID, nil) - check := options.checkExistingWorkspace(context.Background(), db) + options := testCheckExistingWorkspaceOptions(nil) + check := options.checkExistingWorkspace(context.Background(), db, chatID) require.Error(t, check.Err) require.Contains(t, check.Err.Error(), "existing workspace build failed") @@ -935,8 +952,8 @@ func TestCheckExistingWorkspace_ConnectingAgentWaits(t *testing.T) { return nil, func() {}, nil } - options := testCheckExistingWorkspaceOptions(chatID, connFn) - check := options.checkExistingWorkspace(context.Background(), db) + options := testCheckExistingWorkspaceOptions(connFn) + check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.True(t, check.Done) @@ -996,8 +1013,8 @@ func TestCheckExistingWorkspace_DeadAgentAllowsCreation(t *testing.T) { GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{tc.agent}, nil) - options := testCheckExistingWorkspaceOptions(chatID, nil) - check := options.checkExistingWorkspace(context.Background(), db) + options := testCheckExistingWorkspaceOptions(nil) + check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.False(t, check.Done) @@ -1014,11 +1031,20 @@ func TestWaitForBuild_CanceledJob(t *testing.T) { ownerID := uuid.New() orgID := uuid.New() + chatID := uuid.New() templateID := uuid.New() workspaceID := uuid.New() jobID := uuid.New() buildID := uuid.New() + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ID: chatID}, nil) + + db.EXPECT(). + UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()). + Return(database.Chat{ID: chatID}, nil) + db.EXPECT(). GetAuthorizationUserRoles(gomock.Any(), ownerID). Return(database.GetAuthorizationUserRolesRow{ @@ -1067,10 +1093,9 @@ func TestWaitForBuild_CanceledJob(t *testing.T) { }, nil } - tool := CreateWorkspace(orgID, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, orgID, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, - ChatID: uuid.Nil, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), @@ -1111,8 +1136,8 @@ func TestCheckExistingWorkspace_StoppedWorkspace(t *testing.T) { database.WorkspaceTransitionStop, ) - options := testCheckExistingWorkspaceOptions(chatID, nil) - check := options.checkExistingWorkspace(context.Background(), db) + options := testCheckExistingWorkspaceOptions(nil) + check := options.checkExistingWorkspace(context.Background(), db, chatID) require.True(t, check.Done) require.NoError(t, check.Err) @@ -1144,8 +1169,8 @@ func TestCheckExistingWorkspace_DeletedWorkspace(t *testing.T) { Deleted: true, }, nil) - options := testCheckExistingWorkspaceOptions(chatID, nil) - check := options.checkExistingWorkspace(context.Background(), db) + options := testCheckExistingWorkspaceOptions(nil) + check := options.checkExistingWorkspace(context.Background(), db, chatID) require.NoError(t, check.Err) require.False(t, check.Done, "should allow creation for deleted workspace") @@ -1153,11 +1178,9 @@ func TestCheckExistingWorkspace_DeletedWorkspace(t *testing.T) { } func testCheckExistingWorkspaceOptions( - chatID uuid.UUID, agentConnFn AgentConnFunc, ) CreateWorkspaceOptions { return CreateWorkspaceOptions{ - ChatID: chatID, AgentConnFn: agentConnFn, AgentInactiveDisconnectTimeout: 30 * time.Second, } @@ -1266,7 +1289,6 @@ func TestCreateWorkspace_OnChatUpdatedFiresAfterBuild(t *testing.T) { CompletedAt: validNullTime(time.Now()), }, nil) - // GetChatByID — called after waitForBuild for second OnChatUpdated. // GetChatByID — called after waitForBuild for second OnChatUpdated. db.EXPECT(). GetChatByID(gomock.Any(), chatID). @@ -1295,10 +1317,9 @@ func TestCreateWorkspace_OnChatUpdatedFiresAfterBuild(t *testing.T) { }, nil } - tool := CreateWorkspace(uuid.Nil, db, CreateWorkspaceOptions{ + tool := CreateWorkspace(db, uuid.Nil, chatID, CreateWorkspaceOptions{ OwnerID: ownerID, - ChatID: chatID, CreateFn: createFn, WorkspaceMu: &sync.Mutex{}, Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), @@ -1455,9 +1476,8 @@ func TestCreateWorkspace_WithPresetID(t *testing.T) { return nil, func() {}, nil } - tool := CreateWorkspace(s.OrgID, s.DB, CreateWorkspaceOptions{ + tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{ OwnerID: s.OwnerID, - ChatID: s.ChatID, CreateFn: createFn, AgentConnFn: agentConnFn, WorkspaceMu: &sync.Mutex{}, @@ -1487,9 +1507,8 @@ func TestCreateWorkspace_InvalidPresetID(t *testing.T) { s := setupCreateWorkspacePresetTest(t) - tool := CreateWorkspace(s.OrgID, s.DB, CreateWorkspaceOptions{ + tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{ OwnerID: s.OwnerID, - ChatID: s.ChatID, CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { t.Fatal("CreateFn should not be called with invalid preset_id") return codersdk.Workspace{}, nil @@ -1538,9 +1557,8 @@ func TestCreateWorkspace_WithPresetAndParams(t *testing.T) { return nil, func() {}, nil } - tool := CreateWorkspace(s.OrgID, s.DB, CreateWorkspaceOptions{ + tool := CreateWorkspace(s.DB, s.OrgID, s.ChatID, CreateWorkspaceOptions{ OwnerID: s.OwnerID, - ChatID: s.ChatID, CreateFn: createFn, AgentConnFn: agentConnFn, WorkspaceMu: &sync.Mutex{}, diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index df5af5a2f3..3c6d31c1b0 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -35,7 +35,8 @@ type listTemplatesArgs struct { // The agent uses this to discover templates before creating a workspace. // Results are ordered by number of active developers (most popular first) // and paginated at 10 per page. -func ListTemplates(organizationID uuid.UUID, db database.Store, options ListTemplatesOptions) fantasy.AgentTool { +// db must not be nil. +func ListTemplates(db database.Store, organizationID uuid.UUID, options ListTemplatesOptions) fantasy.AgentTool { return fantasy.NewAgentTool( "list_templates", "List available workspace templates. Optionally filter by a "+ @@ -44,10 +45,6 @@ func ListTemplates(organizationID uuid.UUID, db database.Store, options ListTemp "Results are ordered by number of active developers (most popular first). "+ "Returns 10 per page. Use the page parameter to paginate through results.", func(ctx context.Context, args listTemplatesArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - if db == nil { - return fantasy.NewTextErrorResponse("database is not configured"), nil - } - ctx, err := asOwner(ctx, db, options.OwnerID) if err != nil { return fantasy.NewTextErrorResponse(err.Error()), nil diff --git a/coderd/x/chatd/chattool/listtemplates_test.go b/coderd/x/chatd/chattool/listtemplates_test.go index c874430901..0cf25d2c43 100644 --- a/coderd/x/chatd/chattool/listtemplates_test.go +++ b/coderd/x/chatd/chattool/listtemplates_test.go @@ -49,7 +49,7 @@ func TestListTemplates_OrganizationFilter(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) - tool := chattool.ListTemplates(orgA.ID, db, chattool.ListTemplatesOptions{ + tool := chattool.ListTemplates(db, orgA.ID, chattool.ListTemplatesOptions{ OwnerID: user.ID, }) @@ -69,7 +69,7 @@ func TestListTemplates_OrganizationFilter(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) - tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{ + tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{ OwnerID: user.ID, // Pass uuid.Nil to skip org filtering. }) @@ -89,7 +89,7 @@ func TestListTemplates_OrganizationFilter(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) // Tool scoped to orgA, but requesting a template in orgB. - tool := chattool.ReadTemplate(orgA.ID, db, chattool.ReadTemplateOptions{ + tool := chattool.ReadTemplate(db, orgA.ID, chattool.ReadTemplateOptions{ OwnerID: user.ID, }) @@ -105,7 +105,7 @@ func TestListTemplates_OrganizationFilter(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) // Tool scoped to orgA, requesting a template in orgA. - tool := chattool.ReadTemplate(orgA.ID, db, chattool.ReadTemplateOptions{ + tool := chattool.ReadTemplate(db, orgA.ID, chattool.ReadTemplateOptions{ OwnerID: user.ID, }) @@ -147,7 +147,7 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { t.Run("ListTemplates", func(t *testing.T) { t.Run("NoAllowlist", func(t *testing.T) { - tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{ + tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{ OwnerID: user.ID, }) @@ -160,7 +160,7 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { }) t.Run("EmptyAllowlist", func(t *testing.T) { - tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{ + tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{ OwnerID: user.ID, AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{} }, }) @@ -174,7 +174,7 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { }) t.Run("OneMatch", func(t *testing.T) { - tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{ + tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{ OwnerID: user.ID, AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} }, }) @@ -190,7 +190,7 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { }) t.Run("NoMatches", func(t *testing.T) { - tool := chattool.ListTemplates(uuid.Nil, db, chattool.ListTemplatesOptions{ + tool := chattool.ListTemplates(db, uuid.Nil, chattool.ListTemplatesOptions{ OwnerID: user.ID, AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} }, }) @@ -206,7 +206,7 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { t.Run("ReadTemplate", func(t *testing.T) { t.Run("Allowed", func(t *testing.T) { - tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{ + tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{ OwnerID: user.ID, AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} }, }) @@ -221,7 +221,7 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { }) t.Run("Disallowed", func(t *testing.T) { - tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{ + tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{ OwnerID: user.ID, AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} }, }) @@ -233,7 +233,7 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { }) t.Run("NoAllowlist", func(t *testing.T) { - tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{ + tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{ OwnerID: user.ID, }) input := `{"template_id":"` + t2.ID.String() + `"}` @@ -245,8 +245,21 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { t.Run("CreateWorkspace", func(t *testing.T) { t.Run("Allowed", func(t *testing.T) { + // CreateWorkspace requires a real chat row so the existing + // workspace lookup can fall through to creation. + model := seedModelConfig(t, db) + chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "allowed-create", + Status: database.ChatStatusWaiting, + ClientType: database.ChatClientTypeApi, + }) + require.NoError(t, err) + createCalled := false - tool := chattool.CreateWorkspace(org.ID, db, chattool.CreateWorkspaceOptions{ + tool := chattool.CreateWorkspace(db, org.ID, chat.ID, chattool.CreateWorkspaceOptions{ OwnerID: user.ID, AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t1.ID: true} }, @@ -268,10 +281,10 @@ func TestTemplateAllowlistEnforcement(t *testing.T) { }) t.Run("Disallowed", func(t *testing.T) { - createCalled := false - tool := chattool.CreateWorkspace(uuid.Nil, db, chattool.CreateWorkspaceOptions{ + var createCalled bool + tool := chattool.CreateWorkspace(db, org.ID, uuid.New(), chattool.CreateWorkspaceOptions{ OwnerID: user.ID, - AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{uuid.New(): true} }, + AllowedTemplateIDs: func() map[uuid.UUID]bool { return map[uuid.UUID]bool{t2.ID: true} }, CreateFn: func(_ context.Context, _ uuid.UUID, _ codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { createCalled = true t.Fatal("CreateFn should not be called for blocked template") diff --git a/coderd/x/chatd/chattool/readtemplate.go b/coderd/x/chatd/chattool/readtemplate.go index 4048c734a4..09179237ba 100644 --- a/coderd/x/chatd/chattool/readtemplate.go +++ b/coderd/x/chatd/chattool/readtemplate.go @@ -25,7 +25,8 @@ type readTemplateArgs struct { // ReadTemplate returns a tool that retrieves details about a specific // template, including its configurable rich parameters. The agent // uses this after list_templates and before create_workspace. -func ReadTemplate(organizationID uuid.UUID, db database.Store, options ReadTemplateOptions) fantasy.AgentTool { +// db must not be nil. +func ReadTemplate(db database.Store, organizationID uuid.UUID, options ReadTemplateOptions) fantasy.AgentTool { return fantasy.NewAgentTool( "read_template", "Get details about a workspace template, including its "+ @@ -33,10 +34,6 @@ func ReadTemplate(organizationID uuid.UUID, db database.Store, options ReadTempl "after finding a template with list_templates and before "+ "creating a workspace with create_workspace.", func(ctx context.Context, args readTemplateArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - if db == nil { - return fantasy.NewTextErrorResponse("database is not configured"), nil - } - templateIDStr := strings.TrimSpace(args.TemplateID) if templateIDStr == "" { return fantasy.NewTextErrorResponse("template_id is required"), nil diff --git a/coderd/x/chatd/chattool/readtemplate_test.go b/coderd/x/chatd/chattool/readtemplate_test.go index cadeba6ccd..e0ea0d6848 100644 --- a/coderd/x/chatd/chattool/readtemplate_test.go +++ b/coderd/x/chatd/chattool/readtemplate_test.go @@ -62,7 +62,7 @@ func TestReadTemplate_IncludesPresets(t *testing.T) { }) ctx := testutil.Context(t, testutil.WaitShort) - tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{ + tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{ OwnerID: user.ID, }) @@ -162,7 +162,7 @@ func TestReadTemplate_NoPresets(t *testing.T) { }) ctx := testutil.Context(t, testutil.WaitShort) - tool := chattool.ReadTemplate(org.ID, db, chattool.ReadTemplateOptions{ + tool := chattool.ReadTemplate(db, org.ID, chattool.ReadTemplateOptions{ OwnerID: user.ID, }) diff --git a/coderd/x/chatd/chattool/startworkspace.go b/coderd/x/chatd/chattool/startworkspace.go index aca85b9a0e..332beaa128 100644 --- a/coderd/x/chatd/chattool/startworkspace.go +++ b/coderd/x/chatd/chattool/startworkspace.go @@ -26,9 +26,7 @@ type StartWorkspaceFn func( // StartWorkspaceOptions configures the start_workspace tool. type StartWorkspaceOptions struct { - DB database.Store OwnerID uuid.UUID - ChatID uuid.UUID StartFn StartWorkspaceFn AgentConnFn AgentConnFunc WorkspaceMu *sync.Mutex @@ -43,7 +41,8 @@ type startWorkspaceArgs struct { // StartWorkspace returns a tool that starts a stopped workspace // associated with the current chat. The tool is idempotent: if the // workspace is already running or building, it returns immediately. -func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { +// db must not be nil and chatID must not be uuid.Nil. +func StartWorkspace(db database.Store, chatID uuid.UUID, options StartWorkspaceOptions) fantasy.AgentTool { return fantasy.NewAgentTool( "start_workspace", "Start the chat's workspace if it is currently stopped. "+ @@ -63,11 +62,7 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { defer options.WorkspaceMu.Unlock() } - if options.DB == nil || options.ChatID == uuid.Nil { - return fantasy.NewTextErrorResponse("start_workspace is not properly configured"), nil - } - - chat, err := options.DB.GetChatByID(ctx, options.ChatID) + chat, err := db.GetChatByID(ctx, chatID) if err != nil { return fantasy.NewTextErrorResponse( xerrors.Errorf("load chat: %w", err).Error(), @@ -79,7 +74,7 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { ), nil } - ws, err := options.DB.GetWorkspaceByID(ctx, chat.WorkspaceID.UUID) + ws, err := db.GetWorkspaceByID(ctx, chat.WorkspaceID.UUID) if err != nil { return fantasy.NewTextErrorResponse( xerrors.Errorf("load workspace: %w", err).Error(), @@ -91,14 +86,14 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { ), nil } - build, err := options.DB.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID) + build, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID) if err != nil { return fantasy.NewTextErrorResponse( xerrors.Errorf("get latest build: %w", err).Error(), ), nil } - job, err := options.DB.GetProvisionerJobByID(ctx, build.JobID) + job, err := db.GetProvisionerJobByID(ctx, build.JobID) if err != nil { return fantasy.NewTextErrorResponse( xerrors.Errorf("get provisioner job: %w", err).Error(), @@ -111,8 +106,8 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { database.ProvisionerJobStatusRunning: // Publish the build ID to the frontend so it // can start streaming logs immediately. - updatedChat, bindErr := options.DB.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ - ID: options.ChatID, + updatedChat, bindErr := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ + ID: chatID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, BuildID: uuid.NullUUID{ UUID: build.ID, @@ -122,14 +117,14 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { }) if bindErr != nil { options.Logger.Error(ctx, "failed to persist build ID on chat binding", - slog.F("chat_id", options.ChatID), + slog.F("chat_id", chatID), slog.F("build_id", build.ID), slog.Error(bindErr), ) } else if options.OnChatUpdated != nil { options.OnChatUpdated(updatedChat) } - if err := waitForBuild(ctx, options.DB, build.ID); err != nil { + if err := waitForBuild(ctx, db, build.ID); err != nil { // newBuildError returns via toolResponse (IsError: false) // rather than NewTextErrorResponse (IsError: true) so the // JSON result preserves build_id for the frontend's log @@ -141,13 +136,13 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { build.ID, )), nil } - result := waitForAgentAndRespond(ctx, options.DB, options.AgentConnFn, ws, build.ID) + result := waitForAgentAndRespond(ctx, db, options.AgentConnFn, ws, build.ID) // Re-fire after the agent is fully ready so // callers can load instruction files (AGENTS.md). // This must happen after waitForAgentAndRespond — // firing earlier races with agent startup. if options.OnChatUpdated != nil { - if latest, err := options.DB.GetChatByID(ctx, options.ChatID); err == nil { + if latest, err := db.GetChatByID(ctx, chatID); err == nil { options.OnChatUpdated(latest) } } @@ -156,7 +151,7 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { // If the latest successful build is a start // transition, the workspace should be running. if build.Transition == database.WorkspaceTransitionStart { - return toolResponse(waitForAgentAndRespond(ctx, options.DB, options.AgentConnFn, ws, uuid.Nil)), nil + return toolResponse(waitForAgentAndRespond(ctx, db, options.AgentConnFn, ws, uuid.Nil)), nil } // Otherwise it is stopped (or deleted) — proceed // to start it below. @@ -166,7 +161,7 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { } // Set up dbauthz context for the start call. - ownerCtx, ownerErr := asOwner(ctx, options.DB, options.OwnerID) + ownerCtx, ownerErr := asOwner(ctx, db, options.OwnerID) if ownerErr != nil { return fantasy.NewTextErrorResponse(ownerErr.Error()), nil } @@ -198,8 +193,8 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { // Persist the build ID on the chat binding so the // frontend can stream logs without polling. - updatedChat, bindErr := options.DB.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ - ID: options.ChatID, + updatedChat, bindErr := db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ + ID: chatID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, BuildID: uuid.NullUUID{ UUID: startBuild.ID, @@ -209,21 +204,21 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { }) if bindErr != nil { options.Logger.Error(ctx, "failed to persist build ID on chat binding", - slog.F("chat_id", options.ChatID), + slog.F("chat_id", chatID), slog.F("build_id", startBuild.ID), slog.Error(bindErr), ) } else if options.OnChatUpdated != nil { options.OnChatUpdated(updatedChat) } - if err := waitForBuild(ctx, options.DB, startBuild.ID); err != nil { + if err := waitForBuild(ctx, db, startBuild.ID); err != nil { return buildToolResponse(newBuildError( xerrors.Errorf("workspace start build failed: %w", err).Error(), startBuild.ID, )), nil } - result := waitForAgentAndRespond(ctx, options.DB, options.AgentConnFn, ws, startBuild.ID) + result := waitForAgentAndRespond(ctx, db, options.AgentConnFn, ws, startBuild.ID) // If the template version changed, annotate the // response so the model knows an auto-update @@ -241,7 +236,7 @@ func StartWorkspace(options StartWorkspaceOptions) fantasy.AgentTool { // This must happen after waitForAgentAndRespond — // firing earlier races with agent startup. if options.OnChatUpdated != nil { - if latest, err := options.DB.GetChatByID(ctx, options.ChatID); err == nil { + if latest, err := db.GetChatByID(ctx, chatID); err == nil { options.OnChatUpdated(latest) } } diff --git a/coderd/x/chatd/chattool/startworkspace_test.go b/coderd/x/chatd/chattool/startworkspace_test.go index c36aae5eca..bf1e1f966f 100644 --- a/coderd/x/chatd/chattool/startworkspace_test.go +++ b/coderd/x/chatd/chattool/startworkspace_test.go @@ -49,9 +49,7 @@ func TestStartWorkspace(t *testing.T) { Title: "test-no-workspace", }) - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, - ChatID: chat.ID, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called") return codersdk.WorkspaceBuild{}, nil @@ -96,10 +94,8 @@ func TestStartWorkspace(t *testing.T) { return nil, func() {}, nil } - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, AgentConnFn: agentConnFn, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for already-running workspace") @@ -178,10 +174,8 @@ func TestStartWorkspace(t *testing.T) { return nil, func() {}, nil } - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, AgentConnFn: agentConnFn, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for already-running workspace") @@ -231,10 +225,8 @@ func TestStartWorkspace(t *testing.T) { Title: "test-running-no-agent", }) - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { t.Fatal("AgentConnFn should not be called when no agents exist") return nil, func() {}, nil @@ -293,10 +285,8 @@ func TestStartWorkspace(t *testing.T) { Title: "test-running-selection-error", }) - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { t.Fatal("AgentConnFn should not be called when agent selection fails") return nil, func() {}, nil @@ -369,10 +359,8 @@ func TestStartWorkspace(t *testing.T) { return nil, func() {}, nil } - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, StartFn: startFn, AgentConnFn: agentConnFn, WorkspaceMu: &sync.Mutex{}, @@ -432,10 +420,8 @@ func TestStartWorkspace(t *testing.T) { }, nil } - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, StartFn: startFn, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil @@ -496,10 +482,8 @@ func TestStartWorkspace(t *testing.T) { return codersdk.WorkspaceBuild{ID: buildResp.Build.ID}, nil } - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, StartFn: startFn, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil @@ -543,10 +527,8 @@ func TestStartWorkspace(t *testing.T) { Title: "test-start-workspace-manual-update-required", }) - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { return codersdk.WorkspaceBuild{}, httperror.NewResponseError(400, codersdk.Response{ Message: "The workspace needs the template's active version before it can start. Use read_template with this workspace's template_id to inspect the active version's required parameters, then retry start_workspace with a parameters object that supplies any missing or changed values.", @@ -610,10 +592,8 @@ func TestStartWorkspace(t *testing.T) { Title: "test-start-workspace-responder-error-without-validations", }) - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { return codersdk.WorkspaceBuild{}, httperror.NewResponseError(502, codersdk.Response{ Message: "workspace start failed", @@ -677,10 +657,8 @@ func TestStartWorkspace(t *testing.T) { } var onChatUpdatedCalled atomic.Bool - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: wrappedDB, + tool := chattool.StartWorkspace(wrappedDB, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, AgentConnFn: agentConnFn, StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for an in-progress build") @@ -761,10 +739,8 @@ func TestStartWorkspace(t *testing.T) { jobRead := make(chan struct{}, 1) wrappedDB := &jobInterceptStore{Store: db, jobRead: jobRead} - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: wrappedDB, + tool := chattool.StartWorkspace(wrappedDB, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil }, @@ -856,10 +832,8 @@ func TestStartWorkspace(t *testing.T) { jobRead := make(chan struct{}, 2) wrappedDB := &jobInterceptStore{Store: db, jobRead: jobRead} - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: wrappedDB, + tool := chattool.StartWorkspace(wrappedDB, chat.ID, chattool.StartWorkspaceOptions{ OwnerID: user.ID, - ChatID: chat.ID, StartFn: startFn, AgentConnFn: func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil @@ -934,9 +908,7 @@ func TestStartWorkspace(t *testing.T) { Title: "test-deleted-workspace", }) - tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ - DB: db, - ChatID: chat.ID, + tool := chattool.StartWorkspace(db, chat.ID, chattool.StartWorkspaceOptions{ StartFn: func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { t.Fatal("StartFn should not be called for deleted workspace") return codersdk.WorkspaceBuild{}, nil