Files
coder/coderd/exp_chats_internal_test.go
T
Michael Suchacz 57f38b5c24 fix: keep chat attachments while a linking chat exists
Fixes https://linear.app/codercom/issue/CODAGT-616/keep-chat-attachments-while-chats-remain-unarchived

Chat attachments could disappear even though the chat was still available. This happened when a message was saved without recording which attachments it used, or when cleanup deleted attachments before an archived chat itself was removed.

Creating a chat, sending or queuing a message, and editing a message now record both the message and which attachments it uses as one operation. If the chat is already at the 50-attachment limit, the chat change fails without being partially saved.

Concurrent attachment writes serialize the 50-file cap per chat. Cleanup locks candidates and checks again for new links before deleting. If a file becomes unavailable after input validation, create, send, and edit return a clear client error and roll back the chat change.

An attachment stays available while any chat that uses it still exists. After an archived chat reaches the end of its retention period and is deleted, an old attachment that no remaining chat uses can be cleaned up. The retention guide and unavailable-attachment UI text document this lifecycle. This change cannot restore attachments that were already deleted.

The database migration adds two indexes so attachment cleanup stays fast as attachments accumulate.

> This PR was authored by Mux (AI) on Mike's behalf.
2026-08-11 13:53:15 +02:00

468 lines
15 KiB
Go

package coderd
import (
"database/sql"
"encoding/json"
"net/http"
"net/http/httptest"
"reflect"
"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/coderd/util/ptr"
"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
"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 TestWriteChatFileErrorUnavailable(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitShort)
rec := httptest.NewRecorder()
handled := writeChatFileError(ctx, rec, xerrors.Errorf("link files: %w", chatstate.ErrChatFileUnavailable))
require.True(t, handled)
require.Equal(t, http.StatusBadRequest, rec.Code)
var response codersdk.Response
require.NoError(t, json.NewDecoder(rec.Body).Decode(&response))
require.Equal(t, "Chat attachment unavailable.", response.Message)
require.Equal(t, "An attachment is no longer available. Upload it again and retry.", response.Detail)
}
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)
})
}
}
// Every ChatModelCallConfig field must classify a config as non-zero when set,
// or unmarshalChatModelCallConfig hides it from API responses while the stored
// value stays active. Fails when a new field is added without a sample here.
func TestIsZeroChatModelCallConfigCoversEveryField(t *testing.T) {
t.Parallel()
sampled := codersdk.ChatModelCallConfig{
MaxOutputTokens: ptr.Ref(int64(4096)),
Temperature: ptr.Ref(0.7),
TopP: ptr.Ref(0.9),
TopK: ptr.Ref(int64(40)),
PresencePenalty: ptr.Ref(0.1),
FrequencyPenalty: ptr.Ref(0.2),
ReasoningEffort: &codersdk.ChatModelReasoningEffortConfig{
Default: ptr.Ref("medium"),
},
OpenAIConfig: &codersdk.ChatModelOpenAIConfig{
UseResponsesAPI: ptr.Ref(true),
},
ProviderOptions: &codersdk.ChatModelProviderOptions{
OpenAI: &codersdk.ChatModelOpenAIProviderOptions{},
},
}
require.True(t, isZeroChatModelCallConfig(nil))
require.True(t, isZeroChatModelCallConfig(&codersdk.ChatModelCallConfig{}))
sampledValue := reflect.ValueOf(sampled)
for i := 0; i < sampledValue.NumField(); i++ {
field := sampledValue.Type().Field(i)
require.Falsef(t, sampledValue.Field(i).IsZero(),
"field %s needs a non-zero sample value", field.Name)
config := &codersdk.ChatModelCallConfig{}
reflect.ValueOf(config).Elem().Field(i).Set(sampledValue.Field(i))
require.Falsef(t, isZeroChatModelCallConfig(config),
"isZeroChatModelCallConfig ignores field %s", field.Name)
}
}