mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
test(coderd): centralize chat test harness and stabilize flakes (#25171)
Chat tests previously constructed a real `openai` provider with a fake API key and no `BaseURL`, so background title generation hit `api.openai.com` and timed out under `-race`. The same root cause produced several distinct flakes: title regeneration races with synchronous `UpdateChat`/`ProposeChatTitle`, and pagination races against `updated_at` bumps from real-network processing. This moves the fake OpenAI-compatible provider and the chat-settle wait into first-class `coderdtest` capabilities. `coderd.Options.ChatProviderAPIKeys` is the new seam tests use to redirect chat traffic to a local `httptest.Server`. `coderdtest.WaitForChatSettled` replaces per-test waiters and drains tracked chat-daemon work after the chat row leaves `pending`/`running`. The `newChatClient*` constructors funnel through one options builder that installs the fake provider before the coderd test server so cleanup ordering is deterministic. Closes https://github.com/coder/internal/issues/1528 & Closes ENG-2659 Closes https://github.com/coder/internal/issues/1480 & Closes CODAGT-359 Closes https://github.com/coder/internal/issues/1507 & Closes CODAGT-368 Relates to https://github.com/coder/internal/issues/1397 & Relates to CODAGT-374
This commit is contained in:
@@ -0,0 +1,8 @@
|
||||
package coderd
|
||||
|
||||
import "github.com/coder/coder/v2/coderd/x/chatd"
|
||||
|
||||
// ChatDaemonForTest returns the background chat processor for test harnesses.
|
||||
func (api *API) ChatDaemonForTest() *chatd.Server {
|
||||
return api.chatDaemon
|
||||
}
|
||||
+10
-1
@@ -95,6 +95,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/workspacestats"
|
||||
"github.com/coder/coder/v2/coderd/wsbuilder"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/mcpclient"
|
||||
"github.com/coder/coder/v2/coderd/x/gitsync"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
@@ -248,6 +249,9 @@ type Options struct {
|
||||
// ChatSubscribeFn provides cross-replica subscription merging.
|
||||
// Set by enterprise for HA deployments. Nil in AGPL single-replica.
|
||||
ChatSubscribeFn chatd.SubscribeFn
|
||||
// ChatProviderAPIKeys overrides deployment-derived provider keys.
|
||||
// Test harnesses use this to route chat models to local providers.
|
||||
ChatProviderAPIKeys *chatprovider.ProviderAPIKeys
|
||||
|
||||
UpdateAgentMetrics func(ctx context.Context, labels prometheusmetrics.AgentMetricLabels, metrics []*agentproto.Stats_Metric)
|
||||
StatsBatcher workspacestats.Batcher
|
||||
@@ -792,13 +796,18 @@ func New(options *Options) *API {
|
||||
options.Logger.Named("mcp-user-oidc"),
|
||||
)
|
||||
}
|
||||
providerAPIKeys := ChatProviderAPIKeysFromDeploymentValues(options.DeploymentValues)
|
||||
if options.ChatProviderAPIKeys != nil {
|
||||
providerAPIKeys = *options.ChatProviderAPIKeys
|
||||
}
|
||||
|
||||
api.chatDaemon = chatd.New(chatd.Config{
|
||||
Logger: options.Logger.Named("chatd"),
|
||||
Database: options.Database,
|
||||
ReplicaID: api.ID,
|
||||
SubscribeFn: options.ChatSubscribeFn,
|
||||
MaxChatsPerAcquire: int32(maxChatsPerAcquire), //nolint:gosec // maxChatsPerAcquire is clamped to int32 range above.
|
||||
ProviderAPIKeys: ChatProviderAPIKeysFromDeploymentValues(options.DeploymentValues),
|
||||
ProviderAPIKeys: providerAPIKeys,
|
||||
AlwaysEnableDebugLogs: options.DeploymentValues.AI.Chat.DebugLoggingEnabled.Value(),
|
||||
AgentConn: api.agentProvider.AgentConn,
|
||||
AgentInactiveDisconnectTimeout: api.AgentInactiveDisconnectTimeout,
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
package coderdtest
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/coderd"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattest"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
const (
|
||||
// TestChatProviderOpenAICompat is the default provider for chat runtime tests.
|
||||
TestChatProviderOpenAICompat = "openai-compat"
|
||||
// TestChatProviderAPIKey is a non-secret API key for local chat providers.
|
||||
TestChatProviderAPIKey = "test-api-key"
|
||||
// TestChatModelOpenAICompat is the default model for chat runtime tests.
|
||||
TestChatModelOpenAICompat = "gpt-4o-mini"
|
||||
)
|
||||
|
||||
// OpenAICompatProviderAPIKeys returns provider keys that route OpenAI-compatible
|
||||
// chat calls to baseURL.
|
||||
func OpenAICompatProviderAPIKeys(baseURL string) chatprovider.ProviderAPIKeys {
|
||||
return chatprovider.ProviderAPIKeys{
|
||||
ByProvider: map[string]string{
|
||||
TestChatProviderOpenAICompat: TestChatProviderAPIKey,
|
||||
},
|
||||
BaseURLByProvider: map[string]string{
|
||||
TestChatProviderOpenAICompat: baseURL,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// FakeOpenAICompatProviderAPIKeys starts a fake OpenAI-compatible provider and
|
||||
// returns provider keys for coderdtest.Options.
|
||||
func FakeOpenAICompatProviderAPIKeys(t testing.TB) chatprovider.ProviderAPIKeys {
|
||||
t.Helper()
|
||||
return OpenAICompatProviderAPIKeys(chattest.OpenAI(t))
|
||||
}
|
||||
|
||||
// CreateOpenAICompatChatModelConfig creates the default provider and model
|
||||
// config used by chat runtime tests. Tests that create chats should also set
|
||||
// Options.ChatProviderAPIKeys, usually via FakeOpenAICompatProviderAPIKeys, so
|
||||
// background chat work routes to a local provider until coderd closes. baseURL,
|
||||
// when non-empty, is stored on the provider config.
|
||||
func CreateOpenAICompatChatModelConfig(
|
||||
t testing.TB,
|
||||
client *codersdk.ExperimentalClient,
|
||||
baseURL string,
|
||||
) codersdk.ChatModelConfig {
|
||||
t.Helper()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
_, err := client.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
|
||||
Provider: TestChatProviderOpenAICompat,
|
||||
APIKey: TestChatProviderAPIKey,
|
||||
BaseURL: baseURL,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
contextLimit := int64(4096)
|
||||
isDefault := true
|
||||
modelConfig, err := client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
|
||||
Provider: TestChatProviderOpenAICompat,
|
||||
Model: TestChatModelOpenAICompat,
|
||||
ContextLimit: &contextLimit,
|
||||
IsDefault: &isDefault,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return modelConfig
|
||||
}
|
||||
|
||||
// WaitForChatSettled waits for a chat to leave active processing and drains
|
||||
// tracked chat daemon work before returning the final row.
|
||||
func WaitForChatSettled(
|
||||
ctx context.Context,
|
||||
t testing.TB,
|
||||
api *coderd.API,
|
||||
chatID uuid.UUID,
|
||||
) database.Chat {
|
||||
t.Helper()
|
||||
|
||||
require.NotNil(t, api)
|
||||
waitForChatTerminalState(ctx, t, api.Database, chatID)
|
||||
|
||||
server := api.ChatDaemonForTest()
|
||||
require.NotNil(t, server)
|
||||
chatd.WaitUntilIdleForTest(server)
|
||||
|
||||
chat, err := getChatByIDAsSystem(ctx, api.Database, chatID)
|
||||
require.NoError(t, err)
|
||||
return chat
|
||||
}
|
||||
|
||||
func waitForChatTerminalState(
|
||||
ctx context.Context,
|
||||
t testing.TB,
|
||||
db database.Store,
|
||||
chatID uuid.UUID,
|
||||
) {
|
||||
t.Helper()
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
chat, err := getChatByIDAsSystem(ctx, db, chatID)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return chat.Status != database.ChatStatusPending && chat.Status != database.ChatStatusRunning
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
}
|
||||
|
||||
func getChatByIDAsSystem(
|
||||
ctx context.Context,
|
||||
db database.Store,
|
||||
chatID uuid.UUID,
|
||||
) (database.Chat, error) {
|
||||
// Test helper needs system scope to observe chatd-owned status changes.
|
||||
//nolint:gocritic
|
||||
return db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chatID)
|
||||
}
|
||||
@@ -91,6 +91,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/workspaceapps/appurl"
|
||||
"github.com/coder/coder/v2/coderd/workspacestats"
|
||||
"github.com/coder/coder/v2/coderd/wsbuilder"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/agentsdk"
|
||||
"github.com/coder/coder/v2/codersdk/drpcsdk"
|
||||
@@ -151,6 +152,7 @@ type Options struct {
|
||||
// IncludeProvisionerDaemon when true means to start an in-memory provisionerD
|
||||
IncludeProvisionerDaemon bool
|
||||
ChatdInstructionLookupTimeout time.Duration
|
||||
ChatProviderAPIKeys *chatprovider.ProviderAPIKeys
|
||||
ProvisionerDaemonVersion string
|
||||
ProvisionerDaemonTags map[string]string
|
||||
MetricsCacheRefreshInterval time.Duration
|
||||
@@ -584,6 +586,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
|
||||
// agents are not marked as disconnected during slow tests.
|
||||
AgentInactiveDisconnectTimeout: testutil.WaitShort,
|
||||
ChatdInstructionLookupTimeout: options.ChatdInstructionLookupTimeout,
|
||||
ChatProviderAPIKeys: options.ChatProviderAPIKeys,
|
||||
AccessURL: accessURL,
|
||||
AppHostname: options.AppHostname,
|
||||
AppHostnameRegex: appHostnameRegex,
|
||||
|
||||
+109
-102
@@ -41,6 +41,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/x/chatd"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattest"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/serpent"
|
||||
@@ -60,44 +61,73 @@ func chatDeploymentValues(t testing.TB) *codersdk.DeploymentValues {
|
||||
return values
|
||||
}
|
||||
|
||||
func newChatClient(t testing.TB, overrides ...func(*coderdtest.Options)) *codersdk.ExperimentalClient {
|
||||
// newChatTestOptions builds coderdtest options for chat runtime tests. Unless
|
||||
// a test sets ChatProviderAPIKeys explicitly, it installs a fake
|
||||
// OpenAI-compatible provider before coderd starts so background chat work stays
|
||||
// local, and the fake server outlives chatd during cleanup.
|
||||
func newChatTestOptions(
|
||||
t testing.TB,
|
||||
values *codersdk.DeploymentValues,
|
||||
overrides ...func(*coderdtest.Options),
|
||||
) *coderdtest.Options {
|
||||
t.Helper()
|
||||
|
||||
opts := &coderdtest.Options{
|
||||
DeploymentValues: chatDeploymentValues(t),
|
||||
DeploymentValues: values,
|
||||
}
|
||||
for _, override := range overrides {
|
||||
override(opts)
|
||||
}
|
||||
if opts.ChatProviderAPIKeys == nil {
|
||||
providerKeys := coderdtest.FakeOpenAICompatProviderAPIKeys(t)
|
||||
opts.ChatProviderAPIKeys = &providerKeys
|
||||
}
|
||||
return opts
|
||||
}
|
||||
|
||||
func newChatClient(t testing.TB, overrides ...func(*coderdtest.Options)) *codersdk.ExperimentalClient {
|
||||
t.Helper()
|
||||
|
||||
opts := newChatTestOptions(t, chatDeploymentValues(t), overrides...)
|
||||
client := coderdtest.New(t, opts)
|
||||
return codersdk.NewExperimentalClient(client)
|
||||
}
|
||||
|
||||
func newChatClientWithAPI(t testing.TB, overrides ...func(*coderdtest.Options)) (*codersdk.ExperimentalClient, *coderd.API) {
|
||||
t.Helper()
|
||||
|
||||
opts := newChatTestOptions(t, chatDeploymentValues(t), overrides...)
|
||||
client, _, api := coderdtest.NewWithAPI(t, opts)
|
||||
return codersdk.NewExperimentalClient(client), api
|
||||
}
|
||||
|
||||
func newChatClientWithDeploymentValues(
|
||||
t testing.TB,
|
||||
values *codersdk.DeploymentValues,
|
||||
) *codersdk.ExperimentalClient {
|
||||
t.Helper()
|
||||
|
||||
client := coderdtest.New(t, &coderdtest.Options{
|
||||
DeploymentValues: values,
|
||||
})
|
||||
opts := newChatTestOptions(t, values)
|
||||
client := coderdtest.New(t, opts)
|
||||
return codersdk.NewExperimentalClient(client)
|
||||
}
|
||||
|
||||
func newChatClientWithDatabase(t testing.TB, overrides ...func(*coderdtest.Options)) (*codersdk.ExperimentalClient, database.Store) {
|
||||
t.Helper()
|
||||
|
||||
opts := &coderdtest.Options{
|
||||
DeploymentValues: chatDeploymentValues(t),
|
||||
}
|
||||
for _, override := range overrides {
|
||||
override(opts)
|
||||
}
|
||||
opts := newChatTestOptions(t, chatDeploymentValues(t), overrides...)
|
||||
client, db := coderdtest.NewWithDatabase(t, opts)
|
||||
return codersdk.NewExperimentalClient(client), db
|
||||
}
|
||||
|
||||
func newChatClientWithAPIAndDatabase(t testing.TB, overrides ...func(*coderdtest.Options)) (*codersdk.ExperimentalClient, database.Store, *coderd.API) {
|
||||
t.Helper()
|
||||
|
||||
opts := newChatTestOptions(t, chatDeploymentValues(t), overrides...)
|
||||
client, _, api := coderdtest.NewWithAPI(t, opts)
|
||||
return codersdk.NewExperimentalClient(client), api.Database, api
|
||||
}
|
||||
|
||||
type failNextChatSystemPromptStore struct {
|
||||
database.Store
|
||||
|
||||
@@ -1390,14 +1420,14 @@ func TestListChatModels(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
_ = coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
|
||||
models, err := client.ListChatModels(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
var openAIProvider *codersdk.ChatModelProvider
|
||||
for i := range models.Providers {
|
||||
if models.Providers[i].Provider == "openai" {
|
||||
if models.Providers[i].Provider == modelConfig.Provider {
|
||||
openAIProvider = &models.Providers[i]
|
||||
break
|
||||
}
|
||||
@@ -1407,7 +1437,7 @@ func TestListChatModels(t *testing.T) {
|
||||
|
||||
foundModel := false
|
||||
for _, model := range openAIProvider.Models {
|
||||
if model.Provider == "openai" && model.Model == "gpt-4o-mini" {
|
||||
if model.Provider == modelConfig.Provider && model.Model == modelConfig.Model {
|
||||
foundModel = true
|
||||
break
|
||||
}
|
||||
@@ -1433,14 +1463,14 @@ func TestListChatModels(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
_ = coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
|
||||
models, err := client.ListChatModels(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
var openAIProvider *codersdk.ChatModelProvider
|
||||
for i := range models.Providers {
|
||||
if models.Providers[i].Provider == "openai" {
|
||||
if models.Providers[i].Provider == modelConfig.Provider {
|
||||
openAIProvider = &models.Providers[i]
|
||||
break
|
||||
}
|
||||
@@ -1925,14 +1955,14 @@ func TestListChatProviders(t *testing.T) {
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
_ = coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
modelConfig := createChatModelConfig(t, client)
|
||||
|
||||
providers, err := client.ListChatProviders(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
var openAIProvider *codersdk.ChatProviderConfig
|
||||
for i := range providers {
|
||||
if providers[i].Provider == "openai" {
|
||||
if providers[i].Provider == modelConfig.Provider {
|
||||
openAIProvider = &providers[i]
|
||||
break
|
||||
}
|
||||
@@ -3335,8 +3365,8 @@ func TestListChatModelConfigs(t *testing.T) {
|
||||
for _, config := range configs {
|
||||
if config.ID == modelConfig.ID {
|
||||
found = true
|
||||
require.Equal(t, "openai", config.Provider)
|
||||
require.Equal(t, "gpt-4o-mini", config.Model)
|
||||
require.Equal(t, modelConfig.Provider, config.Provider)
|
||||
require.Equal(t, modelConfig.Model, config.Model)
|
||||
require.True(t, config.IsDefault)
|
||||
}
|
||||
}
|
||||
@@ -3395,7 +3425,7 @@ func TestListChatModelConfigs(t *testing.T) {
|
||||
contextLimit := int64(4096)
|
||||
enabled := false
|
||||
_, err := adminClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
|
||||
Provider: "openai",
|
||||
Provider: enabledConfig.Provider,
|
||||
Model: "gpt-4o-disabled",
|
||||
DisplayName: "GPT-4o Disabled",
|
||||
Enabled: &enabled,
|
||||
@@ -3468,8 +3498,8 @@ func TestListChatModelConfigs(t *testing.T) {
|
||||
for _, config := range configs {
|
||||
if config.ID == modelConfig.ID {
|
||||
found = true
|
||||
require.Equal(t, "openai", config.Provider)
|
||||
require.Equal(t, "gpt-4o-mini", config.Model)
|
||||
require.Equal(t, modelConfig.Provider, config.Provider)
|
||||
require.Equal(t, modelConfig.Model, config.Model)
|
||||
}
|
||||
}
|
||||
require.True(t, found)
|
||||
@@ -4310,21 +4340,6 @@ func TestPatchChat(t *testing.T) {
|
||||
return db2sdk.Chat(dbChat, nil, nil)
|
||||
}
|
||||
|
||||
// waitChatSettled polls the chat until its background title-generation
|
||||
// daemon has left the Pending/Running state. Without this, an immediate
|
||||
// UpdateChat can hit a 409 (title regeneration in progress).
|
||||
waitChatSettled := func(ctx context.Context, t *testing.T, client *codersdk.ExperimentalClient, chatID uuid.UUID) {
|
||||
t.Helper()
|
||||
require.Eventually(t, func() bool {
|
||||
c, err := client.GetChat(ctx, chatID)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return c.Status != codersdk.ChatStatusPending &&
|
||||
c.Status != codersdk.ChatStatusRunning
|
||||
}, testutil.WaitShort, testutil.IntervalFast)
|
||||
}
|
||||
|
||||
t.Run("PlanMode", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -4592,13 +4607,13 @@ func TestPatchChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
client, api := newChatClientWithAPI(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
|
||||
chat := createChat(ctx, t, client, firstUser.OrganizationID, "original title")
|
||||
|
||||
waitChatSettled(ctx, t, client, chat.ID)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
err := client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{
|
||||
Title: ptr.Ref("renamed title"),
|
||||
@@ -4613,13 +4628,13 @@ func TestPatchChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
client, api := newChatClientWithAPI(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
|
||||
chat := createChat(ctx, t, client, firstUser.OrganizationID, "before trim")
|
||||
|
||||
waitChatSettled(ctx, t, client, chat.ID)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
err := client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{
|
||||
Title: ptr.Ref(" padded title "),
|
||||
@@ -4713,11 +4728,11 @@ func TestPatchChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
client, api := newChatClientWithAPI(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
chat := createChat(ctx, t, client, firstUser.OrganizationID, "boundary baseline")
|
||||
waitChatSettled(ctx, t, client, chat.ID)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
err := client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{
|
||||
Title: ptr.Ref(tc.title),
|
||||
@@ -4739,17 +4754,19 @@ func TestPatchChat(t *testing.T) {
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t)
|
||||
clientRaw := coderdtest.New(t, &coderdtest.Options{
|
||||
DeploymentValues: chatDeploymentValues(t),
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
providerKeys := coderdtest.FakeOpenAICompatProviderAPIKeys(t)
|
||||
clientRaw, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
|
||||
DeploymentValues: chatDeploymentValues(t),
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
ChatProviderAPIKeys: &providerKeys,
|
||||
})
|
||||
client := codersdk.NewExperimentalClient(clientRaw)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
|
||||
chat := createChat(ctx, t, client, firstUser.OrganizationID, "rename me")
|
||||
waitChatSettled(ctx, t, client, chat.ID)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
past := time.Now().UTC().Add(-2 * time.Hour).Truncate(time.Second)
|
||||
_, err := sqlDB.ExecContext(ctx,
|
||||
@@ -4774,17 +4791,19 @@ func TestPatchChat(t *testing.T) {
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t)
|
||||
clientRaw := coderdtest.New(t, &coderdtest.Options{
|
||||
DeploymentValues: chatDeploymentValues(t),
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
providerKeys := coderdtest.FakeOpenAICompatProviderAPIKeys(t)
|
||||
clientRaw, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{
|
||||
DeploymentValues: chatDeploymentValues(t),
|
||||
Database: db,
|
||||
Pubsub: ps,
|
||||
ChatProviderAPIKeys: &providerKeys,
|
||||
})
|
||||
client := codersdk.NewExperimentalClient(clientRaw)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
|
||||
chat := createChat(ctx, t, client, firstUser.OrganizationID, "steady title")
|
||||
waitChatSettled(ctx, t, client, chat.ID)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
past := time.Now().UTC().Add(-2 * time.Hour).Truncate(time.Second)
|
||||
_, err := sqlDB.ExecContext(ctx,
|
||||
@@ -4808,13 +4827,13 @@ func TestPatchChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
client, api := newChatClientWithAPI(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
|
||||
chat := createChat(ctx, t, client, firstUser.OrganizationID, "announce me")
|
||||
|
||||
waitChatSettled(ctx, t, client, chat.ID)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
conn, err := client.Dial(ctx, "/api/experimental/chats/watch", nil)
|
||||
require.NoError(t, err)
|
||||
@@ -5774,7 +5793,7 @@ func TestSendMessageWithModelOverrideUpdatesLastModelConfigID(t *testing.T) {
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
user := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfigA := createChatModelConfig(t, client)
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-override-"+uuid.NewString())
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, modelConfigA.Provider, "gpt-4o-mini-override-"+uuid.NewString())
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: user.OrganizationID,
|
||||
@@ -5817,7 +5836,7 @@ func TestSendMessageQueuesEffectiveModelConfigID(t *testing.T) {
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
user := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfigA := createChatModelConfig(t, client)
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-queued-"+uuid.NewString())
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, modelConfigA.Provider, "gpt-4o-mini-queued-"+uuid.NewString())
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: user.OrganizationID,
|
||||
@@ -5868,7 +5887,7 @@ func TestQueuedMessageWithoutOverrideCapturesEnqueueTimeModel(t *testing.T) {
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
user := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfigA := createChatModelConfig(t, client)
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-later-"+uuid.NewString())
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, modelConfigA.Provider, "gpt-4o-mini-later-"+uuid.NewString())
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: user.OrganizationID,
|
||||
@@ -5920,7 +5939,7 @@ func TestSubsequentSendWithoutOverrideUsesPersistedModel(t *testing.T) {
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
user := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-persisted-"+uuid.NewString())
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, coderdtest.TestChatProviderOpenAICompat, "gpt-4o-mini-persisted-"+uuid.NewString())
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: user.OrganizationID,
|
||||
@@ -5961,7 +5980,7 @@ func TestWatchChatsStatusChangeCarriesUpdatedLastModelConfigID(t *testing.T) {
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
user := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfigA := createChatModelConfig(t, client)
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-watch-direct-"+uuid.NewString())
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, modelConfigA.Provider, "gpt-4o-mini-watch-direct-"+uuid.NewString())
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: user.OrganizationID,
|
||||
@@ -5994,7 +6013,7 @@ func TestWatchChatsStatusChangeCarriesUpdatedLastModelConfigID(t *testing.T) {
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
user := coderdtest.CreateFirstUser(t, client.Client)
|
||||
modelConfigA := createChatModelConfig(t, client)
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-watch-promote-"+uuid.NewString())
|
||||
modelConfigB := createAdditionalChatModelConfig(t, client, modelConfigA.Provider, "gpt-4o-mini-watch-promote-"+uuid.NewString())
|
||||
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: user.OrganizationID,
|
||||
@@ -7627,9 +7646,9 @@ func TestRegenerateChatTitle(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
client, db, api := newChatClientWithAPIAndDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
_ = createChatModelConfigWithTitleFailure(t, client)
|
||||
|
||||
chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
@@ -7642,16 +7661,7 @@ func TestRegenerateChatTitle(t *testing.T) {
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for background processing triggered by signalWake
|
||||
// to finish before setting the status, otherwise the
|
||||
// processor may update updated_at concurrently.
|
||||
require.Eventually(t, func() bool {
|
||||
c, getErr := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
return c.Status != database.ChatStatusPending && c.Status != database.ChatStatusRunning
|
||||
}, testutil.WaitShort, testutil.IntervalFast)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
_, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{
|
||||
ID: chat.ID,
|
||||
@@ -7724,9 +7734,9 @@ func TestProposeChatTitle(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client, db := newChatClientWithDatabase(t)
|
||||
client, db, api := newChatClientWithAPIAndDatabase(t)
|
||||
firstUser := coderdtest.CreateFirstUser(t, client.Client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
_ = createChatModelConfigWithTitleFailure(t, client)
|
||||
|
||||
chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
|
||||
OrganizationID: firstUser.OrganizationID,
|
||||
@@ -7736,13 +7746,7 @@ func TestProposeChatTitle(t *testing.T) {
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
c, getErr := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
return c.Status != database.ChatStatusPending && c.Status != database.ChatStatusRunning
|
||||
}, testutil.WaitShort, testutil.IntervalFast)
|
||||
coderdtest.WaitForChatSettled(ctx, t, api, chat.ID)
|
||||
|
||||
before, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID)
|
||||
require.NoError(t, err)
|
||||
@@ -9700,26 +9704,29 @@ func TestWatchChatGitAuthz(t *testing.T) {
|
||||
require.Equal(t, http.StatusForbidden, res.StatusCode)
|
||||
}
|
||||
|
||||
func createChatModelConfig(t *testing.T, client *codersdk.ExperimentalClient) codersdk.ChatModelConfig {
|
||||
func createChatModelConfig(t testing.TB, client *codersdk.ExperimentalClient) codersdk.ChatModelConfig {
|
||||
t.Helper()
|
||||
return coderdtest.CreateOpenAICompatChatModelConfig(t, client, "")
|
||||
}
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
_, err := client.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
|
||||
Provider: "openai",
|
||||
APIKey: "test-api-key",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
func createChatModelConfigWithBaseURL(t testing.TB, client *codersdk.ExperimentalClient, baseURL string) codersdk.ChatModelConfig {
|
||||
t.Helper()
|
||||
return coderdtest.CreateOpenAICompatChatModelConfig(t, client, baseURL)
|
||||
}
|
||||
|
||||
contextLimit := int64(4096)
|
||||
isDefault := true
|
||||
modelConfig, err := client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
|
||||
Provider: "openai",
|
||||
Model: "gpt-4o-mini",
|
||||
ContextLimit: &contextLimit,
|
||||
IsDefault: &isDefault,
|
||||
// createChatModelConfigWithTitleFailure provisions a model whose streaming chat
|
||||
// responses succeed, while non-streaming requests fail. The non-streaming path
|
||||
// is how quick title generation requests structured output, so tests can fail
|
||||
// title generation without breaking the main assistant response.
|
||||
func createChatModelConfigWithTitleFailure(t testing.TB, client *codersdk.ExperimentalClient) codersdk.ChatModelConfig {
|
||||
t.Helper()
|
||||
baseURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if req.Stream {
|
||||
return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("Hello from test server.")...)
|
||||
}
|
||||
return chattest.OpenAIErrorResponse(http.StatusUnauthorized, "invalid_api_key", "test title failure")
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return modelConfig
|
||||
return createChatModelConfigWithBaseURL(t, client, baseURL)
|
||||
}
|
||||
|
||||
func createAdditionalChatModelConfig(
|
||||
@@ -10603,11 +10610,11 @@ func TestUserChatPersonalModelOverrides(t *testing.T) {
|
||||
noKeyClient := codersdk.NewExperimentalClient(noKeyClientRaw)
|
||||
|
||||
defaultModelConfig := createChatModelConfig(t, adminClient)
|
||||
provider := enableUserChatProviderKey(t, adminClient, memberClient, "openai")
|
||||
provider := enableUserChatProviderKey(t, adminClient, memberClient, defaultModelConfig.Provider)
|
||||
modelConfig := createAdditionalChatModelConfig(
|
||||
t,
|
||||
adminClient,
|
||||
"openai",
|
||||
defaultModelConfig.Provider,
|
||||
"gpt-4o-personal-"+uuid.NewString(),
|
||||
)
|
||||
err := adminClient.UpdateChatModelOverride(ctx, codersdk.ChatModelOverrideContextGeneral, codersdk.UpdateChatModelOverrideRequest{
|
||||
@@ -10622,7 +10629,7 @@ func TestUserChatPersonalModelOverrides(t *testing.T) {
|
||||
disabledModelConfig := createDisabledChatModelConfig(
|
||||
t,
|
||||
adminClient,
|
||||
"openai",
|
||||
defaultModelConfig.Provider,
|
||||
"gpt-4o-personal-disabled-"+uuid.NewString(),
|
||||
)
|
||||
disabledProvider, err := adminClient.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
|
||||
|
||||
+4
-22
@@ -30,8 +30,10 @@ func mcpDeploymentValues(t testing.TB) *codersdk.DeploymentValues {
|
||||
func newMCPClient(t testing.TB) *codersdk.Client {
|
||||
t.Helper()
|
||||
|
||||
providerKeys := coderdtest.FakeOpenAICompatProviderAPIKeys(t)
|
||||
return coderdtest.New(t, &coderdtest.Options{
|
||||
DeploymentValues: mcpDeploymentValues(t),
|
||||
DeploymentValues: mcpDeploymentValues(t),
|
||||
ChatProviderAPIKeys: &providerKeys,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1409,29 +1411,9 @@ func TestChatWithMCPServerIDs(t *testing.T) {
|
||||
require.Contains(t, fetched.MCPServerIDs, mcpConfig.ID)
|
||||
}
|
||||
|
||||
// createChatModelConfigForMCP sets up a chat provider and model
|
||||
// config so that CreateChat succeeds. This mirrors the helper in
|
||||
// chats_test.go but is defined here to avoid coupling.
|
||||
func createChatModelConfigForMCP(t testing.TB, client *codersdk.ExperimentalClient) codersdk.ChatModelConfig {
|
||||
t.Helper()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
_, err := client.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
|
||||
Provider: "openai",
|
||||
APIKey: "test-api-key",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
contextLimit := int64(4096)
|
||||
isDefault := true
|
||||
modelConfig, err := client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
|
||||
Provider: "openai",
|
||||
Model: "gpt-4o-mini",
|
||||
ContextLimit: &contextLimit,
|
||||
IsDefault: &isDefault,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
return modelConfig
|
||||
return coderdtest.CreateOpenAICompatChatModelConfig(t, client, "")
|
||||
}
|
||||
|
||||
func TestMCPOAuth2DiscoveryEdgeCases(t *testing.T) {
|
||||
|
||||
@@ -10,14 +10,6 @@ import (
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
// WaitUntilIdleForTest waits for background chat work tracked by the server to
|
||||
// finish without shutting the server down. Tests use this to assert final
|
||||
// database state only after asynchronous chat processing has completed.
|
||||
// Close waits for the same tracked work, but also stops the server.
|
||||
func WaitUntilIdleForTest(server *Server) {
|
||||
server.drainInflight()
|
||||
}
|
||||
|
||||
// FinishActiveChatForTest exposes the unexported cleanup TX so tests
|
||||
// can drive the post-run state machine deterministically. Returns the
|
||||
// resulting chat, the promoted message (if any), the synthetic
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
package chatd
|
||||
|
||||
// WaitUntilIdleForTest waits for background chat work tracked by the server to
|
||||
// finish without shutting the server down. Tests use this to assert final
|
||||
// database state only after asynchronous chat processing has completed.
|
||||
// Close waits for the same tracked work, but also stops the server.
|
||||
func WaitUntilIdleForTest(server *Server) {
|
||||
server.drainInflight()
|
||||
}
|
||||
Reference in New Issue
Block a user