diff --git a/.gitignore b/.gitignore index 0732c41daa..abf76379f9 100644 --- a/.gitignore +++ b/.gitignore @@ -103,6 +103,6 @@ PLAN.md # Ignore any dev licenses license.txt --e + # Agent planning documents (local working files). docs/plans/ diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 82bc69f0ef..f8daa8f81a 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1550,6 +1550,7 @@ func Chat(c database.Chat, diffStatus *database.ChatDiffStatus, files []database } chat := codersdk.Chat{ ID: c.ID, + OrganizationID: c.OrganizationID, OwnerID: c.OwnerID, LastModelConfigID: c.LastModelConfigID, Title: c.Title, diff --git a/coderd/database/db2sdk/db2sdk_test.go b/coderd/database/db2sdk/db2sdk_test.go index 38584a6017..d7903d1d54 100644 --- a/coderd/database/db2sdk/db2sdk_test.go +++ b/coderd/database/db2sdk/db2sdk_test.go @@ -530,6 +530,7 @@ func TestChat_AllFieldsPopulated(t *testing.T) { input := database.Chat{ ID: uuid.New(), OwnerID: uuid.New(), + OrganizationID: uuid.New(), WorkspaceID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, BuildID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, AgentID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index c361282c75..0ace3ba527 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4846,7 +4846,7 @@ func (q *querier) InsertAuditLog(ctx context.Context, arg database.InsertAuditLo } func (q *querier) InsertChat(ctx context.Context, arg database.InsertChatParams) (database.Chat, error) { - return insert(q.log, q.auth, rbac.ResourceChat.WithOwner(arg.OwnerID.String()), q.db.InsertChat)(ctx, arg) + return insert(q.log, q.auth, rbac.ResourceChat.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), q.db.InsertChat)(ctx, arg) } func (q *querier) InsertChatFile(ctx context.Context, arg database.InsertChatFileParams) (database.InsertChatFileRow, error) { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index b156a84e2b..333ea10002 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -783,7 +783,7 @@ func (s *MethodTestSuite) TestChats() { }) chat := testutil.Fake(s.T(), faker, database.Chat{OwnerID: arg.OwnerID}) dbm.EXPECT().InsertChat(gomock.Any(), arg).Return(chat, nil).AnyTimes() - check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.OwnerID.String()), policy.ActionCreate).Returns(chat) + check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionCreate).Returns(chat) })) s.Run("InsertChatFile", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { arg := testutil.Fake(s.T(), faker, database.InsertChatFileParams{}) diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index c63e1daeb2..0659fc68d5 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -1670,9 +1670,10 @@ func TestDeleteOldChatFiles(t *testing.T) { // createChat inserts a chat and optionally archives it, then // backdates updated_at to control the "archived since" window. - createChat := func(ctx context.Context, t *testing.T, db database.Store, rawDB *sql.DB, ownerID, modelConfigID uuid.UUID, archived bool, updatedAt time.Time) database.Chat { + createChat := func(ctx context.Context, t *testing.T, db database.Store, rawDB *sql.DB, ownerID, orgID, modelConfigID uuid.UUID, archived bool, updatedAt time.Time) database.Chat { t.Helper() chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: orgID, OwnerID: ownerID, LastModelConfigID: modelConfigID, Title: "test-chat", @@ -1736,7 +1737,7 @@ func TestDeleteOldChatFiles(t *testing.T) { require.NoError(t, err) // Create an old archived chat and an orphaned old file. - oldChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) + oldChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) oldFileID := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour)) done := awaitDoTick(ctx, t, clk) @@ -1766,7 +1767,7 @@ func TestDeleteOldChatFiles(t *testing.T) { require.NoError(t, err) // Old archived chat (31 days) — should be deleted. - oldChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) + oldChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) // Insert a message so we can verify CASCADE. _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: oldChat.ID, @@ -1791,10 +1792,10 @@ func TestDeleteOldChatFiles(t *testing.T) { require.NoError(t, err) // Recently archived chat (10 days) — should be retained. - recentChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, true, now.Add(-10*24*time.Hour)) + recentChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-10*24*time.Hour)) // Active chat — should be retained. - activeChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, false, now) + activeChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) done := awaitDoTick(ctx, t, clk) closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, clk, prometheus.NewRegistry()) @@ -1839,7 +1840,7 @@ func TestDeleteOldChatFiles(t *testing.T) { // File B: 31 days old, in an active chat -> should be retained. fileB := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour)) - activeChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, false, now) + activeChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ ChatID: activeChat.ID, MaxFileLinks: 100, @@ -1887,7 +1888,7 @@ func TestDeleteOldChatFiles(t *testing.T) { // File D: 31 days old, in a chat archived 31 days ago -> should be deleted. fileD := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour)) - oldArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) + oldArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ ChatID: oldArchivedChat.ID, MaxFileLinks: 100, @@ -1901,7 +1902,7 @@ func TestDeleteOldChatFiles(t *testing.T) { // File E: 31 days old, in a chat archived 10 days ago -> should be retained. fileE := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour)) - recentArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, true, now.Add(-10*24*time.Hour)) + recentArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-10*24*time.Hour)) _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ ChatID: recentArchivedChat.ID, MaxFileLinks: 100, @@ -1914,7 +1915,7 @@ func TestDeleteOldChatFiles(t *testing.T) { // File F: 31 days old, in BOTH an active chat AND an old archived chat -> should be retained. fileF := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now.Add(-31*24*time.Hour)) - anotherOldArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) + anotherOldArchivedChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ ChatID: anotherOldArchivedChat.ID, MaxFileLinks: 100, @@ -1925,7 +1926,7 @@ func TestDeleteOldChatFiles(t *testing.T) { now.Add(-31*24*time.Hour), anotherOldArchivedChat.ID) require.NoError(t, err) - activeChatForF := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, false, now) + activeChatForF := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) _, err = db.LinkChatFiles(ctx, database.LinkChatFilesParams{ ChatID: activeChatForF.ID, MaxFileLinks: 100, @@ -1964,7 +1965,7 @@ func TestDeleteOldChatFiles(t *testing.T) { fileB := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now) fileC := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now) - chat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, false, now) + chat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) _, err := db.LinkChatFiles(ctx, database.LinkChatFilesParams{ ChatID: chat.ID, MaxFileLinks: 100, @@ -2009,14 +2010,16 @@ func TestDeleteOldChatFiles(t *testing.T) { // Test parent+child cascade: deleting files should // clean up links for both parent and child chats // independently via FK cascade. - parentChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, false, now) + parentChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) childChat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: deps.org.ID, OwnerID: deps.user.ID, LastModelConfigID: deps.modelConfig.ID, Title: "child-chat", Status: database.ChatStatusWaiting, }) require.NoError(t, err) + // Set root_chat_id to link child to parent. _, err = rawDB.ExecContext(ctx, "UPDATE chats SET root_chat_id = $1 WHERE id = $2", parentChat.ID, childChat.ID) require.NoError(t, err) @@ -2104,7 +2107,7 @@ func TestDeleteOldChatFiles(t *testing.T) { // Create 3 deletable old archived chats. for range 3 { - createChat(ctx, t, db, rawDB, deps.user.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) + createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) } // Delete with limit 2 — should delete 2, leave 1. diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 871f530641..e14330e5a5 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1431,7 +1431,8 @@ CREATE TABLE chats ( pin_order integer DEFAULT 0 NOT NULL, last_read_message_id bigint, last_injected_context jsonb, - dynamic_tools jsonb + dynamic_tools jsonb, + organization_id uuid NOT NULL ); CREATE TABLE connection_logs ( @@ -3789,6 +3790,8 @@ CREATE INDEX idx_chats_labels ON chats USING gin (labels); CREATE INDEX idx_chats_last_model_config_id ON chats USING btree (last_model_config_id); +CREATE INDEX idx_chats_organization_id ON chats USING btree (organization_id); + CREATE INDEX idx_chats_owner ON chats USING btree (owner_id); CREATE INDEX idx_chats_parent_chat_id ON chats USING btree (parent_chat_id); @@ -4104,6 +4107,9 @@ ALTER TABLE ONLY chats ALTER TABLE ONLY chats ADD CONSTRAINT chats_last_model_config_id_fkey FOREIGN KEY (last_model_config_id) REFERENCES chat_model_configs(id); +ALTER TABLE ONLY chats + ADD CONSTRAINT chats_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; + ALTER TABLE ONLY chats ADD CONSTRAINT chats_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE; diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index f682e0d760..18fd2908a9 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -25,6 +25,7 @@ const ( ForeignKeyChatsAgentID ForeignKeyConstraint = "chats_agent_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_agent_id_fkey FOREIGN KEY (agent_id) REFERENCES workspace_agents(id) ON DELETE SET NULL; ForeignKeyChatsBuildID ForeignKeyConstraint = "chats_build_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_build_id_fkey FOREIGN KEY (build_id) REFERENCES workspace_builds(id) ON DELETE SET NULL; ForeignKeyChatsLastModelConfigID ForeignKeyConstraint = "chats_last_model_config_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_last_model_config_id_fkey FOREIGN KEY (last_model_config_id) REFERENCES chat_model_configs(id); + ForeignKeyChatsOrganizationID ForeignKeyConstraint = "chats_organization_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; ForeignKeyChatsOwnerID ForeignKeyConstraint = "chats_owner_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyChatsParentChatID ForeignKeyConstraint = "chats_parent_chat_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_parent_chat_id_fkey FOREIGN KEY (parent_chat_id) REFERENCES chats(id) ON DELETE SET NULL; ForeignKeyChatsRootChatID ForeignKeyConstraint = "chats_root_chat_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_root_chat_id_fkey FOREIGN KEY (root_chat_id) REFERENCES chats(id) ON DELETE SET NULL; diff --git a/coderd/database/migrations/000467_chat_organization_id.down.sql b/coderd/database/migrations/000467_chat_organization_id.down.sql new file mode 100644 index 0000000000..3ba7d3848d --- /dev/null +++ b/coderd/database/migrations/000467_chat_organization_id.down.sql @@ -0,0 +1 @@ +ALTER TABLE chats DROP COLUMN organization_id; diff --git a/coderd/database/migrations/000467_chat_organization_id.up.sql b/coderd/database/migrations/000467_chat_organization_id.up.sql new file mode 100644 index 0000000000..a589219920 --- /dev/null +++ b/coderd/database/migrations/000467_chat_organization_id.up.sql @@ -0,0 +1,20 @@ +-- Step 1: Add nullable column with FK. +ALTER TABLE chats + ADD COLUMN organization_id UUID REFERENCES organizations(id) ON DELETE CASCADE; + +-- Step 2: Backfill from workspace org (primary path). Fall back to +-- user's oldest org membership, then default org for rows where +-- workspace_id was NULLed out by ON DELETE SET NULL or never set. +UPDATE chats c +SET organization_id = COALESCE( + (SELECT w.organization_id FROM workspaces w WHERE w.id = c.workspace_id), + (SELECT om.organization_id FROM organization_members om + WHERE om.user_id = c.owner_id ORDER BY om.created_at ASC LIMIT 1), + (SELECT id FROM organizations WHERE is_default = true LIMIT 1) +); + +-- Step 3: Enforce NOT NULL going forward. +ALTER TABLE chats ALTER COLUMN organization_id SET NOT NULL; + +-- Step 4: Index for efficient lookups by organization. +CREATE INDEX idx_chats_organization_id ON chats (organization_id); diff --git a/coderd/database/modelmethods.go b/coderd/database/modelmethods.go index 8d5ef4906d..b03e3729fd 100644 --- a/coderd/database/modelmethods.go +++ b/coderd/database/modelmethods.go @@ -176,7 +176,7 @@ func (t Task) RBACObject() rbac.Object { } func (c Chat) RBACObject() rbac.Object { - return rbac.ResourceChat.WithID(c.ID).WithOwner(c.OwnerID.String()) + return rbac.ResourceChat.WithID(c.ID).WithOwner(c.OwnerID.String()).InOrg(c.OrganizationID) } func (r GetChatsRow) RBACObject() rbac.Object { diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index 39ec442b39..77ddca13a1 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -799,6 +799,7 @@ func (q *sqlQuerier) GetAuthorizedChats(ctx context.Context, arg GetChatsParams, &i.Chat.LastReadMessageID, &i.Chat.LastInjectedContext, &i.Chat.DynamicTools, + &i.Chat.OrganizationID, &i.HasUnread); err != nil { return nil, err } diff --git a/coderd/database/models.go b/coderd/database/models.go index 14f4ac8457..3a8b88a392 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4246,6 +4246,7 @@ type Chat struct { LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"` DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` } type ChatDiffStatus struct { diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 20b5635faf..c7d24e35b5 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -1258,6 +1258,11 @@ func TestGetAuthorizedChats(t *testing.T) { RBACRoles: pq.StringArray{rbac.RoleAgentsAccess().String()}, }) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: owner.ID, OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: member.ID, OrganizationID: org.ID}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: secondMember.ID, OrganizationID: org.ID}) + // Create FK dependencies: a chat provider and model config. ctx := testutil.Context(t, testutil.WaitMedium) _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ @@ -1286,6 +1291,7 @@ func TestGetAuthorizedChats(t *testing.T) { // Create 3 chats owned by owner. for i := range 3 { _, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: owner.ID, LastModelConfigID: modelCfg.ID, @@ -1297,6 +1303,7 @@ func TestGetAuthorizedChats(t *testing.T) { // Create 2 chats owned by member. for i := range 2 { _, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: member.ID, LastModelConfigID: modelCfg.ID, @@ -1340,8 +1347,8 @@ func TestGetAuthorizedChats(t *testing.T) { require.NoError(t, err) require.Len(t, secondRows, 0) - // Org admin should NOT see other users' chats — chats are - // not org-scoped resources. + // Org admin should NOT see other users' chats when they are + // in a different org than the chat owner. orgs, err := db.GetOrganizations(ctx, database.GetOrganizationsParams{}) require.NoError(t, err) require.NotEmpty(t, orgs) @@ -1359,6 +1366,21 @@ func TestGetAuthorizedChats(t *testing.T) { require.NoError(t, err) require.Len(t, orgAdminRows, 0, "org admin with no chats should see 0 chats") + // Org admin in SAME org should see all chats in that org. + sameOrgAdmin := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: sameOrgAdmin.ID, + OrganizationID: org.ID, + Roles: []string{rbac.RoleOrgAdmin()}, + }) + sameOrgAdminSubject, _, err := httpmw.UserRBACSubject(ctx, db, sameOrgAdmin.ID, rbac.ExpandableScope(rbac.ScopeAll)) + require.NoError(t, err) + preparedSameOrgAdmin, err := authorizer.Prepare(ctx, sameOrgAdminSubject, policy.ActionRead, rbac.ResourceChat.Type) + require.NoError(t, err) + sameOrgAdminRows, err := db.GetAuthorizedChats(ctx, database.GetChatsParams{}, preparedSameOrgAdmin) + require.NoError(t, err) + require.GreaterOrEqual(t, len(sameOrgAdminRows), 5, "same-org admin should see all chats in their org") + // OwnerID filter: member queries their own chats. memberFilterSelf, err := db.GetAuthorizedChats(ctx, database.GetChatsParams{ OwnerID: member.ID, @@ -1417,8 +1439,10 @@ func TestGetAuthorizedChats(t *testing.T) { paginationUser := dbgen.User(t, db, database.User{ RBACRoles: pq.StringArray{rbac.RoleAgentsAccess().String()}, }) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: paginationUser.ID, OrganizationID: org.ID}) for i := range 7 { _, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: paginationUser.ID, LastModelConfigID: modelCfg.ID, @@ -9753,8 +9777,9 @@ func TestInsertChatMessages(t *testing.T) { store, _ := dbtestutil.NewDB(t) ctx := context.Background() - dbgen.Organization(t, store, database.Organization{}) + org := dbgen.Organization(t, store, database.Organization{}) user := dbgen.User(t, store, database.User{}) + dbgen.OrganizationMember(t, store, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) provider := "openai" _, err := store.InsertChatProvider(ctx, database.InsertChatProviderParams{ @@ -9778,6 +9803,7 @@ func TestInsertChatMessages(t *testing.T) { ) chat, err := store.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: modelConfigA.ID, @@ -9921,6 +9947,8 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { // Helper: create a chat model config (required FK for chats). user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) // A chat_providers row is required as a FK for model configs. _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ @@ -9949,6 +9977,7 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { newChat := func(t *testing.T) database.Chat { t.Helper() chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, @@ -10288,12 +10317,13 @@ func TestGetPRInsights(t *testing.T) { // setupChatInfra creates a fresh database with a user, chat provider, // and model config. Returns the store, user ID, and model config ID. - setupChatInfra := func(t *testing.T) (database.Store, uuid.UUID, uuid.UUID) { + setupChatInfra := func(t *testing.T) (database.Store, uuid.UUID, uuid.UUID, uuid.UUID) { t.Helper() store, _ := dbtestutil.NewDB(t) ctx := context.Background() - dbgen.Organization(t, store, database.Organization{}) + org := dbgen.Organization(t, store, database.Organization{}) user := dbgen.User(t, store, database.User{}) + dbgen.OrganizationMember(t, store, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) _, err := store.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: "anthropic", @@ -10318,12 +10348,13 @@ func TestGetPRInsights(t *testing.T) { }) require.NoError(t, err) - return store, user.ID, mc.ID + return store, user.ID, mc.ID, org.ID } - createChat := func(t *testing.T, store database.Store, userID, mcID uuid.UUID, title string) database.Chat { + createChat := func(t *testing.T, store database.Store, userID, mcID, orgID uuid.UUID, title string) database.Chat { t.Helper() chat, err := store.InsertChat(context.Background(), database.InsertChatParams{ + OrganizationID: orgID, Status: database.ChatStatusWaiting, OwnerID: userID, LastModelConfigID: mcID, @@ -10384,12 +10415,12 @@ func TestGetPRInsights(t *testing.T) { t.Run("MultipleChatsSamePR_CostSummed", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) - chatA := createChat(t, store, userID, mcID, "chat-A") + chatA := createChat(t, store, userID, mcID, orgID, "chat-A") insertCostMessage(t, store, chatA.ID, userID, mcID, 5_000_000) // $5 - chatB := createChat(t, store, userID, mcID, "chat-B") + chatB := createChat(t, store, userID, mcID, orgID, "chat-B") insertCostMessage(t, store, chatB.ID, userID, mcID, 3_000_000) // $3 prURL := "https://github.com/org/repo/pull/123" @@ -10420,13 +10451,13 @@ func TestGetPRInsights(t *testing.T) { t.Run("DifferentPRs_NoDuplication", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) - chatA := createChat(t, store, userID, mcID, "chat-A") + chatA := createChat(t, store, userID, mcID, orgID, "chat-A") insertCostMessage(t, store, chatA.ID, userID, mcID, 5_000_000) linkPR(t, store, chatA.ID, "https://github.com/org/repo/pull/1", "merged", "feat: A", 50, 10, 2) - chatB := createChat(t, store, userID, mcID, "chat-B") + chatB := createChat(t, store, userID, mcID, orgID, "chat-B") insertCostMessage(t, store, chatB.ID, userID, mcID, 3_000_000) linkPR(t, store, chatB.ID, "https://github.com/org/repo/pull/2", "open", "feat: B", 80, 30, 4) @@ -10455,9 +10486,10 @@ func TestGetPRInsights(t *testing.T) { // createChildChat creates a chat with ParentChatID and RootChatID // set, simulating a subagent/child chat in a tree. - createChildChat := func(t *testing.T, store database.Store, userID, mcID, parentID, rootID uuid.UUID, title string) database.Chat { + createChildChat := func(t *testing.T, store database.Store, userID, mcID, orgID, parentID, rootID uuid.UUID, title string) database.Chat { t.Helper() chat, err := store.InsertChat(context.Background(), database.InsertChatParams{ + OrganizationID: orgID, Status: database.ChatStatusWaiting, OwnerID: userID, LastModelConfigID: mcID, @@ -10471,11 +10503,11 @@ func TestGetPRInsights(t *testing.T) { t.Run("DuplicatePRUrl_CountedOnce", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) prURL := "https://github.com/org/repo/pull/99" for i := 0; i < 3; i++ { - chat := createChat(t, store, userID, mcID, fmt.Sprintf("chat-%d", i)) + chat := createChat(t, store, userID, mcID, orgID, fmt.Sprintf("chat-%d", i)) insertCostMessage(t, store, chat.ID, userID, mcID, 1_000_000) linkPR(t, store, chat.ID, prURL, "merged", "fix: same PR", 40, 10, 3) } @@ -10500,19 +10532,19 @@ func TestGetPRInsights(t *testing.T) { t.Run("ChildChatCostsIncluded", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // Parent chat with a $5 cost. - parent := createChat(t, store, userID, mcID, "parent-chat") + parent := createChat(t, store, userID, mcID, orgID, "parent-chat") insertCostMessage(t, store, parent.ID, userID, mcID, 5_000_000) // Two child chats (subagents) with $2 each. Only the parent // has a chat_diff_statuses entry, but the children's costs // should be included via the tree join. - child1 := createChildChat(t, store, userID, mcID, parent.ID, parent.ID, "child-1") + child1 := createChildChat(t, store, userID, mcID, orgID, parent.ID, parent.ID, "child-1") insertCostMessage(t, store, child1.ID, userID, mcID, 2_000_000) - child2 := createChildChat(t, store, userID, mcID, parent.ID, parent.ID, "child-2") + child2 := createChildChat(t, store, userID, mcID, orgID, parent.ID, parent.ID, "child-2") insertCostMessage(t, store, child2.ID, userID, mcID, 2_000_000) prURL := "https://github.com/org/repo/pull/42" @@ -10542,19 +10574,19 @@ func TestGetPRInsights(t *testing.T) { t.Run("SiblingPRs_NoCrossContamination", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // Parent chat with $10 orchestration cost. - parent := createChat(t, store, userID, mcID, "parent") + parent := createChat(t, store, userID, mcID, orgID, "parent") insertCostMessage(t, store, parent.ID, userID, mcID, 10_000_000) // Child C1 ($5) creates PR1. - c1 := createChildChat(t, store, userID, mcID, parent.ID, parent.ID, "child-1") + c1 := createChildChat(t, store, userID, mcID, orgID, parent.ID, parent.ID, "child-1") insertCostMessage(t, store, c1.ID, userID, mcID, 5_000_000) linkPR(t, store, c1.ID, "https://github.com/org/repo/pull/10", "merged", "feat: PR1", 50, 10, 2) // Child C2 ($3) creates PR2. - c2 := createChildChat(t, store, userID, mcID, parent.ID, parent.ID, "child-2") + c2 := createChildChat(t, store, userID, mcID, orgID, parent.ID, parent.ID, "child-2") insertCostMessage(t, store, c2.ID, userID, mcID, 3_000_000) linkPR(t, store, c2.ID, "https://github.com/org/repo/pull/11", "open", "feat: PR2", 30, 5, 1) @@ -10585,23 +10617,23 @@ func TestGetPRInsights(t *testing.T) { t.Run("ParentAndChildDifferentPRs_NoCrossContamination", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // Parent P ($10) creates PR1. - parent := createChat(t, store, userID, mcID, "parent") + parent := createChat(t, store, userID, mcID, orgID, "parent") insertCostMessage(t, store, parent.ID, userID, mcID, 10_000_000) linkPR(t, store, parent.ID, "https://github.com/org/repo/pull/20", "merged", "feat: parent PR", 80, 20, 4) // Child C1 ($5) has its own PR2. Because C1 has its own // chat_diff_statuses entry, its cost should NOT be included // under PR1 — it belongs to PR2 only. - c1 := createChildChat(t, store, userID, mcID, parent.ID, parent.ID, "child-1") + c1 := createChildChat(t, store, userID, mcID, orgID, parent.ID, parent.ID, "child-1") insertCostMessage(t, store, c1.ID, userID, mcID, 5_000_000) linkPR(t, store, c1.ID, "https://github.com/org/repo/pull/21", "open", "feat: child PR", 30, 5, 1) // Child C2 ($2) has NO cds entry — pure subagent. // Its cost should be included under PR1 (the parent's PR). - c2 := createChildChat(t, store, userID, mcID, parent.ID, parent.ID, "child-2") + c2 := createChildChat(t, store, userID, mcID, orgID, parent.ID, parent.ID, "child-2") insertCostMessage(t, store, c2.ID, userID, mcID, 2_000_000) // PR1 cost = parent ($10) + C2 ($2) = $12 (C1 excluded) @@ -10630,16 +10662,16 @@ func TestGetPRInsights(t *testing.T) { t.Run("EmptyURLNotCollapsed", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // Two chats with empty-string URLs should be treated as // separate PRs (NULLIF converts '' to NULL, falling back // to c.id::text). - chatX := createChat(t, store, userID, mcID, "chat-X") + chatX := createChat(t, store, userID, mcID, orgID, "chat-X") insertCostMessage(t, store, chatX.ID, userID, mcID, 4_000_000) linkPR(t, store, chatX.ID, "", "open", "draft: X", 10, 2, 1) - chatY := createChat(t, store, userID, mcID, "chat-Y") + chatY := createChat(t, store, userID, mcID, orgID, "chat-Y") insertCostMessage(t, store, chatY.ID, userID, mcID, 6_000_000) linkPR(t, store, chatY.ID, "", "merged", "draft: Y", 20, 5, 2) @@ -10663,14 +10695,14 @@ func TestGetPRInsights(t *testing.T) { t.Run("ParentAndChildSameURL_DedupedWithCombinedCost", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // Parent P ($10) links to a PR. - parent := createChat(t, store, userID, mcID, "parent") + parent := createChat(t, store, userID, mcID, orgID, "parent") insertCostMessage(t, store, parent.ID, userID, mcID, 10_000_000) // Child C ($5) also links to the same PR URL. - child := createChildChat(t, store, userID, mcID, parent.ID, parent.ID, "child") + child := createChildChat(t, store, userID, mcID, orgID, parent.ID, parent.ID, "child") insertCostMessage(t, store, child.ID, userID, mcID, 5_000_000) prURL := "https://github.com/org/repo/pull/50" @@ -10700,11 +10732,11 @@ func TestGetPRInsights(t *testing.T) { t.Run("ZeroCostChat_StillCounted", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // A chat linked to a PR but with NO chat_messages at all. // The PR should still appear with zero cost. - chat := createChat(t, store, userID, mcID, "zero-cost-chat") + chat := createChat(t, store, userID, mcID, orgID, "zero-cost-chat") linkPR(t, store, chat.ID, "https://github.com/org/repo/pull/60", "open", "feat: no messages", 25, 5, 2) summary, err := store.GetPRInsightsSummary(context.Background(), database.GetPRInsightsSummaryParams{ @@ -10728,7 +10760,7 @@ func TestGetPRInsights(t *testing.T) { t.Run("BlankDisplayNameFallsBackToModel", func(t *testing.T) { t.Parallel() - store, userID, _ := setupChatInfra(t) + store, userID, _, orgID := setupChatInfra(t) const modelName = "claude-4.1" emptyDisplayModel, err := store.InsertChatModelConfig(context.Background(), database.InsertChatModelConfigParams{ @@ -10745,7 +10777,7 @@ func TestGetPRInsights(t *testing.T) { }) require.NoError(t, err) - chat := createChat(t, store, userID, emptyDisplayModel.ID, "chat-empty-display-name") + chat := createChat(t, store, userID, emptyDisplayModel.ID, orgID, "chat-empty-display-name") insertCostMessage(t, store, chat.ID, userID, emptyDisplayModel.ID, 1_000_000) linkPR(t, store, chat.ID, "https://github.com/org/repo/pull/72", "merged", "fix: blank display name", 10, 2, 1) @@ -10770,15 +10802,15 @@ func TestGetPRInsights(t *testing.T) { t.Run("MergedCostMicros_OnlyCountsMerged", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // Merged PR with $5 cost. - chatMerged := createChat(t, store, userID, mcID, "chat-merged") + chatMerged := createChat(t, store, userID, mcID, orgID, "chat-merged") insertCostMessage(t, store, chatMerged.ID, userID, mcID, 5_000_000) linkPR(t, store, chatMerged.ID, "https://github.com/org/repo/pull/70", "merged", "fix: merged", 40, 10, 2) // Open PR with $3 cost. - chatOpen := createChat(t, store, userID, mcID, "chat-open") + chatOpen := createChat(t, store, userID, mcID, orgID, "chat-open") insertCostMessage(t, store, chatOpen.ID, userID, mcID, 3_000_000) linkPR(t, store, chatOpen.ID, "https://github.com/org/repo/pull/71", "open", "feat: open", 20, 5, 1) @@ -10796,13 +10828,13 @@ func TestGetPRInsights(t *testing.T) { t.Run("AllPRsReturnedWithSafetyCap", func(t *testing.T) { t.Parallel() - store, userID, mcID := setupChatInfra(t) + store, userID, mcID, orgID := setupChatInfra(t) // Create 25 distinct PRs — more than the old LIMIT 20 — and // verify all are returned. const prCount = 25 for i := range prCount { - chat := createChat(t, store, userID, mcID, fmt.Sprintf("chat-%d", i)) + chat := createChat(t, store, userID, mcID, orgID, fmt.Sprintf("chat-%d", i)) insertCostMessage(t, store, chat.ID, userID, mcID, 1_000_000) linkPR(t, store, chat.ID, fmt.Sprintf("https://github.com/org/repo/pull/%d", 100+i), @@ -10825,11 +10857,13 @@ func TestChatPinOrderQueries(t *testing.T) { t.SkipNow() } - setup := func(t *testing.T) (context.Context, database.Store, uuid.UUID, uuid.UUID) { + setup := func(t *testing.T) (context.Context, database.Store, uuid.UUID, uuid.UUID, uuid.UUID) { t.Helper() db, _ := dbtestutil.NewDB(t) + org := dbgen.Organization(t, db, database.Organization{}) owner := dbgen.User(t, db, database.User{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: owner.ID, OrganizationID: org.ID}) // Use background context for fixture setup so the // timed test context doesn't tick during DB init. @@ -10858,13 +10892,14 @@ func TestChatPinOrderQueries(t *testing.T) { require.NoError(t, err) ctx := testutil.Context(t, testutil.WaitMedium) - return ctx, db, owner.ID, modelCfg.ID + return ctx, db, owner.ID, modelCfg.ID, org.ID } - createChat := func(t *testing.T, ctx context.Context, db database.Store, ownerID, modelCfgID uuid.UUID, title string) database.Chat { + createChat := func(t *testing.T, ctx context.Context, db database.Store, ownerID, modelCfgID, orgID uuid.UUID, title string) database.Chat { t.Helper() chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: orgID, Status: database.ChatStatusWaiting, OwnerID: ownerID, LastModelConfigID: modelCfgID, @@ -10887,13 +10922,13 @@ func TestChatPinOrderQueries(t *testing.T) { t.Run("PinChatByIDAppendsWithinOwner", func(t *testing.T) { t.Parallel() - ctx, db, ownerID, modelCfgID := setup(t) - first := createChat(t, ctx, db, ownerID, modelCfgID, "first") - second := createChat(t, ctx, db, ownerID, modelCfgID, "second") - third := createChat(t, ctx, db, ownerID, modelCfgID, "third") + ctx, db, ownerID, modelCfgID, orgID := setup(t) + first := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "first") + second := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "second") + third := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "third") otherOwner := dbgen.User(t, db, database.User{}) - other := createChat(t, ctx, db, otherOwner.ID, modelCfgID, "other-owner") + other := createChat(t, ctx, db, otherOwner.ID, modelCfgID, orgID, "other-owner") require.NoError(t, db.PinChatByID(ctx, other.ID)) require.NoError(t, db.PinChatByID(ctx, first.ID)) @@ -10911,10 +10946,10 @@ func TestChatPinOrderQueries(t *testing.T) { t.Run("UpdateChatPinOrderShiftsNeighborsAndClamps", func(t *testing.T) { t.Parallel() - ctx, db, ownerID, modelCfgID := setup(t) - first := createChat(t, ctx, db, ownerID, modelCfgID, "first") - second := createChat(t, ctx, db, ownerID, modelCfgID, "second") - third := createChat(t, ctx, db, ownerID, modelCfgID, "third") + ctx, db, ownerID, modelCfgID, orgID := setup(t) + first := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "first") + second := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "second") + third := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "third") for _, chat := range []database.Chat{first, second, third} { require.NoError(t, db.PinChatByID(ctx, chat.ID)) @@ -10944,10 +10979,10 @@ func TestChatPinOrderQueries(t *testing.T) { t.Run("UnpinChatByIDCompactsPinnedChats", func(t *testing.T) { t.Parallel() - ctx, db, ownerID, modelCfgID := setup(t) - first := createChat(t, ctx, db, ownerID, modelCfgID, "first") - second := createChat(t, ctx, db, ownerID, modelCfgID, "second") - third := createChat(t, ctx, db, ownerID, modelCfgID, "third") + ctx, db, ownerID, modelCfgID, orgID := setup(t) + first := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "first") + second := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "second") + third := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "third") for _, chat := range []database.Chat{first, second, third} { require.NoError(t, db.PinChatByID(ctx, chat.ID)) @@ -10964,10 +10999,10 @@ func TestChatPinOrderQueries(t *testing.T) { t.Run("ArchiveClearsPinAndExcludesFromRanking", func(t *testing.T) { t.Parallel() - ctx, db, ownerID, modelCfgID := setup(t) - first := createChat(t, ctx, db, ownerID, modelCfgID, "first") - second := createChat(t, ctx, db, ownerID, modelCfgID, "second") - third := createChat(t, ctx, db, ownerID, modelCfgID, "third") + ctx, db, ownerID, modelCfgID, orgID := setup(t) + first := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "first") + second := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "second") + third := createChat(t, ctx, db, ownerID, modelCfgID, orgID, "third") for _, chat := range []database.Chat{first, second, third} { require.NoError(t, db.PinChatByID(ctx, chat.ID)) @@ -11014,6 +11049,8 @@ func TestChatLabels(t *testing.T) { ctx := testutil.Context(t, testutil.WaitMedium) owner := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: owner.ID, OrganizationID: org.ID}) _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: "openai", @@ -11047,6 +11084,7 @@ func TestChatLabels(t *testing.T) { require.NoError(t, err) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: owner.ID, LastModelConfigID: modelCfg.ID, @@ -11070,6 +11108,7 @@ func TestChatLabels(t *testing.T) { ctx := testutil.Context(t, testutil.WaitMedium) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: owner.ID, LastModelConfigID: modelCfg.ID, @@ -11086,6 +11125,7 @@ func TestChatLabels(t *testing.T) { ctx := testutil.Context(t, testutil.WaitMedium) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: owner.ID, LastModelConfigID: modelCfg.ID, @@ -11127,6 +11167,7 @@ func TestChatLabels(t *testing.T) { require.NoError(t, err) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: owner.ID, LastModelConfigID: modelCfg.ID, @@ -11164,10 +11205,10 @@ func TestChatLabels(t *testing.T) { labelsJSON, err := json.Marshal(tc.labels) require.NoError(t, err) _, err = db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: owner.ID, - LastModelConfigID: modelCfg.ID, - Title: tc.title, + LastModelConfigID: modelCfg.ID, Title: tc.title, Labels: pqtype.NullRawMessage{ RawMessage: labelsJSON, Valid: true, @@ -11224,8 +11265,9 @@ func TestChatHasUnread(t *testing.T) { store, _ := dbtestutil.NewDB(t) ctx := context.Background() - dbgen.Organization(t, store, database.Organization{}) + org := dbgen.Organization(t, store, database.Organization{}) user := dbgen.User(t, store, database.User{}) + dbgen.OrganizationMember(t, store, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) _, err := store.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: "openai", @@ -11251,6 +11293,7 @@ func TestChatHasUnread(t *testing.T) { require.NoError(t, err) chat, err := store.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index b5381ffe62..d383f13a47 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -4239,7 +4239,7 @@ WHERE $3::int ) RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type AcquireChatsParams struct { @@ -4284,6 +4284,7 @@ func (q *sqlQuerier) AcquireChats(ctx context.Context, arg AcquireChatsParams) ( &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ); err != nil { return nil, err } @@ -4422,9 +4423,9 @@ WITH chats AS ( UPDATE chats SET archived = true, pin_order = 0, updated_at = NOW() WHERE id = $1::uuid OR root_chat_id = $1::uuid - RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id FROM chats ORDER BY (id = $1::uuid) DESC, created_at ASC, id ASC ` @@ -4463,6 +4464,7 @@ func (q *sqlQuerier) ArchiveChatByID(ctx context.Context, id uuid.UUID) ([]Chat, &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ); err != nil { return nil, err } @@ -4612,7 +4614,7 @@ func (q *sqlQuerier) DeleteOldChats(ctx context.Context, arg DeleteOldChatsParam } const getActiveChatsByAgentID = `-- name: GetActiveChatsByAgentID :many -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id FROM chats WHERE agent_id = $1::uuid AND archived = false @@ -4657,6 +4659,7 @@ func (q *sqlQuerier) GetActiveChatsByAgentID(ctx context.Context, agentID uuid.U &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ); err != nil { return nil, err } @@ -4673,7 +4676,7 @@ func (q *sqlQuerier) GetActiveChatsByAgentID(ctx context.Context, agentID uuid.U const getChatByID = `-- name: GetChatByID :one SELECT - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id FROM chats WHERE @@ -4708,12 +4711,13 @@ func (q *sqlQuerier) GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } const getChatByIDForUpdate = `-- name: GetChatByIDForUpdate :one -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools FROM chats WHERE id = $1::uuid FOR UPDATE +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id FROM chats WHERE id = $1::uuid FOR UPDATE ` func (q *sqlQuerier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Chat, error) { @@ -4744,6 +4748,7 @@ func (q *sqlQuerier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Ch &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -5798,7 +5803,7 @@ func (q *sqlQuerier) GetChatUsageLimitUserOverride(ctx context.Context, userID u const getChats = `-- name: GetChats :many SELECT - chats.id, chats.owner_id, chats.workspace_id, chats.title, chats.status, chats.worker_id, chats.started_at, chats.heartbeat_at, chats.created_at, chats.updated_at, chats.parent_chat_id, chats.root_chat_id, chats.last_model_config_id, chats.archived, chats.last_error, chats.mode, chats.mcp_server_ids, chats.labels, chats.build_id, chats.agent_id, chats.pin_order, chats.last_read_message_id, chats.last_injected_context, chats.dynamic_tools, + chats.id, chats.owner_id, chats.workspace_id, chats.title, chats.status, chats.worker_id, chats.started_at, chats.heartbeat_at, chats.created_at, chats.updated_at, chats.parent_chat_id, chats.root_chat_id, chats.last_model_config_id, chats.archived, chats.last_error, chats.mode, chats.mcp_server_ids, chats.labels, chats.build_id, chats.agent_id, chats.pin_order, chats.last_read_message_id, chats.last_injected_context, chats.dynamic_tools, chats.organization_id, EXISTS ( SELECT 1 FROM chat_messages cm WHERE cm.chat_id = chats.id @@ -5911,6 +5916,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha &i.Chat.LastReadMessageID, &i.Chat.LastInjectedContext, &i.Chat.DynamicTools, + &i.Chat.OrganizationID, &i.HasUnread, ); err != nil { return nil, err @@ -5927,7 +5933,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha } const getChatsByWorkspaceIDs = `-- name: GetChatsByWorkspaceIDs :many -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id FROM chats WHERE archived = false AND workspace_id = ANY($1::uuid[]) @@ -5968,6 +5974,7 @@ func (q *sqlQuerier) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ); err != nil { return nil, err } @@ -6095,7 +6102,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh const getStaleChats = `-- name: GetStaleChats :many SELECT - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id FROM chats WHERE @@ -6143,6 +6150,7 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ); err != nil { return nil, err } @@ -6199,6 +6207,7 @@ func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUI const insertChat = `-- name: InsertChat :one INSERT INTO chats ( + organization_id, owner_id, workspace_id, build_id, @@ -6220,18 +6229,20 @@ INSERT INTO chats ( $5::uuid, $6::uuid, $7::uuid, - $8::text, - $9::chat_mode, - $10::chat_status, - COALESCE($11::uuid[], '{}'::uuid[]), - COALESCE($12::jsonb, '{}'::jsonb), - $13::jsonb + $8::uuid, + $9::text, + $10::chat_mode, + $11::chat_status, + COALESCE($12::uuid[], '{}'::uuid[]), + COALESCE($13::jsonb, '{}'::jsonb), + $14::jsonb ) RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type InsertChatParams struct { + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` BuildID uuid.NullUUID `db:"build_id" json:"build_id"` @@ -6249,6 +6260,7 @@ type InsertChatParams struct { func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat, error) { row := q.db.QueryRowContext(ctx, insertChat, + arg.OrganizationID, arg.OwnerID, arg.WorkspaceID, arg.BuildID, @@ -6289,6 +6301,7 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -6796,9 +6809,9 @@ WITH chats AS ( archived = false, updated_at = NOW() WHERE id = $1::uuid OR root_chat_id = $1::uuid - RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id FROM chats ORDER BY (id = $1::uuid) DESC, created_at ASC, id ASC ` @@ -6841,6 +6854,7 @@ func (q *sqlQuerier) UnarchiveChatByID(ctx context.Context, id uuid.UUID) ([]Cha &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ); err != nil { return nil, err } @@ -6921,7 +6935,7 @@ UPDATE chats SET updated_at = NOW() WHERE id = $3::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatBuildAgentBindingParams struct { @@ -6958,6 +6972,7 @@ func (q *sqlQuerier) UpdateChatBuildAgentBinding(ctx context.Context, arg Update &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -6971,7 +6986,7 @@ SET WHERE id = $2::uuid RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatByIDParams struct { @@ -7007,6 +7022,7 @@ func (q *sqlQuerier) UpdateChatByID(ctx context.Context, arg UpdateChatByIDParam &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -7065,7 +7081,7 @@ SET WHERE id = $2::uuid RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatLabelsByIDParams struct { @@ -7101,6 +7117,7 @@ func (q *sqlQuerier) UpdateChatLabelsByID(ctx context.Context, arg UpdateChatLab &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -7110,7 +7127,7 @@ UPDATE chats SET last_injected_context = $1::jsonb WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatLastInjectedContextParams struct { @@ -7150,6 +7167,7 @@ func (q *sqlQuerier) UpdateChatLastInjectedContext(ctx context.Context, arg Upda &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -7163,7 +7181,7 @@ SET WHERE id = $2::uuid RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatLastModelConfigByIDParams struct { @@ -7199,6 +7217,7 @@ func (q *sqlQuerier) UpdateChatLastModelConfigByID(ctx context.Context, arg Upda &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -7230,7 +7249,7 @@ SET WHERE id = $2::uuid RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatMCPServerIDsParams struct { @@ -7266,6 +7285,7 @@ func (q *sqlQuerier) UpdateChatMCPServerIDs(ctx context.Context, arg UpdateChatM &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -7401,7 +7421,7 @@ SET WHERE id = $6::uuid RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatStatusParams struct { @@ -7448,6 +7468,7 @@ func (q *sqlQuerier) UpdateChatStatus(ctx context.Context, arg UpdateChatStatusP &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -7465,7 +7486,7 @@ SET WHERE id = $7::uuid RETURNING - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatStatusPreserveUpdatedAtParams struct { @@ -7514,6 +7535,7 @@ func (q *sqlQuerier) UpdateChatStatusPreserveUpdatedAt(ctx context.Context, arg &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } @@ -7525,7 +7547,7 @@ UPDATE chats SET agent_id = $3::uuid, updated_at = NOW() WHERE id = $4::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id ` type UpdateChatWorkspaceBindingParams struct { @@ -7568,6 +7590,7 @@ func (q *sqlQuerier) UpdateChatWorkspaceBinding(ctx context.Context, arg UpdateC &i.LastReadMessageID, &i.LastInjectedContext, &i.DynamicTools, + &i.OrganizationID, ) return i, err } diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index a47d4bcc9a..caa3552b3f 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -392,6 +392,7 @@ LIMIT -- name: InsertChat :one INSERT INTO chats ( + organization_id, owner_id, workspace_id, build_id, @@ -406,6 +407,7 @@ INSERT INTO chats ( labels, dynamic_tools ) VALUES ( + @organization_id::uuid, @owner_id::uuid, sqlc.narg('workspace_id')::uuid, sqlc.narg('build_id')::uuid, diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index dee223cab3..23b86c7856 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -398,6 +398,33 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { return } + // Validate organization membership. + if req.OrganizationID == uuid.Nil { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "organization_id is required.", + }) + return + } + orgMembers, err := api.Database.OrganizationMembers(ctx, database.OrganizationMembersParams{ + OrganizationID: req.OrganizationID, + UserID: apiKey.UserID, + IncludeSystem: false, + GithubUserID: 0, + }) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to validate organization membership.", + Detail: err.Error(), + }) + return + } + if len(orgMembers) == 0 { + httpapi.Write(ctx, rw, http.StatusForbidden, codersdk.Response{ + Message: "You are not a member of the specified organization.", + }) + return + } + // Validate per-chat system prompt length. const maxSystemPromptLen = 10000 if len(req.SystemPrompt) > maxSystemPromptLen { @@ -407,7 +434,6 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { }) return } - contentBlocks, titleSource, fileIDs, inputError := createChatInputFromRequest(ctx, api.Database, req) if inputError != nil { httpapi.Write(ctx, rw, http.StatusBadRequest, *inputError) @@ -528,6 +554,7 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { } chat, err := api.chatDaemon.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: req.OrganizationID, OwnerID: apiKey.UserID, WorkspaceID: workspaceSelection.WorkspaceID, Title: title, @@ -1847,6 +1874,18 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { chat := httpmw.ChatParam(r) chatID := chat.ID + // Gate message sending behind the same agents-access check + // used by postChats. Sending a message triggers AI/LLM + // inference, so it should require the same authorization as + // chat creation. This is a handler-level band-aid; the + // structural fix is to make agents-access org-aware so + // dbauthz enforces this at the RBAC layer. + // See: https://github.com/coder/coder/issues/24250 + if !api.Authorize(r, policy.ActionCreate, rbac.ResourceChat.WithOwner(apiKey.UserID.String())) { + httpapi.Forbidden(rw) + return + } + if api.chatDaemon == nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Chat processor is unavailable.", @@ -2091,6 +2130,15 @@ func (api *API) promoteChatQueuedMessage(rw http.ResponseWriter, r *http.Request chat := httpmw.ChatParam(r) chatID := chat.ID + // Gate queued-message promotion behind agents-access. + // Promoting a queued message triggers AI/LLM inference, + // same as sending a new message. + // See: https://github.com/coder/coder/issues/24250 + if !api.Authorize(r, policy.ActionCreate, rbac.ResourceChat.WithOwner(apiKey.UserID.String())) { + httpapi.Forbidden(rw) + return + } + queuedMessageIDStr := chi.URLParam(r, "queuedMessage") queuedMessageID, err := strconv.ParseInt(queuedMessageIDStr, 10, 64) if err != nil { @@ -2913,6 +2961,12 @@ func (api *API) validateCreateChatWorkspaceSelection( Valid: true, } + if workspace.OrganizationID != req.OrganizationID { + return selection, http.StatusBadRequest, &codersdk.Response{ + Message: "Workspace does not belong to the specified organization.", + } + } + if !api.Authorize(r, policy.ActionSSH, workspace) { return selection, http.StatusBadRequest, &codersdk.Response{ Message: "Workspace not found or you do not have access to this resource", @@ -5802,6 +5856,15 @@ func (api *API) postChatToolResults(rw http.ResponseWriter, r *http.Request) { chat := httpmw.ChatParam(r) apiKey := httpmw.APIKey(r) + // Gate tool-result submission behind agents-access. + // Submitting tool results resumes AI/LLM inference on + // a chat in requires_action state. + // See: https://github.com/coder/coder/issues/24250 + if !api.Authorize(r, policy.ActionCreate, rbac.ResourceChat.WithOwner(apiKey.UserID.String())) { + httpapi.Forbidden(rw) + return + } + // Cap the raw request body to prevent excessive memory use. r.Body = http.MaxBytesReader(rw, r.Body, int64(2*maxSystemPromptLenBytes)) var req codersdk.SubmitToolResultsRequest diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index e5ea8b7e16..10313fe705 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -220,7 +220,7 @@ func TestPostChats(t *testing.T) { memberClient := codersdk.NewExperimentalClient(memberClientRaw) chat, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ - Content: []codersdk.ChatInputPart{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, Text: "hello from chats route tests", @@ -274,6 +274,7 @@ func TestPostChats(t *testing.T) { memberClient := codersdk.NewExperimentalClient(memberClientRaw) _, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -289,10 +290,11 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -314,10 +316,11 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + user := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -366,10 +369,11 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + user := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -398,12 +402,12 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + user := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) longPrompt := strings.Repeat("a", 10001) _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ - Content: []codersdk.ChatInputPart{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, Text: "hello", @@ -429,6 +433,7 @@ func TestPostChats(t *testing.T) { }).WithAgent().Do() _, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -466,6 +471,7 @@ func TestPostChats(t *testing.T) { }).WithAgent().Do() _, err := orgAdminClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -487,10 +493,11 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) workspaceID := uuid.New() _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -521,6 +528,7 @@ func TestPostChats(t *testing.T) { }).WithAgent().Do() chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -540,9 +548,10 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -559,10 +568,11 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ - Content: nil, + OrganizationID: firstUser.OrganizationID, + Content: nil, }) sdkErr := requireSDKError(t, err, http.StatusBadRequest) require.Equal(t, "Content is required.", sdkErr.Message) @@ -574,9 +584,10 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -594,9 +605,10 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartType("image"), @@ -619,6 +631,7 @@ func TestPostChats(t *testing.T) { wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) existingChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -629,6 +642,7 @@ func TestPostChats(t *testing.T) { insertAssistantCostMessage(ctx, t, db, existingChat.ID, modelConfig.ID, 100) _, err = client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "over limit", @@ -636,6 +650,94 @@ func TestPostChats(t *testing.T) { }) requireChatUsageLimitExceededError(t, err, 100, 100, wantResetsAt) }) + + t.Run("NilOrganizationID", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID, rbac.RoleAgentsAccess()) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + + _, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: uuid.Nil, + Content: []codersdk.ChatInputPart{ + { + Type: codersdk.ChatInputPartTypeText, + Text: "hello", + }, + }, + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "organization_id is required.", sdkErr.Message) + }) + + t.Run("NonMemberOrganization", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID, rbac.RoleAgentsAccess()) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + + // Create a second organization via the database since the + // API endpoint is enterprise-only. + secondOrg := dbgen.Organization(t, db, database.Organization{}) + + _, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: secondOrg.ID, + Content: []codersdk.ChatInputPart{ + { + Type: codersdk.ChatInputPartTypeText, + Text: "hello", + }, + }, + }) + sdkErr := requireSDKError(t, err, http.StatusForbidden) + require.Equal(t, "You are not a member of the specified organization.", sdkErr.Message) + }) + + t.Run("CrossOrgWorkspaceMismatch", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + workspaceBuild := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + }).WithAgent().Do() + + // Create a second organization and add the admin as a member + // so the request passes the membership check but fails on + // the workspace org mismatch. + secondOrg := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + OrganizationID: secondOrg.ID, + UserID: firstUser.UserID, + }) + + _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: secondOrg.ID, + Content: []codersdk.ChatInputPart{ + { + Type: codersdk.ChatInputPartTypeText, + Text: "hello", + }, + }, + WorkspaceID: &workspaceBuild.Workspace.ID, + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "Workspace does not belong to the specified organization.", sdkErr.Message) + }) } func TestListChats(t *testing.T) { @@ -650,6 +752,7 @@ func TestListChats(t *testing.T) { modelConfig := createChatModelConfig(t, client) firstChatA, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -660,6 +763,7 @@ func TestListChats(t *testing.T) { require.NoError(t, err) firstChatB, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -672,6 +776,7 @@ func TestListChats(t *testing.T) { memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID, rbac.RoleAgentsAccess()) memberClient := codersdk.NewExperimentalClient(memberClientRaw) memberDBChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, @@ -738,21 +843,25 @@ func TestListChats(t *testing.T) { require.Equal(t, memberChats[0].ID, memberChats[0].DiffStatus.ChatID) }) - t.Run("MemberWithoutAgentsAccess", func(t *testing.T) { + t.Run("OrgMemberWithoutAgentsAccessCanAccessOwnChats", func(t *testing.T) { t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) // Create a member without agents-access and insert a chat - // owned by them via system context. This verifies the - // RBAC filter actually excludes results rather than - // returning empty because no chats exist. + // owned by them via system context. With org-scoped chats, + // org members get full CRUD on their own chats through + // OrgMemberPermissions, without needing agents-access. + // The agents-access role only gates chat creation (postChats) + // and message sending (postChatMessages). Metadata operations + // like archive/pin/label and reading are not gated. + // See: https://github.com/coder/coder/issues/24250 memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) memberClient := codersdk.NewExperimentalClient(memberClientRaw) _, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, @@ -762,7 +871,13 @@ func TestListChats(t *testing.T) { chats, err := memberClient.ListChats(ctx, nil) require.NoError(t, err) - require.Empty(t, chats) + require.Len(t, chats, 1) + + // Verify member without agents-access can update own chat. + err = memberClient.UpdateChat(ctx, chats[0].ID, codersdk.UpdateChatRequest{ + Title: ptr.Ref("new title"), + }) + require.NoError(t, err) }) t.Run("Unauthenticated", func(t *testing.T) { @@ -781,7 +896,7 @@ func TestListChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, _ := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) // Create 5 chats. @@ -789,6 +904,7 @@ func TestListChats(t *testing.T) { createdChats := make([]codersdk.Chat, 0, totalChats) for i := 0; i < totalChats; i++ { chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -883,12 +999,13 @@ func TestListChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, _ := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) // Create the chat that will later be pinned. It gets the // earliest updated_at because it is inserted first. pinnedChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "pinned-chat", @@ -902,6 +1019,7 @@ func TestListChats(t *testing.T) { fillerChats := make([]codersdk.Chat, 0, fillerCount) for i := range fillerCount { c, createErr := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: fmt.Sprintf("filler-%d", i), @@ -965,7 +1083,7 @@ func TestListChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, _ := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) // Create 5 chats: 2 will be pinned, 3 unpinned. @@ -973,6 +1091,7 @@ func TestListChats(t *testing.T) { createdChats := make([]codersdk.Chat, 0, totalChats) for i := range totalChats { c, createErr := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: fmt.Sprintf("cursor-pin-chat-%d", i), @@ -1287,7 +1406,7 @@ func TestWatchChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) conn, err := client.Dial(ctx, "/api/experimental/chats/watch", nil) @@ -1295,6 +1414,7 @@ func TestWatchChats(t *testing.T) { defer conn.Close(websocket.StatusNormalClosure, "done") createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -1327,7 +1447,7 @@ func TestWatchChats(t *testing.T) { // API → pubsub → websocket pipeline. ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) conn, err := client.Dial(ctx, "/api/experimental/chats/watch", nil) @@ -1335,6 +1455,7 @@ func TestWatchChats(t *testing.T) { defer conn.Close(websocket.StatusNormalClosure, "done") createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -1383,6 +1504,7 @@ func TestWatchChats(t *testing.T) { // Insert a chat and a diff status row. chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -1479,6 +1601,7 @@ func TestWatchChats(t *testing.T) { modelConfig := createChatModelConfig(t, client) parentChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -1489,22 +1612,22 @@ func TestWatchChats(t *testing.T) { require.NoError(t, err) childOne, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "watch child 1", - ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "watch child 1", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) require.NoError(t, err) childTwo, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "watch child 2", - ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "watch child 2", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) require.NoError(t, err) @@ -3256,6 +3379,7 @@ func TestGetChat(t *testing.T) { modelConfig := createChatModelConfig(t, client) createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3302,6 +3426,7 @@ func TestGetChat(t *testing.T) { _ = createChatModelConfig(t, client) createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3332,9 +3457,9 @@ func TestGetChat(t *testing.T) { // Create a chat with a text + file part. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ - {Type: codersdk.ChatInputPartTypeText, Text: "check file hydration"}, - {Type: codersdk.ChatInputPartTypeFile, FileID: uploadResp.ID}, + {Type: codersdk.ChatInputPartTypeText, Text: "check file hydration"}, {Type: codersdk.ChatInputPartTypeFile, FileID: uploadResp.ID}, }, }) require.NoError(t, err) @@ -3366,6 +3491,7 @@ func TestGetChat(t *testing.T) { // Create a chat via the API so all metadata is set up. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ {Type: codersdk.ChatInputPartTypeText, Text: "tool file test"}, }, @@ -3473,10 +3599,11 @@ func TestArchiveChat(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chatToArchive, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3487,6 +3614,7 @@ func TestArchiveChat(t *testing.T) { require.NoError(t, err) chatToKeep, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3548,6 +3676,7 @@ func TestArchiveChat(t *testing.T) { // Create a parent chat via the API. parentChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3559,6 +3688,7 @@ func TestArchiveChat(t *testing.T) { // Insert child chats directly via the database. child1, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -3569,6 +3699,7 @@ func TestArchiveChat(t *testing.T) { require.NoError(t, err) child2, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -3612,10 +3743,11 @@ func TestUnarchiveChat(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3666,6 +3798,7 @@ func TestUnarchiveChat(t *testing.T) { modelConfig := createChatModelConfig(t, client) parentChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3676,22 +3809,22 @@ func TestUnarchiveChat(t *testing.T) { require.NoError(t, err) child1, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "child 1", - ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "child 1", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) require.NoError(t, err) child2, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "child 2", - ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "child 2", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) require.NoError(t, err) @@ -3754,10 +3887,11 @@ func TestUnarchiveChat(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3786,10 +3920,11 @@ func TestUnarchiveChat(t *testing.T) { func TestChatPinOrder(t *testing.T) { t.Parallel() - createChat := func(ctx context.Context, t *testing.T, client *codersdk.ExperimentalClient, title string) codersdk.Chat { + createChat := func(ctx context.Context, t *testing.T, client *codersdk.ExperimentalClient, orgID uuid.UUID, title string) codersdk.Chat { t.Helper() chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: orgID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3814,12 +3949,12 @@ func TestChatPinOrder(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - first := createChat(ctx, t, client, "first pinned chat") - second := createChat(ctx, t, client, "second pinned chat") - third := createChat(ctx, t, client, "third pinned chat") + first := createChat(ctx, t, client, firstUser.OrganizationID, "first pinned chat") + second := createChat(ctx, t, client, firstUser.OrganizationID, "second pinned chat") + third := createChat(ctx, t, client, firstUser.OrganizationID, "third pinned chat") err := client.UpdateChat(ctx, first.ID, codersdk.UpdateChatRequest{PinOrder: ptr.Ref(int32(1))}) require.NoError(t, err) @@ -3861,11 +3996,11 @@ func TestChatPinOrder(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - first := createChat(ctx, t, client, "pinned then archived") - second := createChat(ctx, t, client, "stays pinned") + first := createChat(ctx, t, client, firstUser.OrganizationID, "pinned then archived") + second := createChat(ctx, t, client, firstUser.OrganizationID, "stays pinned") // Pin both. err := client.UpdateChat(ctx, first.ID, codersdk.UpdateChatRequest{PinOrder: ptr.Ref(int32(1))}) @@ -3891,10 +4026,10 @@ func TestChatPinOrder(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - chat := createChat(ctx, t, client, "negative pin order") + chat := createChat(ctx, t, client, firstUser.OrganizationID, "negative pin order") err := client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{PinOrder: ptr.Ref(int32(-1))}) sdkErr := requireSDKError(t, err, http.StatusBadRequest) require.Equal(t, "Pin order must be non-negative.", sdkErr.Message) @@ -3912,10 +4047,11 @@ func TestPostChatMessages(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -3997,15 +4133,51 @@ func TestPostChatMessages(t *testing.T) { } }) + t.Run("MemberWithoutAgentsAccess", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + // Create a member without agents-access and insert a + // chat owned by them via system context. Even though + // the member can read the chat through org membership, + // sending messages should be gated by agents-access + // because it triggers AI/LLM inference. + memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, + Status: database.ChatStatusWaiting, + OwnerID: member.ID, + LastModelConfigID: modelConfig.ID, + Title: "member chat", + }) + require.NoError(t, err) + + _, err = memberClient.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + Content: []codersdk.ChatInputPart{ + { + Type: codersdk.ChatInputPartTypeText, + Text: "this should fail", + }, + }, + }) + requireSDKError(t, err, http.StatusForbidden) + }) + t.Run("EmptyText", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -4033,10 +4205,11 @@ func TestPostChatMessages(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "initial message for usage-limit test", @@ -4080,10 +4253,11 @@ func TestChatMessageWithFileReferences(t *testing.T) { t.Parallel() // createChat is a helper that creates a chat so we can post messages to it. - createChatForTest := func(t *testing.T, client *codersdk.ExperimentalClient) codersdk.Chat { + createChatForTest := func(t *testing.T, client *codersdk.ExperimentalClient, orgID uuid.UUID) codersdk.Chat { t.Helper() ctx := testutil.Context(t, testutil.WaitLong) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: orgID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "initial message", @@ -4098,9 +4272,9 @@ func TestChatMessageWithFileReferences(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - chat := createChatForTest(t, client) + chat := createChatForTest(t, client, firstUser.OrganizationID) created, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -4160,9 +4334,9 @@ func TestChatMessageWithFileReferences(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - chat := createChatForTest(t, client) + chat := createChatForTest(t, client, firstUser.OrganizationID) created, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -4213,9 +4387,9 @@ func TestChatMessageWithFileReferences(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - chat := createChatForTest(t, client) + chat := createChatForTest(t, client, firstUser.OrganizationID) created, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -4266,9 +4440,9 @@ func TestChatMessageWithFileReferences(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - chat := createChatForTest(t, client) + chat := createChatForTest(t, client, firstUser.OrganizationID) created, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -4319,9 +4493,9 @@ func TestChatMessageWithFileReferences(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - chat := createChatForTest(t, client) + chat := createChatForTest(t, client, firstUser.OrganizationID) created, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{ @@ -4432,9 +4606,9 @@ func TestChatMessageWithFileReferences(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) - chat := createChatForTest(t, client) + chat := createChatForTest(t, client, firstUser.OrganizationID) _, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -4454,11 +4628,12 @@ func TestChatMessageWithFileReferences(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) // File references should also work in the initial CreateChat call. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeFileReference, FileName: "bug.py", @@ -4494,6 +4669,7 @@ func TestChatMessageWithFiles(t *testing.T) { // Create a chat with text first. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -4538,6 +4714,7 @@ func TestChatMessageWithFiles(t *testing.T) { // Create a chat with text first. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -4598,6 +4775,7 @@ func TestChatMessageWithFiles(t *testing.T) { // Create a new chat with only a file part. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeFile, @@ -4624,11 +4802,12 @@ func TestChatMessageWithFiles(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) // Create a chat with text first. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -4662,6 +4841,7 @@ func TestChatMessageWithFiles(t *testing.T) { // Create a text-only chat (no files initially). chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ {Type: codersdk.ChatInputPartTypeText, Text: "no files yet"}, }, @@ -4705,9 +4885,9 @@ func TestChatMessageWithFiles(t *testing.T) { // Create a chat with a file. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ - {Type: codersdk.ChatInputPartTypeText, Text: "first mention"}, - {Type: codersdk.ChatInputPartTypeFile, FileID: uploadResp.ID}, + {Type: codersdk.ChatInputPartTypeText, Text: "first mention"}, {Type: codersdk.ChatInputPartTypeFile, FileID: uploadResp.ID}, }, }) require.NoError(t, err) @@ -4754,7 +4934,7 @@ func TestChatMessageWithFiles(t *testing.T) { for _, fid := range fileIDs { parts = append(parts, codersdk.ChatInputPart{Type: codersdk.ChatInputPartTypeFile, FileID: fid}) } - chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{Content: parts}) + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{OrganizationID: firstUser.OrganizationID, Content: parts}) require.NoError(t, err) require.Empty(t, chat.Warnings, "creating a chat at exactly the cap should not warn") require.Len(t, chat.Files, codersdk.MaxChatFileIDs, "all files should be linked on creation") @@ -4819,7 +4999,7 @@ func TestChatMessageWithFiles(t *testing.T) { for _, fid := range fileIDs { parts = append(parts, codersdk.ChatInputPart{Type: codersdk.ChatInputPartTypeFile, FileID: fid}) } - chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{Content: parts}) + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{OrganizationID: firstUser.OrganizationID, Content: parts}) require.NoError(t, err, "chat creation should succeed even when cap is exceeded") require.NotEmpty(t, chat.Warnings, "response should warn about unlinked files") require.Contains(t, chat.Warnings[0], "file linking skipped") @@ -4844,10 +5024,11 @@ func TestPatchChatMessage(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -4930,6 +5111,7 @@ func TestPatchChatMessage(t *testing.T) { // Create a chat with a text + file part. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5016,10 +5198,11 @@ func TestPatchChatMessage(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "hello before edit", @@ -5056,10 +5239,11 @@ func TestPatchChatMessage(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5086,10 +5270,11 @@ func TestPatchChatMessage(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5130,6 +5315,7 @@ func TestPatchChatMessage(t *testing.T) { // Create a text-only chat. chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ {Type: codersdk.ChatInputPartTypeText, Text: "before file edit"}, }, @@ -5190,7 +5376,7 @@ func TestPatchChatMessage(t *testing.T) { require.NoError(t, err) parts = append(parts, codersdk.ChatInputPart{Type: codersdk.ChatInputPartTypeFile, FileID: up.ID}) } - chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{Content: parts}) + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{OrganizationID: firstUser.OrganizationID, Content: parts}) require.NoError(t, err) require.Empty(t, chat.Warnings, "all files should link on create") @@ -5235,11 +5421,12 @@ func TestStreamChat(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) const initialMessage = "stream chat route initial message" chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5314,6 +5501,7 @@ func TestInterruptChat(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5394,6 +5582,7 @@ func TestRegenerateChatTitle(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5414,6 +5603,7 @@ func TestRegenerateChatTitle(t *testing.T) { _ = createChatModelConfig(t, client) createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5434,10 +5624,11 @@ func TestRegenerateChatTitle(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "chat for unauthenticated regeneration", @@ -5455,10 +5646,11 @@ func TestRegenerateChatTitle(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "chat over usage limit", @@ -5503,6 +5695,7 @@ func TestRegenerateChatTitle(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5544,6 +5737,7 @@ func TestRegenerateChatTitle(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5590,10 +5784,11 @@ func TestRegenerateChatTitle(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5660,6 +5855,7 @@ func TestGetChatDiffStatus(t *testing.T) { modelConfig := createChatModelConfig(t, client) noCachedStatusChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5673,6 +5869,7 @@ func TestGetChatDiffStatus(t *testing.T) { require.Nil(t, noCachedChat.DiffStatus) cachedStatusChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5742,6 +5939,7 @@ func TestGetChatDiffStatus(t *testing.T) { _ = createChatModelConfig(t, client) createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5780,6 +5978,7 @@ func TestGetChatDiffContents(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5817,10 +6016,11 @@ func TestGetChatDiffContents(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5849,6 +6049,7 @@ func TestGetChatDiffContents(t *testing.T) { _ = createChatModelConfig(t, client) createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5877,6 +6078,7 @@ func TestDeleteChatQueuedMessage(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5929,6 +6131,7 @@ func TestDeleteChatQueuedMessage(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -5964,6 +6167,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -6035,6 +6239,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { enableDailyChatUsageLimit(ctx, t, db, 100) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -6110,6 +6315,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -6131,6 +6337,52 @@ func TestPromoteChatQueuedMessage(t *testing.T) { require.Equal(t, "Invalid queued message ID.", sdkErr.Message) require.Contains(t, sdkErr.Detail, "invalid syntax") }) + + t.Run("MemberWithoutAgentsAccess", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + // Create a member without agents-access. Even though the + // member owns the chat, promoting a queued message should + // be gated by agents-access because it triggers inference. + memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, + Status: database.ChatStatusWaiting, + OwnerID: member.ID, + LastModelConfigID: modelConfig.ID, + Title: "promote queued no agents access", + }) + require.NoError(t, err) + + queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("queued message no agents access"), + }) + require.NoError(t, err) + queuedMessage, err := db.InsertChatQueuedMessage( + dbauthz.AsSystemRestricted(ctx), + database.InsertChatQueuedMessageParams{ + ChatID: chat.ID, + Content: queuedContent, + }, + ) + require.NoError(t, err) + + promoteRes, err := memberClient.Request( + ctx, + http.MethodPost, + fmt.Sprintf("/api/experimental/chats/%s/queue/%d/promote", chat.ID, queuedMessage.ID), + nil, + ) + require.NoError(t, err) + defer promoteRes.Body.Close() + require.Equal(t, http.StatusForbidden, promoteRes.StatusCode) + }) } func TestChatUsageLimitOverrideRoutes(t *testing.T) { @@ -6693,6 +6945,7 @@ func seedChatCostFixture(t *testing.T) chatCostTestFixture { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, @@ -6814,6 +7067,7 @@ func TestChatCostSummary_AdminDrilldown(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, @@ -6883,6 +7137,7 @@ func TestChatCostUsers(t *testing.T) { modelConfig := createChatModelConfig(t, client) adminChat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, @@ -6911,6 +7166,7 @@ func TestChatCostUsers(t *testing.T) { require.NoError(t, err) memberChat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, @@ -6995,6 +7251,7 @@ func TestChatCostSummary_DateRange(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, @@ -7061,6 +7318,7 @@ func TestChatCostSummary_UnpricedMessages(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: firstUser.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, @@ -7171,10 +7429,11 @@ func TestWatchChatDesktop(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -7251,6 +7510,7 @@ func TestChatSystemPrompt(t *testing.T) { t.Helper() chat, err := adminClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -7383,7 +7643,7 @@ func TestChatSystemPrompt(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) legacyClient, legacyDB := newChatClientWithDatabase(t) - _ = coderdtest.CreateFirstUser(t, legacyClient.Client) + firstUser := coderdtest.CreateFirstUser(t, legacyClient.Client) _ = createChatModelConfig(t, legacyClient) require.NoError(t, legacyDB.UpsertChatSystemPrompt(dbauthz.AsSystemRestricted(ctx), "Legacy custom instructions")) @@ -7395,6 +7655,7 @@ func TestChatSystemPrompt(t *testing.T) { require.Equal(t, chatd.DefaultSystemPrompt, resp.DefaultSystemPrompt) chat, err := legacyClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: fmt.Sprintf("legacy custom prompt %s", t.Name()), @@ -7522,7 +7783,7 @@ func TestChatSystemPrompt(t *testing.T) { Pubsub: pubsub, DeploymentValues: chatDeploymentValues(t), })) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) err := client.UpdateChatSystemPrompt(ctx, codersdk.UpdateChatSystemPromptRequest{ @@ -7533,6 +7794,7 @@ func TestChatSystemPrompt(t *testing.T) { store.failNextGetChatSystemPromptConfig.Store(true) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: fmt.Sprintf("config-read fallback %s", t.Name()), @@ -7568,7 +7830,7 @@ func TestChatSystemPrompt(t *testing.T) { Pubsub: pubsub, DeploymentValues: chatDeploymentValues(t), })) - _ = coderdtest.CreateFirstUser(t, client.Client) + firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) err := client.UpdateChatSystemPrompt(ctx, codersdk.UpdateChatSystemPromptRequest{ @@ -7581,6 +7843,7 @@ func TestChatSystemPrompt(t *testing.T) { // include_default=false, so chat creation falls back to the built-in default. store.failNextGetChatSystemPromptConfig.Store(true) chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: fmt.Sprintf("config-read fallback %s", t.Name()), @@ -8184,6 +8447,7 @@ func TestGetChatsByWorkspace(t *testing.T) { // Helper to insert a chat linked to a workspace. insertChat := func(ctx context.Context, title string, workspaceID uuid.UUID) database.Chat { chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -8310,6 +8574,7 @@ func TestSubmitToolResults(t *testing.T) { t *testing.T, db database.Store, ownerID uuid.UUID, + organizationID uuid.UUID, modelConfigID uuid.UUID, dynamicToolName string, toolCallIDs []string, @@ -8326,11 +8591,11 @@ func TestSubmitToolResults(t *testing.T) { require.NoError(t, err) chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: organizationID, Status: database.ChatStatusWaiting, OwnerID: ownerID, LastModelConfigID: modelConfigID, - Title: "tool-results-test", - DynamicTools: pqtype.NullRawMessage{RawMessage: dtJSON, Valid: true}, + Title: "tool-results-test", DynamicTools: pqtype.NullRawMessage{RawMessage: dtJSON, Valid: true}, }) require.NoError(t, err) @@ -8390,7 +8655,7 @@ func TestSubmitToolResults(t *testing.T) { const toolName = "my_dynamic_tool" toolCallIDs := []string{"call_abc", "call_def"} - chat := setupRequiresAction(ctx, t, db, user.UserID, modelConfig.ID, toolName, toolCallIDs) + chat := setupRequiresAction(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, toolName, toolCallIDs) err := client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ Results: []codersdk.ToolResult{ @@ -8433,6 +8698,7 @@ func TestSubmitToolResults(t *testing.T) { // Create a chat that is NOT in requires_action status. chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OrganizationID: user.OrganizationID, Status: database.ChatStatusWaiting, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, @@ -8459,7 +8725,7 @@ func TestSubmitToolResults(t *testing.T) { const toolName = "my_dynamic_tool" toolCallIDs := []string{"call_one", "call_two"} - chat := setupRequiresAction(ctx, t, db, user.UserID, modelConfig.ID, toolName, toolCallIDs) + chat := setupRequiresAction(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, toolName, toolCallIDs) // Submit only one of the two required results. err := client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ @@ -8481,7 +8747,7 @@ func TestSubmitToolResults(t *testing.T) { const toolName = "my_dynamic_tool" toolCallIDs := []string{"call_real"} - chat := setupRequiresAction(ctx, t, db, user.UserID, modelConfig.ID, toolName, toolCallIDs) + chat := setupRequiresAction(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, toolName, toolCallIDs) // Submit a result with a wrong tool_call_id. err := client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ @@ -8503,7 +8769,7 @@ func TestSubmitToolResults(t *testing.T) { const toolName = "my_dynamic_tool" toolCallIDs := []string{"call_json"} - chat := setupRequiresAction(ctx, t, db, user.UserID, modelConfig.ID, toolName, toolCallIDs) + chat := setupRequiresAction(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, toolName, toolCallIDs) // We must bypass the SDK client because json.RawMessage // rejects invalid JSON during json.Marshal. A raw HTTP @@ -8533,7 +8799,7 @@ func TestSubmitToolResults(t *testing.T) { const toolName = "my_dynamic_tool" toolCallIDs := []string{"call_dup1", "call_dup2"} - chat := setupRequiresAction(ctx, t, db, user.UserID, modelConfig.ID, toolName, toolCallIDs) + chat := setupRequiresAction(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, toolName, toolCallIDs) err := client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ Results: []codersdk.ToolResult{ @@ -8556,7 +8822,7 @@ func TestSubmitToolResults(t *testing.T) { const toolName = "my_dynamic_tool" toolCallIDs := []string{"call_empty"} - chat := setupRequiresAction(ctx, t, db, user.UserID, modelConfig.ID, toolName, toolCallIDs) + chat := setupRequiresAction(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, toolName, toolCallIDs) err := client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ Results: []codersdk.ToolResult{}, @@ -8575,7 +8841,7 @@ func TestSubmitToolResults(t *testing.T) { const toolName = "my_dynamic_tool" toolCallIDs := []string{"call_other"} - chat := setupRequiresAction(ctx, t, db, user.UserID, modelConfig.ID, toolName, toolCallIDs) + chat := setupRequiresAction(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, toolName, toolCallIDs) // Create a second user and try to submit tool results // to user A's chat. @@ -8592,6 +8858,33 @@ func TestSubmitToolResults(t *testing.T) { }) requireSDKError(t, err, http.StatusNotFound) }) + + t.Run("MemberWithoutAgentsAccess", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + // Create a member without agents-access. Even though the + // member owns the chat, submitting tool results should be + // gated by agents-access because it triggers inference. + memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + + const toolName = "my_dynamic_tool" + toolCallIDs := []string{"call_noaccess"} + + chat := setupRequiresAction(ctx, t, db, member.ID, firstUser.OrganizationID, modelConfig.ID, toolName, toolCallIDs) + + err := memberClient.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ + Results: []codersdk.ToolResult{ + {ToolCallID: "call_noaccess", Output: json.RawMessage(`"should fail"`)}, + }, + }) + requireSDKError(t, err, http.StatusForbidden) + }) } func TestPostChats_DynamicToolValidation(t *testing.T) { @@ -8602,7 +8895,7 @@ func TestPostChats_DynamicToolValidation(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + user := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) tools := make([]codersdk.DynamicTool, 251) @@ -8613,6 +8906,7 @@ func TestPostChats_DynamicToolValidation(t *testing.T) { } _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "hello", @@ -8628,10 +8922,11 @@ func TestPostChats_DynamicToolValidation(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + user := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "hello", @@ -8649,10 +8944,11 @@ func TestPostChats_DynamicToolValidation(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newChatClient(t) - _ = coderdtest.CreateFirstUser(t, client.Client) + user := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "hello", diff --git a/coderd/httpmw/chatparam_test.go b/coderd/httpmw/chatparam_test.go index 8585b9462c..0c758f15bd 100644 --- a/coderd/httpmw/chatparam_test.go +++ b/coderd/httpmw/chatparam_test.go @@ -35,7 +35,7 @@ func TestChatParam(t *testing.T) { return r, user } - insertChat := func(t *testing.T, db database.Store, ownerID uuid.UUID) database.Chat { + insertChat := func(t *testing.T, db database.Store, ownerID, organizationID uuid.UUID) database.Chat { t.Helper() _, err := db.InsertChatProvider(context.Background(), database.InsertChatProviderParams{ @@ -63,6 +63,7 @@ func TestChatParam(t *testing.T) { require.NoError(t, err) chat, err := db.InsertChat(context.Background(), database.InsertChatParams{ + OrganizationID: organizationID, Status: database.ChatStatusWaiting, OwnerID: ownerID, WorkspaceID: uuid.NullUUID{}, @@ -147,7 +148,8 @@ func TestChatParam(t *testing.T) { }) r, user := setupAuthentication(db) - chat := insertChat(t, db, user.ID) + org := dbgen.Organization(t, db, database.Organization{}) + chat := insertChat(t, db, user.ID, org.ID) chi.RouteContext(r.Context()).URLParams.Add("chat", chat.ID.String()) rw := httptest.NewRecorder() diff --git a/coderd/mcp_test.go b/coderd/mcp_test.go index 5b791adf11..85ee74ec95 100644 --- a/coderd/mcp_test.go +++ b/coderd/mcp_test.go @@ -1276,7 +1276,7 @@ func TestChatWithMCPServerIDs(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client := newMCPClient(t) - _ = coderdtest.CreateFirstUser(t, client) + firstUser := coderdtest.CreateFirstUser(t, client) expClient := codersdk.NewExperimentalClient(client) @@ -1288,6 +1288,7 @@ func TestChatWithMCPServerIDs(t *testing.T) { // Create a chat referencing the MCP server. chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, diff --git a/coderd/rbac/authz.go b/coderd/rbac/authz.go index 264970928b..4cd05d2855 100644 --- a/coderd/rbac/authz.go +++ b/coderd/rbac/authz.go @@ -689,11 +689,11 @@ func ConfigWithoutACL() regosql.ConvertConfig { } // ConfigChats is the configuration for converting rego to SQL when -// the target table is "chats", which has no organization_id or ACL +// the target table is "chats", which has no ACL // columns. func ConfigChats() regosql.ConvertConfig { return regosql.ConvertConfig{ - VariableConverter: regosql.ChatConverter(), + VariableConverter: regosql.NoACLConverter(), } } diff --git a/coderd/rbac/regosql/compile_test.go b/coderd/rbac/regosql/compile_test.go index 63f9302d3a..e42e3f70e8 100644 --- a/coderd/rbac/regosql/compile_test.go +++ b/coderd/rbac/regosql/compile_test.go @@ -287,16 +287,15 @@ neq(input.object.owner, ""); Queries: []string{ `"me" = input.object.owner; input.object.owner != ""; input.object.org_owner = ""`, }, - ExpectedSQL: p(p("'me' = owner_id :: text") + " AND " + p("owner_id :: text != ''") + " AND " + p("'' = ''")), - VariableConverter: regosql.ChatConverter(), + ExpectedSQL: p(p("'me' = owner_id :: text") + " AND " + p("owner_id :: text != ''") + " AND " + p("organization_id :: text = ''")), + VariableConverter: regosql.NoACLConverter(), }, { - Name: "ChatOrgScopedNeverMatches", + Name: "ChatOrgScopedMatches", Queries: []string{ `input.object.org_owner = "org-id"`, }, - ExpectedSQL: p("'' = 'org-id'"), - VariableConverter: regosql.ChatConverter(), + ExpectedSQL: p("organization_id :: text = 'org-id'"), VariableConverter: regosql.NoACLConverter(), }, { Name: "AuditLogUUID", diff --git a/coderd/rbac/regosql/configs.go b/coderd/rbac/regosql/configs.go index 2066d93473..22302a5296 100644 --- a/coderd/rbac/regosql/configs.go +++ b/coderd/rbac/regosql/configs.go @@ -126,30 +126,6 @@ func NoACLConverter() *sqltypes.VariableConverter { return matcher } -// ChatConverter should be used for the chats table, which has no -// organization_id, group_acl, or user_acl columns. -func ChatConverter() *sqltypes.VariableConverter { - matcher := sqltypes.NewVariableConverter().RegisterMatcher( - resourceIDMatcher(), - // The chats table has no organization_id column. Map org_owner - // to a literal empty string so that: - // - User-level ownership checks (org_owner = '') activate correctly. - // - Org-scoped permissions never match (org_owner will never equal - // a real org UUID), which is intentional since chats are not - // org-scoped resources. - // Note: custom org roles that include "chat" permissions will - // silently have no effect because of this mapping. - sqltypes.StringVarMatcher("''", []string{"input", "object", "org_owner"}), - userOwnerMatcher(), - ) - matcher.RegisterMatcher( - sqltypes.AlwaysFalse(groupACLMatcher(matcher)), - sqltypes.AlwaysFalse(userACLMatcher(matcher)), - ) - - return matcher -} - func DefaultVariableConverter() *sqltypes.VariableConverter { matcher := sqltypes.NewVariableConverter().RegisterMatcher( resourceIDMatcher(), diff --git a/coderd/rbac/roles.go b/coderd/rbac/roles.go index b68878e3cd..f5805bf7fa 100644 --- a/coderd/rbac/roles.go +++ b/coderd/rbac/roles.go @@ -437,17 +437,22 @@ func ReloadBuiltinRoles(opts *RoleOptions) { return auditorRole }, + // templateAdmin grants all actions on templates, files, + // provisioner daemons, and prebuilt workspaces. templateAdmin: func(_ uuid.UUID) Role { return templateAdminRole }, + // userAdmin grants all actions on users, groups, roles, + // and organization membership. userAdmin: func(_ uuid.UUID) Role { return userAdminRole }, // agentsAccess grants all actions on chat resources owned - // by the user. Without this role, members cannot create - // or interact with chats. + // by the user. Without this role, members can still read, + // update, and delete their own chats via org membership, + // but cannot create chats or trigger AI inference. agentsAccess: func(_ uuid.UUID) Role { return agentsAccessRole }, diff --git a/coderd/rbac/roles_test.go b/coderd/rbac/roles_test.go index 1cc7e2cade..bbd229621a 100644 --- a/coderd/rbac/roles_test.go +++ b/coderd/rbac/roles_test.go @@ -205,6 +205,34 @@ func TestRolePermissions(t *testing.T) { orgTemplateAdmin := authSubject{Name: "org_template_admin", Actor: rbac.Subject{ID: userAdminID.String(), Roles: rbac.RoleIdentifiers{rbac.RoleMember(), rbac.ScopedRoleOrgTemplateAdmin(orgID)}, Scope: rbac.ScopeAll}.WithCachedASTValue()} orgAdminBanWorkspace := authSubject{Name: "org_admin_workspace_ban", Actor: rbac.Subject{ID: adminID.String(), Roles: rbac.RoleIdentifiers{rbac.RoleMember(), rbac.ScopedRoleOrgAdmin(orgID), rbac.ScopedRoleOrgWorkspaceCreationBan(orgID)}, Scope: rbac.ScopeAll}.WithCachedASTValue()} agentsAccessUser := authSubject{Name: "chat_access", Actor: rbac.Subject{ID: currentUser.String(), Roles: rbac.RoleIdentifiers{rbac.RoleMember(), rbac.RoleAgentsAccess()}, Scope: rbac.ScopeAll}.WithCachedASTValue()} + orgMemberMe := func() authSubject { + memberRole, err := rbac.RoleByName(rbac.RoleMember()) + require.NoError(t, err) + perms := rbac.OrgMemberPermissions(rbac.OrgSettings{ + ShareableWorkspaceOwners: rbac.ShareableWorkspaceOwnersEveryone, + }) + return authSubject{ + Name: "org_member_me", + Actor: rbac.Subject{ + ID: currentUser.String(), + Roles: rbac.Roles{ + memberRole, + { + Identifier: rbac.ScopedRoleOrgMember(orgID), + Site: []rbac.Permission{}, + User: []rbac.Permission{}, + ByOrgID: map[string]rbac.OrgPermissions{ + orgID.String(): { + Org: perms.Org, + Member: perms.Member, + }, + }, + }, + }, + Scope: rbac.ScopeAll, + }.WithCachedASTValue(), + } + }() setOrgNotMe := authSubjectSet{orgAdmin, orgAuditor, orgUserAdmin, orgTemplateAdmin} otherOrgAdmin := authSubject{Name: "org_admin_other", Actor: rbac.Subject{ID: uuid.NewString(), Roles: rbac.RoleIdentifiers{rbac.RoleMember(), rbac.ScopedRoleOrgAdmin(otherOrg)}, Scope: rbac.ScopeAll}.WithCachedASTValue()} @@ -1070,16 +1098,10 @@ func TestRolePermissions(t *testing.T) { { Name: "ChatUsage", Actions: []policy.Action{policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete}, - Resource: rbac.ResourceChat.WithOwner(currentUser.String()), + Resource: rbac.ResourceChat.WithID(uuid.New()).InOrg(orgID).WithOwner(currentUser.String()), AuthorizeMap: map[bool][]hasAuthSubjects{ - true: {owner, agentsAccessUser}, - false: { - memberMe, - orgAdmin, otherOrgAdmin, - orgAuditor, otherOrgAuditor, - templateAdmin, orgTemplateAdmin, otherOrgTemplateAdmin, - userAdmin, orgUserAdmin, otherOrgUserAdmin, - }, + true: {owner, orgAdmin, orgMemberMe}, + false: {setOtherOrg, memberMe, agentsAccessUser, userAdmin, templateAdmin, orgTemplateAdmin, orgUserAdmin, orgAuditor}, }, }, } diff --git a/coderd/telemetry/telemetry_test.go b/coderd/telemetry/telemetry_test.go index 56d22f5172..a07f1c74cf 100644 --- a/coderd/telemetry/telemetry_test.go +++ b/coderd/telemetry/telemetry_test.go @@ -1646,6 +1646,7 @@ func TestChatsTelemetry(t *testing.T) { }) rootChat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, Title: "Root Chat", @@ -1657,6 +1658,7 @@ func TestChatsTelemetry(t *testing.T) { // Create a child chat (has parent + root). childChat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: modelCfg2.ID, Title: "Child Chat", diff --git a/coderd/workspaceagents_active_chat_internal_test.go b/coderd/workspaceagents_active_chat_internal_test.go index 68e3beeda1..4af46f6605 100644 --- a/coderd/workspaceagents_active_chat_internal_test.go +++ b/coderd/workspaceagents_active_chat_internal_test.go @@ -35,6 +35,7 @@ func TestActiveAgentChatDefinitionsAgree(t *testing.T) { for _, archived := range []bool{false, true} { for _, status := range database.AllChatStatusValues() { chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: status, OwnerID: owner.ID, LastModelConfigID: modelConfig.ID, diff --git a/coderd/workspaceagents_chat_context_test.go b/coderd/workspaceagents_chat_context_test.go index d3f67fb8bc..aa7d474f00 100644 --- a/coderd/workspaceagents_chat_context_test.go +++ b/coderd/workspaceagents_chat_context_test.go @@ -157,7 +157,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) for _, step := range tc.steps { resp, err := setup.agentClient.AddChatContext(ctx, step.req) @@ -219,7 +219,7 @@ func TestAgentChatContext(t *testing.T) { }, ) require.NoError(t, err) - chat := createAgentChatContextChat(ctx, t, baseDB, user.UserID, originalModel.ID, workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, baseDB, user.OrganizationID, user.UserID, originalModel.ID, workspace.Agents[0].ID, t.Name()) interceptDB.beforeInTx = func() { _, err := baseDB.UpdateChatLastModelConfigByID( @@ -260,7 +260,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) skillPart := codersdk.ChatMessagePart{ Type: codersdk.ChatMessagePartTypeSkill, @@ -322,7 +322,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) skillPart := codersdk.ChatMessagePart{ Type: codersdk.ChatMessagePartTypeSkill, @@ -421,7 +421,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -478,7 +478,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -540,7 +540,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("assistant reply"), @@ -606,7 +606,7 @@ func TestAgentChatContext(t *testing.T) { OwnerID: user.UserID, }).WithAgent().Do() - chat := createAgentChatContextChat(ctx, t, db, user.UserID, model.ID, firstWorkspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, db, user.OrganizationID, user.UserID, model.ID, firstWorkspace.Agents[0].ID, t.Name()) secondAgentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(secondWorkspace.AgentToken)) _, err := secondAgentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ @@ -628,7 +628,7 @@ func TestAgentChatContext(t *testing.T) { setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: chat.ID, @@ -683,7 +683,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: chat.ID, @@ -705,7 +705,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) largeContent := strings.Repeat("a", maxContextFileBytes+100) resp, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ @@ -740,7 +740,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) visible := strings.Repeat("a", maxContextFileBytes-1) content := visible + strings.Repeat("\u200b", 100) + "z" @@ -782,8 +782,8 @@ func TestAgentChatContext(t *testing.T) { setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - ownerChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") - foreignChat := createAgentChatContextChat(ctx, t, setup.db, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") + ownerChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") + foreignChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") resp, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -806,8 +806,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - rootChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") - childChat := createAgentChatContextChildChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") + rootChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") + childChat := createAgentChatContextChildChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") resp, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -830,8 +830,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat1") - createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat2") + createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat1") + createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat2") _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -850,8 +850,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - rootChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") - childChat := createAgentChatContextChildChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") + rootChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") + childChat := createAgentChatContextChildChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: rootChat.ID, @@ -878,8 +878,8 @@ func TestAgentChatContext(t *testing.T) { setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - ownerChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") - _ = createAgentChatContextChat(ctx, t, setup.db, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") + ownerChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") + _ = createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: ownerChat.ID, @@ -904,7 +904,7 @@ func TestAgentChatContext(t *testing.T) { setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.ClearChatContext(ctx, agentsdk.ClearChatContextRequest{ChatID: chat.ID}) sdkErr := requireSDKError(t, err, http.StatusForbidden) @@ -917,7 +917,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, @@ -1023,6 +1023,7 @@ func createAgentChatContextChat( ctx context.Context, t testing.TB, db database.Store, + orgID uuid.UUID, ownerID uuid.UUID, modelConfigID uuid.UUID, agentID uuid.UUID, @@ -1032,6 +1033,7 @@ func createAgentChatContextChat( chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ Status: database.ChatStatusWaiting, + OrganizationID: orgID, OwnerID: ownerID, LastModelConfigID: modelConfigID, Title: title, @@ -1046,6 +1048,7 @@ func createAgentChatContextChildChat( ctx context.Context, t testing.TB, db database.Store, + orgID uuid.UUID, ownerID uuid.UUID, modelConfigID uuid.UUID, agentID uuid.UUID, @@ -1056,6 +1059,7 @@ func createAgentChatContextChildChat( chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ Status: database.ChatStatusWaiting, + OrganizationID: orgID, OwnerID: ownerID, LastModelConfigID: modelConfigID, Title: title, diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 19efbbc7a4..454d8eb3a6 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -775,6 +775,7 @@ func (e *UsageLimitExceededError) Error() string { // CreateOptions controls chat creation in the shared chat mutation path. type CreateOptions struct { + OrganizationID uuid.UUID OwnerID uuid.UUID WorkspaceID uuid.NullUUID BuildID uuid.NullUUID @@ -852,6 +853,9 @@ type PromoteQueuedResult struct { // CreateChat creates a chat, inserts optional system prompt and initial user // message, and moves the chat into pending status. func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.Chat, error) { + if opts.OrganizationID == uuid.Nil { + return database.Chat{}, xerrors.New("organization_id is required") + } if opts.OwnerID == uuid.Nil { return database.Chat{}, xerrors.New("owner_id is required") } @@ -883,6 +887,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C } insertedChat, err := tx.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: opts.OrganizationID, OwnerID: opts.OwnerID, WorkspaceID: opts.WorkspaceID, BuildID: opts.BuildID, @@ -5009,6 +5014,7 @@ func (p *Server) runChat( chattool.ListTemplates(chattool.ListTemplatesOptions{ DB: p.db, OwnerID: chat.OwnerID, + OrganizationID: chat.OrganizationID, AllowedTemplateIDs: p.chatTemplateAllowlist, }), chattool.ReadTemplate(chattool.ReadTemplateOptions{ @@ -5019,6 +5025,7 @@ func (p *Server) runChat( chattool.CreateWorkspace(chattool.CreateWorkspaceOptions{ DB: p.db, OwnerID: chat.OwnerID, + OrganizationID: chat.OrganizationID, ChatID: chat.ID, CreateFn: p.createWorkspaceFn, AgentConnFn: chattool.AgentConnFunc(p.agentConnFn), diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index c290eb33f2..6b31fa8ad7 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -62,9 +62,10 @@ func TestInterruptChatBroadcastsStatusAcrossInstances(t *testing.T) { replicaB := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replicaA.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "interrupt-me", ModelConfigID: model.ID, @@ -185,6 +186,7 @@ func TestSubagentChatExcludesWorkspaceProvisioningTools(t *testing.T) { // Create a root chat whose first model call will spawn a subagent. chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -271,9 +273,10 @@ func TestInterruptChatClearsWorkerInDatabase(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "db-transition", ModelConfigID: model.ID, @@ -307,10 +310,11 @@ func TestArchiveChatMovesPendingChatToWaiting(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, + OrganizationID: org.ID, Title: "archive-pending", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -372,11 +376,12 @@ func TestArchiveChatInterruptsActiveProcessing(t *testing.T) { }) server := newActiveTestServer(t, db, ps) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, + OrganizationID: org.ID, Title: "archive-interrupt", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -481,9 +486,10 @@ func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "heartbeat-ownership", ModelConfigID: model.ID, @@ -528,9 +534,10 @@ func TestSendMessageQueueBehaviorQueuesWhenBusy(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "queue-when-busy", ModelConfigID: model.ID, @@ -579,9 +586,10 @@ func TestSendMessageQueuesWhenWaitingWithQueuedBacklog(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "queue-when-waiting-with-backlog", ModelConfigID: model.ID, @@ -645,9 +653,10 @@ func TestSendMessageInterruptBehaviorQueuesAndInterruptsWhenBusy(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "interrupt-when-busy", ModelConfigID: model.ID, @@ -709,9 +718,10 @@ func TestEditMessageUpdatesAndTruncatesAndClearsQueue(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "edit-message", ModelConfigID: model.ID, @@ -818,9 +828,8 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) - org := dbgen.Organization(t, db, database.Organization{}) tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ OrganizationID: org.ID, CreatedBy: user.ID, @@ -837,6 +846,7 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, Title: "test-with-workspace", @@ -870,9 +880,10 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "test-without-workspace", ModelConfigID: model.ID, @@ -906,7 +917,7 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ Enabled: true, @@ -916,6 +927,7 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { require.NoError(t, err) existingChat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, Title: "existing-limit-chat", @@ -959,6 +971,7 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { require.Len(t, beforeChats, 1) _, err = replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "over-limit", ModelConfigID: model.ID, @@ -988,7 +1001,7 @@ func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ Enabled: true, @@ -998,6 +1011,7 @@ func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing require.NoError(t, err) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "queued-limit-reached", ModelConfigID: model.ID, @@ -1166,10 +1180,11 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { }) acquireTrap.MustWait(ctx).MustRelease(ctx) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "interrupt-autopromote-limit", ModelConfigID: model.ID, @@ -1207,6 +1222,7 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { require.NotNil(t, laterQueuedResult.QueuedMessage) spendChat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{}, @@ -1289,7 +1305,7 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ Enabled: true, @@ -1299,6 +1315,7 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { require.NoError(t, err) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "edit-limit-reached", ModelConfigID: model.ID, @@ -1370,9 +1387,10 @@ func TestEditMessageRejectsMissingMessage(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "missing-edited-message", ModelConfigID: model.ID, @@ -1396,9 +1414,10 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "non-user-edited-message", ModelConfigID: model.ID, @@ -1448,7 +1467,7 @@ func TestRecoverStaleChatsPeriodically(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Use a very short stale threshold so the periodic recovery // kicks in quickly during the test. @@ -1458,6 +1477,7 @@ func TestRecoverStaleChatsPeriodically(t *testing.T) { // to running with a heartbeat in the past. deadWorkerID := uuid.New() chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, Title: "stale-recovery-periodic", @@ -1504,6 +1524,7 @@ func TestRecoverStaleChatsPeriodically(t *testing.T) { // This tests the periodic recovery, not just the startup one. deadWorkerID2 := uuid.New() chat2, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, Title: "stale-recovery-periodic-2", @@ -1537,7 +1558,7 @@ func TestRecoverStaleRequiresActionChat(t *testing.T) { db, ps, rawDB := dbtestutil.NewDBWithSQLDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Use a very short stale threshold so the periodic recovery // kicks in quickly during the test. @@ -1547,6 +1568,7 @@ func TestRecoverStaleRequiresActionChat(t *testing.T) { // client that disappeared while the chat was waiting for // dynamic tool results. chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, Title: "stale-requires-action", @@ -1601,12 +1623,13 @@ func TestNewReplicaRecoversStaleChatFromDeadReplica(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Simulate a chat left running by a dead replica with a stale // heartbeat (well beyond the stale threshold). deadReplicaID := uuid.New() chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, Title: "orphaned-chat", @@ -1645,11 +1668,12 @@ func TestWaitingChatsAreNotRecoveredAsStale(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Create a chat in waiting status — this should NOT be touched // by stale recovery. chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, Title: "waiting-chat", @@ -1690,9 +1714,10 @@ func TestUpdateChatStatusPersistsLastError(t *testing.T) { _ = newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, Title: "error-persisted", @@ -1746,9 +1771,10 @@ func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "status-snapshot", ModelConfigID: model.ID, @@ -1817,7 +1843,7 @@ func TestPersistToolResultWithBinaryData(t *testing.T) { // /chat/completions endpoint, where the mock server supports // streaming tool calls. The default "openai" provider routes to // /responses which only handles text deltas in the mock. - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) ctrl := gomock.NewController(t) @@ -1865,10 +1891,11 @@ func TestPersistToolResultWithBinaryData(t *testing.T) { }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "binary-tool-result", - ModelConfigID: model.ID, - WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "binary-tool-result", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Read /home/coder/binary_file.bin."), }, @@ -1988,7 +2015,7 @@ func TestDynamicToolCallPausesAndResumes(t *testing.T) { ) }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) // Dynamic tools do not need a workspace connection, but the // chatd server always builds workspace tools. Use an active @@ -2012,9 +2039,10 @@ func TestDynamicToolCallPausesAndResumes(t *testing.T) { require.NoError(t, err) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "dynamic-tool-pause-resume", - ModelConfigID: model.ID, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "dynamic-tool-pause-resume", + ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Please call the dynamic tool."), }, @@ -2189,7 +2217,7 @@ func TestDynamicToolCallMixedWithBuiltIn(t *testing.T) { ) }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) server := newActiveTestServer(t, db, ps) // Create a chat with a dynamic tool. @@ -2207,9 +2235,10 @@ func TestDynamicToolCallMixedWithBuiltIn(t *testing.T) { require.NoError(t, err) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "mixed-builtin-dynamic", - ModelConfigID: model.ID, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "mixed-builtin-dynamic", + ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Call both tools."), }, @@ -2327,7 +2356,7 @@ func TestSubmitToolResultsConcurrency(t *testing.T) { ) }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) server := newActiveTestServer(t, db, ps) // Create a chat with a dynamic tool. @@ -2345,9 +2374,10 @@ func TestSubmitToolResultsConcurrency(t *testing.T) { require.NoError(t, err) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "concurrency-tool-results", - ModelConfigID: model.ID, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "concurrency-tool-results", + ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Please call the dynamic tool."), }, @@ -2470,9 +2500,10 @@ func TestSubscribeNoPubsubNoDuplicateMessageParts(t *testing.T) { replica := newTestServer(t, db, nil, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "no-dup-parts", ModelConfigID: model.ID, @@ -2517,10 +2548,11 @@ func TestSubscribeAfterMessageID(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Create a chat — this inserts one initial "user" message. chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "after-id-test", ModelConfigID: model.ID, @@ -2694,6 +2726,7 @@ func TestCreateWorkspaceTool_EndToEnd(t *testing.T) { require.NoError(t, err) chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -2865,6 +2898,7 @@ func TestStartWorkspaceTool_EndToEnd(t *testing.T) { // Create a chat with the stopped workspace pre-associated. chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -2986,15 +3020,16 @@ func TestStoppedWorkspaceWithPersistedAgentBindingDoesNotBlockChat(t *testing.T) ) }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) inactive := newTestServer(t, db, ps, uuid.New()) chat, err := inactive.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "stopped-workspace-regression", - ModelConfigID: model.ID, - WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "stopped-workspace-regression", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Run echo hi in the workspace."), }, @@ -3126,7 +3161,7 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { if !req.Stream { return chattest.OpenAINonStreamingResponse("ok") @@ -3144,11 +3179,6 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { // Create a workspace with a full build chain so we can verify // both last_used_at (dormancy) and deadline (autostop) bumps. - org := dbgen.Organization(t, db, database.Organization{}) - _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user.ID, - OrganizationID: org.ID, - }) tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ OrganizationID: org.ID, CreatedBy: user.ID, @@ -3223,6 +3253,7 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { // the chatd server processes everything under that role. chatCtx := dbauthz.AsChatd(ctx) chat, err := server.CreateChat(chatCtx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "usage-tracking-test", ModelConfigID: model.ID, @@ -3304,7 +3335,7 @@ func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { if !req.Stream { return chattest.OpenAINonStreamingResponse("ok") @@ -3342,6 +3373,7 @@ func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { // Create a chat WITHOUT linking a workspace. chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "no-workspace-test", ModelConfigID: model.ID, @@ -3459,23 +3491,29 @@ func seedChatDependencies( ctx context.Context, t *testing.T, db database.Store, -) (database.User, database.ChatModelConfig) { +) (database.User, database.Organization, database.ChatModelConfig) { t.Helper() return seedChatDependenciesWithProvider(ctx, t, db, "openai", "") } -// seedChatDependenciesWithProvider creates a user, chat provider, and -// model config for the given provider type and base URL. +// seedChatDependenciesWithProvider creates a user, organization, +// chat provider, and model config for the given provider type and +// base URL. func seedChatDependenciesWithProvider( ctx context.Context, t *testing.T, db database.Store, provider string, baseURL string, -) (database.User, database.ChatModelConfig) { +) (database.User, database.Organization, database.ChatModelConfig) { t.Helper() user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: provider, DisplayName: provider, @@ -3499,7 +3537,7 @@ func seedChatDependenciesWithProvider( Options: json.RawMessage(`{}`), }) require.NoError(t, err) - return user, model + return user, org, model } func seedChatDependenciesWithProviderPolicy( @@ -3512,10 +3550,15 @@ func seedChatDependenciesWithProviderPolicy( centralAPIKeyEnabled bool, allowUserAPIKey bool, allowCentralAPIKeyFallback bool, -) (database.User, database.ChatProvider, database.ChatModelConfig) { +) (database.User, database.Organization, database.ChatProvider, database.ChatModelConfig) { t.Helper() user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) providerConfig, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: provider, DisplayName: provider, @@ -3543,7 +3586,7 @@ func seedChatDependenciesWithProviderPolicy( }) require.NoError(t, err) - return user, providerConfig, model + return user, org, providerConfig, model } func waitForTerminalChatStatusEvent( @@ -3712,10 +3755,11 @@ func TestInterruptChatDoesNotSendWebPushNotification(t *testing.T) { require.NoError(t, server.Close()) }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "interrupt-no-push", ModelConfigID: model.ID, @@ -3823,10 +3867,11 @@ func TestSuccessfulChatSendsWebPushWithNavigationData(t *testing.T) { require.NoError(t, server.Close()) }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "push-nav-test", ModelConfigID: model.ID, @@ -3907,10 +3952,11 @@ func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) require.NoError(t, serverA.Close()) }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := serverA.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "shutdown-retry", ModelConfigID: model.ID, @@ -4011,10 +4057,11 @@ func TestSuccessfulChatSendsWebPushWithSummary(t *testing.T) { require.NoError(t, server.Close()) }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) _, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "summary-push-test", ModelConfigID: model.ID, @@ -4071,10 +4118,11 @@ func TestSuccessfulChatSendsWebPushFallbackWithoutSummaryForEmptyAssistantText(t require.NoError(t, server.Close()) }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) _, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "empty-summary-push-test", ModelConfigID: model.ID, @@ -4232,7 +4280,7 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { }) // Seed the DB: user, openai-compat provider, model config. - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) // Add an Anthropic provider pointing to our mock server. _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ @@ -4290,10 +4338,11 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { // Create a root chat with a workspace so the child inherits it. chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "computer-use-detection", - ModelConfigID: model.ID, - WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "computer-use-detection", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Use the desktop to check the UI"), }, @@ -4437,10 +4486,11 @@ func TestInterruptChatPersistsPartialResponse(t *testing.T) { require.NoError(t, server.Close()) }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "interrupt-persist-test", ModelConfigID: model.ID, @@ -4549,7 +4599,7 @@ func TestProcessChat_UserProviderKey_Success(t *testing.T) { ) }) - user, provider, model := seedChatDependenciesWithProviderPolicy( + user, org, provider, model := seedChatDependenciesWithProviderPolicy( ctx, t, db, @@ -4569,9 +4619,10 @@ func TestProcessChat_UserProviderKey_Success(t *testing.T) { creator := newTestServer(t, db, ps, uuid.New()) chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "user-provider-key-success", - ModelConfigID: model.ID, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "user-provider-key-success", + ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("say hello"), }, @@ -4614,7 +4665,7 @@ func TestProcessChat_UserProviderKey_MissingKeyError(t *testing.T) { ) }) - user, _, model := seedChatDependenciesWithProviderPolicy( + user, org, _, model := seedChatDependenciesWithProviderPolicy( ctx, t, db, @@ -4628,9 +4679,10 @@ func TestProcessChat_UserProviderKey_MissingKeyError(t *testing.T) { creator := newTestServer(t, db, ps, uuid.New()) chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "user-provider-key-missing", - ModelConfigID: model.ID, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "user-provider-key-missing", + ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("say hello"), }, @@ -4678,16 +4730,17 @@ func TestProcessChatPanicRecovery(t *testing.T) { }) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) // Pass the panic wrapper to the server, but use the real // database for seeding so those operations don't panic. server := newActiveTestServer(t, panicWrapper, ps) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "panic-recovery", - ModelConfigID: model.ID, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "panic-recovery", + ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("hello"), }, @@ -4813,7 +4866,7 @@ func TestMCPServerToolInvocation(t *testing.T) { ) }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) // Seed the MCP server config in the database. This must // happen after seedChatDependencies so user.ID exists for @@ -4855,10 +4908,12 @@ func TestMCPServerToolInvocation(t *testing.T) { }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "mcp-tool-test", ModelConfigID: model.ID, - WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, - MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "mcp-tool-test", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Echo something via MCP."), }, @@ -5045,7 +5100,7 @@ func TestMCPServerOAuth2TokenRefresh(t *testing.T) { ) }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) // Seed the MCP server config with OAuth2 auth pointing to our // mock token endpoint. @@ -5098,11 +5153,12 @@ func TestMCPServerOAuth2TokenRefresh(t *testing.T) { }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "oauth2-refresh-test", - ModelConfigID: model.ID, - WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, - MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "oauth2-refresh-test", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Echo something via the authed MCP."), }, @@ -5180,7 +5236,7 @@ func TestMCPServerOAuth2TokenRefreshFailureGraceful(t *testing.T) { ) }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ DisplayName: "Broken MCP", @@ -5212,10 +5268,11 @@ func TestMCPServerOAuth2TokenRefreshFailureGraceful(t *testing.T) { server := newActiveTestServer(t, db, ps) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "graceful-degradation-test", - ModelConfigID: model.ID, - MCPServerIDs: []uuid.UUID{mcpConfig.ID}, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "graceful-degradation-test", + ModelConfigID: model.ID, + MCPServerIDs: []uuid.UUID{mcpConfig.ID}, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Hello, just reply."), }, @@ -5299,14 +5356,9 @@ func TestChatTemplateAllowlistEnforcement(t *testing.T) { } }) - user, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) // Create two templates the user can see. - org := dbgen.Organization(t, db, database.Organization{}) - _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user.ID, - OrganizationID: org.ID, - }) tplAllowed = dbgen.Template(t, db, database.Template{ OrganizationID: org.ID, CreatedBy: user.ID, @@ -5340,9 +5392,10 @@ func TestChatTemplateAllowlistEnforcement(t *testing.T) { }) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - Title: "allowlist-test", - ModelConfigID: model.ID, + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "allowlist-test", + ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("Test allowlist enforcement"), }, @@ -5448,11 +5501,12 @@ func TestSignalWakeImmediateAcquisition(t *testing.T) { cfg.InFlightChatStaleAfter = testutil.WaitSuperLong }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // CreateChat sets status=pending and calls signalWake(). chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "wake-test", ModelConfigID: model.ID, @@ -5510,11 +5564,12 @@ func TestSignalWakeSendMessage(t *testing.T) { cfg.InFlightChatStaleAfter = testutil.WaitSuperLong }) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // CreateChat triggers wake -> processes first turn. chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "wake-send-test", ModelConfigID: model.ID, @@ -5645,7 +5700,8 @@ func TestAgentContextFilesAndSkillsLoadedIntoChat(t *testing.T) { workspaceID := workspace.ID chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ - WorkspaceID: &workspaceID, + OrganizationID: user.OrganizationID, + WorkspaceID: &workspaceID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, diff --git a/coderd/x/chatd/chatprompt/chatprompt_test.go b/coderd/x/chatd/chatprompt/chatprompt_test.go index 4507fc533d..73815fcb33 100644 --- a/coderd/x/chatd/chatprompt/chatprompt_test.go +++ b/coderd/x/chatd/chatprompt/chatprompt_test.go @@ -1482,7 +1482,14 @@ func TestNulEscapeRoundTrip(t *testing.T) { }) require.NoError(t, err) + org := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: model.ID, @@ -1942,6 +1949,11 @@ func TestMediaToolResultRoundTrip(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: "anthropic", @@ -1981,6 +1993,7 @@ func TestMediaToolResultRoundTrip(t *testing.T) { t.Helper() chat, chatErr := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: model.ID, diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index 72d4241c1d..c3c7ca75e3 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -62,6 +62,7 @@ type AgentConnFunc func( type CreateWorkspaceOptions struct { DB database.Store OwnerID uuid.UUID + OrganizationID uuid.UUID ChatID uuid.UUID CreateFn CreateWorkspaceFn AgentConnFn AgentConnFunc @@ -142,6 +143,24 @@ func CreateWorkspace(options CreateWorkspaceOptions) fantasy.AgentTool { ctx = ownerCtx } + // Verify the template belongs to the same org as the + // chat. Without this check the tool could silently + // bind a cross-org workspace to the chat. + if options.DB != nil && options.OrganizationID != uuid.Nil { + tmpl, tmplErr := options.DB.GetTemplateByID(ctx, templateID) + if tmplErr != nil { + return fantasy.NewTextErrorResponse( + xerrors.Errorf("look up template: %w", tmplErr).Error(), + ), nil + } + if tmpl.OrganizationID != options.OrganizationID { + return fantasy.NewTextErrorResponse( + "template belongs to a different organization than this chat; " + + "use list_templates to find templates in the correct organization", + ), nil + } + } + var ttlMs *int64 if options.DB != nil { raw, err := options.DB.GetChatWorkspaceTTL(ctx) @@ -168,7 +187,12 @@ func CreateWorkspace(options CreateWorkspaceOptions) fantasy.AgentTool { TTLMillis: ttlMs, } - // Resolve workspace name. + // Resolve workspace name. This does a second + // GetTemplateByID when no name is provided; the first + // is the org-validation check above. Consolidating + // them would couple the security gate to the + // name-fallback path, and the cost is negligible next + // to the workspace build that follows. name := strings.TrimSpace(args.Name) if name == "" { seed := "workspace" diff --git a/coderd/x/chatd/chattool/createworkspace_test.go b/coderd/x/chatd/chattool/createworkspace_test.go index 09fab451cd..89ac86eb11 100644 --- a/coderd/x/chatd/chattool/createworkspace_test.go +++ b/coderd/x/chatd/chattool/createworkspace_test.go @@ -507,6 +507,69 @@ func TestCreateWorkspace_GlobalTTL(t *testing.T) { } } +func TestCreateWorkspace_RejectsCrossOrgTemplate(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + + ownerID := uuid.New() + chatOrgID := uuid.New() + templateOrgID := uuid.New() // Different org. + templateID := uuid.New() + + chatID := uuid.New() + + // Chat exists but has no workspace binding. + db.EXPECT(). + GetChatByID(gomock.Any(), chatID). + Return(database.Chat{ + ID: chatID, + WorkspaceID: uuid.NullUUID{}, + }, nil) + + db.EXPECT(). + GetAuthorizationUserRoles(gomock.Any(), ownerID). + Return(database.GetAuthorizationUserRolesRow{ + ID: ownerID, + Roles: []string{}, + Groups: []string{}, + Status: database.UserStatusActive, + }, nil) + + db.EXPECT(). + GetTemplateByID(gomock.Any(), templateID). + Return(database.Template{ + ID: templateID, + OrganizationID: templateOrgID, + Name: "wrong-org-template", + }, nil) + + createCalled := false + tool := CreateWorkspace(CreateWorkspaceOptions{ + DB: db, + OwnerID: ownerID, + OrganizationID: chatOrgID, + ChatID: chatID, + CreateFn: func(context.Context, uuid.UUID, codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { + createCalled = true + return codersdk.Workspace{}, nil + }, + WorkspaceMu: &sync.Mutex{}, + Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + }) + + input := fmt.Sprintf(`{"template_id":%q}`, templateID.String()) + resp, err := tool.Run(context.Background(), fantasy.ToolCall{ + ID: "call-1", + Name: "create_workspace", + Input: input, + }) + require.NoError(t, err) + require.False(t, createCalled, "CreateFn must not be called for cross-org template") + require.Contains(t, resp.Content, "organization") +} + func TestCheckExistingWorkspace_ConnectedAgent(t *testing.T) { t.Parallel() ctrl := gomock.NewController(t) diff --git a/coderd/x/chatd/chattool/listtemplates.go b/coderd/x/chatd/chattool/listtemplates.go index 9c81412497..2ba3e5f58b 100644 --- a/coderd/x/chatd/chattool/listtemplates.go +++ b/coderd/x/chatd/chattool/listtemplates.go @@ -24,6 +24,7 @@ const listTemplatesPageSize = 10 type ListTemplatesOptions struct { DB database.Store OwnerID uuid.UUID + OrganizationID uuid.UUID AllowedTemplateIDs func() map[uuid.UUID]bool } @@ -55,7 +56,8 @@ func ListTemplates(options ListTemplatesOptions) fantasy.AgentTool { } filterParams := database.GetTemplatesWithFilterParams{ - Deleted: false, + Deleted: false, + OrganizationID: options.OrganizationID, Deprecated: sql.NullBool{ Bool: false, Valid: true, @@ -121,8 +123,9 @@ func ListTemplates(options ListTemplatesOptions) fantasy.AgentTool { items := make([]map[string]any, 0, len(pageTemplates)) for _, t := range pageTemplates { item := map[string]any{ - "id": t.ID.String(), - "name": t.Name, + "id": t.ID.String(), + "name": t.Name, + "organization_id": t.OrganizationID.String(), } if display := strings.TrimSpace(t.DisplayName); display != "" { item["display_name"] = display diff --git a/coderd/x/chatd/chattool/startworkspace_test.go b/coderd/x/chatd/chattool/startworkspace_test.go index 19961be1a6..33cd8089f6 100644 --- a/coderd/x/chatd/chattool/startworkspace_test.go +++ b/coderd/x/chatd/chattool/startworkspace_test.go @@ -35,8 +35,14 @@ func TestStartWorkspace(t *testing.T) { user := dbgen.User(t, db, database.User{}) modelCfg := seedModelConfig(ctx, t, db, user.ID) + org := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, @@ -80,6 +86,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -161,6 +168,7 @@ func TestStartWorkspace(t *testing.T) { require.NotEqual(t, uuid.Nil, preferredAgentID) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -221,6 +229,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -284,6 +293,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -341,6 +351,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -412,6 +423,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -507,6 +519,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -588,6 +601,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -685,6 +699,7 @@ func TestStartWorkspace(t *testing.T) { ws := wsResp.Workspace chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, diff --git a/coderd/x/chatd/integration_test.go b/coderd/x/chatd/integration_test.go index f0e80dde32..185d36b216 100644 --- a/coderd/x/chatd/integration_test.go +++ b/coderd/x/chatd/integration_test.go @@ -41,7 +41,7 @@ func TestAnthropicWebSearchRoundTrip(t *testing.T) { client := coderdtest.New(t, &coderdtest.Options{ DeploymentValues: deploymentValues, }) - _ = coderdtest.CreateFirstUser(t, client) + user := coderdtest.CreateFirstUser(t, client) expClient := codersdk.NewExperimentalClient(client) // Configure an Anthropic provider with the real API key. @@ -73,6 +73,7 @@ func TestAnthropicWebSearchRoundTrip(t *testing.T) { // --- Step 1: Send a message that triggers web_search --- t.Log("Creating chat with web search query...") chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -301,7 +302,7 @@ func TestOpenAIReasoningRoundTrip(t *testing.T) { client := coderdtest.New(t, &coderdtest.Options{ DeploymentValues: deploymentValues, }) - _ = coderdtest.CreateFirstUser(t, client) + user := coderdtest.CreateFirstUser(t, client) expClient := codersdk.NewExperimentalClient(client) // Configure an OpenAI provider with the real API key. @@ -337,6 +338,7 @@ func TestOpenAIReasoningRoundTrip(t *testing.T) { // --- Step 1: Send a message that triggers reasoning --- t.Log("Creating chat with reasoning query...") chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -455,7 +457,7 @@ func TestOpenAIReasoningRoundTripStoreFalse(t *testing.T) { client := coderdtest.New(t, &coderdtest.Options{ DeploymentValues: deploymentValues, }) - _ = coderdtest.CreateFirstUser(t, client) + user := coderdtest.CreateFirstUser(t, client) expClient := codersdk.NewExperimentalClient(client) // Configure an OpenAI provider with the real API key. @@ -490,6 +492,7 @@ func TestOpenAIReasoningRoundTripStoreFalse(t *testing.T) { // --- Step 1: Send a message that triggers reasoning --- t.Log("Creating chat with reasoning query...") chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, diff --git a/coderd/x/chatd/recording_internal_test.go b/coderd/x/chatd/recording_internal_test.go index 67ea48dc2f..7f0acc5edd 100644 --- a/coderd/x/chatd/recording_internal_test.go +++ b/coderd/x/chatd/recording_internal_test.go @@ -70,6 +70,7 @@ func createComputerUseParentChild( t *testing.T, server *Server, user database.User, + org database.Organization, model database.ChatModelConfig, workspace database.WorkspaceTable, agent database.WorkspaceAgent, @@ -80,6 +81,7 @@ func createComputerUseParentChild( // Insert the parent chat directly via DB to avoid triggering // the server's background processing. parent, err := server.db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true}, @@ -93,6 +95,7 @@ func createComputerUseParentChild( // the server's background processing (which would try to run // the chat without an LLM and get stuck). child, err = server.db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true}, @@ -155,7 +158,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create the server WITHOUT agentConnFn so the background @@ -163,7 +166,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) parent, child := createComputerUseParentChild( - ctx, t, server, user, model, workspace, agent, + ctx, t, server, user, org, model, workspace, agent, "parent-recording", "computer-use-child", ) @@ -235,13 +238,13 @@ func TestWaitAgentComputerUseRecordingWithThumbnail(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) parent, child := createComputerUseParentChild( - ctx, t, server, user, model, workspace, agent, + ctx, t, server, user, org, model, workspace, agent, "parent-recording-thumb", "computer-use-child-thumb", ) @@ -316,12 +319,12 @@ func TestWaitAgentNonComputerUseNoRecording(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) // Create parent and regular (non-computer_use) child. - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) // Add an assistant message so the report is extracted. insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Done.") @@ -366,7 +369,7 @@ func TestWaitAgentRecordingStartFails(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create the server WITHOUT agentConnFn so the background @@ -375,7 +378,7 @@ func TestWaitAgentRecordingStartFails(t *testing.T) { // Create parent + computer_use child. parent, child := createComputerUseParentChild( - ctx, t, server, user, model, workspace, agent, + ctx, t, server, user, org, model, workspace, agent, "parent-start-fail", "computer-use-start-fail", ) @@ -419,7 +422,7 @@ func TestWaitAgentRecordingStopFails(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create the server WITHOUT agentConnFn so the background @@ -428,7 +431,7 @@ func TestWaitAgentRecordingStopFails(t *testing.T) { // Create parent + computer_use child. parent, child := createComputerUseParentChild( - ctx, t, server, user, model, workspace, agent, + ctx, t, server, user, org, model, workspace, agent, "parent-stop-fail", "computer-use-stop-fail", ) @@ -480,12 +483,12 @@ func TestWaitAgentTimeoutLeavesRecordingRunning(t *testing.T) { // Use the mock clock server; don't set agentConnFn yet. server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create parent + computer_use child. _, child := createComputerUseParentChild( - ctx, t, server, user, model, workspace, agent, + ctx, t, server, user, org, model, workspace, agent, "parent-timeout", "computer-use-timeout", ) @@ -564,7 +567,7 @@ func TestStopAndStoreRecording_Oversized(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -612,7 +615,7 @@ func TestStopAndStoreRecording_OversizedThumbnail(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -675,7 +678,7 @@ func TestStopAndStoreRecording_DuplicatePartsIgnored(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -717,7 +720,7 @@ func TestStopAndStoreRecording_Empty(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -747,7 +750,7 @@ func TestStopAndStoreRecording_WithThumbnail(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -800,7 +803,7 @@ func TestStopAndStoreRecording_VideoOnly(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -842,7 +845,7 @@ func TestStopAndStoreRecording_DownloadFailure(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -874,7 +877,7 @@ func TestStopAndStoreRecording_UnknownPartIgnored(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -927,7 +930,7 @@ func TestStopAndStoreRecording_MalformedContentType(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -962,7 +965,7 @@ func TestStopAndStoreRecording_MissingBoundary(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(ctx, t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index f0f5211f0e..300ce08aeb 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -460,6 +460,7 @@ func (p *Server) createChildSubagentChatWithOptions( } insertedChat, err := tx.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: parent.OrganizationID, OwnerID: parent.OwnerID, WorkspaceID: parent.WorkspaceID, BuildID: parent.BuildID, diff --git a/coderd/x/chatd/subagent_context_internal_test.go b/coderd/x/chatd/subagent_context_internal_test.go index a2f908cb33..e0bf95f84b 100644 --- a/coderd/x/chatd/subagent_context_internal_test.go +++ b/coderd/x/chatd/subagent_context_internal_test.go @@ -148,9 +148,10 @@ func createParentChatWithInheritedContext( ) database.Chat { t.Helper() - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-with-context", ModelConfigID: model.ID, @@ -328,9 +329,10 @@ func createParentChatWithRotatedInheritedContext( ) database.Chat { t.Helper() - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-with-rotated-context", ModelConfigID: model.ID, diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 144addbd08..999c552be8 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -111,16 +111,22 @@ func newInternalTestServerWithClock( } // seedInternalChatDeps inserts an OpenAI provider and model config -// into the database and returns the created user and model. This -// deliberately does NOT create an Anthropic provider. +// into the database and returns the created user, organization, +// and model. This deliberately does NOT create an Anthropic +// provider. func seedInternalChatDeps( ctx context.Context, t *testing.T, db database.Store, -) (database.User, database.ChatModelConfig) { +) (database.User, database.Organization, database.ChatModelConfig) { t.Helper() user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: "openai", DisplayName: "OpenAI", @@ -147,7 +153,7 @@ func seedInternalChatDeps( }) require.NoError(t, err) - return user, model + return user, org, model } func seedWorkspaceBinding( @@ -212,11 +218,12 @@ func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) workspace, build, agent := seedWorkspaceBinding(t, db, user.ID) parent, err := server.CreateChat(ctx, CreateOptions{ - OwnerID: user.ID, + OrganizationID: org.ID, + OwnerID: user.ID, WorkspaceID: uuid.NullUUID{ UUID: workspace.ID, Valid: true, @@ -243,6 +250,7 @@ func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) { childChat, err := db.GetChatByID(ctx, child.ID) require.NoError(t, err) + require.Equal(t, parentChat.OrganizationID, childChat.OrganizationID) require.Equal(t, parentChat.WorkspaceID, childChat.WorkspaceID) require.Equal(t, parentChat.BuildID, childChat.BuildID) require.Equal(t, parentChat.AgentID, childChat.AgentID) @@ -257,10 +265,11 @@ func TestSpawnComputerUseAgent_NoAnthropicProvider(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) // Create a root parent chat. parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-no-anthropic", ModelConfigID: model.ID, @@ -288,10 +297,11 @@ func TestSpawnComputerUseAgent_NotAvailableForChildChats(t *testing.T) { }) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) // Create a root parent chat. parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "root-parent", ModelConfigID: model.ID, @@ -301,7 +311,8 @@ func TestSpawnComputerUseAgent_NotAvailableForChildChats(t *testing.T) { // Create a child chat under the parent. child, err := server.CreateChat(ctx, CreateOptions{ - OwnerID: user.ID, + OrganizationID: org.ID, + OwnerID: user.ID, ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -347,8 +358,9 @@ func TestSpawnComputerUseAgent_DesktopDisabled(t *testing.T) { }) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-desktop-disabled", ModelConfigID: model.ID, @@ -374,7 +386,7 @@ func TestSpawnComputerUseAgent_UsesComputerUseModelNotParent(t *testing.T) { }) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) workspace, build, agent := seedWorkspaceBinding(t, db, user.ID) // The parent uses an OpenAI model. @@ -382,7 +394,8 @@ func TestSpawnComputerUseAgent_UsesComputerUseModelNotParent(t *testing.T) { "seed helper must create an OpenAI model") parent, err := server.CreateChat(ctx, CreateOptions{ - OwnerID: user.ID, + OrganizationID: org.ID, + OwnerID: user.ID, WorkspaceID: uuid.NullUUID{ UUID: workspace.ID, Valid: true, @@ -455,7 +468,7 @@ func TestCreateChildSubagentChat_InheritsMCPServerIDs(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) // Insert two MCP server configs so we can verify both are // inherited by the child chat. @@ -493,6 +506,7 @@ func TestCreateChildSubagentChat_InheritsMCPServerIDs(t *testing.T) { // Create a parent chat with MCP servers. parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-with-mcp", ModelConfigID: model.ID, @@ -533,7 +547,7 @@ func TestSpawnComputerUseAgent_InheritsMCPServerIDs(t *testing.T) { }) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) // Insert an MCP server config. mcpCfg, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ @@ -555,6 +569,7 @@ func TestSpawnComputerUseAgent_InheritsMCPServerIDs(t *testing.T) { // Create a parent chat with MCP servers. parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-cu-mcp", ModelConfigID: model.ID, @@ -602,10 +617,11 @@ func TestCreateChildSubagentChat_NoMCPServersStaysEmpty(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) // Create a parent chat without any MCP servers. parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-no-mcp", ModelConfigID: model.ID, @@ -638,10 +654,11 @@ func TestIsSubagentDescendant(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) // Build a chain: root -> child -> grandchild. root, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "root", ModelConfigID: model.ID, @@ -650,7 +667,8 @@ func TestIsSubagentDescendant(t *testing.T) { require.NoError(t, err) child, err := server.CreateChat(ctx, CreateOptions{ - OwnerID: user.ID, + OrganizationID: org.ID, + OwnerID: user.ID, ParentChatID: uuid.NullUUID{ UUID: root.ID, Valid: true, @@ -666,7 +684,8 @@ func TestIsSubagentDescendant(t *testing.T) { require.NoError(t, err) grandchild, err := server.CreateChat(ctx, CreateOptions{ - OwnerID: user.ID, + OrganizationID: org.ID, + OwnerID: user.ID, ParentChatID: uuid.NullUUID{ UUID: child.ID, Valid: true, @@ -683,6 +702,7 @@ func TestIsSubagentDescendant(t *testing.T) { // Build a separate, unrelated chain. unrelated, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "unrelated-root", ModelConfigID: model.ID, @@ -691,7 +711,8 @@ func TestIsSubagentDescendant(t *testing.T) { require.NoError(t, err) unrelatedChild, err := server.CreateChat(ctx, CreateOptions{ - OwnerID: user.ID, + OrganizationID: org.ID, + OwnerID: user.ID, ParentChatID: uuid.NullUUID{ UUID: unrelated.ID, Valid: true, @@ -774,11 +795,13 @@ func createParentChildChats( t *testing.T, server *Server, user database.User, + org database.Organization, model database.ChatModelConfig, ) (parent database.Chat, child database.Chat) { t.Helper() parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent-" + t.Name(), ModelConfigID: model.ID, @@ -787,7 +810,8 @@ func createParentChildChats( require.NoError(t, err) child, err = server.CreateChat(ctx, CreateOptions{ - OwnerID: user.ID, + OrganizationID: org.ID, + OwnerID: user.ID, ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -876,15 +900,16 @@ func TestAwaitSubagentCompletion(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) t.Run("NotDescendant", func(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) - parent, _ := createParentChildChats(ctx, t, server, user, model) + parent, _ := createParentChildChats(ctx, t, server, user, org, model) unrelated, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "unrelated", ModelConfigID: model.ID, @@ -902,7 +927,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") insertAssistantMessage(ctx, t, db, child.ID, model.ID, "task complete") @@ -920,7 +945,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) setChatStatus(ctx, t, db, child.ID, database.ChatStatusError, "something broke") insertAssistantMessage(ctx, t, db, child.ID, model.ID, "partial work done") @@ -936,7 +961,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) setChatStatus(ctx, t, db, child.ID, database.ChatStatusError, "crash") @@ -956,9 +981,9 @@ func TestAwaitSubagentCompletion(t *testing.T) { mClock := quartz.NewMock(t) server := newInternalTestServerWithClock(t, db, nil, chatprovider.ProviderAPIKeys{}, mClock) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) // Set the trap BEFORE starting the goroutine so we // deterministically catch the ticker creation. @@ -1002,9 +1027,9 @@ func TestAwaitSubagentCompletion(t *testing.T) { mClock := quartz.NewMock(t) server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) // signalWake from CreateChat may trigger immediate processing. // Wait for it to settle, then reset chats to the state we need. @@ -1089,7 +1114,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) // signalWake from CreateChat may trigger immediate processing. // Wait for it to settle, then set the terminal state we need. @@ -1113,9 +1138,9 @@ func TestAwaitSubagentCompletion(t *testing.T) { mClock := quartz.NewMock(t) server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock) ctx := chatdTestContext(t) - user, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(ctx, t, db) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) // Trap the timeout timer to know when the function // has entered its poll loop. @@ -1149,7 +1174,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) // signalWake from CreateChat triggers background // processing. drainInflight waits for in-flight goroutines @@ -1181,7 +1206,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) - parent, child := createParentChildChats(ctx, t, server, user, model) + parent, child := createParentChildChats(ctx, t, server, user, org, model) // Pre-complete the child so it returns immediately. setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") diff --git a/coderd/x/chatd/subagent_test.go b/coderd/x/chatd/subagent_test.go index 69ed10e6b5..6a4cb6f50a 100644 --- a/coderd/x/chatd/subagent_test.go +++ b/coderd/x/chatd/subagent_test.go @@ -21,10 +21,11 @@ func TestSpawnComputerUseAgent_CreatesChildWithChatMode(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Create a parent chat. parent, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent", ModelConfigID: model.ID, @@ -37,7 +38,8 @@ func TestSpawnComputerUseAgent_CreatesChildWithChatMode(t *testing.T) { prompt := "Use the desktop to open Firefox" child, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: parent.OwnerID, + OrganizationID: org.ID, + OwnerID: parent.OwnerID, ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -75,9 +77,10 @@ func TestSpawnComputerUseAgent_SystemPromptFormat(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) parent, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent", ModelConfigID: model.ID, @@ -89,7 +92,8 @@ func TestSpawnComputerUseAgent_SystemPromptFormat(t *testing.T) { systemPrompt := "Computer use instructions\n\n" + prompt child, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: parent.OwnerID, + OrganizationID: org.ID, + OwnerID: parent.OwnerID, ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -132,9 +136,10 @@ func TestSpawnComputerUseAgent_ChildIsListedUnderParent(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) parent, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "parent", ModelConfigID: model.ID, @@ -145,7 +150,8 @@ func TestSpawnComputerUseAgent_ChildIsListedUnderParent(t *testing.T) { prompt := "Check the UI layout" child, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: parent.OwnerID, + OrganizationID: org.ID, + OwnerID: parent.OwnerID, ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -175,10 +181,11 @@ func TestSpawnComputerUseAgent_RootChatIDPropagation(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Create a root parent chat (no parent of its own). parent, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "root-parent", ModelConfigID: model.ID, @@ -189,7 +196,8 @@ func TestSpawnComputerUseAgent_RootChatIDPropagation(t *testing.T) { prompt := "Take a screenshot" child, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: parent.OwnerID, + OrganizationID: org.ID, + OwnerID: parent.OwnerID, ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, diff --git a/coderd/x/gitsync/worker_test.go b/coderd/x/gitsync/worker_test.go index ce89a3f77f..a1b664d879 100644 --- a/coderd/x/gitsync/worker_test.go +++ b/coderd/x/gitsync/worker_test.go @@ -954,8 +954,9 @@ func TestWorker(t *testing.T) { // 1. Real database store. db, _ := dbtestutil.NewDB(t) - // 2. Create a user (FK for chats). + // 2. Create a user and an organization (FKs for chats). user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) // 3. Set up FK chain: chat_providers -> chat_model_configs -> chats. _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ @@ -978,6 +979,7 @@ func TestWorker(t *testing.T) { require.NoError(t, err) chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, diff --git a/codersdk/chats.go b/codersdk/chats.go index e73f32d2bf..6d1c4a2175 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -55,6 +55,7 @@ const ( // Chat represents a chat session with an AI agent. type Chat struct { ID uuid.UUID `json:"id" format:"uuid"` + OrganizationID uuid.UUID `json:"organization_id" format:"uuid"` OwnerID uuid.UUID `json:"owner_id" format:"uuid"` WorkspaceID *uuid.UUID `json:"workspace_id,omitempty" format:"uuid"` BuildID *uuid.UUID `json:"build_id,omitempty" format:"uuid"` @@ -381,12 +382,13 @@ type ToolResult struct { // CreateChatRequest is the request to create a new chat. type CreateChatRequest struct { - Content []ChatInputPart `json:"content"` - SystemPrompt string `json:"system_prompt,omitempty"` - WorkspaceID *uuid.UUID `json:"workspace_id,omitempty" format:"uuid"` - ModelConfigID *uuid.UUID `json:"model_config_id,omitempty" format:"uuid"` - MCPServerIDs []uuid.UUID `json:"mcp_server_ids,omitempty" format:"uuid"` - Labels map[string]string `json:"labels,omitempty"` + OrganizationID uuid.UUID `json:"organization_id" format:"uuid"` + Content []ChatInputPart `json:"content"` + SystemPrompt string `json:"system_prompt,omitempty"` + WorkspaceID *uuid.UUID `json:"workspace_id,omitempty" format:"uuid"` + ModelConfigID *uuid.UUID `json:"model_config_id,omitempty" format:"uuid"` + MCPServerIDs []uuid.UUID `json:"mcp_server_ids,omitempty" format:"uuid"` + Labels map[string]string `json:"labels,omitempty"` // UnsafeDynamicTools declares client-executed tools that the // LLM can invoke. This API is highly experimental and highly // subject to change. diff --git a/docs/ai-coder/agents/chats-api.md b/docs/ai-coder/agents/chats-api.md index 55d78b5ab5..666b1a60ea 100644 --- a/docs/ai-coder/agents/chats-api.md +++ b/docs/ai-coder/agents/chats-api.md @@ -26,6 +26,7 @@ curl -X POST https://coder.example.com/api/experimental/chats \ -H "Coder-Session-Token: $CODER_SESSION_TOKEN" \ -H "Content-Type: application/json" \ -d '{ + "organization_id": "", "content": [ {"type": "text", "text": "hello world"} ] @@ -37,6 +38,7 @@ The response is the newly created `Chat` object: ```json { "id": "a1b2c3d4-...", + "organization_id": "...", "owner_id": "...", "workspace_id": null, "build_id": null, @@ -82,6 +84,7 @@ A typical integration follows three steps: | Field | Type | Required | Description | |-------------------|---------------------|----------|-------------------------------------------------| | `content` | `ChatInputPart[]` | yes | The user's prompt as one or more content parts. | +| `organization_id` | `uuid` | yes | The organization this chat belongs to. | | `workspace_id` | `uuid` | no | Pin the chat to a specific workspace. | | `model_config_id` | `uuid` | no | Override the default model configuration. | | `mcp_server_ids` | `uuid[]` | no | Attach MCP servers to this chat. | diff --git a/enterprise/coderd/exp_chats_test.go b/enterprise/coderd/exp_chats_test.go index 36609d8f4f..bfb146038b 100644 --- a/enterprise/coderd/exp_chats_test.go +++ b/enterprise/coderd/exp_chats_test.go @@ -14,6 +14,8 @@ import ( "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/rbac" + "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" @@ -30,7 +32,7 @@ func TestChatStreamRelay(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, pubsub := dbtestutil.NewDB(t) - firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{ + firstClient, firstUser := coderdenttest.New(t, &coderdenttest.Options{ Options: &coderdtest.Options{ Database: db, Pubsub: pubsub, @@ -93,6 +95,7 @@ func TestChatStreamRelay(t *testing.T) { // Create a chat on the first replica chat, err := codersdk.NewExperimentalClient(firstClient).CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "Test chat for relay", @@ -186,7 +189,7 @@ func TestChatStreamRelay(t *testing.T) { certificates := []tls.Certificate{testutil.GenerateTLSCertificate(t, "localhost")} db, pubsub := dbtestutil.NewDB(t) - firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{ + firstClient, firstUser := coderdenttest.New(t, &coderdenttest.Options{ Options: &coderdtest.Options{ Database: db, Pubsub: pubsub, @@ -281,6 +284,7 @@ func TestChatStreamRelay(t *testing.T) { // Create a chat on the first replica. chat, err := codersdk.NewExperimentalClient(firstClient).CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "Test chat for TLS relay", @@ -380,7 +384,7 @@ func TestChatStreamRelay(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, pubsub := dbtestutil.NewDB(t) - firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{ + firstClient, firstUser := coderdenttest.New(t, &coderdenttest.Options{ Options: &coderdtest.Options{ Database: db, Pubsub: pubsub, @@ -450,6 +454,7 @@ func TestChatStreamRelay(t *testing.T) { require.NoError(t, err) chat, err := codersdk.NewExperimentalClient(firstClient).CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "Test cookie-only relay", @@ -547,7 +552,7 @@ func TestChatStreamRelay(t *testing.T) { dv.HTTPCookies.EnableHostPrefix = true dv.HTTPCookies.Secure = true }) - firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{ + firstClient, firstUser := coderdenttest.New(t, &coderdenttest.Options{ Options: &coderdtest.Options{ Database: db, Pubsub: pubsub, @@ -621,6 +626,7 @@ func TestChatStreamRelay(t *testing.T) { require.NoError(t, err) chat, err := codersdk.NewExperimentalClient(firstClient).CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "Test host-prefix relay", @@ -706,7 +712,7 @@ func TestChatStreamRelay(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, pubsub := dbtestutil.NewDB(t) - firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{ + firstClient, firstUser := coderdenttest.New(t, &coderdenttest.Options{ Options: &coderdtest.Options{ Database: db, Pubsub: pubsub, @@ -768,6 +774,7 @@ func TestChatStreamRelay(t *testing.T) { // Create a chat on the first replica. chat, err := codersdk.NewExperimentalClient(firstClient).CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, Text: "Test chat for buffered relay", @@ -1087,3 +1094,81 @@ func (p cookieOnlySessionTokenProvider) SetDialOption(opts *websocket.DialOption } opts.HTTPHeader.Set("Cookie", cookieName+"="+p.token) } + +func TestCreateChatNonDefaultOrg(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + client, firstUser := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: func() *codersdk.DeploymentValues { + v := coderdtest.DeploymentValues(t) + v.Experiments = []string{string(codersdk.ExperimentAgents)} + return v + }(), + }, + LicenseOptions: &coderdenttest.LicenseOptions{ + Features: license.Features{ + codersdk.FeatureMultipleOrganizations: 1, + }, + }, + }) + expClient := codersdk.NewExperimentalClient(client) + + // Set up a chat provider and model config. + provider, err := expClient.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{ + Provider: "openai", + DisplayName: "OpenAI", + APIKey: "test-key", + BaseURL: "https://example.com", + }) + require.NoError(t, err) + _, err = expClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{ + Provider: provider.Provider, + Model: "gpt-4o-mini", + DisplayName: "Test Model", + IsDefault: ptr.Ref(true), + ContextLimit: ptr.Ref(int64(1000)), + CompressionThreshold: ptr.Ref(int32(70)), + }) + require.NoError(t, err) + + // Create a second (non-default) org via the API. + secondOrg := coderdenttest.CreateOrganization(t, client, coderdenttest.CreateOrganizationOptions{}) + + // Create a member in the default org, then add them to the second org. + memberClientRaw, member := coderdtest.CreateAnotherUser( + t, client, firstUser.OrganizationID, rbac.RoleAgentsAccess(), + ) + _, err = client.PostOrganizationMember(ctx, secondOrg.ID, member.Username) + require.NoError(t, err) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + + // Create a chat in the non-default org. + chat, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: secondOrg.ID, + Content: []codersdk.ChatInputPart{ + { + Type: codersdk.ChatInputPartTypeText, + Text: "hello from non-default org", + }, + }, + }) + require.NoError(t, err) + require.Equal(t, secondOrg.ID, chat.OrganizationID) + require.Equal(t, member.ID, chat.OwnerID) + + // Verify the chat is visible when listing. + chats, err := memberClient.ListChats(ctx, nil) + require.NoError(t, err) + var found bool + for _, c := range chats { + if c.ID == chat.ID { + found = true + require.Equal(t, secondOrg.ID, c.OrganizationID) + break + } + } + require.True(t, found, "chat should be visible in list") +} diff --git a/enterprise/coderd/x/chatd/chatd_test.go b/enterprise/coderd/x/chatd/chatd_test.go index 4b1ca1939a..fa44ea9d8a 100644 --- a/enterprise/coderd/x/chatd/chatd_test.go +++ b/enterprise/coderd/x/chatd/chatd_test.go @@ -86,13 +86,13 @@ func newActiveWorkerServer( return server } -// seedChatDependencies creates a user and chat model config in the -// database for use in relay tests. +// seedChatDependencies creates a user, organization, and chat model +// config in the database for use in relay tests. func seedChatDependencies( ctx context.Context, t *testing.T, db database.Store, -) (database.User, database.ChatModelConfig) { +) (database.User, database.Organization, database.ChatModelConfig) { t.Helper() safetyNet := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { @@ -103,6 +103,11 @@ func seedChatDependencies( t.Cleanup(safetyNet.Close) user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ Provider: "openai", DisplayName: "OpenAI", @@ -127,13 +132,14 @@ func seedChatDependencies( Options: json.RawMessage(`{}`), }) require.NoError(t, err) - return user, model + return user, org, model } func seedWaitingChat( ctx context.Context, t *testing.T, db database.Store, + orgID uuid.UUID, user database.User, model database.ChatModelConfig, title string, @@ -141,6 +147,7 @@ func seedWaitingChat( t.Helper() chat, err := db.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: orgID, Status: database.ChatStatusWaiting, OwnerID: user.ID, LastModelConfigID: model.ID, @@ -155,6 +162,7 @@ func seedRemoteRunningChat( ctx context.Context, t *testing.T, db database.Store, + orgID uuid.UUID, user database.User, model database.ChatModelConfig, workerID uuid.UUID, @@ -162,7 +170,7 @@ func seedRemoteRunningChat( ) database.Chat { t.Helper() - chat := seedWaitingChat(ctx, t, db, user, model, title) + chat := seedWaitingChat(ctx, t, db, orgID, user, model, title) now := time.Now() chat, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, @@ -247,9 +255,9 @@ func TestSubscribeRelayReconnectsOnDrop(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) - chat := seedRemoteRunningChat(ctx, t, db, user, model, workerID, "relay-reconnect") + chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-reconnect") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -325,11 +333,11 @@ func TestSubscribeRelayAsyncDoesNotBlock(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Seed a waiting chat so Subscribe does not trigger a synchronous // relay. - chat := seedWaitingChat(ctx, t, db, user, model, "relay-async-nonblock") + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-async-nonblock") // Subscribe before the chat is marked running so the relay opens // via pubsub notification (openRelayAsync path). @@ -427,9 +435,9 @@ func TestSubscribeRelaySnapshotDelivered(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) - chat := seedRemoteRunningChat(ctx, t, db, user, model, workerID, "relay-snapshot") + chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-snapshot") initialSnapshot, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -515,10 +523,11 @@ func TestSubscribeRetryEventAcrossInstances(t *testing.T) { }, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := worker.CreateChat(ctx, osschatd.CreateOptions{ + OrganizationID: org.ID, OwnerID: user.ID, Title: "retry-across-instances", ModelConfigID: model.ID, @@ -651,11 +660,11 @@ func TestSubscribeRelayStaleDialDiscardedAfterInterrupt(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Seed the chat in waiting state so Subscribe does not try an initial // relay. - chat := seedWaitingChat(ctx, t, db, user, model, "stale-dial-test") + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "stale-dial-test") // Subscribe while chat is in "waiting" state — no relay opened. _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) @@ -803,11 +812,11 @@ func TestSubscribeCancelDuringInFlightDial(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Seed the chat in waiting state so Subscribe does not open a // synchronous relay. - chat := seedWaitingChat(ctx, t, db, user, model, "cancel-inflight-dial") + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "cancel-inflight-dial") _, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -889,10 +898,10 @@ func TestSubscribeRelayRunningToRunningSwitch(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Seed the chat in waiting state so Subscribe does not open a relay. - chat := seedWaitingChat(ctx, t, db, user, model, "running-to-running") + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "running-to-running") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -997,11 +1006,11 @@ func TestSubscribeRelayFailedDialRetries(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) // Seed the chat in waiting state so Subscribe does not open a // synchronous relay dial. - chat := seedWaitingChat(ctx, t, db, user, model, "failed-dial-retry") + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "failed-dial-retry") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -1093,12 +1102,13 @@ func TestSubscribeRunningLocalWorkerClosesRelay(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat := seedRemoteRunningChat( ctx, t, db, + org.ID, user, model, remoteWorkerID, @@ -1192,12 +1202,13 @@ func TestSubscribeRelayMultipleReconnects(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) chat := seedRemoteRunningChat( ctx, t, db, + org.ID, user, model, workerID, @@ -1334,13 +1345,13 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) { }, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // Create the chat in waiting state so the subscriber sees it // before the worker picks it up (avoids the synchronous relay // path in Subscribe). - chat := seedWaitingChat(ctx, t, db, user, model, "fast-completion-relay-race") + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "fast-completion-relay-race") // Subscribe from the subscriber replica while the chat is idle. // No relay is opened because the chat is in waiting state. @@ -1496,11 +1507,11 @@ func TestSubscribeRelayEstablishedMidStream(t *testing.T) { // call) involves multiple DB round-trips that can be slow under // load. ctx := testutil.Context(t, testutil.WaitSuperLong) - user, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(ctx, t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // Create the chat in waiting state. - chat := seedWaitingChat(ctx, t, db, user, model, "mid-stream-relay") + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "mid-stream-relay") // Subscribe from the subscriber replica while the chat is idle. _, events, subCancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 148b23a9b3..fccde9fedc 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -80,6 +80,7 @@ const makeChat = ( overrides?: Partial, ): TypesGen.Chat => ({ id, + organization_id: "test-org-id", owner_id: "owner-1", last_model_config_id: "model-1", mcp_server_ids: [], diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 1b8db2f6e1..e55158ecb1 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1186,6 +1186,7 @@ export interface ChangePasswordWithOneTimePasscodeRequest { */ export interface Chat { readonly id: string; + readonly organization_id: string; readonly owner_id: string; readonly workspace_id?: string; readonly build_id?: string; @@ -2488,6 +2489,7 @@ export interface CreateChatProviderConfigRequest { * CreateChatRequest is the request to create a new chat. */ export interface CreateChatRequest { + readonly organization_id: string; readonly content: readonly ChatInputPart[]; readonly system_prompt?: string; readonly workspace_id?: string; diff --git a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx index b47fe44902..af610a67f1 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -121,6 +121,7 @@ const mockModelConfigs: TypesGen.ChatModelConfig[] = [ ]; const baseChatFields = { + organization_id: "test-org-id", owner_id: "owner-id", workspace_id: mockWorkspace.id, last_model_config_id: MODEL_CONFIG_ID, diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index bd0a5c90e8..ced1b0506f 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -40,6 +40,7 @@ const oneWeekAgo = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000).toISOString(); const buildChat = (overrides: Partial = {}): TypesGen.Chat => ({ id: AGENT_ID, + organization_id: "test-org-id", owner_id: "owner-1", title: "Help me refactor", status: "completed", diff --git a/site/src/pages/AgentsPage/AgentCreatePage.tsx b/site/src/pages/AgentsPage/AgentCreatePage.tsx index 5fa2ad57e8..b0f332e973 100644 --- a/site/src/pages/AgentsPage/AgentCreatePage.tsx +++ b/site/src/pages/AgentsPage/AgentCreatePage.tsx @@ -44,6 +44,7 @@ const AgentCreatePage: FC = () => { workspaceId, model, mcpServerIds, + organizationId, }: CreateChatOptions) => { const modelConfigID = model || nilUUID; const content: TypesGen.ChatInputPart[] = []; @@ -56,6 +57,7 @@ const AgentCreatePage: FC = () => { } } const createdChat = await createMutation.mutateAsync({ + organization_id: organizationId, content, workspace_id: workspaceId, model_config_id: modelConfigID, diff --git a/site/src/pages/AgentsPage/AgentsPageView.stories.tsx b/site/src/pages/AgentsPage/AgentsPageView.stories.tsx index 043e04d8f6..afdb33d2b9 100644 --- a/site/src/pages/AgentsPage/AgentsPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentsPageView.stories.tsx @@ -128,6 +128,7 @@ const todayTimestamp = new Date().toISOString(); const buildChat = (overrides: Partial = {}): Chat => ({ id: "chat-default", + organization_id: "test-org-id", owner_id: "owner-1", title: "Agent", status: "completed", diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx index 14e6ee73ab..4a304e1eeb 100644 --- a/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx +++ b/site/src/pages/AgentsPage/components/AgentCreateForm.stories.tsx @@ -1,6 +1,11 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; -import { expect, fn, userEvent, waitFor, within } from "storybook/test"; -import { MockWorkspace } from "#/testHelpers/entities"; +import { expect, fn, screen, userEvent, waitFor, within } from "storybook/test"; +import { ConfirmDialog } from "#/components/Dialogs/ConfirmDialog/ConfirmDialog"; +import { + MockDefaultOrganization, + MockOrganization2, + MockWorkspace, +} from "#/testHelpers/entities"; import { withDashboardProvider } from "#/testHelpers/storybook"; import { AgentCreateForm } from "./AgentCreateForm"; @@ -310,6 +315,51 @@ export const ForbiddenErrorWithRole: Story = { }, }; +export const WithOrganizationPicker: Story = { + parameters: { + showOrganizations: true, + organizations: [MockDefaultOrganization, MockOrganization2], + }, +}; + +/** + * Standalone story for the org-change confirmation dialog. Renders + * the ConfirmDialog directly in its open state, following the same + * pattern as DeleteConfirmationDialog in AgentsPageView.stories. + */ +export const OrgChangeConfirmation: Story = { + render: () => ( + + ), + play: async () => { + const dialog = await screen.findByRole("dialog"); + await expect(dialog).toBeInTheDocument(); + await expect( + within(dialog).getByText("Change organization?"), + ).toBeInTheDocument(); + await expect( + within(dialog).getByText( + "Changing organization will remove your current attachments.", + ), + ).toBeInTheDocument(); + await expect( + within(dialog).getByRole("button", { name: /continue/i }), + ).toBeInTheDocument(); + await expect( + within(dialog).getByRole("button", { name: /cancel/i }), + ).toBeInTheDocument(); + }, +}; + export const ForbiddenNoAgentsRole: Story = { args: { ...defaultArgs, diff --git a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx index 73bc0a6ef3..19ec6c2358 100644 --- a/site/src/pages/AgentsPage/components/AgentCreateForm.tsx +++ b/site/src/pages/AgentsPage/components/AgentCreateForm.tsx @@ -6,6 +6,9 @@ import type * as TypesGen from "#/api/typesGenerated"; import { Alert, AlertDescription } from "#/components/Alert/Alert"; import { ErrorAlert } from "#/components/Alert/ErrorAlert"; import { Button } from "#/components/Button/Button"; +import { ConfirmDialog } from "#/components/Dialogs/ConfirmDialog/ConfirmDialog"; +import { Label } from "#/components/Label/Label"; +import { OrganizationAutocomplete } from "#/components/OrganizationAutocomplete/OrganizationAutocomplete"; import { useDashboard } from "#/modules/dashboard/useDashboard"; import { docs } from "#/utils/docs"; import { useFileAttachments } from "../hooks/useFileAttachments"; @@ -42,6 +45,7 @@ export type CreateChatOptions = { workspaceId?: string; model?: string; mcpServerIds?: string[]; + organizationId: string; }; /** @@ -144,7 +148,7 @@ export const AgentCreateForm: FC = ({ workspacesError, isWorkspacesLoading, }) => { - const { organizations } = useDashboard(); + const { organizations, showOrganizations } = useDashboard(); const { initialInputValue, initialEditorState, @@ -184,9 +188,40 @@ export const AgentCreateForm: FC = ({ : preferredModelID; const [selectedWorkspaceId, setSelectedWorkspaceId] = useState( () => { - return localStorage.getItem(selectedWorkspaceIdStorageKey) || null; + const stored = localStorage.getItem(selectedWorkspaceIdStorageKey); + if (!stored) return null; + + // If workspaces haven't loaded yet, keep the stored value. + // It will be re-validated once the list arrives via + // filteredWorkspaces clearing the selection if stale. + if (workspaceOptions.length === 0) return stored; + + // Validate the stored workspace still exists and belongs + // to the initial org. Without this, a workspace from a + // previously selected org persists across sessions and + // gets submitted even though it's hidden from the picker. + const workspace = workspaceOptions.find((ws) => ws.id === stored); + if (!workspace) { + localStorage.removeItem(selectedWorkspaceIdStorageKey); + return null; + } + if ( + showOrganizations && + organizations[0] && + workspace.organization_id !== organizations[0].id + ) { + localStorage.removeItem(selectedWorkspaceIdStorageKey); + return null; + } + return stored; }, ); + const [selectedOrg, setSelectedOrg] = useState( + organizations[0] ?? null, + ); + const [pendingOrgChange, setPendingOrgChange] = + useState(null); + const organizationId = selectedOrg?.id ?? ""; const hasModelOptions = modelOptions.length > 0; const hasConfiguredModels = hasConfiguredModelsInCatalog(modelCatalog); const hasUserFixableModelProviders = hasUserFixableProviders(modelCatalog); @@ -225,6 +260,7 @@ export const AgentCreateForm: FC = ({ // the shared input component re-rendering on every change. const selectedWorkspaceIdRef = useRef(selectedWorkspaceId); const selectedModelRef = useRef(selectedModel); + const organizationIdRef = useRef(organizationId); const [userMCPServerIds, setUserMCPServerIds] = useState( null, ); @@ -243,6 +279,7 @@ export const AgentCreateForm: FC = ({ selectedWorkspaceIdRef.current = selectedWorkspaceId; selectedModelRef.current = selectedModel; selectedMCPServerIdsRef.current = effectiveMCPServerIds; + organizationIdRef.current = organizationId; }); const handleWorkspaceChange = (value: string | null) => { if (value === null) { @@ -266,6 +303,7 @@ export const AgentCreateForm: FC = ({ fileIDs, workspaceId: selectedWorkspaceIdRef.current ?? undefined, model: selectedModelRef.current || undefined, + organizationId: organizationIdRef.current, mcpServerIds: selectedMCPServerIdsRef.current.length > 0 ? [...selectedMCPServerIdsRef.current] @@ -288,7 +326,7 @@ export const AgentCreateForm: FC = ({ handleAttach, handleRemoveAttachment, resetAttachments, - } = useFileAttachments(organizations[0]?.id, { persist: true }); + } = useFileAttachments(organizationId || undefined, { persist: true }); const handleSendWithAttachments = async (message: string) => { const fileIds: string[] = []; @@ -319,81 +357,133 @@ export const AgentCreateForm: FC = ({ const isForbidden = !canCreateChat; + // Filter workspaces by the selected organization. We use + // client-side filtering of the full "owner:me" fetch rather + // than re-querying with an org filter because it avoids + // extra loading/error states on org change. The full list is + // already small (user's own workspaces) and limit: 0 + // guarantees completeness. If workspace counts grow large + // enough to warrant pagination, this should switch to a + // server-side organization: query filter. + const filteredWorkspaces = + showOrganizations && selectedOrg + ? workspaceOptions.filter((ws) => ws.organization_id === selectedOrg.id) + : workspaceOptions; + return ( -
-
- {isForbidden ? ( - - ) : createError ? ( - isApiError(createError) && - createError.response?.status === 409 && - isUsageLimitData(createError.response.data) ? ( - - View Usage - - } + <> +
+
+ {isForbidden ? ( + + ) : createError ? ( + isApiError(createError) && + createError.response?.status === 409 && + isUsageLimitData(createError.response.data) ? ( + + View Usage + + } + > + + {formatUsageLimitMessage(createError.response.data)} + + + ) : ( + + ) + ) : null} + {workspacesError != null && } + {showOrganizations && ( +
+ + { + const orgChanged = newOrg?.id !== selectedOrg?.id; + if (orgChanged && attachments.length > 0) { + setPendingOrgChange(newOrg); + return; + } + if (orgChanged) { + handleWorkspaceChange(null); + } + setSelectedOrg(newOrg); + }} + /> +
+ )} + { + setUserMCPServerIds(ids); + saveMCPSelection(ids); + }} + onMCPAuthComplete={onMCPAuthComplete} + workspaceOptions={filteredWorkspaces} + selectedWorkspaceId={selectedWorkspaceId} + onWorkspaceChange={handleWorkspaceChange} + isWorkspaceLoading={isWorkspacesLoading} + /> + {modelSelectorHelp ? ( +
+ {modelSelectorHelp} +
+ ) : null} +

+ - - {formatUsageLimitMessage(createError.response.data)} - - - ) : ( - - ) - ) : null} - {workspacesError != null && } - { - setUserMCPServerIds(ids); - saveMCPSelection(ids); - }} - onMCPAuthComplete={onMCPAuthComplete} - workspaceOptions={workspaceOptions} - selectedWorkspaceId={selectedWorkspaceId} - onWorkspaceChange={handleWorkspaceChange} - isWorkspaceLoading={isWorkspacesLoading} - /> - {modelSelectorHelp ? ( -

- ) : null} -

- - Introductory access - {" "} - to Coder Agents through September 2026 -

+ Introductory access + {" "} + to Coder Agents through September 2026 +

+
-
+ { + resetAttachments(); + handleWorkspaceChange(null); + setSelectedOrg(pendingOrgChange); + setPendingOrgChange(null); + }} + onClose={() => setPendingOrgChange(null)} + /> + ); }; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 50c4d74a18..eb264d5d9a 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -205,6 +205,7 @@ const createTestQueryClient = (): QueryClient => const makeChat = (chatID: string): TypesGen.Chat => ({ id: chatID, + organization_id: "test-org-id", owner_id: "owner-1", last_model_config_id: "model-1", mcp_server_ids: [], diff --git a/site/src/pages/AgentsPage/components/ChatTopBar.stories.tsx b/site/src/pages/AgentsPage/components/ChatTopBar.stories.tsx index 07df88f10d..244ab9d27a 100644 --- a/site/src/pages/AgentsPage/components/ChatTopBar.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatTopBar.stories.tsx @@ -57,6 +57,7 @@ export const WithParentChat: Story = { args: { parentChat: { id: "parent-chat-1", + organization_id: "test-org-id", owner_id: "owner-id", last_model_config_id: "model-config-1", mcp_server_ids: [], diff --git a/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.stories.tsx b/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.stories.tsx index 22f2d6b9c0..cb5455ca23 100644 --- a/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.stories.tsx +++ b/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.stories.tsx @@ -40,6 +40,7 @@ const oneWeekAgo = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000).toISOString(); const buildChat = (overrides: Partial = {}): Chat => ({ id: "chat-default", + organization_id: "test-org-id", owner_id: "owner-1", title: "Agent", status: "completed", diff --git a/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.test.tsx b/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.test.tsx index 3cf5e2809c..2e7040b71d 100644 --- a/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.test.tsx +++ b/site/src/pages/AgentsPage/components/Sidebar/AgentsSidebar.test.tsx @@ -53,6 +53,7 @@ const oneWeekAgo = new Date(Date.now() - 7 * 24 * 60 * 60 * 1000).toISOString(); const buildChat = (overrides: Partial = {}): Chat => ({ id: "chat-default", + organization_id: "test-org-id", owner_id: "owner-1", title: "Agent", status: "completed",