coder/coderd/exp_chats_internal_test.go

package coderd

import (
	"database/sql"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
	"golang.org/x/xerrors"

	"cdr.dev/slog/v3"
	"cdr.dev/slog/v3/sloggers/slogtest"
	"github.com/coder/coder/v2/coderd/database"
	"github.com/coder/coder/v2/coderd/database/dbauthz"
	"github.com/coder/coder/v2/coderd/database/dbmock"
	"github.com/coder/coder/v2/coderd/httpmw"
	"github.com/coder/coder/v2/codersdk"
	"github.com/coder/coder/v2/testutil"
)

// ExtractChatParam authorizes the read, then GetAIBridgeChatCost authorizes it
// again. A denial on the second check means the ACL changed in between (a
// read-authz race). Assert it surfaces as 404, not 500.
func TestGetChatCostSurfacesReadAuthzRace(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	dbm := dbmock.NewMockStore(ctrl)
	chat := database.Chat{
		ID:             uuid.New(),
		OrganizationID: uuid.New(),
		OwnerID:        uuid.New(),
	}

	dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil)
	dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), chat.ID).Return(
		database.GetAIBridgeChatCostRow{},
		dbauthz.NotAuthorizedError{Err: sql.ErrNoRows},
	)

	api := &API{Options: &Options{Database: dbm}}
	rtr := chi.NewRouter()
	rtr.With(httpmw.ExtractChatParam(dbm)).Get("/chats/{chat}/cost", api.getChatCost)

	req := httptest.NewRequest(http.MethodGet, "/chats/"+chat.ID.String()+"/cost", nil)
	rec := httptest.NewRecorder()
	rtr.ServeHTTP(rec, req)
	resp := rec.Result()
	defer resp.Body.Close()

	require.Equal(t, http.StatusNotFound, resp.StatusCode)
}

// AI Gateway attributes a subagent's requests to the chat that spawned it, so
// a subagent request must be answered with its root chat's tree cost.
func TestGetChatCostQueriesRootChat(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	dbm := dbmock.NewMockStore(ctrl)
	rootID := uuid.New()
	child := database.Chat{
		ID:             uuid.New(),
		OrganizationID: uuid.New(),
		OwnerID:        uuid.New(),
		ParentChatID:   uuid.NullUUID{UUID: rootID, Valid: true},
		RootChatID:     uuid.NullUUID{UUID: rootID, Valid: true},
	}

	dbm.EXPECT().GetChatByID(gomock.Any(), child.ID).Return(child, nil)
	dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), rootID).Return(
		database.GetAIBridgeChatCostRow{
			TotalCostMicros:      250,
			RequestCount:         2,
			UnpricedRequestCount: 1,
		},
		nil,
	)

	api := &API{Options: &Options{Database: dbm}}
	rtr := chi.NewRouter()
	rtr.With(httpmw.ExtractChatParam(dbm)).Get("/chats/{chat}/cost", api.getChatCost)

	req := httptest.NewRequest(http.MethodGet, "/chats/"+child.ID.String()+"/cost", nil)
	rec := httptest.NewRecorder()
	rtr.ServeHTTP(rec, req)
	resp := rec.Result()
	defer resp.Body.Close()

	require.Equal(t, http.StatusOK, resp.StatusCode)
	var cost codersdk.ChatCost
	require.NoError(t, json.NewDecoder(resp.Body).Decode(&cost))
	require.Equal(t, child.ID, cost.ChatID)
	require.Equal(t, int64(250), cost.TotalCostMicros)
	require.Equal(t, int64(2), cost.RequestCount)
	require.Equal(t, int64(1), cost.UnpricedRequestCount)
}

func TestGetChatCostFallsBackToParentChat(t *testing.T) {
	t.Parallel()

	dbm := dbmock.NewMockStore(gomock.NewController(t))
	parentID := uuid.New()
	// chats.parent_chat_id and chats.root_chat_id are both ON DELETE SET NULL,
	// so deleting a root leaves descendants with only a parent.
	child := database.Chat{
		ID:           uuid.New(),
		OwnerID:      uuid.New(),
		ParentChatID: uuid.NullUUID{UUID: parentID, Valid: true},
	}

	dbm.EXPECT().GetChatByID(gomock.Any(), child.ID).Return(child, nil)
	dbm.EXPECT().GetAIBridgeChatCost(gomock.Any(), parentID).Return(
		database.GetAIBridgeChatCostRow{TotalCostMicros: 125, RequestCount: 1},
		nil,
	)

	api := &API{Options: &Options{Database: dbm}}
	rtr := chi.NewRouter()
	rtr.With(httpmw.ExtractChatParam(dbm)).Get("/chats/{chat}/cost", api.getChatCost)

	req := httptest.NewRequest(http.MethodGet, "/chats/"+child.ID.String()+"/cost", nil)
	rec := httptest.NewRecorder()
	rtr.ServeHTTP(rec, req)
	resp := rec.Result()
	defer resp.Body.Close()

	require.Equal(t, http.StatusOK, resp.StatusCode)
	var cost codersdk.ChatCost
	require.NoError(t, json.NewDecoder(resp.Body).Decode(&cost))
	require.Equal(t, int64(125), cost.TotalCostMicros)
}

func TestEnrichMissingChatAgentIDs(t *testing.T) {
	t.Parallel()
	newAPI := func(t *testing.T) (*API, *dbmock.MockStore) {
		t.Helper()
		mDB := dbmock.NewMockStore(gomock.NewController(t))
		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
		return &API{Options: &Options{Database: mDB, Logger: logger}}, mDB
	}
	workspaceID, otherWorkspaceID := uuid.New(), uuid.New()
	rootAgentID, otherAgentID := uuid.New(), uuid.New()
	row := func(workspaceID, id uuid.UUID, parentID uuid.NullUUID, name string) database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow {
		return database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{
			WorkspaceID: workspaceID,
			WorkspaceAgent: database.WorkspaceAgent{
				ID:       id,
				ParentID: parentID,
				Name:     name,
			},
		}
	}
	t.Run("batch selection and shared workspace", func(t *testing.T) {
		t.Parallel()
		api, mDB := newAPI(t)
		mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), gomock.Any()).DoAndReturn(func(_ any, ids []uuid.UUID) ([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow, error) {
			require.ElementsMatch(t, []uuid.UUID{workspaceID, otherWorkspaceID}, ids)
			return []database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{
				row(workspaceID, uuid.New(), uuid.NullUUID{UUID: rootAgentID, Valid: true}, "sub"), row(workspaceID, rootAgentID, uuid.NullUUID{}, "root"), row(otherWorkspaceID, otherAgentID, uuid.NullUUID{}, "root"),
			}, nil
		}).Times(1)
		chats := []codersdk.Chat{{WorkspaceID: &workspaceID, Children: []codersdk.Chat{{WorkspaceID: &workspaceID}}}, {WorkspaceID: &otherWorkspaceID}}
		api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
		require.Equal(t, rootAgentID, *chats[0].AgentID)
		require.Equal(t, rootAgentID, *chats[0].Children[0].AgentID)
		require.Equal(t, otherAgentID, *chats[1].AgentID)
	})
	t.Run("query error", func(t *testing.T) {
		t.Parallel()
		api, mDB := newAPI(t)
		mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), gomock.Any()).Return(nil, xerrors.New("boom"))
		chats := []codersdk.Chat{{WorkspaceID: &workspaceID}, {WorkspaceID: &otherWorkspaceID}}
		api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
		require.Nil(t, chats[0].AgentID)
		require.Nil(t, chats[1].AgentID)
	})
	t.Run("selection error and skips bound or unbound", func(t *testing.T) {
		t.Parallel()
		api, mDB := newAPI(t)
		mDB.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceIDs(gomock.Any(), []uuid.UUID{workspaceID}).Return([]database.GetWorkspaceAgentsInLatestBuildByWorkspaceIDsRow{row(workspaceID, uuid.New(), uuid.NullUUID{UUID: rootAgentID, Valid: true}, "sub")}, nil)
		bound := otherAgentID
		chats := []codersdk.Chat{{}, {WorkspaceID: &workspaceID}, {WorkspaceID: &workspaceID, AgentID: &bound}}
		api.enrichChatWithWorkspaceAgentIDs(testutil.Context(t, testutil.WaitShort), chats)
		require.Nil(t, chats[1].AgentID)
		require.Equal(t, bound, *chats[2].AgentID)
	})
}

func TestValidateChatModelProviderOptions_AnthropicThinkingDisplay(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name    string
		display string
		wantErr string
	}{
		{name: "Summarized", display: "summarized"},
		{name: "Omitted", display: " omitted "},
		{name: "Empty", display: " "},
		{
			name:    "Invalid",
			display: "summrized",
			wantErr: "provider_options.anthropic.thinking_display must be one of summarized, omitted",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			display := tt.display
			err := validateChatModelProviderOptions(&codersdk.ChatModelProviderOptions{
				Anthropic: &codersdk.ChatModelAnthropicProviderOptions{
					ThinkingDisplay: &display,
				},
			})
			if tt.wantErr != "" {
				require.EqualError(t, err, tt.wantErr)
				return
			}
			require.NoError(t, err)
		})
	}
}

func TestValidateChatModelConfigProviderModel(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name       string
		model      string
		provider   database.AIProvider
		wantErr    bool
		wantDetail string
	}{
		{
			name:  "OpenRouterNameWithOpenAITypeAndSlashModel",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name: "openrouter",
				Type: database.AIProviderTypeOpenai,
			},
			wantErr:    true,
			wantDetail: "Change the AI provider type to openrouter or openai-compat.",
		},
		{
			name:  "OpenRouterNameWithWhitespaceAndCase",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name: " OpenRouter ",
				Type: database.AIProviderTypeOpenai,
			},
			wantErr:    true,
			wantDetail: "Change the AI provider type to openrouter or openai-compat.",
		},
		{
			name:  "OpenRouterHostWithOpenAITypeAndSlashModel",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name:    "private-relay",
				Type:    database.AIProviderTypeOpenai,
				BaseUrl: "https://openrouter.ai/api/v1",
			},
			wantErr:    true,
			wantDetail: "Change the AI provider type to openrouter or openai-compat.",
		},
		{
			name:  "OpenRouterHostWithPort",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name:    "private-relay",
				Type:    database.AIProviderTypeOpenai,
				BaseUrl: "https://openrouter.ai:443/api/v1",
			},
			wantErr:    true,
			wantDetail: "Change the AI provider type to openrouter or openai-compat.",
		},
		{
			name:  "OpenRouterSubdomainWithOpenAIType",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name:    "private-relay",
				Type:    database.AIProviderTypeOpenai,
				BaseUrl: "https://api.openrouter.ai/v1",
			},
			wantErr:    true,
			wantDetail: "Change the AI provider type to openrouter or openai-compat.",
		},
		{
			name:  "OpenRouterTypeAllowsSlashModel",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name: "openrouter",
				Type: database.AIProviderTypeOpenrouter,
			},
		},
		{
			name:  "OpenAICompatTypeAllowsSlashModel",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name: "openrouter",
				Type: database.AIProviderTypeOpenaiCompat,
			},
		},
		{
			name:  "PrivateOpenAIProxyAllowsSlashModel",
			model: "anthropic/claude-opus-4.6",
			provider: database.AIProvider{
				Name:    "private-relay",
				Type:    database.AIProviderTypeOpenai,
				BaseUrl: "https://llm-relay.internal/v1",
			},
		},
		{
			name:  "OpenRouterNameWithPlainModelAllowed",
			model: "gpt-4.1",
			provider: database.AIProvider{
				Name: "openrouter",
				Type: database.AIProviderTypeOpenai,
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			got := validateChatModelConfigProviderModel(tt.provider, tt.model)
			if tt.wantErr {
				require.NotNil(t, got)
				require.Contains(t, got.Response.Detail, tt.wantDetail)
				return
			}
			require.Nil(t, got)
		})
	}
}

func TestRewriteChatStartWorkspaceManualUpdateResponse(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name           string
		resp           codersdk.Response
		fallbackDetail string
		wantDetail     string
	}{
		{
			name: "NoValidationsAndEmptyDetail",
			resp: codersdk.Response{
				Message: "missing required parameter",
			},
			fallbackDetail: "wrapped missing required parameter",
			wantDetail:     "missing required parameter",
		},
		{
			name: "NoValidationsAndExistingDetail",
			resp: codersdk.Response{
				Message: "missing required parameter",
				Detail:  "region must be set before the workspace can start",
			},
			fallbackDetail: "wrapped missing required parameter",
			wantDetail:     "missing required parameter: region must be set before the workspace can start",
		},
		{
			name: "ValidationsAndEmptyDetail",
			resp: codersdk.Response{
				Message: "missing required parameter",
				Validations: []codersdk.ValidationError{{
					Field:  "region",
					Detail: "region must be set before the workspace can start",
				}},
			},
			fallbackDetail: "wrapped missing required parameter",
			wantDetail:     "wrapped missing required parameter",
		},
		{
			name: "ValidationsAndExistingDetail",
			resp: codersdk.Response{
				Message: "missing required parameter",
				Detail:  "region must be set before the workspace can start",
				Validations: []codersdk.ValidationError{{
					Field:  "region",
					Detail: "region must be set before the workspace can start",
				}},
			},
			fallbackDetail: "wrapped missing required parameter",
			wantDetail:     "region must be set before the workspace can start",
		},
	}

	const retryInstructions = "Use read_template before retrying start_workspace."
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			got := rewriteChatStartWorkspaceManualUpdateResponse(tt.resp, tt.fallbackDetail, retryInstructions)
			require.Equal(t, retryInstructions, got.Message)
			require.Equal(t, tt.wantDetail, got.Detail)
			require.Equal(t, tt.resp.Validations, got.Validations)
		})
	}
}