mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
Michael Suchacz
6b8f820493
## Stack Context This stack makes AI Gateway data and budgets the source of truth for AI spend controls. 1. Re-back the per-chat cost endpoint with AI Gateway data (#27328, merged). 2. Remove native chat usage limits (#27329, merged). 3. **This PR, now based on `main`:** remove native chat cost tracking and its dedicated admin UI. ## Summary Removes native per-message price calculation, model pricing fields, cost persistence, aggregate cost queries, and admin cost API types. It also deletes the Analytics and Spend pages plus their legacy redirects. The AI Gateway-backed per-chat cost row and compact budget indicators remain. The spend documentation is renamed to `spend-management.md` and updated for the remaining surfaces, group budget APIs, CSV export, upgrade handling for native pricing and cost history, and the absence of a deployment-wide spend dashboard. The per-chat cost API documents that data follows AI Gateway retention and reports zero after all matching requests are purged. No schema is dropped in this release. `chat_messages.total_cost_micros` remains nullable and unwritten so replicas from the previous release can continue inserting messages during rolling upgrades. #27600 tracks removal after the compatibility window. > Mux prepared this PR on Mike's behalf.
452 lines
15 KiB
Go
452 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/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)
|
|
})
|
|
}
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
}
|