mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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.
468 lines
15 KiB
Go
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)
|
|
}
|
|
}
|