diff --git a/codersdk/toolsdk/chats.go b/codersdk/toolsdk/chats.go index 8647e4bed44..0f94b185124 100644 --- a/codersdk/toolsdk/chats.go +++ b/codersdk/toolsdk/chats.go @@ -97,7 +97,7 @@ The chat runs asynchronously. Poll coder_get_chat for status and read the transc }, "organization_id": map[string]any{ "type": "string", - "description": "Optional organization UUID. Defaults to the authenticated user's first organization.", + "description": "Optional organization UUID. Defaults to the organization of the authenticated user's most recently updated chat. If the user has never created a chat, defaults only when they belong to one organization.", }, "model_config_id": map[string]any{ "type": "string", @@ -125,15 +125,11 @@ The chat runs asynchronously. Poll coder_get_chat for status and read the transc return ChatToolStatus{}, xerrors.New("organization_id must be a valid UUID") } } else { - me, err := deps.coderClient.User(ctx, codersdk.Me) + var err error + orgID, err = defaultCreateChatOrganization(ctx, deps) if err != nil { return ChatToolStatus{}, err } - // Admins can remove a user's only organization membership. - if len(me.OrganizationIDs) == 0 { - return ChatToolStatus{}, xerrors.New("authenticated user belongs to no organization; pass organization_id explicitly") - } - orgID = me.OrganizationIDs[0] } var modelConfigID *uuid.UUID if args.ModelConfigID != "" { @@ -159,6 +155,44 @@ The chat runs asynchronously. Poll coder_get_chat for status and read the transc }, } +func defaultCreateChatOrganization(ctx context.Context, deps Deps) (uuid.UUID, error) { + me, err := deps.coderClient.User(ctx, codersdk.Me) + if err != nil { + return uuid.Nil, err + } + // Admins can remove a user's only organization membership. + if len(me.OrganizationIDs) == 0 { + return uuid.Nil, xerrors.New("authenticated user belongs to no organization; pass organization_id explicitly") + } + if len(me.OrganizationIDs) == 1 { + return me.OrganizationIDs[0], nil + } + + expClient := codersdk.NewExperimentalClient(deps.coderClient) + chats, err := expClient.ListChats(ctx, &codersdk.ListChatsOptions{ + Source: codersdk.ChatListSourceCreatedByMe, + Pagination: codersdk.Pagination{ + Limit: 100, + }, + }) + if err != nil { + return uuid.Nil, xerrors.Errorf("list chats to determine organization: %w", err) + } + // Pinned chats sort before recently updated chats. If the user has 100 + // pinned chats, this batch may not contain their latest chat, which is an + // acceptable tradeoff for keeping organization selection to one request. + var latest *codersdk.Chat + for i := range chats { + if latest == nil || chats[i].UpdatedAt.After(latest.UpdatedAt) { + latest = &chats[i] + } + } + if latest != nil { + return latest.OrganizationID, nil + } + return uuid.Nil, xerrors.New("organization_id is required because the authenticated user belongs to multiple organizations and has not created a chat yet") +} + type GetChatArgs struct { ChatID string `json:"chat_id"` } diff --git a/codersdk/toolsdk/chats_test.go b/codersdk/toolsdk/chats_test.go index c6a36158af8..e50ab4db7c1 100644 --- a/codersdk/toolsdk/chats_test.go +++ b/codersdk/toolsdk/chats_test.go @@ -122,8 +122,9 @@ func TestChatTools(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) created, err := testTool(t, toolsdk.CreateChat, tb, toolsdk.CreateChatArgs{ - Prompt: "Say hello.", - Labels: map[string]string{"purpose": "toolsdk-test"}, + Prompt: "Say hello.", + OrganizationID: firstUser.OrganizationID.String(), + Labels: map[string]string{"purpose": "toolsdk-test"}, }) require.NoError(t, err) chatID, err := uuid.Parse(created.ID) @@ -736,8 +737,9 @@ func TestChatTools(t *testing.T) { blockingModelConfig := coderdtest.CreateOpenAICompatChatModel(t, expClient, blockingURL) created, err := testTool(t, toolsdk.CreateChat, tb, toolsdk.CreateChatArgs{ - Prompt: "Block forever.", - ModelConfigID: blockingModelConfig.ID.String(), + Prompt: "Block forever.", + OrganizationID: firstUser.OrganizationID.String(), + ModelConfigID: blockingModelConfig.ID.String(), }) require.NoError(t, err) require.Equal(t, codersdk.ChatStatusRunning, created.Status) @@ -794,6 +796,59 @@ func TestChatTools(t *testing.T) { require.Contains(t, auditorIDs, defaultModelConfig.ID.String()) }) + t.Run("CreateChatUsesLastUpdatedOrganization", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + model := coderdtest.CreateOpenAICompatChatModel(t, expClient, "") + organization := dbgen.Organization(t, api.Database, database.Organization{}) + dbgen.OrganizationMember(t, api.Database, database.OrganizationMember{ + OrganizationID: organization.ID, + UserID: firstUser.UserID, + }) + defaultOrgChat := dbgen.Chat(t, api.Database, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: model.ID, + }) + dbgen.Chat(t, api.Database, database.Chat{ + OrganizationID: organization.ID, + OwnerID: firstUser.UserID, + LastModelConfigID: model.ID, + }) + + err := expClient.UpdateChat(ctx, defaultOrgChat.ID, codersdk.UpdateChatRequest{ + Archived: ptr.Ref(true), + }) + require.NoError(t, err) + err = expClient.UpdateChat(ctx, defaultOrgChat.ID, codersdk.UpdateChatRequest{ + Archived: ptr.Ref(false), + }) + require.NoError(t, err) + + created, err := testTool(t, toolsdk.CreateChat, tb, toolsdk.CreateChatArgs{ + Prompt: "Reuse the last updated organization.", + ModelConfigID: model.ID.String(), + }) + require.NoError(t, err) + chat, err := expClient.GetChat(ctx, uuid.MustParse(created.ID)) + require.NoError(t, err) + require.Equal(t, firstUser.OrganizationID, chat.OrganizationID) + }) + + t.Run("CreateChatRequiresOrganizationForNewMultiOrgUser", func(t *testing.T) { + organization := dbgen.Organization(t, api.Database, database.Organization{}) + memberClient, member := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID) + dbgen.OrganizationMember(t, api.Database, database.OrganizationMember{ + OrganizationID: organization.ID, + UserID: member.ID, + }) + memberDeps, err := toolsdk.NewDeps(memberClient) + require.NoError(t, err) + + _, err = testTool(t, toolsdk.CreateChat, memberDeps, toolsdk.CreateChatArgs{Prompt: "hi"}) + require.ErrorContains(t, err, "belongs to multiple organizations") + require.ErrorContains(t, err, "organization_id is required") + }) + t.Run("CreateChatZeroOrgUser", func(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong)