diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 551e62de84..e50edfbfea 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -29,6 +29,7 @@ import ( "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/rbac/rolestore" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/cryptorand" "github.com/coder/coder/v2/provisionerd/proto" @@ -75,6 +76,166 @@ func AuditLog(t testing.TB, db database.Store, seed database.AuditLog) database. return log } +func Chat(t testing.TB, db database.Store, seed database.Chat) database.Chat { + t.Helper() + + var labels pqtype.NullRawMessage + if seed.Labels != nil { + raw, err := json.Marshal(seed.Labels) + require.NoError(t, err, "marshal chat labels") + labels = pqtype.NullRawMessage{RawMessage: raw, Valid: true} + } + + chat, err := db.InsertChat(genCtx, database.InsertChatParams{ + OrganizationID: takeFirst(seed.OrganizationID, uuid.New()), + OwnerID: takeFirst(seed.OwnerID, uuid.New()), + WorkspaceID: seed.WorkspaceID, + BuildID: seed.BuildID, + AgentID: seed.AgentID, + ParentChatID: seed.ParentChatID, + RootChatID: seed.RootChatID, + LastModelConfigID: takeFirst(seed.LastModelConfigID, uuid.New()), + Title: takeFirst(seed.Title, testutil.GetRandomName(t)), + Mode: seed.Mode, + PlanMode: seed.PlanMode, + Status: takeFirst(seed.Status, database.ChatStatusWaiting), + MCPServerIDs: seed.MCPServerIDs, + Labels: labels, + DynamicTools: seed.DynamicTools, + ClientType: takeFirst(seed.ClientType, database.ChatClientTypeUi), + }) + require.NoError(t, err, "insert chat") + return chat +} + +func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) database.ChatMessage { + t.Helper() + + content := "[]" + if seed.Content.Valid { + content = string(seed.Content.RawMessage) + } + + msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{ + ChatID: seed.ChatID, + CreatedBy: []uuid.UUID{seed.CreatedBy.UUID}, + ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID}, + Role: []database.ChatMessageRole{takeFirst(seed.Role, database.ChatMessageRoleUser)}, + Content: []string{content}, + ContentVersion: []int16{takeFirst(seed.ContentVersion, chatprompt.CurrentContentVersion)}, + Visibility: []database.ChatMessageVisibility{takeFirst(seed.Visibility, database.ChatMessageVisibilityBoth)}, + InputTokens: []int64{seed.InputTokens.Int64}, + OutputTokens: []int64{seed.OutputTokens.Int64}, + TotalTokens: []int64{seed.TotalTokens.Int64}, + ReasoningTokens: []int64{seed.ReasoningTokens.Int64}, + CacheCreationTokens: []int64{seed.CacheCreationTokens.Int64}, + CacheReadTokens: []int64{seed.CacheReadTokens.Int64}, + ContextLimit: []int64{seed.ContextLimit.Int64}, + Compressed: []bool{seed.Compressed}, + TotalCostMicros: []int64{seed.TotalCostMicros.Int64}, + RuntimeMs: []int64{seed.RuntimeMs.Int64}, + ProviderResponseID: []string{seed.ProviderResponseID.String}, + }) + require.NoError(t, err, "insert chat message") + require.Len(t, msgs, 1) + return msgs[0] +} + +const ( + // Match the default OpenAI test model's effective context settings. + defaultChatModelContextLimit int64 = 128000 + defaultChatModelCompressionThreshold int32 = 70 +) + +func ChatModelConfig(t testing.TB, db database.Store, seed database.ChatModelConfig, munge ...func(*database.InsertChatModelConfigParams)) database.ChatModelConfig { + t.Helper() + params := database.InsertChatModelConfigParams{ + Provider: takeFirst(seed.Provider, "openai"), + Model: takeFirst(seed.Model, "gpt-4o-mini"), + DisplayName: takeFirst(seed.DisplayName, "Test Model"), + CreatedBy: seed.CreatedBy, + UpdatedBy: seed.UpdatedBy, + Enabled: takeFirst(seed.Enabled, true), + IsDefault: seed.IsDefault, + ContextLimit: takeFirst(seed.ContextLimit, defaultChatModelContextLimit), + CompressionThreshold: takeFirst(seed.CompressionThreshold, defaultChatModelCompressionThreshold), + Options: takeFirstSlice(seed.Options, json.RawMessage(`{}`)), + } + for _, fn := range munge { + fn(¶ms) + } + cfg, err := db.InsertChatModelConfig(genCtx, params) + require.NoError(t, err, "insert chat model config") + return cfg +} + +func ChatProvider(t testing.TB, db database.Store, seed database.ChatProvider, munge ...func(*database.InsertChatProviderParams)) database.ChatProvider { + t.Helper() + params := database.InsertChatProviderParams{ + Provider: takeFirst(seed.Provider, "openai"), + DisplayName: takeFirst(seed.DisplayName, seed.Provider, "openai"), + APIKey: takeFirst(seed.APIKey, "test-key"), + BaseUrl: seed.BaseUrl, + ApiKeyKeyID: seed.ApiKeyKeyID, + CreatedBy: seed.CreatedBy, + Enabled: takeFirst(seed.Enabled, true), + CentralApiKeyEnabled: takeFirst(seed.CentralApiKeyEnabled, true), + AllowUserApiKey: seed.AllowUserApiKey, + AllowCentralApiKeyFallback: seed.AllowCentralApiKeyFallback, + } + for _, fn := range munge { + fn(¶ms) + } + provider, err := db.InsertChatProvider(genCtx, params) + require.NoError(t, err, "insert chat provider") + return provider +} + +func MCPServerConfig(t testing.TB, db database.Store, seed database.MCPServerConfig) database.MCPServerConfig { + t.Helper() + + // CreatedBy and UpdatedBy are user FKs, so default fixtures create a user. + createdBy := seed.CreatedBy.UUID + if createdBy == uuid.Nil { + createdBy = User(t, db, database.User{}).ID + } + updatedBy := seed.UpdatedBy.UUID + if updatedBy == uuid.Nil { + updatedBy = createdBy + } + + cfg, err := db.InsertMCPServerConfig(genCtx, database.InsertMCPServerConfigParams{ + DisplayName: takeFirst(seed.DisplayName, "Test MCP Server"), + Slug: takeFirst(seed.Slug, testutil.GetRandomName(t)), + Description: seed.Description, + IconURL: seed.IconURL, + Transport: takeFirst(seed.Transport, "streamable_http"), + Url: takeFirst(seed.Url, "https://mcp.example.com"), + AuthType: takeFirst(seed.AuthType, "none"), + OAuth2ClientID: seed.OAuth2ClientID, + OAuth2ClientSecret: seed.OAuth2ClientSecret, + OAuth2ClientSecretKeyID: seed.OAuth2ClientSecretKeyID, + OAuth2AuthURL: seed.OAuth2AuthURL, + OAuth2TokenURL: seed.OAuth2TokenURL, + OAuth2Scopes: seed.OAuth2Scopes, + APIKeyHeader: seed.APIKeyHeader, + APIKeyValue: seed.APIKeyValue, + APIKeyValueKeyID: seed.APIKeyValueKeyID, + CustomHeaders: seed.CustomHeaders, + CustomHeadersKeyID: seed.CustomHeadersKeyID, + ToolAllowList: takeFirstSlice(seed.ToolAllowList, []string{}), + ToolDenyList: takeFirstSlice(seed.ToolDenyList, []string{}), + Availability: takeFirst(seed.Availability, "default_off"), + Enabled: takeFirst(seed.Enabled, true), + ModelIntent: seed.ModelIntent, + AllowInPlanMode: seed.AllowInPlanMode, + CreatedBy: createdBy, + UpdatedBy: updatedBy, + }) + require.NoError(t, err, "insert MCP server config") + return cfg +} + func ConnectionLog(t testing.TB, db database.Store, seed database.UpsertConnectionLogParams) database.ConnectionLog { arg := database.UpsertConnectionLogParams{ ID: takeFirst(seed.ID, uuid.New()), diff --git a/coderd/database/dbgen/dbgen_test.go b/coderd/database/dbgen/dbgen_test.go index bd2e4ae36c..a07a9c5881 100644 --- a/coderd/database/dbgen/dbgen_test.go +++ b/coderd/database/dbgen/dbgen_test.go @@ -2,14 +2,18 @@ package dbgen_test import ( "context" + "database/sql" + "encoding/json" "testing" "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" ) func TestGenerator(t *testing.T) { @@ -252,6 +256,191 @@ func TestGenerator(t *testing.T) { require.Len(t, actual, 1) require.Equal(t, exp, actual[0]) }) + + t.Run("ChatProvider", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + + // Defaults. + p := dbgen.ChatProvider(t, db, database.ChatProvider{}) + require.NotEqual(t, uuid.Nil, p.ID) + require.Equal(t, "openai", p.Provider) + require.Equal(t, "openai", p.DisplayName) + require.True(t, p.Enabled) + require.True(t, p.CentralApiKeyEnabled) + require.Equal(t, "test-key", p.APIKey) + + // Overrides. + p2 := dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "anthropic", + DisplayName: "Claude", + APIKey: "sk-custom", + }) + require.Equal(t, "anthropic", p2.Provider) + require.Equal(t, "Claude", p2.DisplayName) + require.Equal(t, "sk-custom", p2.APIKey) + + p3 := dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openrouter", + }, func(params *database.InsertChatProviderParams) { + params.APIKey = "" + }) + require.Empty(t, p3.APIKey) + }) + + t.Run("ChatModelConfig", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + _ = dbgen.ChatProvider(t, db, database.ChatProvider{}) + + // Defaults. + cfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + require.NotEqual(t, uuid.Nil, cfg.ID) + require.Equal(t, "openai", cfg.Provider) + require.Equal(t, "gpt-4o-mini", cfg.Model) + require.Equal(t, "Test Model", cfg.DisplayName) + require.True(t, cfg.Enabled) + require.Equal(t, int64(128000), cfg.ContextLimit) + require.Equal(t, int32(70), cfg.CompressionThreshold) + + // Overrides. + _ = dbgen.ChatProvider(t, db, database.ChatProvider{Provider: "anthropic"}) + cfg2 := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "anthropic", + Model: "claude-4", + ContextLimit: 200000, + }) + require.Equal(t, "anthropic", cfg2.Provider) + require.Equal(t, "claude-4", cfg2.Model) + require.Equal(t, int64(200000), cfg2.ContextLimit) + }) + + t.Run("Chat", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + u := dbgen.User(t, db, database.User{}) + o := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: u.ID, + OrganizationID: o.ID, + }) + p := dbgen.ChatProvider(t, db, database.ChatProvider{}) + m := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{Provider: p.Provider}) + + // Defaults. + chat := dbgen.Chat(t, db, database.Chat{ + OwnerID: u.ID, + OrganizationID: o.ID, + LastModelConfigID: m.ID, + }) + require.NotEqual(t, uuid.Nil, chat.ID) + require.Equal(t, database.ChatStatusWaiting, chat.Status) + require.Equal(t, database.ChatClientTypeUi, chat.ClientType) + require.NotEmpty(t, chat.Title) + + // Overrides. + chat2 := dbgen.Chat(t, db, database.Chat{ + OwnerID: u.ID, + OrganizationID: o.ID, + LastModelConfigID: m.ID, + Title: "custom-title", + Status: database.ChatStatusRunning, + }) + require.Equal(t, "custom-title", chat2.Title) + require.Equal(t, database.ChatStatusRunning, chat2.Status) + }) + + t.Run("ChatMessage", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + u := dbgen.User(t, db, database.User{}) + o := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: u.ID, + OrganizationID: o.ID, + }) + p := dbgen.ChatProvider(t, db, database.ChatProvider{}) + m := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{Provider: p.Provider}) + chat := dbgen.Chat(t, db, database.Chat{ + OwnerID: u.ID, + OrganizationID: o.ID, + LastModelConfigID: m.ID, + }) + + // Defaults. + msg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + }) + require.NotZero(t, msg.ID) + require.Equal(t, database.ChatMessageRoleUser, msg.Role) + require.Equal(t, database.ChatMessageVisibilityBoth, msg.Visibility) + require.Equal(t, chatprompt.CurrentContentVersion, msg.ContentVersion) + + // Overrides. + rawContent := json.RawMessage(`[{"type":"text","text":"hello"}]`) + msg2 := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{ + RawMessage: rawContent, + Valid: true, + }, + InputTokens: sql.NullInt64{Int64: 11, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 22, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 33, Valid: true}, + ReasoningTokens: sql.NullInt64{Int64: 44, Valid: true}, + CacheCreationTokens: sql.NullInt64{Int64: 55, Valid: true}, + CacheReadTokens: sql.NullInt64{Int64: 66, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 77, Valid: true}, + Compressed: true, + TotalCostMicros: sql.NullInt64{Int64: 88, Valid: true}, + ProviderResponseID: sql.NullString{String: "resp-123", Valid: true}, + }) + require.Equal(t, database.ChatMessageRoleAssistant, msg2.Role) + require.True(t, msg2.Content.Valid) + require.JSONEq(t, string(rawContent), string(msg2.Content.RawMessage)) + require.Equal(t, sql.NullInt64{Int64: 11, Valid: true}, msg2.InputTokens) + require.Equal(t, sql.NullInt64{Int64: 22, Valid: true}, msg2.OutputTokens) + require.Equal(t, sql.NullInt64{Int64: 33, Valid: true}, msg2.TotalTokens) + require.Equal(t, sql.NullInt64{Int64: 44, Valid: true}, msg2.ReasoningTokens) + require.Equal(t, sql.NullInt64{Int64: 55, Valid: true}, msg2.CacheCreationTokens) + require.Equal(t, sql.NullInt64{Int64: 66, Valid: true}, msg2.CacheReadTokens) + require.Equal(t, sql.NullInt64{Int64: 77, Valid: true}, msg2.ContextLimit) + require.True(t, msg2.Compressed) + require.Equal(t, sql.NullInt64{Int64: 88, Valid: true}, msg2.TotalCostMicros) + require.Equal(t, sql.NullString{String: "resp-123", Valid: true}, msg2.ProviderResponseID) + }) + + t.Run("MCPServerConfig", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + + // Defaults. + cfg := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{}) + require.NotEqual(t, uuid.Nil, cfg.ID) + require.Equal(t, "streamable_http", cfg.Transport) + require.Equal(t, "none", cfg.AuthType) + require.Equal(t, "default_off", cfg.Availability) + require.True(t, cfg.Enabled) + require.Empty(t, cfg.ToolAllowList) + require.Empty(t, cfg.ToolDenyList) + require.NotEmpty(t, cfg.Slug) + require.NotEmpty(t, cfg.Url) + + // Overrides. + cfg2 := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Custom MCP", + Slug: "custom-mcp", + Url: "https://custom.example.com", + AuthType: "oauth2", + AllowInPlanMode: true, + }) + require.Equal(t, "Custom MCP", cfg2.DisplayName) + require.Equal(t, "custom-mcp", cfg2.Slug) + require.Equal(t, "https://custom.example.com", cfg2.Url) + require.Equal(t, "oauth2", cfg2.AuthType) + require.True(t, cfg2.AllowInPlanMode) + }) } func must[T any](value T, err error) T { diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index 166ed12bc0..9d02aba6e1 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -1839,20 +1839,17 @@ func TestDeleteOldChatFiles(t *testing.T) { // backdates updated_at to control the "archived since" window. createChat := func(ctx context.Context, t *testing.T, db database.Store, rawDB *sql.DB, ownerID, orgID, modelConfigID uuid.UUID, archived bool, updatedAt time.Time) database.Chat { t.Helper() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: orgID, OwnerID: ownerID, LastModelConfigID: modelConfigID, Title: "test-chat", - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) if archived { - _, err = db.ArchiveChatByID(ctx, chat.ID) + _, err := db.ArchiveChatByID(ctx, chat.ID) require.NoError(t, err) } - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET updated_at = $1 WHERE id = $2", updatedAt, chat.ID) + _, err := rawDB.ExecContext(ctx, "UPDATE chats SET updated_at = $1 WHERE id = $2", updatedAt, chat.ID) require.NoError(t, err) return chat } @@ -1863,25 +1860,20 @@ func TestDeleteOldChatFiles(t *testing.T) { org database.Organization modelConfig database.ChatModelConfig } - setupChatDeps := func(ctx context.Context, t *testing.T, db database.Store) chatDeps { + setupChatDeps := func(t *testing.T, db database.Store) chatDeps { t.Helper() user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", }) - require.NoError(t, err) - mc, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ + mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ Provider: "openai", Model: "test-model", ContextLimit: 8192, - Options: json.RawMessage("{}"), }) - require.NoError(t, err) return chatDeps{user: user, org: org, modelConfig: mc} } @@ -1898,7 +1890,7 @@ func TestDeleteOldChatFiles(t *testing.T) { db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - deps := setupChatDeps(ctx, t, db) + deps := setupChatDeps(t, db) // Disable retention. err := db.UpsertChatRetentionDays(ctx, int32(0)) @@ -1929,7 +1921,7 @@ func TestDeleteOldChatFiles(t *testing.T) { db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - deps := setupChatDeps(ctx, t, db) + deps := setupChatDeps(t, db) err := db.UpsertChatRetentionDays(ctx, int32(30)) require.NoError(t, err) @@ -1937,27 +1929,12 @@ func TestDeleteOldChatFiles(t *testing.T) { // Old archived chat (31 days) — should be deleted. oldChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-31*24*time.Hour)) // Insert a message so we can verify CASCADE. - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: oldChat.ID, - CreatedBy: []uuid.UUID{deps.user.ID}, - ModelConfigID: []uuid.UUID{deps.modelConfig.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - Content: []string{`[{"type":"text","text":"hello"}]`}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: oldChat.ID, + CreatedBy: uuid.NullUUID{UUID: deps.user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: deps.modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleUser, }) - require.NoError(t, err) // Recently archived chat (10 days) — should be retained. recentChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, true, now.Add(-10*24*time.Hour)) @@ -1998,7 +1975,7 @@ func TestDeleteOldChatFiles(t *testing.T) { db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - deps := setupChatDeps(ctx, t, db) + deps := setupChatDeps(t, db) err := db.UpsertChatRetentionDays(ctx, int32(30)) require.NoError(t, err) @@ -2049,7 +2026,7 @@ func TestDeleteOldChatFiles(t *testing.T) { db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - deps := setupChatDeps(ctx, t, db) + deps := setupChatDeps(t, db) err := db.UpsertChatRetentionDays(ctx, int32(30)) require.NoError(t, err) @@ -2126,7 +2103,7 @@ func TestDeleteOldChatFiles(t *testing.T) { // file purge should show only surviving files. ctx := testutil.Context(t, testutil.WaitLong) db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) - deps := setupChatDeps(ctx, t, db) + deps := setupChatDeps(t, db) // Create a chat with three attached files. fileA := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now) @@ -2179,19 +2156,13 @@ func TestDeleteOldChatFiles(t *testing.T) { // clean up links for both parent and child chats // independently via FK cascade. parentChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) - childChat, err := db.InsertChat(ctx, database.InsertChatParams{ + childChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: deps.org.ID, OwnerID: deps.user.ID, LastModelConfigID: deps.modelConfig.ID, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, Title: "child-chat", - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) - - // Set root_chat_id to link child to parent. - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET root_chat_id = $1 WHERE id = $2", parentChat.ID, childChat.ID) - require.NoError(t, err) // Attach different files to parent and child. parentFileKeep := createChatFile(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, now) @@ -2243,7 +2214,7 @@ func TestDeleteOldChatFiles(t *testing.T) { run: func(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) - deps := setupChatDeps(ctx, t, db) + deps := setupChatDeps(t, db) // Create 3 deletable orphaned files (all 31 days old). for range 3 { @@ -2272,7 +2243,7 @@ func TestDeleteOldChatFiles(t *testing.T) { run: func(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) - deps := setupChatDeps(ctx, t, db) + deps := setupChatDeps(t, db) // Create 3 deletable old archived chats. for range 3 { @@ -2307,25 +2278,20 @@ func TestDeleteOldChatFiles(t *testing.T) { // helpers for TestAutoArchiveInactiveChats. Kept scoped to the // test so they don't leak into the package surface area. -func archiveTestDeps(ctx context.Context, t *testing.T, db database.Store) chatAutoArchiveDeps { +func archiveTestDeps(t *testing.T, db database.Store) chatAutoArchiveDeps { t.Helper() user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", }) - require.NoError(t, err) - mc, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ + mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ Provider: "openai", Model: "test-model", ContextLimit: 8192, - Options: json.RawMessage("{}"), }) - require.NoError(t, err) return chatAutoArchiveDeps{user: user, org: org, modelConfig: mc} } @@ -2361,7 +2327,7 @@ func newArchiveHarness(t *testing.T, now time.Time) *archiveHarness { db: db, rawDB: rawDB, logger: logger, - deps: archiveTestDeps(ctx, t, db), + deps: archiveTestDeps(t, db), } } @@ -2370,16 +2336,13 @@ func newArchiveHarness(t *testing.T, now time.Time) *archiveHarness { // digest contents. func createArchiveChat(ctx context.Context, t *testing.T, db database.Store, rawDB *sql.DB, deps chatAutoArchiveDeps, title string, createdAt time.Time) database.Chat { t.Helper() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: deps.org.ID, OwnerID: deps.user.ID, LastModelConfigID: deps.modelConfig.ID, Title: title, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET created_at = $1, updated_at = $1 WHERE id = $2", createdAt, chat.ID) + _, err := rawDB.ExecContext(ctx, "UPDATE chats SET created_at = $1, updated_at = $1 WHERE id = $2", createdAt, chat.ID) require.NoError(t, err) return chat } @@ -2389,29 +2352,13 @@ func createArchiveChat(ctx context.Context, t *testing.T, db database.Store, raw // auto-archive query's LATERAL subquery. func insertTextMessage(ctx context.Context, t *testing.T, db database.Store, rawDB *sql.DB, chatID, userID, modelConfigID uuid.UUID, createdAt time.Time) { t.Helper() - msgs, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{userID}, - ModelConfigID: []uuid.UUID{modelConfigID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - Content: []string{`[{"type":"text","text":"hello"}]`}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, + msg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleUser, }) - require.NoError(t, err) - require.Len(t, msgs, 1) - _, err = rawDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1 WHERE id = $2", createdAt, msgs[0].ID) + _, err := rawDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1 WHERE id = $2", createdAt, msg.ID) require.NoError(t, err) } diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 0f96dcc172..30ae724ff7 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -1260,54 +1260,37 @@ func TestGetAuthorizedChats(t *testing.T) { dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: secondMember.ID, OrganizationID: org.ID, Roles: []string{rbac.RoleAgentsAccess()}}) // Create FK dependencies: a chat provider and model config. - ctx := testutil.Context(t, testutil.WaitMedium) - _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - APIKey: "test-key", - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", }) - require.NoError(t, err) - - modelCfg, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ + modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ Provider: "openai", Model: "test-model", - DisplayName: "Test Model", CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, - Enabled: true, IsDefault: true, - ContextLimit: 128000, CompressionThreshold: 80, - Options: json.RawMessage(`{}`), }) - require.NoError(t, err) // Create 3 chats owned by owner. for i := range 3 { - _, err := db.InsertChat(ctx, database.InsertChatParams{ + dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: owner.ID, LastModelConfigID: modelCfg.ID, Title: fmt.Sprintf("owner chat %d", i+1), }) - require.NoError(t, err) } // Create 2 chats owned by member. for i := range 2 { - _, err := db.InsertChat(ctx, database.InsertChatParams{ + dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelCfg.ID, Title: fmt.Sprintf("member chat %d", i+1), }) - require.NoError(t, err) } t.Run("sqlQuerier", func(t *testing.T) { @@ -1437,15 +1420,12 @@ func TestGetAuthorizedChats(t *testing.T) { paginationUser := dbgen.User(t, db, database.User{}) dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: paginationUser.ID, OrganizationID: org.ID, Roles: []string{rbac.RoleAgentsAccess()}}) for i := range 7 { - _, err := db.InsertChat(ctx, database.InsertChatParams{ + dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: paginationUser.ID, LastModelConfigID: modelCfg.ID, Title: fmt.Sprintf("pagination chat %d", i+1), }) - require.NoError(t, err) } pagUserSubject, _, err := httpmw.UserRBACSubject(ctx, db, paginationUser.ID, rbac.ExpandableScope(rbac.ScopeAll)) diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 25fea235c2..8d1b370d52 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -217,7 +217,6 @@ func enableDailyChatUsageLimit( } func insertAssistantCostMessage( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -231,26 +230,13 @@ func insertAssistantCostMessage( }) require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfigID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{totalCostMicros}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + TotalCostMicros: sql.NullInt64{Int64: totalCostMicros, Valid: true}, }) - require.NoError(t, err) } func TestPostChats(t *testing.T) { @@ -690,19 +676,15 @@ func TestPostChats(t *testing.T) { modelConfig := createChatModelConfig(t, client) wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - existingChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + existingChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "existing-limit-chat", }) - require.NoError(t, err) + insertAssistantCostMessage(t, db, existingChat.ID, modelConfig.ID, 100) - insertAssistantCostMessage(ctx, t, db, existingChat.ID, modelConfig.ID, 100) - - _, err = client.CreateChat(ctx, codersdk.CreateChatRequest{ + _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: user.OrganizationID, Content: []codersdk.ChatInputPart{{ Type: codersdk.ChatInputPartTypeText, @@ -909,15 +891,12 @@ func TestListChats(t *testing.T) { memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID, rbac.ScopedRoleAgentsAccess(firstUser.OrganizationID)) memberClient := codersdk.NewExperimentalClient(memberClientRaw) - memberDBChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + memberDBChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, Title: "member chat only", }) - require.NoError(t, err) chats, err := client.ListChats(ctx, nil) require.NoError(t, err) @@ -992,15 +971,12 @@ func TestListChats(t *testing.T) { // a specific chat returns 404 (dbauthz wraps as not found). memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) memberClient := codersdk.NewExperimentalClient(memberClientRaw) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, Title: "member chat", }) - require.NoError(t, err) // Listing chats returns empty because the SQL auth // filter excludes chats the member cannot read. @@ -1333,29 +1309,23 @@ func TestListChats(t *testing.T) { require.NoError(t, err) // Insert child chats directly via the database. - child1, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child1 := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "child one", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) - child2, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child2 := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "child two", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) // Also create a standalone root chat to verify it still appears. standalone, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ @@ -1437,17 +1407,14 @@ func TestListChats(t *testing.T) { }) require.NoError(t, err) for j := range 2 { - _, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + _ = dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: fmt.Sprintf("child %d-%d", i, j), ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, }) - require.NoError(t, err) } } @@ -1791,19 +1758,15 @@ func TestWatchChats(t *testing.T) { modelConfig := createChatModelConfig(t, client) // Insert a chat and a diff status row. - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "diff status watch test", }) - require.NoError(t, err) - refreshedAt := time.Now().UTC().Truncate(time.Second) staleAt := refreshedAt.Add(time.Hour) - _, err = db.UpsertChatDiffStatusReference( + _, err := db.UpsertChatDiffStatusReference( dbauthz.AsSystemRestricted(ctx), database.UpsertChatDiffStatusReferenceParams{ ChatID: chat.ID, @@ -1924,27 +1887,23 @@ func TestWatchChats(t *testing.T) { }) require.NoError(t, err) - childOne, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + childOne := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "watch child 1", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "watch child 1", + ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) - childTwo, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + childTwo := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "watch child 2", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "watch child 2", + ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) conn, err := client.Dial(ctx, "/api/experimental/chats/watch", nil) require.NoError(t, err) @@ -2744,7 +2703,7 @@ func TestDeleteChatProvider(t *testing.T) { require.NoError(t, err) require.Equal(t, configToDelete.ID, chat.LastModelConfigID) - insertAssistantCostMessage(ctx, t, db, chat.ID, configToDelete.ID, 500) + insertAssistantCostMessage(t, db, chat.ID, configToDelete.ID, 500) _, err = client.UpsertUserChatProviderKey(ctx, providerToDelete.ID, codersdk.CreateUserChatProviderKeyRequest{ APIKey: "user-delete-key", @@ -2851,7 +2810,7 @@ func TestDeleteChatProvider(t *testing.T) { require.NoError(t, err) require.Equal(t, config.ID, chat.LastModelConfigID) - insertAssistantCostMessage(ctx, t, db, chat.ID, config.ID, 250) + insertAssistantCostMessage(t, db, chat.ID, config.ID, 250) err = client.DeleteChatProvider(ctx, provider.ID) require.NoError(t, err) @@ -3516,19 +3475,16 @@ func TestListChatModelConfigs(t *testing.T) { require.NoError(t, err) legacyOptions := json.RawMessage(`{"input_price_per_million_tokens":0.15,"output_price_per_million_tokens":0.6,"cache_read_price_per_million_tokens":0.03,"cache_write_price_per_million_tokens":0.3}`) - storedConfig, err := db.InsertChatModelConfig(dbauthz.AsSystemRestricted(ctx), database.InsertChatModelConfigParams{ + storedConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ Provider: "openai", Model: "gpt-4o-mini-legacy", DisplayName: "GPT-4o Mini Legacy", CreatedBy: uuid.NullUUID{UUID: firstUser.UserID, Valid: true}, UpdatedBy: uuid.NullUUID{UUID: firstUser.UserID, Valid: true}, - Enabled: true, - IsDefault: false, ContextLimit: 4096, CompressionThreshold: 80, Options: legacyOptions, }) - require.NoError(t, err) configs, err := client.ListChatModelConfigs(ctx) require.NoError(t, err) @@ -4320,17 +4276,14 @@ func TestGetChat(t *testing.T) { }) require.NoError(t, err) - child, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "child for getChat", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) // Fetching the root chat should embed its children. result, err := client.GetChat(ctx, parentChat.ID) @@ -4399,15 +4352,12 @@ func TestPatchChat(t *testing.T) { ) codersdk.Chat { t.Helper() - dbChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + dbChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: orgID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: ownerID, LastModelConfigID: modelConfigID, Title: title, }) - require.NoError(t, err) return db2sdk.Chat(dbChat, nil, nil) } @@ -5049,29 +4999,23 @@ func TestArchiveChat(t *testing.T) { require.NoError(t, err) // Insert child chats directly via the database. - child1, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child1 := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "child 1", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) - child2, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child2 := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "child 2", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) // Archive the parent via the API. err = client.UpdateChat(ctx, parentChat.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) @@ -5143,17 +5087,14 @@ func TestArchiveChat(t *testing.T) { require.NoError(t, err) // Insert a child chat directly via the database. - child, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "child", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) // Individual child archive is permitted and leaves the // parent active; the invariant is one-way. @@ -5266,27 +5207,23 @@ func TestUnarchiveChat(t *testing.T) { }) require.NoError(t, err) - child1, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child1 := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "child 1", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "child 1", + ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) - child2, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child2 := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, - Title: "child 2", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + Title: "child 2", + ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) err = client.UpdateChat(ctx, parentChat.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) require.NoError(t, err) @@ -5389,17 +5326,15 @@ func TestUnarchiveChat(t *testing.T) { // Insert a child directly via the database, then archive the // parent so the whole family is archived (cascade). - child, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "child", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) + err = client.UpdateChat(ctx, parentChat.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) require.NoError(t, err) @@ -5438,17 +5373,15 @@ func TestUnarchiveChat(t *testing.T) { // Simulate legacy lone-archived child (from before the // child-archive gate existed) by inserting it directly // with archived=true while the parent is not archived. - child, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "legacy child", ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) + _, err = db.ArchiveChatByID(dbauthz.AsSystemRestricted(ctx), child.ID) require.NoError(t, err) @@ -5606,19 +5539,18 @@ func TestChatPinOrder(t *testing.T) { parentChat := createChat(ctx, t, client, firstUser.OrganizationID, "parent chat") - child, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + child := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "child chat", Status: database.ChatStatusCompleted, - ClientType: database.ChatClientTypeUi, ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - require.NoError(t, err) - err = client.UpdateChat(ctx, child.ID, codersdk.UpdateChatRequest{PinOrder: ptr.Ref(int32(1))}) + err := client.UpdateChat(ctx, child.ID, codersdk.UpdateChatRequest{PinOrder: ptr.Ref(int32(1))}) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) require.Equal(t, "Cannot pin a child chat.", sdkErr.Message) @@ -5736,17 +5668,14 @@ func TestPostChatMessages(t *testing.T) { // before the handler can check agents-access. memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) memberClient := codersdk.NewExperimentalClient(memberClientRaw) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, Title: "member chat", }) - require.NoError(t, err) - _, err = memberClient.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + _, err := memberClient.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -5807,7 +5736,7 @@ func TestPostChatMessages(t *testing.T) { require.NoError(t, err) wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - insertAssistantCostMessage(ctx, t, db, chat.ID, modelConfig.ID, 100) + insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) _, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -5897,15 +5826,12 @@ func TestSendMessageWithModelOverrideUpdatesLastModelConfigID(t *testing.T) { modelConfigA := createChatModelConfig(t, client) modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-override-"+uuid.NewString()) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfigA.ID, Title: "mid-chat model switch direct send", }) - require.NoError(t, err) resp, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -5943,17 +5869,14 @@ func TestSendMessageQueuesEffectiveModelConfigID(t *testing.T) { modelConfigA := createChatModelConfig(t, client) modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-queued-"+uuid.NewString()) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfigA.ID, Title: "mid-chat model switch queued send", }) - require.NoError(t, err) - _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ + _, err := db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRunning, WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, @@ -5997,17 +5920,14 @@ func TestQueuedMessageWithoutOverrideCapturesEnqueueTimeModel(t *testing.T) { modelConfigA := createChatModelConfig(t, client) modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-later-"+uuid.NewString()) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfigA.ID, Title: "capture queued enqueue-time model", }) - require.NoError(t, err) - _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ + _, err := db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRunning, WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, @@ -6052,15 +5972,12 @@ func TestSubsequentSendWithoutOverrideUsesPersistedModel(t *testing.T) { _ = createChatModelConfig(t, client) modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-persisted-"+uuid.NewString()) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfigB.ID, Title: "subsequent send uses persisted model", }) - require.NoError(t, err) resp, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -6096,15 +6013,12 @@ func TestWatchChatsStatusChangeCarriesUpdatedLastModelConfigID(t *testing.T) { modelConfigA := createChatModelConfig(t, client) modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-watch-direct-"+uuid.NewString()) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfigA.ID, Title: "watch direct model switch", }) - require.NoError(t, err) conn, err := client.Dial(ctx, "/api/experimental/chats/watch", nil) require.NoError(t, err) @@ -6132,17 +6046,14 @@ func TestWatchChatsStatusChangeCarriesUpdatedLastModelConfigID(t *testing.T) { modelConfigA := createChatModelConfig(t, client) modelConfigB := createAdditionalChatModelConfig(t, client, "openai", "gpt-4o-mini-watch-promote-"+uuid.NewString()) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfigA.ID, Title: "watch queued promotion model switch", }) - require.NoError(t, err) - _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ + _, err := db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRunning, WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, @@ -7167,7 +7078,7 @@ func TestPatchChatMessage(t *testing.T) { require.NotZero(t, userMessageID) wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - insertAssistantCostMessage(ctx, t, db, chat.ID, modelConfig.ID, 100) + insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) _, err = client.EditChatMessage(ctx, chat.ID, userMessageID, codersdk.EditChatMessageRequest{ Content: []codersdk.ChatInputPart{{ @@ -7488,17 +7399,15 @@ func TestInterruptChat(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "interrupt route test", }) - require.NoError(t, err) runningWorkerID := uuid.New() + var err error chat, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRunning, @@ -7506,6 +7415,7 @@ func TestInterruptChat(t *testing.T) { StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, }) + require.NoError(t, err) require.Equal(t, database.ChatStatusRunning, chat.Status) require.True(t, chat.WorkerID.Valid) @@ -7570,17 +7480,14 @@ func TestRegenerateChatTitle(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "chat with update denied", }) - require.NoError(t, err) - _, err = client.RegenerateChatTitle(ctx, chat.ID) + _, err := client.RegenerateChatTitle(ctx, chat.ID) requireSDKError(t, err, http.StatusNotFound) }) @@ -7649,7 +7556,7 @@ func TestRegenerateChatTitle(t *testing.T) { require.NoError(t, err) wantResetsAt := enableDailyChatUsageLimit(ctx, t, db, 100) - insertAssistantCostMessage(ctx, t, db, chat.ID, modelConfig.ID, 100) + insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, @@ -7684,23 +7591,21 @@ func TestRegenerateChatTitle(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "chat with lock held", }) - require.NoError(t, err) - _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ + _, err := db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusCompleted, WorkerID: uuid.NullUUID{UUID: uuid.MustParse("00000000-0000-0000-0000-000000000001"), Valid: true}, StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - LastError: sql.NullString{}, + + LastError: sql.NullString{}, }) require.NoError(t, err) @@ -7727,23 +7632,22 @@ func TestRegenerateChatTitle(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "pending chat without worker", }) - require.NoError(t, err) + var err error chat, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusPending, WorkerID: uuid.NullUUID{}, StartedAt: sql.NullTime{}, HeartbeatAt: sql.NullTime{}, - LastError: sql.NullString{}, + + LastError: sql.NullString{}, }) require.NoError(t, err) @@ -7856,17 +7760,15 @@ func TestProposeChatTitle(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "chat with update denied", }) - require.NoError(t, err) - _, err = client.ProposeChatTitle(ctx, chat.ID) + _, err := client.ProposeChatTitle(ctx, chat.ID) + requireSDKError(t, err, http.StatusNotFound) }) @@ -7932,30 +7834,24 @@ func TestGetChatDiffStatus(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - noCachedStatusChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + noCachedStatusChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "get diff status route no cache", }) - require.NoError(t, err) noCachedChat, err := client.GetChat(ctx, noCachedStatusChat.ID) require.NoError(t, err) require.Equal(t, noCachedStatusChat.ID, noCachedChat.ID) require.Nil(t, noCachedChat.DiffStatus) - cachedStatusChat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + cachedStatusChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "get diff status route cached", }) - require.NoError(t, err) refreshedAt := time.Now().UTC().Truncate(time.Second) staleAt := refreshedAt.Add(time.Hour) @@ -8057,17 +7953,14 @@ func TestGetChatDiffContents(t *testing.T) { db := api.Database user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "diff contents with cached repository reference", }) - require.NoError(t, err) - _, err = db.UpsertChatDiffStatusReference( + _, err := db.UpsertChatDiffStatusReference( dbauthz.AsSystemRestricted(ctx), database.UpsertChatDiffStatusReferenceParams{ ChatID: chat.ID, @@ -8158,15 +8051,12 @@ func TestDeleteChatQueuedMessage(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "delete queued message route test", }) - require.NoError(t, err) deleteContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued message for delete route"), @@ -8212,15 +8102,12 @@ func TestDeleteChatQueuedMessage(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "delete queued invalid id", }) - require.NoError(t, err) invalidRes, err := client.Request( ctx, @@ -8229,6 +8116,7 @@ func TestDeleteChatQueuedMessage(t *testing.T) { nil, ) require.NoError(t, err) + defer invalidRes.Body.Close() err = codersdk.ReadBodyAsError(invalidRes) @@ -8249,15 +8137,12 @@ func TestPromoteChatQueuedMessage(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "promote queued message route test", }) - require.NoError(t, err) const queuedText = "queued message for promote route" queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ @@ -8322,17 +8207,15 @@ func TestPromoteChatQueuedMessage(t *testing.T) { modelConfig := createChatModelConfig(t, client) enableDailyChatUsageLimit(ctx, t, db, 100) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "promote queued usage limit", }) - require.NoError(t, err) const queuedText = "queued message for promote route" + queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText(queuedText), }) @@ -8346,7 +8229,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { ) require.NoError(t, err) - insertAssistantCostMessage(ctx, t, db, chat.ID, modelConfig.ID, 100) + insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, @@ -8399,15 +8282,12 @@ func TestPromoteChatQueuedMessage(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "promote queued invalid id", }) - require.NoError(t, err) invalidRes, err := client.Request( ctx, @@ -8438,15 +8318,12 @@ func TestPromoteChatQueuedMessage(t *testing.T) { // before the handler can check agents-access. memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) memberClient := codersdk.NewExperimentalClient(memberClientRaw) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, Title: "promote queued no agents access", }) - require.NoError(t, err) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued message no agents access"), @@ -8480,15 +8357,12 @@ func TestPromoteChatQueuedMessage(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "promote queued archived", }) - require.NoError(t, err) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), @@ -9114,42 +8988,38 @@ func (f chatCostTestFixture) safeOptions() codersdk.ChatCostSummaryOptions { func seedChatCostFixture(t *testing.T) chatCostTestFixture { t.Helper() - ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "test chat", }) - require.NoError(t, err) - results, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil, uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID, modelConfig.ID}, - Role: []database.ChatMessageRole{"assistant", "assistant"}, - Content: []string{"null", "null"}, - ContentVersion: []int16{0, 0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth}, - InputTokens: []int64{100, 100}, - OutputTokens: []int64{50, 50}, - TotalTokens: []int64{0, 0}, - ReasoningTokens: []int64{0, 0}, - CacheCreationTokens: []int64{0, 0}, - CacheReadTokens: []int64{0, 0}, - ContextLimit: []int64{0, 0}, - Compressed: []bool{false, false}, - TotalCostMicros: []int64{500, 500}, - RuntimeMs: []int64{1500, 2500}, + msg1 := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 100, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 1500, Valid: true}, }) - require.NoError(t, err) + msg2 := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 100, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 2500, Valid: true}, + }) + results := []database.ChatMessage{msg1, msg2} require.Len(t, results, 2) + earliestCreatedAt := results[0].CreatedAt latestCreatedAt := results[0].CreatedAt for _, msg := range results { @@ -9235,44 +9105,28 @@ func TestChatCostSummary_AfterModelDeletion(t *testing.T) { func TestChatCostSummary_AdminDrilldown(t *testing.T) { t.Parallel() - seedCtx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID) memberClient := codersdk.NewExperimentalClient(memberClientRaw) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, Title: "member chat", }) - require.NoError(t, err) - results, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{"assistant"}, - Content: []string{"null"}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{200}, - OutputTokens: []int64{100}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{750}, - RuntimeMs: []int64{0}, + message := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 200, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 100, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 750, Valid: true}, }) - require.NoError(t, err) - message := results[0] + options := codersdk.ChatCostSummaryOptions{ // Pad the DB-assigned timestamp so the query window cannot race it. StartDate: message.CreatedAt.Add(-time.Minute), @@ -9313,65 +9167,35 @@ func TestChatCostUsers(t *testing.T) { require.NoError(t, err) modelConfig := createChatModelConfig(t, client) - adminChat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + adminChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "admin chat", }) - require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ - ChatID: adminChat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{"assistant"}, - Content: []string{"null"}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{100}, - OutputTokens: []int64{50}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{300}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: adminChat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 100, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 300, Valid: true}, }) - require.NoError(t, err) - memberChat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + memberChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, Title: "member chat", }) - require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ - ChatID: memberChat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{"assistant"}, - Content: []string{"null"}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{200}, - OutputTokens: []int64{100}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{800}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: memberChat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 200, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 100, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 800, Valid: true}, }) - require.NoError(t, err) t.Run("AdminCanListUsers", func(t *testing.T) { t.Parallel() @@ -9424,41 +9248,25 @@ func TestChatCostUsers(t *testing.T) { func TestChatCostSummary_DateRange(t *testing.T) { t.Parallel() - seedCtx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "date range test", }) - require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{"assistant"}, - Content: []string{"null"}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{100}, - OutputTokens: []int64{50}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{500}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 100, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true}, }) - require.NoError(t, err) now := time.Now() @@ -9497,59 +9305,29 @@ func TestChatCostSummary_UnpricedMessages(t *testing.T) { firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "unpriced test", }) - require.NoError(t, err) - pricedResults, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{"assistant"}, - Content: []string{"null"}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{100}, - OutputTokens: []int64{50}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{500}, - RuntimeMs: []int64{0}, + pricedMessage := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 100, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true}, }) - require.NoError(t, err) - pricedMessage := pricedResults[0] - unpricedResults, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - Role: []database.ChatMessageRole{"assistant"}, - Content: []string{"null"}, - ContentVersion: []int16{0}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{200}, - OutputTokens: []int64{75}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + unpricedMessage := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + InputTokens: sql.NullInt64{Int64: 200, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 75, Valid: true}, }) - require.NoError(t, err) - unpricedMessage := unpricedResults[0] earliestCreatedAt := pricedMessage.CreatedAt latestCreatedAt := pricedMessage.CreatedAt @@ -10861,15 +10639,12 @@ func TestChatDebugRuns(t *testing.T) { memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID, rbac.ScopedRoleAgentsAccess(firstUser.OrganizationID)) memberClient := codersdk.NewExperimentalClient(memberClientRaw) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: member.ID, LastModelConfigID: modelConfig.ID, Title: "debug-runs-list", }) - require.NoError(t, err) base := time.Now().UTC().Add(-time.Hour).Round(time.Second) older := seedChatDebugRun(ctx, t, db, chat.ID, base) @@ -10892,15 +10667,12 @@ func TestChatDebugRuns(t *testing.T) { firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-runs-cap", }) - require.NoError(t, err) base := time.Now().UTC().Add(-24 * time.Hour).Round(time.Second) // Seed 101 runs with monotonically increasing started_at. The @@ -10931,15 +10703,12 @@ func TestChatDebugRuns(t *testing.T) { firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-runs-empty", }) - require.NoError(t, err) // Guard against a regression from `make([]..., 0, n)` to // `var summaries []...`, which would silently serialize as @@ -10970,22 +10739,20 @@ func TestChatDebugRuns(t *testing.T) { modelConfig := createChatModelConfig(t, client) // Chat owned by the first (admin) user. - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-runs-other-owner", }) - require.NoError(t, err) seedChatDebugRun(ctx, t, db, chat.ID, time.Now().UTC()) otherClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, firstUser.OrganizationID, rbac.ScopedRoleAgentsAccess(firstUser.OrganizationID)) otherClient := codersdk.NewExperimentalClient(otherClientRaw) - _, err = otherClient.GetChatDebugRuns(ctx, chat.ID) + _, err := otherClient.GetChatDebugRuns(ctx, chat.ID) + requireSDKError(t, err, http.StatusNotFound) }) } @@ -11001,15 +10768,12 @@ func TestChatDebugRun(t *testing.T) { firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-run-detail", }) - require.NoError(t, err) run := seedChatDebugRun(ctx, t, db, chat.ID, time.Now().UTC()) firstStep := seedChatDebugStep(ctx, t, db, run, 1) @@ -11037,15 +10801,12 @@ func TestChatDebugRun(t *testing.T) { firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-run-empty", }) - require.NoError(t, err) run := seedChatDebugRun(ctx, t, db, chat.ID, time.Now().UTC()) got, err := client.GetChatDebugRun(ctx, chat.ID, run.ID) @@ -11063,15 +10824,12 @@ func TestChatDebugRun(t *testing.T) { firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-run-bad-uuid", }) - require.NoError(t, err) // Issue a raw request with a non-UUID run ID to exercise the // handler's parser path. @@ -11090,17 +10848,15 @@ func TestChatDebugRun(t *testing.T) { firstUser := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-run-missing", }) - require.NoError(t, err) - _, err = client.GetChatDebugRun(ctx, chat.ID, uuid.New()) + _, err := client.GetChatDebugRun(ctx, chat.ID, uuid.New()) + requireSDKError(t, err, http.StatusNotFound) }) @@ -11114,28 +10870,23 @@ func TestChatDebugRun(t *testing.T) { // Two chats owned by the same user. A run on chat A must not // be addressable through chat B's URL. - chatA, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chatA := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-run-chat-a", }) - require.NoError(t, err) - chatB, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chatB := dbgen.Chat(t, db, database.Chat{ OrganizationID: firstUser.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: firstUser.UserID, LastModelConfigID: modelConfig.ID, Title: "debug-run-chat-b", }) - require.NoError(t, err) runOnA := seedChatDebugRun(ctx, t, db, chatA.ID, time.Now().UTC()) - _, err = client.GetChatDebugRun(ctx, chatB.ID, runOnA.ID) + _, err := client.GetChatDebugRun(ctx, chatB.ID, runOnA.ID) + requireSDKError(t, err, http.StatusNotFound) }) } @@ -11954,16 +11705,13 @@ func TestGetChatsByWorkspace(t *testing.T) { // Helper to insert a chat linked to a workspace. insertChat := func(ctx context.Context, title string, workspaceID uuid.UUID) database.Chat { - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: title, WorkspaceID: uuid.NullUUID{UUID: workspaceID, Valid: true}, }) - require.NoError(t, err) return chat } @@ -12099,15 +11847,13 @@ func TestSubmitToolResults(t *testing.T) { dtJSON, err := json.Marshal(dynamicTools) require.NoError(t, err) - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: organizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: ownerID, LastModelConfigID: modelConfigID, - Title: "tool-results-test", DynamicTools: pqtype.NullRawMessage{RawMessage: dtJSON, Valid: true}, + Title: "tool-results-test", + DynamicTools: pqtype.NullRawMessage{RawMessage: dtJSON, Valid: true}, }) - require.NoError(t, err) // Build assistant message with tool-call parts. parts := make([]codersdk.ChatMessagePart, 0, len(toolCallIDs)) @@ -12122,26 +11868,12 @@ func TestSubmitToolResults(t *testing.T) { content, err := chatprompt.MarshalParts(parts) require.NoError(t, err) - _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelConfigID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(content.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: content, }) - require.NoError(t, err) // Transition to requires_action. chat, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ @@ -12207,17 +11939,14 @@ func TestSubmitToolResults(t *testing.T) { modelConfig := createChatModelConfig(t, client) // Create a chat that is NOT in requires_action status. - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "wrong-status-test", }) - require.NoError(t, err) - err = client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ + err := client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ Results: []codersdk.ToolResult{ {ToolCallID: "call_xyz", Output: json.RawMessage(`"nope"`)}, }, @@ -12561,7 +12290,6 @@ func TestGetChatMessages_Pagination(t *testing.T) { // the chat and the inserted message IDs in the order they were // persisted (ascending). Callers use these IDs as cursor values. seedChat := func( - ctx context.Context, t *testing.T, db database.Store, ownerID uuid.UUID, @@ -12571,71 +12299,28 @@ func TestGetChatMessages_Pagination(t *testing.T) { ) (database.Chat, []int64) { t.Helper() - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: organizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: ownerID, LastModelConfigID: modelConfigID, Title: "pagination-test", }) - require.NoError(t, err) - createdBy := make([]uuid.UUID, count) - modelIDs := make([]uuid.UUID, count) - roles := make([]database.ChatMessageRole, count) - contents := make([]string, count) - contentVersions := make([]int16, count) - visibility := make([]database.ChatMessageVisibility, count) - inputTokens := make([]int64, count) - outputTokens := make([]int64, count) - totalTokens := make([]int64, count) - reasoningTokens := make([]int64, count) - cacheCreationTokens := make([]int64, count) - cacheReadTokens := make([]int64, count) - contextLimit := make([]int64, count) - compressed := make([]bool, count) - totalCost := make([]int64, count) - runtime := make([]int64, count) + ids := make([]int64, count) for i := range count { - part, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText(fmt.Sprintf("msg %d", i)), }) require.NoError(t, err) - createdBy[i] = ownerID - modelIDs[i] = modelConfigID - roles[i] = database.ChatMessageRoleUser - contents[i] = string(part.RawMessage) - contentVersions[i] = chatprompt.CurrentContentVersion - visibility[i] = database.ChatMessageVisibilityBoth - } - - results, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: createdBy, - ModelConfigID: modelIDs, - Role: roles, - Content: contents, - ContentVersion: contentVersions, - Visibility: visibility, - InputTokens: inputTokens, - OutputTokens: outputTokens, - TotalTokens: totalTokens, - ReasoningTokens: reasoningTokens, - CacheCreationTokens: cacheCreationTokens, - CacheReadTokens: cacheReadTokens, - ContextLimit: contextLimit, - Compressed: compressed, - TotalCostMicros: totalCost, - RuntimeMs: runtime, - }) - require.NoError(t, err) - require.Len(t, results, count) - - ids := make([]int64, count) - for i, m := range results { - ids[i] = m.ID + message := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: ownerID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: content, + }) + ids[i] = message.ID } return chat, ids } @@ -12670,7 +12355,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) seedQueuedMessage(ctx, t, db, chat.ID) resp, err := client.GetChatMessages(ctx, chat.ID, nil) @@ -12695,7 +12380,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) seedQueuedMessage(ctx, t, db, chat.ID) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ @@ -12721,7 +12406,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) seedQueuedMessage(ctx, t, db, chat.ID) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ @@ -12749,7 +12434,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) seedQueuedMessage(ctx, t, db, chat.ID) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ @@ -12776,7 +12461,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) // Seed a queued message so the Empty assertion below verifies // the cursor suppresses queued rows, not just that none exist. seedQueuedMessage(ctx, t, db, chat.ID) @@ -12808,7 +12493,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, _ := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 1) + chat, _ := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 1) res, err := client.Request( ctx, @@ -12839,7 +12524,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, _ := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 1) + chat, _ := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 1) res, err := client.Request( ctx, @@ -12870,7 +12555,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 3) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 3) // Seed a queued message to prove the cursor path suppresses // it even when nothing else comes back. seedQueuedMessage(ctx, t, db, chat.ID) @@ -12890,12 +12575,11 @@ func TestGetChatMessages_Pagination(t *testing.T) { t.Run("AfterIDGreaterThanOrEqualBeforeIDReturns400", func(t *testing.T) { t.Parallel() - ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, 3) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 3) // Transposed cursors: after >= before. Fail loudly rather // than return an empty page indistinguishable from @@ -12940,7 +12624,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { const pageSize = 25 // Seed burstSize+1 rows; ids[0] is the "already acknowledged" // message the client saw before the burst. - chat, ids := seedChat(ctx, t, db, user.UserID, user.OrganizationID, modelConfig.ID, burstSize+1) + chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, burstSize+1) var seen []int64 cursor := ids[0] diff --git a/coderd/httpmw/chatparam_test.go b/coderd/httpmw/chatparam_test.go index 4e40ff4c27..c83355c4cb 100644 --- a/coderd/httpmw/chatparam_test.go +++ b/coderd/httpmw/chatparam_test.go @@ -2,7 +2,6 @@ package httpmw_test import ( "context" - "database/sql" "net/http" "net/http/httptest" "testing" @@ -38,42 +37,22 @@ func TestChatParam(t *testing.T) { insertChat := func(t *testing.T, db database.Store, ownerID, organizationID uuid.UUID) database.Chat { t.Helper() - _, err := db.InsertChatProvider(context.Background(), database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - APIKey: "test-api-key", - BaseUrl: "https://api.openai.com/v1", - ApiKeyKeyID: sql.NullString{}, - CreatedBy: uuid.NullUUID{UUID: ownerID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + APIKey: "test-api-key", + BaseUrl: "https://api.openai.com/v1", + CreatedBy: uuid.NullUUID{UUID: ownerID, Valid: true}, }) - require.NoError(t, err) - modelConfig, err := db.InsertChatModelConfig(context.Background(), database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o-mini", - DisplayName: "Test model", - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: []byte("{}"), + modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + IsDefault: true, }) - require.NoError(t, err) - chat, err := db.InsertChat(context.Background(), database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: organizationID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: ownerID, - WorkspaceID: uuid.NullUUID{}, - ParentChatID: uuid.NullUUID{}, - RootChatID: uuid.NullUUID{}, LastModelConfigID: modelConfig.ID, Title: "Test chat", }) - require.NoError(t, err) return chat } diff --git a/coderd/telemetry/telemetry_test.go b/coderd/telemetry/telemetry_test.go index f0fd563ef4..eb3ee365a6 100644 --- a/coderd/telemetry/telemetry_test.go +++ b/coderd/telemetry/telemetry_test.go @@ -16,6 +16,7 @@ import ( "github.com/go-chi/chi/v5" "github.com/google/go-cmp/cmp" "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/goleak" @@ -1559,60 +1560,39 @@ func TestChatsTelemetry(t *testing.T) { user := dbgen.User(t, db, database.User{}) // Create chat providers (required FK for model configs). - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "anthropic", - DisplayName: "Anthropic", - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "anthropic", + DisplayName: "Anthropic", }) - require.NoError(t, err) - _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", }) - require.NoError(t, err) // Create a model config. - modelCfg, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "anthropic", - Model: "claude-sonnet-4-20250514", - DisplayName: "Claude Sonnet", - Enabled: true, - IsDefault: true, - ContextLimit: 200000, - CompressionThreshold: 70, - Options: json.RawMessage("{}"), + modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + DisplayName: "Claude Sonnet", + IsDefault: true, + ContextLimit: 200000, }) - require.NoError(t, err) // Create a second model config to test full dump. - modelCfg2, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o", - DisplayName: "GPT-4o", - Enabled: true, - IsDefault: false, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage("{}"), + modelCfg2 := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + Model: "gpt-4o", + DisplayName: "GPT-4o", }) - require.NoError(t, err) // Create a soft-deleted model config — should NOT appear in telemetry. - deletedCfg, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "anthropic", - Model: "claude-deleted", - DisplayName: "Deleted Model", - Enabled: true, - IsDefault: false, - ContextLimit: 100000, - CompressionThreshold: 70, - Options: json.RawMessage("{}"), + deletedCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "anthropic", + Model: "claude-deleted", + DisplayName: "Deleted Model", + ContextLimit: 100000, }) - require.NoError(t, err) - err = db.DeleteChatModelConfigByID(ctx, deletedCfg.ID) + err := db.DeleteChatModelConfigByID(ctx, deletedCfg.ID) require.NoError(t, err) // Create a root chat with a workspace. @@ -1645,30 +1625,26 @@ func TestChatsTelemetry(t *testing.T) { JobID: job.ID, }) - rootChat, err := db.InsertChat(ctx, database.InsertChatParams{ + rootChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, Title: "Root Chat", Status: database.ChatStatusRunning, - ClientType: database.ChatClientTypeUi, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, Mode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true}, }) - require.NoError(t, err) // Create a child chat (has parent + root). - childChat, err := db.InsertChat(ctx, database.InsertChatParams{ + childChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: modelCfg2.ID, Title: "Child Chat", Status: database.ChatStatusCompleted, - ClientType: database.ChatClientTypeUi, ParentChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true}, }) - require.NoError(t, err) // Associate a PR with the root chat so PullRequestState is populated. rootChatNow := dbtime.Now() @@ -1681,76 +1657,118 @@ func TestChatsTelemetry(t *testing.T) { require.NoError(t, err) // Insert messages for root chat: 2 user, 2 assistant, 1 tool. - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ ChatID: rootChat.ID, - CreatedBy: []uuid.UUID{user.ID, uuid.Nil, user.ID, uuid.Nil, uuid.Nil}, - ModelConfigID: []uuid.UUID{modelCfg.ID, modelCfg.ID, modelCfg.ID, modelCfg.ID, modelCfg.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleAssistant, database.ChatMessageRoleUser, database.ChatMessageRoleAssistant, database.ChatMessageRoleTool}, - Content: []string{`[{"type":"text","text":"hello"}]`, `[{"type":"text","text":"hi"}]`, `[{"type":"text","text":"help"}]`, `[{"type":"text","text":"sure"}]`, `[{"type":"text","text":"result"}]`}, - ContentVersion: []int16{1, 1, 1, 1, 1}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth}, - InputTokens: []int64{100, 200, 150, 300, 0}, - OutputTokens: []int64{0, 50, 0, 100, 0}, - TotalTokens: []int64{100, 250, 150, 400, 0}, - ReasoningTokens: []int64{0, 10, 0, 20, 0}, - CacheCreationTokens: []int64{50, 0, 30, 0, 0}, - CacheReadTokens: []int64{0, 25, 0, 40, 0}, - ContextLimit: []int64{200000, 200000, 200000, 200000, 200000}, - Compressed: []bool{false, false, false, false, false}, - TotalCostMicros: []int64{1000, 2000, 1500, 3000, 0}, - RuntimeMs: []int64{0, 500, 0, 800, 100}, - ProviderResponseID: []string{"", "resp-1", "", "resp-2", ""}, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"hello"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 100, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 100, Valid: true}, + CacheCreationTokens: sql.NullInt64{Int64: 50, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 200000, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 1000, Valid: true}, + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: rootChat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"hi"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 200, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 250, Valid: true}, + ReasoningTokens: sql.NullInt64{Int64: 10, Valid: true}, + CacheReadTokens: sql.NullInt64{Int64: 25, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 200000, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 2000, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 500, Valid: true}, + ProviderResponseID: sql.NullString{String: "resp-1", Valid: true}, + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: rootChat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"help"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 150, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 150, Valid: true}, + CacheCreationTokens: sql.NullInt64{Int64: 30, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 200000, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 1500, Valid: true}, + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: rootChat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"sure"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 300, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 100, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 400, Valid: true}, + ReasoningTokens: sql.NullInt64{Int64: 20, Valid: true}, + CacheReadTokens: sql.NullInt64{Int64: 40, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 200000, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 3000, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 800, Valid: true}, + ProviderResponseID: sql.NullString{String: "resp-2", Valid: true}, + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: rootChat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleTool, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"result"}]`), Valid: true}, + ContextLimit: sql.NullInt64{Int64: 200000, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 100, Valid: true}, }) - require.NoError(t, err) // Insert messages for child chat: 1 user, 1 assistant (compressed). - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ ChatID: childChat.ID, - CreatedBy: []uuid.UUID{user.ID, uuid.Nil}, - ModelConfigID: []uuid.UUID{modelCfg2.ID, modelCfg2.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleAssistant}, - Content: []string{`[{"type":"text","text":"q"}]`, `[{"type":"text","text":"a"}]`}, - ContentVersion: []int16{1, 1}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth}, - InputTokens: []int64{500, 600}, - OutputTokens: []int64{0, 200}, - TotalTokens: []int64{500, 800}, - ReasoningTokens: []int64{0, 50}, - CacheCreationTokens: []int64{100, 0}, - CacheReadTokens: []int64{0, 75}, - ContextLimit: []int64{128000, 128000}, - Compressed: []bool{false, true}, - TotalCostMicros: []int64{5000, 8000}, - RuntimeMs: []int64{0, 1200}, - ProviderResponseID: []string{"", "resp-3"}, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg2.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"q"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 500, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 500, Valid: true}, + CacheCreationTokens: sql.NullInt64{Int64: 100, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 128000, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 5000, Valid: true}, + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: childChat.ID, + ModelConfigID: uuid.NullUUID{UUID: modelCfg2.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"a"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 600, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 200, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 800, Valid: true}, + ReasoningTokens: sql.NullInt64{Int64: 50, Valid: true}, + CacheReadTokens: sql.NullInt64{Int64: 75, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 128000, Valid: true}, + Compressed: true, + TotalCostMicros: sql.NullInt64{Int64: 8000, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 1200, Valid: true}, + ProviderResponseID: sql.NullString{String: "resp-3", Valid: true}, }) - require.NoError(t, err) // Insert a soft-deleted message on root chat with large token values. // This acts as "poison" — if the deleted filter is missing, totals // will be inflated and assertions below will fail. - poisonMsgs, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + poisonMsg := dbgen.ChatMessage(t, db, database.ChatMessage{ ChatID: rootChat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelCfg.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - Content: []string{`[{"type":"text","text":"poison"}]`}, - ContentVersion: []int16{1}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{999999}, - OutputTokens: []int64{999999}, - TotalTokens: []int64{999999}, - ReasoningTokens: []int64{999999}, - CacheCreationTokens: []int64{999999}, - CacheReadTokens: []int64{999999}, - ContextLimit: []int64{200000}, - Compressed: []bool{false}, - TotalCostMicros: []int64{999999}, - RuntimeMs: []int64{999999}, - ProviderResponseID: []string{""}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"poison"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 999999, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 999999, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 999999, Valid: true}, + ReasoningTokens: sql.NullInt64{Int64: 999999, Valid: true}, + CacheCreationTokens: sql.NullInt64{Int64: 999999, Valid: true}, + CacheReadTokens: sql.NullInt64{Int64: 999999, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 200000, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: 999999, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 999999, Valid: true}, }) - require.NoError(t, err) - err = db.SoftDeleteChatMessageByID(ctx, poisonMsgs[0].ID) + err = db.SoftDeleteChatMessageByID(ctx, poisonMsg.ID) require.NoError(t, err) _, snapshot := collectSnapshot(ctx, t, db, nil) @@ -1890,40 +1908,31 @@ func TestChatDiffStatusSummaryTelemetry(t *testing.T) { org, err := db.GetDefaultOrganization(ctx) require.NoError(t, err) - _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "anthropic", - DisplayName: "Anthropic", - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "anthropic", + DisplayName: "Anthropic", }) - require.NoError(t, err) - modelCfg, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "anthropic", - Model: "claude-sonnet-4-20250514", - DisplayName: "Claude Sonnet", - Enabled: true, - IsDefault: true, - ContextLimit: 200000, - CompressionThreshold: 70, - Options: json.RawMessage("{}"), + modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + DisplayName: "Claude Sonnet", + IsDefault: true, + ContextLimit: 200000, }) - require.NoError(t, err) // Helper to create a chat and upsert its diff status. insertChatWithDiffStatus := func(prURL, state string) uuid.UUID { t.Helper() - chat, chatErr := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, Title: "Chat " + state, Status: database.ChatStatusCompleted, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, chatErr) now := dbtime.Now() - _, chatErr = db.UpsertChatDiffStatus(ctx, database.UpsertChatDiffStatusParams{ + _, chatErr := db.UpsertChatDiffStatus(ctx, database.UpsertChatDiffStatusParams{ ChatID: chat.ID, Url: sql.NullString{String: prURL, Valid: prURL != ""}, PullRequestState: sql.NullString{String: state, Valid: true}, @@ -1945,15 +1954,13 @@ func TestChatDiffStatusSummaryTelemetry(t *testing.T) { // Insert a chat with NULL pull_request_state (no PR yet). // This should be excluded from all counts. - noPRChat, err := db.InsertChat(ctx, database.InsertChatParams{ + noPRChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, Title: "Chat no PR", Status: database.ChatStatusRunning, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) now := dbtime.Now() _, err = db.UpsertChatDiffStatus(ctx, database.UpsertChatDiffStatusParams{ ChatID: noPRChat.ID, diff --git a/coderd/workspaceagents_active_chat_internal_test.go b/coderd/workspaceagents_active_chat_internal_test.go index 4834bbedf2..c2d8291f8e 100644 --- a/coderd/workspaceagents_active_chat_internal_test.go +++ b/coderd/workspaceagents_active_chat_internal_test.go @@ -29,21 +29,19 @@ func TestActiveAgentChatDefinitionsAgree(t *testing.T) { OrganizationID: org.ID, OwnerID: owner.ID, }).WithAgent().Do() - modelConfig := insertAgentChatTestModelConfig(ctx, t, db, owner.ID) + modelConfig := insertAgentChatTestModelConfig(t, db, owner.ID) insertedChats := make([]database.Chat, 0, len(database.AllChatStatusValues())*2) for _, archived := range []bool{false, true} { for _, status := range database.AllChatStatusValues() { - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, Status: status, - ClientType: database.ChatClientTypeUi, OwnerID: owner.ID, LastModelConfigID: modelConfig.ID, Title: fmt.Sprintf("%s-archived-%t", status, archived), AgentID: uuid.NullUUID{UUID: workspace.Agents[0].ID, Valid: true}, }) - require.NoError(t, err) if archived { _, err = db.ArchiveChatByID(ctx, chat.ID) diff --git a/coderd/workspaceagents_chat_context_internal_test.go b/coderd/workspaceagents_chat_context_internal_test.go index 377c79466c..5a2c8e25be 100644 --- a/coderd/workspaceagents_chat_context_internal_test.go +++ b/coderd/workspaceagents_chat_context_internal_test.go @@ -2,7 +2,6 @@ package coderd import ( "context" - "database/sql" "encoding/json" "testing" "time" @@ -14,7 +13,7 @@ import ( "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/dbgen" "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/codersdk" ) @@ -89,40 +88,25 @@ func TestUpdateAgentChatLastInjectedContextFromMessagesUsesMessageIDTieBreaker(t } func insertAgentChatTestModelConfig( - ctx context.Context, t testing.TB, db database.Store, userID uuid.UUID, ) database.ChatModelConfig { t.Helper() - sysCtx := dbauthz.AsSystemRestricted(ctx) createdBy := uuid.NullUUID{UUID: userID, Valid: true} - _, err := db.InsertChatProvider(sysCtx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - APIKey: "test-api-key", - ApiKeyKeyID: sql.NullString{}, - CreatedBy: createdBy, - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", + APIKey: "test-api-key", + CreatedBy: createdBy, }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(sysCtx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: createdBy, - UpdatedBy: createdBy, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + return dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + CreatedBy: createdBy, + UpdatedBy: createdBy, + IsDefault: true, }) - require.NoError(t, err) - - return model } diff --git a/coderd/workspaceagents_chat_context_test.go b/coderd/workspaceagents_chat_context_test.go index c6893e2ffd..2067fe3ff4 100644 --- a/coderd/workspaceagents_chat_context_test.go +++ b/coderd/workspaceagents_chat_context_test.go @@ -2,6 +2,7 @@ package coderd_test import ( "context" + "database/sql" "encoding/json" "net/http" "strings" @@ -16,6 +17,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbfake" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/x/chatd" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" @@ -156,8 +158,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) for _, step := range tc.steps { resp, err := setup.agentClient.AddChatContext(ctx, step.req) @@ -202,24 +204,17 @@ func TestAgentChatContext(t *testing.T) { }).WithAgent().Do() agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(workspace.AgentToken)) - originalModel := coderd.InsertAgentChatTestModelConfig(ctx, t, baseDB, user.UserID) - updatedModel, err := baseDB.InsertChatModelConfig( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatModelConfigParams{ - Provider: originalModel.Provider, - Model: "gpt-4o-mini-updated", - DisplayName: "Updated Test Model", - CreatedBy: uuid.NullUUID{UUID: user.UserID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.UserID, Valid: true}, - Enabled: true, - IsDefault: false, - ContextLimit: originalModel.ContextLimit, - CompressionThreshold: originalModel.CompressionThreshold, - Options: json.RawMessage(`{}`), - }, - ) - require.NoError(t, err) - chat := createAgentChatContextChat(ctx, t, baseDB, user.OrganizationID, user.UserID, originalModel.ID, workspace.Agents[0].ID, t.Name()) + originalModel := coderd.InsertAgentChatTestModelConfig(t, baseDB, user.UserID) + updatedModel := dbgen.ChatModelConfig(t, baseDB, database.ChatModelConfig{ + Provider: originalModel.Provider, + Model: "gpt-4o-mini-updated", + DisplayName: "Updated Test Model", + CreatedBy: uuid.NullUUID{UUID: user.UserID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.UserID, Valid: true}, + ContextLimit: originalModel.ContextLimit, + CompressionThreshold: originalModel.CompressionThreshold, + }) + chat := createAgentChatContextChat(t, baseDB, user.OrganizationID, user.UserID, originalModel.ID, workspace.Agents[0].ID, t.Name()) interceptDB.beforeInTx = func() { _, err := baseDB.UpdateChatLastModelConfigByID( @@ -259,8 +254,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) skillPart := codersdk.ChatMessagePart{ Type: codersdk.ChatMessagePartTypeSkill, @@ -321,8 +316,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) skillPart := codersdk.ChatMessagePart{ Type: codersdk.ChatMessagePartTypeSkill, @@ -363,40 +358,30 @@ func TestAgentChatContext(t *testing.T) { codersdk.ChatMessageText("compressed summary"), }) require.NoError(t, err) - summaryParams := chatd.BuildSingleChatMessageInsertParams( - chat.ID, - database.ChatMessageRoleUser, - summaryContent, - database.ChatMessageVisibilityModel, - chat.LastModelConfigID, - chatprompt.CurrentContentVersion, - setup.user.UserID, - ) - summaryParams.Compressed[0] = true - _, err = setup.db.InsertChatMessages( - dbauthz.AsSystemRestricted(ctx), - summaryParams, - ) - require.NoError(t, err) + _ = dbgen.ChatMessage(t, setup.db, database.ChatMessage{ + ChatID: chat.ID, + Role: database.ChatMessageRoleUser, + Content: summaryContent, + Visibility: database.ChatMessageVisibilityModel, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: setup.user.UserID, Valid: true}, + Compressed: true, + }) regularContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("keep this user message"), }) require.NoError(t, err) - _, err = setup.db.InsertChatMessages( - dbauthz.AsSystemRestricted(ctx), - chatd.BuildSingleChatMessageInsertParams( - chat.ID, - database.ChatMessageRoleUser, - regularContent, - database.ChatMessageVisibilityBoth, - chat.LastModelConfigID, - chatprompt.CurrentContentVersion, - setup.user.UserID, - ), - ) - require.NoError(t, err) - + _ = dbgen.ChatMessage(t, setup.db, database.ChatMessage{ + ChatID: chat.ID, + Role: database.ChatMessageRoleUser, + Content: regularContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: setup.user.UserID, Valid: true}, + }) resp, err := setup.agentClient.ClearChatContext(ctx, agentsdk.ClearChatContextRequest{}) require.NoError(t, err) require.Equal(t, chat.ID, resp.ChatID) @@ -420,8 +405,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -436,20 +421,15 @@ func TestAgentChatContext(t *testing.T) { codersdk.ChatMessageText("keep this user message"), }) require.NoError(t, err) - _, err = setup.db.InsertChatMessages( - dbauthz.AsSystemRestricted(ctx), - chatd.BuildSingleChatMessageInsertParams( - chat.ID, - database.ChatMessageRoleUser, - regularContent, - database.ChatMessageVisibilityBoth, - chat.LastModelConfigID, - chatprompt.CurrentContentVersion, - setup.user.UserID, - ), - ) - require.NoError(t, err) - + _ = dbgen.ChatMessage(t, setup.db, database.ChatMessage{ + ChatID: chat.ID, + Role: database.ChatMessageRoleUser, + Content: regularContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: setup.user.UserID, Valid: true}, + }) resp, err := setup.agentClient.ClearChatContext(ctx, agentsdk.ClearChatContextRequest{}) require.NoError(t, err) require.Equal(t, chat.ID, resp.ChatID) @@ -477,8 +457,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -493,21 +473,15 @@ func TestAgentChatContext(t *testing.T) { codersdk.ChatMessageText("assistant reply"), }) require.NoError(t, err) - assistantParams := chatd.BuildSingleChatMessageInsertParams( - chat.ID, - database.ChatMessageRoleAssistant, - assistantContent, - database.ChatMessageVisibilityBoth, - chat.LastModelConfigID, - chatprompt.CurrentContentVersion, - uuid.Nil, - ) - assistantParams.ProviderResponseID[0] = "resp-123" - _, err = setup.db.InsertChatMessages( - dbauthz.AsSystemRestricted(ctx), - assistantParams, - ) - require.NoError(t, err) + _ = dbgen.ChatMessage(t, setup.db, database.ChatMessage{ + ChatID: chat.ID, + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + ProviderResponseID: sql.NullString{String: "resp-123", Valid: true}, + }) messages := requireAgentChatContextMessages(ctx, t, setup.db, chat.ID) require.Len(t, messages, 2) @@ -539,29 +513,22 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("assistant reply"), }) require.NoError(t, err) - assistantParams := chatd.BuildSingleChatMessageInsertParams( - chat.ID, - database.ChatMessageRoleAssistant, - assistantContent, - database.ChatMessageVisibilityBoth, - chat.LastModelConfigID, - chatprompt.CurrentContentVersion, - uuid.Nil, - ) - assistantParams.ProviderResponseID[0] = "resp-123" - _, err = setup.db.InsertChatMessages( - dbauthz.AsSystemRestricted(ctx), - assistantParams, - ) - require.NoError(t, err) - + _ = dbgen.ChatMessage(t, setup.db, database.ChatMessage{ + ChatID: chat.ID, + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + ProviderResponseID: sql.NullString{String: "resp-123", Valid: true}, + }) resp, err := setup.agentClient.ClearChatContext(ctx, agentsdk.ClearChatContextRequest{ChatID: chat.ID}) require.NoError(t, err) require.Equal(t, chat.ID, resp.ChatID) @@ -595,7 +562,7 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := coderdtest.NewWithDatabase(t, nil) user := coderdtest.CreateFirstUser(t, client) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, db, user.UserID) + model := coderd.InsertAgentChatTestModelConfig(t, db, user.UserID) firstWorkspace := dbfake.WorkspaceBuild(t, db, database.WorkspaceTable{ OrganizationID: user.OrganizationID, @@ -606,7 +573,7 @@ func TestAgentChatContext(t *testing.T) { OwnerID: user.UserID, }).WithAgent().Do() - chat := createAgentChatContextChat(ctx, t, db, user.OrganizationID, user.UserID, model.ID, firstWorkspace.Agents[0].ID, t.Name()) + chat := createAgentChatContextChat(t, db, user.OrganizationID, user.UserID, model.ID, firstWorkspace.Agents[0].ID, t.Name()) secondAgentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(secondWorkspace.AgentToken)) _, err := secondAgentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ @@ -627,8 +594,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: chat.ID, @@ -682,8 +649,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: chat.ID, @@ -704,8 +671,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) largeContent := strings.Repeat("a", maxContextFileBytes+100) resp, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ @@ -739,8 +706,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) visible := strings.Repeat("a", maxContextFileBytes-1) content := visible + strings.Repeat("\u200b", 100) + "z" @@ -781,9 +748,9 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - ownerChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") - foreignChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + ownerChat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") + foreignChat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") resp, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -805,9 +772,9 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - rootChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") - childChat := createAgentChatContextChildChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + rootChat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") + childChat := createAgentChatContextChildChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") resp, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -829,9 +796,9 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat1") - createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat2") + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat1") + createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-chat2") _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ Parts: []codersdk.ChatMessagePart{{ @@ -849,9 +816,9 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - rootChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") - childChat := createAgentChatContextChildChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + rootChat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-root") + childChat := createAgentChatContextChildChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, rootChat.ID, t.Name()+"-child") _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: rootChat.ID, @@ -877,9 +844,9 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - ownerChat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") - _ = createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + ownerChat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-owner") + _ = createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()+"-foreign") _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ ChatID: ownerChat.ID, @@ -903,8 +870,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) _, otherUser := coderdtest.CreateAnotherUser(t, setup.client, setup.user.OrganizationID) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, otherUser.ID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.agentClient.ClearChatContext(ctx, agentsdk.ClearChatContextRequest{ChatID: chat.ID}) sdkErr := requireSDKError(t, err, http.StatusForbidden) @@ -916,8 +883,8 @@ func TestAgentChatContext(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) setup := newAgentChatContextTestSetup(t) - model := coderd.InsertAgentChatTestModelConfig(ctx, t, setup.db, setup.user.UserID) - chat := createAgentChatContextChat(ctx, t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) _, err := setup.db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, @@ -1020,7 +987,6 @@ func newAgentChatContextTestSetup(t *testing.T) agentChatContextTestSetup { } func createAgentChatContextChat( - ctx context.Context, t testing.TB, db database.Store, orgID uuid.UUID, @@ -1031,22 +997,16 @@ func createAgentChatContextChat( ) database.Chat { t.Helper() - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, + return dbgen.Chat(t, db, database.Chat{ OrganizationID: orgID, OwnerID: ownerID, LastModelConfigID: modelConfigID, Title: title, AgentID: uuid.NullUUID{UUID: agentID, Valid: true}, }) - require.NoError(t, err) - - return chat } func createAgentChatContextChildChat( - ctx context.Context, t testing.TB, db database.Store, orgID uuid.UUID, @@ -1058,9 +1018,7 @@ func createAgentChatContextChildChat( ) database.Chat { t.Helper() - chat, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, + return dbgen.Chat(t, db, database.Chat{ OrganizationID: orgID, OwnerID: ownerID, LastModelConfigID: modelConfigID, @@ -1069,9 +1027,6 @@ func createAgentChatContextChildChat( ParentChatID: uuid.NullUUID{UUID: parentChatID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parentChatID, Valid: true}, }) - require.NoError(t, err) - - return chat } func requireAgentChatContextParts(t testing.TB, raw json.RawMessage) []codersdk.ChatMessagePart { diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 1f77d8f039..3a9d84b6b0 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -4847,13 +4847,17 @@ func TestAutoPromote_InsertFailureSkipsStatusUpdate(t *testing.T) { heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), } - // Block model resolution until the running status has been - // published. Returning ErrInterrupted makes processChat enter the - // waiting-state auto-promotion path deterministically. + // Hold model resolution until the interrupt has canceled the chat + // context. Returning ErrInterrupted keeps processChat on the + // interrupted path regardless of whether the cache singleflight sees + // the caller cancellation or the DB fetch result first. modelBlocked := make(chan struct{}) + modelRelease := make(chan struct{}) + var modelBlockedOnce sync.Once db.EXPECT().GetChatModelConfigByID(gomock.Any(), gomock.Any()).DoAndReturn( - func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { - <-modelBlocked + func(_ context.Context, _ uuid.UUID) (database.ChatModelConfig, error) { + modelBlockedOnce.Do(func() { close(modelBlocked) }) + <-modelRelease return database.ChatModelConfig{}, chatloop.ErrInterrupted }, ).AnyTimes() @@ -4916,7 +4920,20 @@ func TestAutoPromote_InsertFailureSkipsStatusUpdate(t *testing.T) { t.Fatal("timed out waiting for running status") } - close(modelBlocked) + select { + case <-modelBlocked: + case <-ctx.Done(): + t.Fatal("timed out waiting for model resolution") + } + + // Publish an interrupt so processChat exits runChat. + interruptMsg, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ + Status: string(database.ChatStatusWaiting), + }) + require.NoError(t, err) + err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chatID), interruptMsg) + require.NoError(t, err) + close(modelRelease) select { case <-processDone: diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index d20f980bdf..f7ff9c3780 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -23,6 +23,7 @@ import ( mcpgo "github.com/mark3labs/mcp-go/mcp" mcpserver "github.com/mark3labs/mcp-go/server" "github.com/prometheus/client_golang/prometheus" + "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/xerrors" @@ -153,7 +154,7 @@ func TestInterruptChatBroadcastsStatusAcrossInstances(t *testing.T) { replicaB := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replicaA.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -759,13 +760,12 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { ) }) - user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) webSearchEnabled := true storeEnabled := true // OpenAI only serializes web_search through the Responses API. // Store=true routes there only for supported Responses models. webSearchModel := insertChatModelConfigWithCallConfig( - ctx, t, db, user.ID, @@ -780,49 +780,33 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { }, }, ) - mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "External Snapshot MCP", - Slug: "external-snapshot-mcp", - Url: externalMCPServer.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "External Snapshot MCP", + Slug: "external-snapshot-mcp", + Url: externalMCPServer.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - _, err = db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Second MCP", - Slug: "second-mcp", - Url: secondMCPServer.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Second MCP", + Slug: "second-mcp", + Url: secondMCPServer.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) - rootChat, err := db.InsertChat(ctx, database.InsertChatParams{ + rootChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, LastModelConfigID: webSearchModel.ID, Title: "root", - Status: database.ChatStatusWaiting, ClientType: database.ChatClientTypeApi, }) - require.NoError(t, err) - exploreChat, err := db.InsertChat(ctx, database.InsertChatParams{ + exploreChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -839,8 +823,17 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { MCPServerIDs: []uuid.UUID{mcpConfig.ID}, ClientType: database.ChatClientTypeApi, }) - require.NoError(t, err) - insertUserTextMessage(ctx, t, db, exploreChat.ID, user.ID, webSearchModel.ID, "inspect the codebase") + + dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: exploreChat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: webSearchModel.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`[{"type":"text","text":"inspect the codebase"}]`), + Valid: true, + }, + }) ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) @@ -936,21 +929,14 @@ func TestRootExploreChatStaysBuiltinOnlyAtRuntime(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) - mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Root Explore Runtime MCP", - Slug: "root-explore-runtime-mcp", - Url: externalMCPServer.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Root Explore Runtime MCP", + Slug: "root-explore-runtime-mcp", + Url: externalMCPServer.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) server := newActiveTestServer(t, db, ps) @@ -1016,13 +1002,12 @@ func TestRootExploreChatExcludesWebSearchProviderToolAtRuntime(t *testing.T) { ) }) - user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) webSearchEnabled := true storeEnabled := true // OpenAI only serializes web_search through the Responses API. // Store=true routes there only for supported Responses models. webSearchModel := insertChatModelConfigWithCallConfig( - ctx, t, db, user.ID, @@ -1145,35 +1130,21 @@ func TestExploreChatSendMessageCannotMutateMCPSnapshot(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) - parentConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Runtime Parent MCP", - Slug: "runtime-parent-mcp", - Url: parentTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + parentConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Runtime Parent MCP", + Slug: "runtime-parent-mcp", + Url: parentTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - injectedConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Runtime Injected MCP", - Slug: "runtime-injected-mcp", - Url: injectedTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + injectedConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Runtime Injected MCP", + Slug: "runtime-injected-mcp", + Url: injectedTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) server := newActiveTestServer(t, db, ps) @@ -1322,55 +1293,34 @@ func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - approvedConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ + approvedConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ DisplayName: "Plan Approved MCP", Slug: "plan-approved-mcp", Url: echoTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, AllowInPlanMode: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - blockedConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Plan Blocked MCP", - Slug: "plan-blocked-mcp", - Url: echoTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - AllowInPlanMode: false, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + blockedConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Plan Blocked MCP", + Slug: "plan-blocked-mcp", + Url: echoTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - filteredConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ + filteredConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ DisplayName: "Plan Filtered MCP", Slug: "plan-filtered-mcp", Url: filteredTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, AllowInPlanMode: true, ToolAllowList: []string{"visible"}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) ctrl := gomock.NewController(t) @@ -1477,7 +1427,7 @@ func TestInterruptChatClearsWorkerInDatabase(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -1514,7 +1464,7 @@ func TestArchiveChatMovesPendingChatToWaiting(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, @@ -1560,7 +1510,7 @@ func TestUnarchiveChildChat(t *testing.T) { db, ps := dbtestutil.NewDB(t) replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) parent, child := insertParentWithArchivedChild(ctx, t, db, user, org, model) @@ -1581,7 +1531,7 @@ func TestUnarchiveChildChat(t *testing.T) { db, ps := dbtestutil.NewDB(t) replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) parent, child := insertParentWithArchivedChild(ctx, t, db, user, org, model) _, err := db.ArchiveChatByID(ctx, parent.ID) @@ -1601,9 +1551,9 @@ func TestUnarchiveChildChat(t *testing.T) { db, ps := dbtestutil.NewDB(t) replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) - _, child := insertParentWithActiveChild(ctx, t, db, user, org, model) + _, child := insertParentWithActiveChild(t, db, user, org, model) require.NoError(t, replica.UnarchiveChat(ctx, child)) @@ -1617,7 +1567,6 @@ func TestUnarchiveChildChat(t *testing.T) { // child chat linked to it. Both are returned in their initial // (active) state. func insertParentWithActiveChild( - ctx context.Context, t *testing.T, db database.Store, user database.User, @@ -1625,27 +1574,20 @@ func insertParentWithActiveChild( model database.ChatModelConfig, ) (parent database.Chat, child database.Chat) { t.Helper() - var err error - parent, err = db.InsertChat(ctx, database.InsertChatParams{ + parent = dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, LastModelConfigID: model.ID, Title: "parent", }) - require.NoError(t, err) - child, err = db.InsertChat(ctx, database.InsertChatParams{ + child = dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, LastModelConfigID: model.ID, Title: "child", ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, }) - require.NoError(t, err) return parent, child } @@ -1661,7 +1603,7 @@ func insertParentWithArchivedChild( model database.ChatModelConfig, ) (parent database.Chat, child database.Chat) { t.Helper() - parent, child = insertParentWithActiveChild(ctx, t, db, user, org, model) + parent, child = insertParentWithActiveChild(t, db, user, org, model) _, err := db.ArchiveChatByID(ctx, child.ID) require.NoError(t, err) child, err = db.GetChatByID(ctx, child.ID) @@ -1701,7 +1643,7 @@ func TestArchiveChatInterruptsActiveProcessing(t *testing.T) { }) server := newActiveTestServer(t, db, ps) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -1811,7 +1753,7 @@ func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -1859,7 +1801,7 @@ func TestSendMessageQueueBehaviorQueuesWhenBusy(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -1928,7 +1870,7 @@ func TestPlanTurnPromptContract(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) planModeInstructions := "Ask about deployment sequencing before finalizing the plan." err := db.UpsertChatPlanModeInstructions(dbauthz.AsSystemRestricted(ctx), planModeInstructions) require.NoError(t, err) @@ -1984,7 +1926,7 @@ func TestSendMessageQueuesWhenWaitingWithQueuedBacklog(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -2051,20 +1993,18 @@ func TestSendMessageRejectsInvalidQueuedModelConfigID(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfig := seedChatDependencies(ctx, t, db) + user, org, modelConfig := seedChatDependencies(t, db) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, Status: database.ChatStatusPending, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: modelConfig.ID, Title: "reject invalid queued model config", }) - require.NoError(t, err) invalidModelConfigID := uuid.New() - _, err = replica.SendMessage(ctx, chatd.SendMessageOptions{ + _, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, ModelConfigID: invalidModelConfigID, @@ -2083,7 +2023,7 @@ func TestSendMessageInterruptBehaviorQueuesAndInterruptsWhenBusy(t *testing.T) { replica := newStartedTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -2150,7 +2090,7 @@ func TestEditMessageUpdatesAndTruncatesAndClearsQueue(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -2258,7 +2198,7 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ OrganizationID: org.ID, @@ -2310,7 +2250,7 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -2347,7 +2287,7 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ Enabled: true, @@ -2356,41 +2296,26 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { }) require.NoError(t, err) - existingChat, err := db.InsertChat(ctx, database.InsertChatParams{ + existingChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, Title: "existing-limit-chat", LastModelConfigID: model.ID, }) - require.NoError(t, err) assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("assistant"), }) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: existingChat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{100}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: existingChat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, + TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, }) - require.NoError(t, err) beforeChats, err := db.GetChats(ctx, database.GetChatsParams{ OwnerID: user.ID, @@ -2432,7 +2357,7 @@ func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing replica := newStartedTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ Enabled: true, @@ -2478,26 +2403,14 @@ func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing }) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{100}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, + TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, }) - require.NoError(t, err) chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, @@ -2537,9 +2450,8 @@ func TestPromoteQueuedMessageUsesQueuedModelConfigID(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfigA := seedChatDependencies(ctx, t, db) + user, org, modelConfigA := seedChatDependencies(t, db) modelConfigB := insertChatModelConfigWithCallConfig( - ctx, t, db, user.ID, @@ -2548,15 +2460,12 @@ func TestPromoteQueuedMessageUsesQueuedModelConfigID(t *testing.T) { codersdk.ChatModelCallConfig{}, ) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: modelConfigA.ID, Title: "promote queued uses stored model", }) - require.NoError(t, err) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("queued with model b")}) require.NoError(t, err) @@ -2598,9 +2507,8 @@ func TestPromoteQueuedMessageReloadsChatWhenModelConfigChangesDuringPending(t *t replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfigA := seedChatDependencies(ctx, t, db) + user, org, modelConfigA := seedChatDependencies(t, db) modelConfigB := insertChatModelConfigWithCallConfig( - ctx, t, db, user.ID, @@ -2628,15 +2536,13 @@ func TestPromoteQueuedMessageReloadsChatWhenModelConfigChangesDuringPending(t *t require.NoError(t, err) defer cancelWatch() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, Status: database.ChatStatusPending, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: modelConfigA.ID, Title: "promote queued reloads pending chat", }) - require.NoError(t, err) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("queued with new model")}) require.NoError(t, err) @@ -2728,9 +2634,8 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { // signalWake. cfg.PendingChatAcquireInterval = time.Hour }) - user, org, modelConfigA := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, modelConfigA := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) modelConfigB := insertChatModelConfigWithCallConfig( - ctx, t, db, user.ID, @@ -2739,7 +2644,6 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { codersdk.ChatModelCallConfig{}, ) modelConfigC := insertChatModelConfigWithCallConfig( - ctx, t, db, user.ID, @@ -2871,7 +2775,7 @@ func testAutoPromoteQueuedMessageFallback(t *testing.T, queuedModelConfigID uuid // trigger the next processing run. cfg.PendingChatAcquireInterval = time.Hour }) - user, org, modelConfig := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, modelConfig := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, @@ -2937,16 +2841,13 @@ func TestPromoteQueuedMessageFallsBackForLegacyQueuedRows(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfigA := seedChatDependencies(ctx, t, db) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + user, org, modelConfigA := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: modelConfigA.ID, Title: "promote queued legacy fallback", }) - require.NoError(t, err) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("legacy queued row")}) require.NoError(t, err) @@ -2977,17 +2878,14 @@ func TestPromoteQueuedMessageFallsBackForInvalidQueuedModelConfigID(t *testing.T replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfig := seedChatDependencies(ctx, t, db) + user, org, modelConfig := seedChatDependencies(t, db) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: modelConfig.ID, Title: "promote queued invalid fallback", }) - require.NoError(t, err) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("invalid queued model")}) require.NoError(t, err) @@ -3108,7 +3006,7 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { cfg.InFlightChatStaleAfter = testutil.WaitSuperLong }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -3144,45 +3042,26 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { require.True(t, laterQueuedResult.Queued) require.NotNil(t, laterQueuedResult.QueuedMessage) - spendChat, err := db.InsertChat(ctx, database.InsertChatParams{ + spendChat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, - WorkspaceID: uuid.NullUUID{}, - ParentChatID: uuid.NullUUID{}, - RootChatID: uuid.NullUUID{}, LastModelConfigID: model.ID, Title: "other-spend", - Mode: database.NullChatMode{}, }) - require.NoError(t, err) assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("spent elsewhere"), }) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: spendChat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{100}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: spendChat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, + TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, }) - require.NoError(t, err) close(allowSecondRequestFinish) testutil.TryReceive(ctx, t, thirdRequestStarted) @@ -3227,7 +3106,7 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ Enabled: true, @@ -3258,26 +3137,14 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { }) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{100}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, + TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, }) - require.NoError(t, err) _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, @@ -3309,7 +3176,7 @@ func TestEditMessageRejectsMissingMessage(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -3336,7 +3203,7 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -3352,27 +3219,13 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { }) require.NoError(t, err) - assistantMessages, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + assistantMessage := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: assistantContent, }) - require.NoError(t, err) - assistantMessage := assistantMessages[0] _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, @@ -3397,7 +3250,7 @@ func TestEditMessageDebugCleanupDeletesPreEditRuns(t *testing.T) { replica := newDebugEnabledTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -3503,7 +3356,7 @@ func TestEditMessageDebugCleanupPreservesRecentRuns(t *testing.T) { replica := newDebugEnabledTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -3582,7 +3435,7 @@ func TestArchiveChatDebugCleanupDeletesPreArchiveRuns(t *testing.T) { replica := newDebugEnabledTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -3662,7 +3515,7 @@ func TestRecoverStaleChatsPeriodically(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Use a very short stale threshold so the periodic recovery // kicks in quickly during the test. @@ -3671,17 +3524,14 @@ func TestRecoverStaleChatsPeriodically(t *testing.T) { // Create a chat and simulate a dead worker by setting the chat // to running with a heartbeat in the past. deadWorkerID := uuid.New() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, Title: "stale-recovery-periodic", LastModelConfigID: model.ID, }) - require.NoError(t, err) - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRunning, WorkerID: uuid.NullUUID{UUID: deadWorkerID, Valid: true}, @@ -3720,15 +3570,12 @@ func TestRecoverStaleChatsPeriodically(t *testing.T) { // Now simulate a second stale chat appearing AFTER startup. // This tests the periodic recovery, not just the startup one. deadWorkerID2 := uuid.New() - chat2, err := db.InsertChat(ctx, database.InsertChatParams{ + chat2 := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, Title: "stale-recovery-periodic-2", LastModelConfigID: model.ID, }) - require.NoError(t, err) _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat2.ID, @@ -3756,7 +3603,7 @@ func TestRecoverStaleRequiresActionChat(t *testing.T) { db, ps, rawDB := dbtestutil.NewDBWithSQLDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Use a very short stale threshold so the periodic recovery // kicks in quickly during the test. @@ -3765,17 +3612,14 @@ func TestRecoverStaleRequiresActionChat(t *testing.T) { // Create a chat and set it to requires_action to simulate a // client that disappeared while the chat was waiting for // dynamic tool results. - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, Title: "stale-requires-action", LastModelConfigID: model.ID, }) - require.NoError(t, err) - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRequiresAction, }) @@ -3823,23 +3667,20 @@ func TestNewReplicaRecoversStaleChatFromDeadReplica(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Simulate a chat left running by a dead replica with a stale // heartbeat (well beyond the stale threshold). deadReplicaID := uuid.New() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, Title: "orphaned-chat", LastModelConfigID: model.ID, }) - require.NoError(t, err) // Set the heartbeat far in the past so it's definitely stale. - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRunning, WorkerID: uuid.NullUUID{UUID: deadReplicaID, Valid: true}, @@ -3869,19 +3710,16 @@ func TestWaitingChatsAreNotRecoveredAsStale(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Create a chat in waiting status — this should NOT be touched // by stale recovery. - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, Title: "waiting-chat", LastModelConfigID: model.ID, }) - require.NoError(t, err) // Start a replica with a short stale threshold. logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) @@ -3917,21 +3755,18 @@ func TestUpdateChatStatusPersistsLastError(t *testing.T) { _ = newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, Title: "error-persisted", LastModelConfigID: model.ID, }) - require.NoError(t, err) // Simulate a chat that failed with an error. errorMessage := "stream response: status 500: internal server error" - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + chat, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusError, WorkerID: uuid.NullUUID{}, @@ -3975,7 +3810,7 @@ func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -4044,7 +3879,7 @@ func TestPersistToolResultWithBinaryData(t *testing.T) { // /chat/completions endpoint, where the mock server supports // streaming tool calls. The default "openai" provider routes to // /responses which only handles text deltas in the mock. - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) ctrl := gomock.NewController(t) @@ -4216,7 +4051,7 @@ func TestDynamicToolCallPausesAndResumes(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) // Dynamic tools do not need a workspace connection, but the // chatd server always builds workspace tools. Use an active @@ -4390,7 +4225,7 @@ func TestDynamicToolNamedProposePlanRemainsAvailableOutsidePlanMode(t *testing.T ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) server := newActiveTestServer(t, db, ps) dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ @@ -4503,7 +4338,7 @@ func TestDynamicToolCallMixedWithBuiltIn(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) server := newActiveTestServer(t, db, ps) // Create a chat with a dynamic tool. @@ -4642,7 +4477,7 @@ func TestSubmitToolResultsConcurrency(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) server := newActiveTestServer(t, db, ps) // Create a chat with a dynamic tool. @@ -4786,7 +4621,7 @@ func TestSubscribeNoPubsubNoDuplicateMessageParts(t *testing.T) { replica := newStartedTestServer(t, db, nil, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -4834,7 +4669,7 @@ func TestSubscribeAfterMessageID(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Create a chat — this inserts one initial "user" message. chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ @@ -4853,53 +4688,26 @@ func TestSubscribeAfterMessageID(t *testing.T) { }) require.NoError(t, err) - msg2Results, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(secondContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + msg2 := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + ContentVersion: chatprompt.CurrentContentVersion, + Content: secondContent, }) - require.NoError(t, err) - msg2 := msg2Results[0] thirdContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("third"), }) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(thirdContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, + Content: thirdContent, }) - require.NoError(t, err) // Control: Subscribe with afterMessageID=0 returns ALL messages. allSnapshot, _, cancelAll, ok := replica.Subscribe(ctx, chat.ID, nil, 0) @@ -5304,7 +5112,7 @@ func TestStoppedWorkspaceWithPersistedAgentBindingDoesNotBlockChat(t *testing.T) ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) inactive := newTestServer(t, db, ps, uuid.New()) @@ -5445,7 +5253,7 @@ func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { if !req.Stream { return chattest.OpenAINonStreamingResponse("ok") @@ -5620,7 +5428,7 @@ func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { if !req.Stream { return chattest.OpenAINonStreamingResponse("ok") @@ -5749,7 +5557,7 @@ func TestPassiveServerDoesNotProcess(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, ps := dbtestutil.NewDB(t, dbtestutil.WithDumpOnFailure()) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) server := newTestServer(t, db, ps, uuid.New()) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -5913,7 +5721,6 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { return tt.response() }) user, org, model := seedChatDependenciesWithProvider( - ctx, t, db, "openai", @@ -5931,7 +5738,7 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { require.NoError(t, server.Close()) }) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, Status: database.ChatStatusCompleted, ClientType: database.ChatClientTypeUi, @@ -5939,9 +5746,7 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { Title: "original title", LastModelConfigID: model.ID, }) - require.NoError(t, err) - messages := insertUserTextMessage( - ctx, + message := insertUserTextMessage( t, db, chat.ID, @@ -5950,7 +5755,7 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { "summarize debug title generation", model.ContextLimit, ) - require.Len(t, messages, 1) + require.NotEqual(t, uuid.Nil, message.ID) gotTitle, err := server.ProposeChatTitle(ctx, chat) if tt.wantErr { @@ -5971,7 +5776,7 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { require.Equal(t, string(tt.wantDebugStatus), runs[0].Status) require.True(t, runs[0].FinishedAt.Valid) require.True(t, runs[0].HistoryTipMessageID.Valid) - require.Equal(t, messages[0].ID, runs[0].HistoryTipMessageID.Int64) + require.Equal(t, message.ID, runs[0].HistoryTipMessageID.Int64) } if !tt.wantErr { var usageMessages int @@ -5988,20 +5793,18 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { } func seedChatDependencies( - ctx context.Context, t *testing.T, db database.Store, ) (database.User, database.Organization, database.ChatModelConfig) { t.Helper() openAIURL := chattest.OpenAI(t) - return seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) + return seedChatDependenciesWithProvider(t, db, "openai", openAIURL) } // seedChatDependenciesWithProvider creates a user, organization, // chat provider, and model config for the given provider type and // base URL. func seedChatDependenciesWithProvider( - ctx context.Context, t *testing.T, db database.Store, provider string, @@ -6015,34 +5818,19 @@ func seedChatDependenciesWithProvider( UserID: user.ID, OrganizationID: org.ID, }) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: provider, - DisplayName: provider, - APIKey: "test-key", - BaseUrl: baseURL, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: provider, + DisplayName: provider, + BaseUrl: baseURL, }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: provider, - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: provider, + IsDefault: true, }) - require.NoError(t, err) return user, org, model } func seedChatDependenciesWithProviderPolicy( - ctx context.Context, t *testing.T, db database.Store, provider string, @@ -6060,32 +5848,23 @@ func seedChatDependenciesWithProviderPolicy( UserID: user.ID, OrganizationID: org.ID, }) - providerConfig, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: provider, - DisplayName: provider, - APIKey: apiKey, - BaseUrl: baseURL, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: centralAPIKeyEnabled, - AllowUserApiKey: allowUserAPIKey, - AllowCentralApiKeyFallback: allowCentralAPIKeyFallback, + providerConfig := dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: provider, + DisplayName: provider, + BaseUrl: baseURL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + Enabled: true, + }, func(p *database.InsertChatProviderParams) { + p.APIKey = apiKey + p.CentralApiKeyEnabled = centralAPIKeyEnabled + p.AllowUserApiKey = allowUserAPIKey + p.AllowCentralApiKeyFallback = allowCentralAPIKeyFallback }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: provider, - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: provider, + IsDefault: true, }) - require.NoError(t, err) return user, org, providerConfig, model } @@ -6143,7 +5922,6 @@ func waitForTerminalChat( } func insertChatModelConfigWithCallConfig( - ctx context.Context, t *testing.T, db database.Store, userID uuid.UUID, @@ -6156,24 +5934,17 @@ func insertChatModelConfigWithCallConfig( options, err := json.Marshal(callConfig) require.NoError(t, err) - modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: provider, - Model: model, - DisplayName: model, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - Enabled: true, - IsDefault: false, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: options, + return dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: provider, + Model: model, + DisplayName: model, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + Options: options, }) - require.NoError(t, err) - return modelConfig } func insertUserTextMessage( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -6181,7 +5952,7 @@ func insertUserTextMessage( modelConfigID uuid.UUID, text string, contextLimit ...int64, -) []database.ChatMessage { +) database.ChatMessage { t.Helper() require.LessOrEqual(t, len(contextLimit), 1) @@ -6192,28 +5963,14 @@ func insertUserTextMessage( content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) require.NoError(t, err) - messages, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{userID}, - ModelConfigID: []uuid.UUID{modelConfigID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - Content: []string{string(content.RawMessage)}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{contextLimitValue}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, + return dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: content.RawMessage, Valid: true}, + ContextLimit: sql.NullInt64{Int64: contextLimitValue, Valid: contextLimitValue != 0}, }) - require.NoError(t, err) - return messages } // seedWorkspaceWithAgent creates a full workspace chain with a connected @@ -6331,7 +6088,7 @@ func TestInterruptChatDoesNotSendWebPushNotification(t *testing.T) { require.NoError(t, server.Close()) }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -6444,7 +6201,7 @@ func TestSuccessfulChatSendsWebPushWithNavigationData(t *testing.T) { require.NoError(t, server.Close()) }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -6530,7 +6287,7 @@ func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) require.NoError(t, serverA.Close()) }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := serverA.CreateChat(ctx, chatd.CreateOptions{ @@ -6637,7 +6394,7 @@ func TestSuccessfulChatSendsWebPushWithSummary(t *testing.T) { require.NoError(t, server.Close()) }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) _, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -6699,7 +6456,7 @@ func TestSuccessfulChatSendsWebPushFallbackWithoutSummaryForEmptyAssistantText(t require.NoError(t, server.Close()) }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) _, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -6861,21 +6618,17 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { }) // Seed the DB: user, openai-compat provider, model config. - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) // Add an Anthropic provider pointing to our mock server. - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "anthropic", - DisplayName: "Anthropic", - APIKey: "test-anthropic-key", - BaseUrl: anthropicSrv.URL, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "anthropic", + DisplayName: "Anthropic", + APIKey: "test-anthropic-key", + BaseUrl: anthropicSrv.URL, }) - require.NoError(t, err) - err = db.UpsertChatDesktopEnabled(ctx, true) + err := db.UpsertChatDesktopEnabled(ctx, true) require.NoError(t, err) // Build workspace + agent records so getWorkspaceConn can @@ -7066,7 +6819,7 @@ func TestInterruptChatPersistsPartialResponse(t *testing.T) { require.NoError(t, server.Close()) }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -7180,7 +6933,6 @@ func TestProcessChat_UserProviderKey_Success(t *testing.T) { }) user, org, provider, model := seedChatDependenciesWithProviderPolicy( - ctx, t, db, "openai-compat", @@ -7246,7 +6998,6 @@ func TestProcessChat_UserProviderKey_MissingKeyError(t *testing.T) { }) user, org, _, model := seedChatDependenciesWithProviderPolicy( - ctx, t, db, "openai-compat", @@ -7310,7 +7061,7 @@ func TestProcessChatPanicRecovery(t *testing.T) { }) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) // Pass the panic wrapper to the server, but use the real // database for seeding so those operations don't panic. @@ -7446,25 +7197,18 @@ func TestMCPServerToolInvocation(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) // Seed the MCP server config in the database. This must // happen after seedChatDependencies so user.ID exists for // the foreign key. - mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Test MCP", - Slug: "test-mcp", - Url: mcpTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "Test MCP", + Slug: "test-mcp", + Url: mcpTS.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) @@ -7630,23 +7374,16 @@ func TestPlanModeRootChatApprovedExternalMCPToolInvocation(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ DisplayName: "Plan Mode MCP", Slug: "plan-mode-mcp", Url: mcpTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, AllowInPlanMode: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) server := newActiveTestServer(t, db, ps) @@ -7745,23 +7482,16 @@ func TestPlanModeRootChatApprovedExternalMCPWorkflowCanReachProposePlan(t *testi } }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ DisplayName: "Plan Workflow MCP", Slug: "plan-workflow-mcp", Url: mcpTS.URL, - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, AllowInPlanMode: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) ctrl := gomock.NewController(t) @@ -7960,29 +7690,23 @@ func TestMCPServerOAuth2TokenRefresh(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) // Seed the MCP server config with OAuth2 auth pointing to our // mock token endpoint. - mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ DisplayName: "Authed MCP", Slug: "authed-mcp", Url: mcpTS.URL, - Transport: "streamable_http", AuthType: "oauth2", OAuth2ClientID: "test-client-id", OAuth2TokenURL: tokenSrv.URL, - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) // Seed an expired OAuth2 token with a valid refresh_token. - _, err = db.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{ + _, err := db.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{ MCPServerConfigID: mcpConfig.ID, UserID: user.ID, AccessToken: "old-expired-access-token", @@ -8096,26 +7820,19 @@ func TestMCPServerOAuth2TokenRefreshFailureGraceful(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ + mcpConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ DisplayName: "Broken MCP", Slug: "broken-mcp", Url: "http://127.0.0.1:0/does-not-exist", - Transport: "streamable_http", AuthType: "oauth2", OAuth2ClientID: "test-client-id", OAuth2TokenURL: tokenSrv.URL, - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - - _, err = db.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{ + _, err := db.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{ MCPServerConfigID: mcpConfig.ID, UserID: user.ID, AccessToken: "old-expired-token", @@ -8123,6 +7840,7 @@ func TestMCPServerOAuth2TokenRefreshFailureGraceful(t *testing.T) { TokenType: "Bearer", Expiry: sql.NullTime{Time: time.Now().Add(-1 * time.Hour), Valid: true}, }) + require.NoError(t, err) server := newActiveTestServer(t, db, ps) @@ -8216,7 +7934,7 @@ func TestChatTemplateAllowlistEnforcement(t *testing.T) { } }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) // Create two templates the user can see. tplAllowed = dbgen.Template(t, db, database.Template{ @@ -8361,7 +8079,7 @@ func TestSignalWakeImmediateAcquisition(t *testing.T) { cfg.InFlightChatStaleAfter = testutil.WaitSuperLong }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // CreateChat sets status=pending and calls signalWake(). @@ -8424,7 +8142,7 @@ func TestSignalWakeSendMessage(t *testing.T) { cfg.InFlightChatStaleAfter = testutil.WaitSuperLong }) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // CreateChat triggers wake -> processes first turn. @@ -8640,7 +8358,7 @@ func TestSendMessageRejectsArchivedChat(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, @@ -8669,7 +8387,7 @@ func TestEditMessageRejectsArchivedChat(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, @@ -8705,7 +8423,7 @@ func TestPromoteQueuedRejectsArchivedChat(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, @@ -8762,7 +8480,7 @@ func TestSubmitToolResultsRejectsArchivedChat(t *testing.T) { replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, @@ -8803,20 +8521,17 @@ func TestAcquireChatsSkipsArchivedPendingChat(t *testing.T) { _ = newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) - archivedChat, err := db.InsertChat(ctx, database.InsertChatParams{ + archivedChat := dbgen.Chat(t, db, database.Chat{ OwnerID: user.ID, OrganizationID: org.ID, Title: "acquire-skip-archived", LastModelConfigID: model.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) // Archive the chat, then force it to pending. - _, err = db.ArchiveChatByID(ctx, archivedChat.ID) + _, err := db.ArchiveChatByID(ctx, archivedChat.ID) require.NoError(t, err) _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ @@ -8827,15 +8542,13 @@ func TestAcquireChatsSkipsArchivedPendingChat(t *testing.T) { // Insert a second, non-archived pending chat so the result // slice is non-empty and the assertion is not vacuously true. - activeChat, err := db.InsertChat(ctx, database.InsertChatParams{ + activeChat := dbgen.Chat(t, db, database.Chat{ OwnerID: user.ID, OrganizationID: org.ID, Title: "acquire-active", LastModelConfigID: model.ID, Status: database.ChatStatusPending, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) now := time.Now() acquired, err := db.AcquireChats(ctx, database.AcquireChatsParams{ @@ -8877,7 +8590,7 @@ func TestAdvisorGating_Disabled(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{ Enabled: false, MaxUsesPerRun: 3, @@ -8974,7 +8687,7 @@ func TestAdvisorGating_RootChat(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{ Enabled: true, MaxUsesPerRun: 3, @@ -9108,7 +8821,7 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { } }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{ Enabled: true, MaxUsesPerRun: 3, @@ -9225,7 +8938,7 @@ func TestAdvisorGating_ChildChat(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{ Enabled: true, MaxUsesPerRun: 3, @@ -9235,7 +8948,7 @@ func TestAdvisorGating_ChildChat(t *testing.T) { // Seed the parent chat directly in the database so the test server // never executes the root turn. That keeps this test focused on the // child-chat gating path without depending on subagent wiring. - parent, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + parent := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, Status: database.ChatStatusWaiting, @@ -9243,7 +8956,6 @@ func TestAdvisorGating_ChildChat(t *testing.T) { LastModelConfigID: model.ID, Title: "advisor-root-parent", }) - require.NoError(t, err) server := newActiveTestServer(t, db, ps) @@ -9317,7 +9029,7 @@ func TestAdvisorGating_PlanMode(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{ Enabled: true, MaxUsesPerRun: 3, @@ -9396,7 +9108,7 @@ func TestAdvisorGating_ExploreSubagent(t *testing.T) { ) }) - user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{ Enabled: true, MaxUsesPerRun: 3, @@ -9529,14 +9241,14 @@ func TestAdvisorChainMode_SnapshotKeepsFullHistory(t *testing.T) { ) }) - user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) storeEnabled := true // The OpenAI Responses API is the only provider code path where // chain mode activates. Store=true is the switch that routes this // provider/model through the Responses API and lets // IsResponsesStoreEnabled return true. responsesModel := insertChatModelConfigWithCallConfig( - ctx, t, db, user.ID, "openai", "gpt-4o", + t, db, user.ID, "openai", "gpt-4o", codersdk.ChatModelCallConfig{ ProviderOptions: &codersdk.ChatModelProviderOptions{ OpenAI: &codersdk.ChatModelOpenAIProviderOptions{ diff --git a/coderd/x/chatd/chatdebug/service_test.go b/coderd/x/chatd/chatdebug/service_test.go index d8a25d8b8c..358ff0e36b 100644 --- a/coderd/x/chatd/chatdebug/service_test.go +++ b/coderd/x/chatd/chatdebug/service_test.go @@ -39,7 +39,7 @@ func TestService_IsEnabled(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _, _ := dbtestutil.NewDBWithSQLDB(t) - _, owner, chat, model := seedChat(ctx, t, db) + _, owner, chat, model := seedChat(t, db) require.NotEqual(t, uuid.Nil, model.ID) svc := chatdebug.NewService(db, testutil.Logger(t), nil) @@ -77,7 +77,7 @@ func TestService_IsEnabled_AlwaysEnable(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _, _ := dbtestutil.NewDBWithSQLDB(t) - _, owner, chat, model := seedChat(ctx, t, db) + _, owner, chat, model := seedChat(t, db) require.NotEqual(t, uuid.Nil, model.ID) svc := chatdebug.NewService(db, testutil.Logger(t), nil, chatdebug.WithAlwaysEnable(true)) @@ -98,11 +98,11 @@ func TestService_CreateRun(t *testing.T) { t.Parallel() fixture := newFixture(t) - rootChat := insertChat(fixture.ctx, t, fixture.db, fixture.org.ID, fixture.owner.ID, fixture.model.ID) - parentChat := insertChat(fixture.ctx, t, fixture.db, fixture.org.ID, fixture.owner.ID, fixture.model.ID) - triggerMsg := insertMessage(fixture.ctx, t, fixture.db, fixture.chat.ID, + rootChat := insertChat(t, fixture.db, fixture.org.ID, fixture.owner.ID, fixture.model.ID) + parentChat := insertChat(t, fixture.db, fixture.org.ID, fixture.owner.ID, fixture.model.ID) + triggerMsg := insertMessage(t, fixture.db, fixture.chat.ID, fixture.owner.ID, fixture.model.ID, database.ChatMessageRoleUser, "trigger") - historyTipMsg := insertMessage(fixture.ctx, t, fixture.db, fixture.chat.ID, + historyTipMsg := insertMessage(t, fixture.db, fixture.chat.ID, fixture.owner.ID, fixture.model.ID, database.ChatMessageRoleAssistant, "history-tip") @@ -279,7 +279,7 @@ func TestService_CreateStep(t *testing.T) { fixture := newFixture(t) run := createRun(t, fixture) - historyTipMsg := insertMessage(fixture.ctx, t, fixture.db, fixture.chat.ID, + historyTipMsg := insertMessage(t, fixture.db, fixture.chat.ID, fixture.owner.ID, fixture.model.ID, database.ChatMessageRoleAssistant, "history-tip") @@ -424,7 +424,7 @@ func TestService_CreateStep_ChatIDMismatchReportsNotFound(t *testing.T) { // attach a step to the existing run using the wrong chat_id. // The insert's locked_run WHERE fails on chat_id, producing // sql.ErrNoRows; classifyMissingRun must report not-found. - otherChat := insertChat(fixture.ctx, t, fixture.db, fixture.org.ID, + otherChat := insertChat(t, fixture.db, fixture.org.ID, fixture.owner.ID, fixture.model.ID) _, err := fixture.svc.CreateStep(fixture.ctx, chatdebug.CreateStepParams{ @@ -454,7 +454,7 @@ func TestService_UpdateStep(t *testing.T) { }) require.NoError(t, err) - assistantMsg := insertMessage(fixture.ctx, t, fixture.db, fixture.chat.ID, + assistantMsg := insertMessage(t, fixture.db, fixture.chat.ID, fixture.owner.ID, fixture.model.ID, database.ChatMessageRoleAssistant, "assistant") finishedAt := time.Now().UTC().Round(time.Microsecond) @@ -598,12 +598,12 @@ func TestService_DeleteAfterMessageID(t *testing.T) { t.Parallel() fixture := newFixture(t) - low := insertMessage(fixture.ctx, t, fixture.db, fixture.chat.ID, fixture.owner.ID, + low := insertMessage(t, fixture.db, fixture.chat.ID, fixture.owner.ID, fixture.model.ID, database.ChatMessageRoleAssistant, "low") - threshold := insertMessage(fixture.ctx, t, fixture.db, fixture.chat.ID, + threshold := insertMessage(t, fixture.db, fixture.chat.ID, fixture.owner.ID, fixture.model.ID, database.ChatMessageRoleAssistant, "threshold") - high := insertMessage(fixture.ctx, t, fixture.db, fixture.chat.ID, fixture.owner.ID, + high := insertMessage(t, fixture.db, fixture.chat.ID, fixture.owner.ID, fixture.model.ID, database.ChatMessageRoleAssistant, "high") require.Less(t, low.ID, threshold.ID) require.Less(t, threshold.ID, high.ID) @@ -685,7 +685,7 @@ func TestService_FinalizeStale(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) - _, owner, chat, model := seedChat(ctx, t, db) + _, owner, chat, model := seedChat(t, db) require.NotEqual(t, uuid.Nil, owner.ID) staleTime := time.Now().Add(-10 * time.Minute).UTC().Round(time.Microsecond) @@ -733,7 +733,7 @@ func TestService_FinalizeStale_BroadcastsFinalizeEvent(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) - _, owner, chat, model := seedChat(ctx, t, db) + _, owner, chat, model := seedChat(t, db) require.NotEqual(t, uuid.Nil, owner.ID) staleTime := time.Now().Add(-10 * time.Minute).UTC().Round(time.Microsecond) @@ -796,7 +796,7 @@ func TestService_FinalizeStale_NoChangesDoesNotBroadcast(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) - _, owner, chat, _ := seedChat(ctx, t, db) + _, owner, chat, _ := seedChat(t, db) require.NotEqual(t, uuid.Nil, owner.ID) memoryPubsub := dbpubsub.NewInMemory() @@ -1018,7 +1018,7 @@ func TestService_PublishesEvents(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) - _, owner, chat, model := seedChat(ctx, t, db) + _, owner, chat, model := seedChat(t, db) require.NotEqual(t, uuid.Nil, owner.ID) memoryPubsub := dbpubsub.NewInMemory() @@ -1069,7 +1069,7 @@ func newFixture(t *testing.T) testFixture { ctx := testutil.Context(t, testutil.WaitLong) db, _ := dbtestutil.NewDB(t) - org, owner, chat, model := seedChat(ctx, t, db) + org, owner, chat, model := seedChat(t, db) return testFixture{ ctx: ctx, db: db, @@ -1082,7 +1082,6 @@ func newFixture(t *testing.T) testFixture { } func seedChat( - ctx context.Context, t *testing.T, db database.Store, ) (database.Organization, database.User, database.Chat, database.ChatModelConfig) { @@ -1091,38 +1090,21 @@ func seedChat( org := dbgen.Organization(t, db, database.Organization{}) owner := dbgen.User(t, db, database.User{}) providerName := "openai" - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: providerName, - DisplayName: "OpenAI", - APIKey: "test-key", - CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: providerName, + DisplayName: "OpenAI", }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(ctx, - database.InsertChatModelConfigParams{ - Provider: providerName, - Model: "model-" + uuid.NewString(), - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), - }, - ) - require.NoError(t, err) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "model-" + uuid.NewString(), + IsDefault: true, + }) - chat := insertChat(ctx, t, db, org.ID, owner.ID, model.ID) + chat := insertChat(t, db, org.ID, owner.ID, model.ID) return org, owner, chat, model } func insertChat( - ctx context.Context, t *testing.T, db database.Store, orgID uuid.UUID, @@ -1131,20 +1113,16 @@ func insertChat( ) database.Chat { t.Helper() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: orgID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: ownerID, LastModelConfigID: modelID, Title: "chat-" + uuid.NewString(), }) - require.NoError(t, err) return chat } func insertMessage( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -1160,29 +1138,16 @@ func insertMessage( }) require.NoError(t, err) - messages, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{createdBy}, - ModelConfigID: []uuid.UUID{modelID}, - Role: []database.ChatMessageRole{role}, - Content: []string{string(parts.RawMessage)}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, + msg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + Role: role, + Content: parts, + ContentVersion: chatprompt.CurrentContentVersion, + ProviderResponseID: sql.NullString{}, }) - require.NoError(t, err) - require.Len(t, messages, 1) - return messages[0] + return msg } func createRun(t *testing.T, fixture testFixture) database.ChatDebugRun { diff --git a/coderd/x/chatd/chatprompt/chatprompt_test.go b/coderd/x/chatd/chatprompt/chatprompt_test.go index 7d06ec0224..c9180eb7fe 100644 --- a/coderd/x/chatd/chatprompt/chatprompt_test.go +++ b/coderd/x/chatd/chatprompt/chatprompt_test.go @@ -1815,35 +1815,16 @@ func TestNulEscapeRoundTrip(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitShort) // Seed minimal dependencies for the DB round-trip path: // user, provider, model config, chat. user := dbgen.User(t, db, database.User{}) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "openai", - APIKey: "test-key", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, - }) - require.NoError(t, err) + dbgen.ChatProvider(t, db, database.ChatProvider{}) - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + IsDefault: true, }) - require.NoError(t, err) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ @@ -1851,15 +1832,12 @@ func TestNulEscapeRoundTrip(t *testing.T) { OrganizationID: org.ID, }) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: model.ID, Title: "nul-roundtrip-test", }) - require.NoError(t, err) textTests := []struct { name string @@ -1945,31 +1923,17 @@ func TestNulEscapeRoundTrip(t *testing.T) { // Full DB round-trip: write to PostgreSQL jsonb, read // back, and verify the value survives storage. ctx := testutil.Context(t, testutil.WaitShort) - dbMsgs, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{user.ID}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - Content: []string{string(encoded.RawMessage)}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + dbMsg := dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: encoded, + ContentVersion: chatprompt.CurrentContentVersion, }) - require.NoError(t, err) - require.Len(t, dbMsgs, 1) - readBack, err := db.GetChatMessageByID(ctx, dbMsgs[0].ID) + readBack, err := db.GetChatMessageByID(ctx, dbMsg.ID) require.NoError(t, err) - dbDecoded, err := chatprompt.ParseContent(readBack) require.NoError(t, err) require.Len(t, dbDecoded, 1) @@ -2392,29 +2356,16 @@ func TestMediaToolResultRoundTrip(t *testing.T) { OrganizationID: org.ID, }) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "anthropic", - DisplayName: "anthropic", - APIKey: "test-key", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "anthropic", }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "anthropic", - Model: "test-model", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 200000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "anthropic", + Model: "test-model", + IsDefault: true, + ContextLimit: 200000, }) - require.NoError(t, err) // Small base64 payload standing in for a real screenshot. const imageData = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVQI12NgAAIABQAB" @@ -2429,15 +2380,12 @@ func TestMediaToolResultRoundTrip(t *testing.T) { ) database.Chat { t.Helper() - chat, chatErr := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: model.ID, Title: "media-roundtrip-" + callID, }) - require.NoError(t, chatErr) // Assistant message with the tool call. callPart := codersdk.ChatMessageToolCall(callID, toolName, json.RawMessage(`{}`)) @@ -2448,26 +2396,22 @@ func TestMediaToolResultRoundTrip(t *testing.T) { resultEncoded, encErr := chatprompt.MarshalParts(resultParts) require.NoError(t, encErr) - _, insertErr := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{user.ID, user.ID}, - ModelConfigID: []uuid.UUID{model.ID, model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant, database.ChatMessageRoleTool}, - Content: []string{string(assistantEncoded.RawMessage), string(resultEncoded.RawMessage)}, - ContentVersion: []int16{chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0, 0}, - OutputTokens: []int64{0, 0}, - TotalTokens: []int64{0, 0}, - ReasoningTokens: []int64{0, 0}, - CacheCreationTokens: []int64{0, 0}, - CacheReadTokens: []int64{0, 0}, - ContextLimit: []int64{0, 0}, - Compressed: []bool{false, false}, - TotalCostMicros: []int64{0, 0}, - RuntimeMs: []int64{0, 0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: assistantEncoded, + ContentVersion: chatprompt.CurrentContentVersion, + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chat.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleTool, + Content: resultEncoded, + ContentVersion: chatprompt.CurrentContentVersion, }) - require.NoError(t, insertErr) return chat } diff --git a/coderd/x/chatd/chattool/startworkspace_test.go b/coderd/x/chatd/chattool/startworkspace_test.go index 4c073e82ec..c36aae5eca 100644 --- a/coderd/x/chatd/chattool/startworkspace_test.go +++ b/coderd/x/chatd/chattool/startworkspace_test.go @@ -35,22 +35,19 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, OrganizationID: org.ID, }) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, Title: "test-no-workspace", }) - require.NoError(t, err) tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ DB: db, @@ -73,7 +70,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -87,16 +84,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-already-running", }) - require.NoError(t, err) agentConnFn := func(_ context.Context, _ uuid.UUID) (workspacesdk.AgentConn, func(), error) { return nil, func() {}, nil @@ -132,7 +126,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -170,16 +164,13 @@ func TestStartWorkspace(t *testing.T) { } require.NotEqual(t, uuid.Nil, preferredAgentID) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-running-preferred-agent", }) - require.NoError(t, err) var connectedAgentID uuid.UUID agentConnFn := func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { @@ -216,7 +207,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -232,16 +223,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-running-no-agent", }) - require.NoError(t, err) tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ DB: db, @@ -275,7 +263,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -297,16 +285,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-running-selection-error", }) - require.NoError(t, err) tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ DB: db, @@ -341,7 +326,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -356,16 +341,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-stopped-workspace", }) - require.NoError(t, err) var startCalled bool var startBuildID uuid.UUID @@ -415,7 +397,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -429,16 +411,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-stopped-workspace-auto-update", }) - require.NoError(t, err) startFn := func(_ context.Context, _ uuid.UUID, wsID uuid.UUID, req codersdk.CreateWorkspaceBuildRequest) (codersdk.WorkspaceBuild, error) { require.Equal(t, codersdk.WorkspaceTransitionStart, req.Transition) @@ -480,7 +459,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -494,16 +473,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-workspace-passes-parameters", - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) expectedParams := []codersdk.WorkspaceBuildParameter{ {Name: "region", Value: "us-east-1"}, @@ -545,7 +521,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -559,16 +535,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-workspace-manual-update-required", }) - require.NoError(t, err) tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ DB: db, @@ -615,7 +588,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -629,16 +602,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-workspace-responder-error-without-validations", - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ DB: db, @@ -671,7 +641,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -686,16 +656,13 @@ func TestStartWorkspace(t *testing.T) { }).Starting().Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-in-progress-build", }) - require.NoError(t, err) // Wrap the DB so we know exactly when the tool reads // the job status. The interceptor signals AFTER the @@ -768,7 +735,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -783,16 +750,13 @@ func TestStartWorkspace(t *testing.T) { }).Starting().Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-failed-build", }) - require.NoError(t, err) jobRead := make(chan struct{}, 1) wrappedDB := &jobInterceptStore{Store: db, jobRead: jobRead} @@ -851,7 +815,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -866,16 +830,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-start-triggered-build-failure", }) - require.NoError(t, err) // StartFn creates a real in-progress build via dbfake. var startBuildJobID uuid.UUID @@ -949,7 +910,7 @@ func TestStartWorkspace(t *testing.T) { db, _ := dbtestutil.NewDB(t) user := dbgen.User(t, db, database.User{}) - modelCfg := seedModelConfig(ctx, t, db, user.ID) + modelCfg := seedModelConfig(t, db) org := dbgen.Organization(t, db, database.Organization{}) _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -965,16 +926,13 @@ func TestStartWorkspace(t *testing.T) { }).Do() ws := wsResp.Workspace - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, LastModelConfigID: modelCfg.ID, Title: "test-deleted-workspace", }) - require.NoError(t, err) tool := chattool.StartWorkspace(chattool.StartWorkspaceOptions{ DB: db, @@ -994,39 +952,15 @@ func TestStartWorkspace(t *testing.T) { // seedModelConfig inserts a provider and model config for testing. func seedModelConfig( - ctx context.Context, t *testing.T, db database.Store, - userID uuid.UUID, ) database.ChatModelConfig { t.Helper() - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - APIKey: "test-key", - BaseUrl: "", - ApiKeyKeyID: sql.NullString{}, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + dbgen.ChatProvider(t, db, database.ChatProvider{}) + return dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + IsDefault: true, }) - require.NoError(t, err) - - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), - }) - require.NoError(t, err) - return model } // jobInterceptStore wraps a database.Store and signals a diff --git a/coderd/x/chatd/integration_responses_test.go b/coderd/x/chatd/integration_responses_test.go index 822bb7269c..adb4c1f708 100644 --- a/coderd/x/chatd/integration_responses_test.go +++ b/coderd/x/chatd/integration_responses_test.go @@ -13,7 +13,9 @@ import ( "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "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/chattest" @@ -63,9 +65,9 @@ func TestOpenAIResponsesNoStaleWebSearchReplay(t *testing.T) { } }) - user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) - model := insertOpenAIResponsesModelConfig(ctx, t, db, user.ID, false, true) - server := newActiveTestServer(t, db, ps) + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model := insertOpenAIResponsesModelConfig(t, db, user.ID, false, true) + server := newOpenAIResponsesTestServer(t, db, ps) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -145,10 +147,10 @@ func TestOpenAIResponsesFullReplayPairsReasoningAndWebSearch(t *testing.T) { } }) - user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) - firstModel := insertOpenAIResponsesModelConfig(ctx, t, db, user.ID, true, true) - secondModel := insertOpenAIResponsesModelConfig(ctx, t, db, user.ID, true, true) - server := newActiveTestServer(t, db, ps) + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + firstModel := insertOpenAIResponsesModelConfig(t, db, user.ID, true, true) + secondModel := insertOpenAIResponsesModelConfig(t, db, user.ID, true, true) + server := newOpenAIResponsesTestServer(t, db, ps) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -205,9 +207,9 @@ func TestOpenAIResponsesChainModeSkipsWhenLocalCallPending(t *testing.T) { return resp }) - user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) - model := insertOpenAIResponsesModelConfig(ctx, t, db, user.ID, true, false) - chat := insertOpenAIResponsesChat(ctx, t, db, org.ID, user.ID, model.ID, "local-pending") + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model := insertOpenAIResponsesModelConfig(t, db, user.ID, true, false) + chat := insertOpenAIResponsesChat(t, db, org.ID, user.ID, model.ID, "local-pending") callID := fmt.Sprintf("call_local_%d", time.Now().UnixNano()) localCall := codersdk.ChatMessageToolCall( @@ -229,7 +231,7 @@ func TestOpenAIResponsesChainModeSkipsWhenLocalCallPending(t *testing.T) { }, ) - server := newActiveTestServer(t, db, ps) + server := newOpenAIResponsesTestServer(t, db, ps) _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, @@ -272,9 +274,9 @@ func TestOpenAIResponsesChainModeStillFiresForProviderExecutedOnly(t *testing.T) return resp }) - user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL) - model := insertOpenAIResponsesModelConfig(ctx, t, db, user.ID, true, true) - chat := insertOpenAIResponsesChat(ctx, t, db, org.ID, user.ID, model.ID, "provider-only") + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model := insertOpenAIResponsesModelConfig(t, db, user.ID, true, true) + chat := insertOpenAIResponsesChat(t, db, org.ID, user.ID, model.ID, "provider-only") const ( previousResponseID = "resp_provider_only_prior" @@ -311,7 +313,7 @@ func TestOpenAIResponsesChainModeStillFiresForProviderExecutedOnly(t *testing.T) }, ) - server := newActiveTestServer(t, db, ps) + server := newOpenAIResponsesTestServer(t, db, ps) _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, @@ -382,8 +384,23 @@ type persistedResponsesMessage struct { providerResponseID string } +func newOpenAIResponsesTestServer( + t *testing.T, + db database.Store, + ps dbpubsub.Pubsub, +) *chatd.Server { + t.Helper() + return newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + // Let CreateChat and SendMessage publish their pending status + // before wake-driven processing starts. The responses tests are + // not exercising periodic polling, and PostgreSQL can otherwise + // deliver that stale pending notification after processChat + // subscribes to control events. + cfg.PendingChatAcquireInterval = testutil.WaitLong + }) +} + func insertOpenAIResponsesModelConfig( - ctx context.Context, t *testing.T, db database.Store, userID uuid.UUID, @@ -392,7 +409,6 @@ func insertOpenAIResponsesModelConfig( ) database.ChatModelConfig { t.Helper() return insertChatModelConfigWithCallConfig( - ctx, t, db, userID, @@ -410,7 +426,6 @@ func insertOpenAIResponsesModelConfig( } func insertOpenAIResponsesChat( - ctx context.Context, t *testing.T, db database.Store, organizationID uuid.UUID, @@ -419,7 +434,7 @@ func insertOpenAIResponsesChat( titlePrefix string, ) database.Chat { t.Helper() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + return dbgen.Chat(t, db, database.Chat{ OrganizationID: organizationID, OwnerID: ownerID, LastModelConfigID: modelConfigID, @@ -428,8 +443,6 @@ func insertOpenAIResponsesChat( MCPServerIDs: []uuid.UUID{}, ClientType: database.ChatClientTypeApi, }) - require.NoError(t, err) - return chat } func insertOpenAIResponsesMessages( @@ -464,6 +477,8 @@ func insertOpenAIResponsesMessages( params.RuntimeMs = append(params.RuntimeMs, 0) params.ProviderResponseID = append(params.ProviderResponseID, message.providerResponseID) } + // Keep this raw because dbgen.ChatMessage inserts one message at a time, + // while this helper needs to preserve variadic batch insert behavior. _, err := db.InsertChatMessages(ctx, params) require.NoError(t, err) } diff --git a/coderd/x/chatd/recording_internal_test.go b/coderd/x/chatd/recording_internal_test.go index d3c852c1df..24bdf3cf76 100644 --- a/coderd/x/chatd/recording_internal_test.go +++ b/coderd/x/chatd/recording_internal_test.go @@ -20,6 +20,7 @@ import ( "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" "github.com/coder/coder/v2/codersdk" @@ -84,7 +85,6 @@ func validRecordingJPEG(extra int, fill byte) []byte { // background processing (which would try to call the LLM and // use the agent connection mock). func createComputerUseParentChild( - ctx context.Context, t *testing.T, server *Server, user database.User, @@ -98,7 +98,7 @@ func createComputerUseParentChild( // Insert the parent chat directly via DB to avoid triggering // the server's background processing. - parent, err := server.db.InsertChat(ctx, database.InsertChatParams{ + parent = dbgen.Chat(t, server.db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, @@ -106,14 +106,12 @@ func createComputerUseParentChild( LastModelConfigID: model.ID, Title: parentTitle, Status: database.ChatStatusPending, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) // Insert the child chat directly via DB to avoid triggering // the server's background processing (which would try to run // the chat without an LLM and get stuck). - child, err = server.db.InsertChat(ctx, database.InsertChatParams{ + child = dbgen.Chat(t, server.db, database.Chat{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, @@ -124,9 +122,7 @@ func createComputerUseParentChild( Title: childTitle, Mode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true}, Status: database.ChatStatusPending, - ClientType: database.ChatClientTypeUi, }) - require.NoError(t, err) return parent, child } @@ -178,7 +174,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create the server WITHOUT agentConnFn so the background @@ -186,7 +182,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) parent, child := createComputerUseParentChild( - ctx, t, server, user, org, model, workspace, agent, + t, server, user, org, model, workspace, agent, "parent-recording", "computer-use-child", ) @@ -201,7 +197,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) { } // Add an assistant message so the report is extracted. - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "I opened Firefox.") + insertAssistantMessage(t, db, child.ID, model.ID, "I opened Firefox.") // Set child to waiting (terminal success state). setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") @@ -268,13 +264,13 @@ func TestWaitAgentComputerUseRecordingWithThumbnail(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) parent, child := createComputerUseParentChild( - ctx, t, server, user, org, model, workspace, agent, + t, server, user, org, model, workspace, agent, "parent-recording-thumb", "computer-use-child-thumb", ) @@ -285,7 +281,7 @@ func TestWaitAgentComputerUseRecordingWithThumbnail(t *testing.T) { return mockConn, func() {}, nil } - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "I opened Firefox and took a screenshot.") + insertAssistantMessage(t, db, child.ID, model.ID, "I opened Firefox and took a screenshot.") setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") @@ -360,7 +356,7 @@ func TestWaitAgentNonComputerUseNoRecording(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -368,7 +364,7 @@ func TestWaitAgentNonComputerUseNoRecording(t *testing.T) { parent, child := createParentChildChats(ctx, t, server, user, org, model) // Add an assistant message so the report is extracted. - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Done.") + insertAssistantMessage(t, db, child.ID, model.ID, "Done.") // Wait for background processing triggered by CreateChat to // settle before setting up the mock agent connection. @@ -411,7 +407,7 @@ func TestWaitAgentRecordingStartFails(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create the server WITHOUT agentConnFn so the background @@ -420,7 +416,7 @@ func TestWaitAgentRecordingStartFails(t *testing.T) { // Create parent + computer_use child. parent, child := createComputerUseParentChild( - ctx, t, server, user, org, model, workspace, agent, + t, server, user, org, model, workspace, agent, "parent-start-fail", "computer-use-start-fail", ) @@ -429,7 +425,7 @@ func TestWaitAgentRecordingStartFails(t *testing.T) { return mockConn, func() {}, nil } - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Opened the browser.") + insertAssistantMessage(t, db, child.ID, model.ID, "Opened the browser.") setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") // StartDesktopRecording fails. StopDesktopRecording must NOT @@ -465,7 +461,7 @@ func TestWaitAgentRecordingStopFails(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create the server WITHOUT agentConnFn so the background @@ -474,7 +470,7 @@ func TestWaitAgentRecordingStopFails(t *testing.T) { // Create parent + computer_use child. parent, child := createComputerUseParentChild( - ctx, t, server, user, org, model, workspace, agent, + t, server, user, org, model, workspace, agent, "parent-stop-fail", "computer-use-stop-fail", ) @@ -483,7 +479,7 @@ func TestWaitAgentRecordingStopFails(t *testing.T) { return mockConn, func() {}, nil } - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Checked settings.") + insertAssistantMessage(t, db, child.ID, model.ID, "Checked settings.") setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") // Start succeeds, stop fails. @@ -526,12 +522,12 @@ func TestWaitAgentTimeoutLeavesRecordingRunning(t *testing.T) { // Use the mock clock server; don't set agentConnFn yet. server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) // Create parent + computer_use child. _, child := createComputerUseParentChild( - ctx, t, server, user, org, model, workspace, agent, + t, server, user, org, model, workspace, agent, "parent-timeout", "computer-use-timeout", ) @@ -610,7 +606,7 @@ func TestStopAndStoreRecording_Oversized(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -659,7 +655,7 @@ func TestStopAndStoreRecording_OversizedThumbnail(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -723,7 +719,7 @@ func TestStopAndStoreRecording_DuplicatePartsIgnored(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -766,7 +762,7 @@ func TestStopAndStoreRecording_Empty(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -796,7 +792,7 @@ func TestStopAndStoreRecording_LinkFailureRollsBackInsert(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -851,7 +847,7 @@ func TestStopAndStoreRecording_WithThumbnail(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -905,7 +901,7 @@ func TestStopAndStoreRecording_VideoOnly(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -947,7 +943,7 @@ func TestStopAndStoreRecording_MismatchedVideoBytesSkipped(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -984,7 +980,7 @@ func TestStopAndStoreRecording_DownloadFailure(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -1017,7 +1013,7 @@ func TestStopAndStoreRecording_UnknownPartIgnored(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -1071,7 +1067,7 @@ func TestStopAndStoreRecording_MalformedContentType(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -1107,7 +1103,7 @@ func TestStopAndStoreRecording_MissingBoundary(t *testing.T) { ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) diff --git a/coderd/x/chatd/subagent_context_internal_test.go b/coderd/x/chatd/subagent_context_internal_test.go index 6d5a6e4513..dc60e3330f 100644 --- a/coderd/x/chatd/subagent_context_internal_test.go +++ b/coderd/x/chatd/subagent_context_internal_test.go @@ -12,6 +12,7 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" @@ -148,7 +149,7 @@ func createParentChatWithInheritedContext( ) database.Chat { t.Helper() - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, @@ -182,26 +183,14 @@ func createParentChatWithInheritedContext( content, err := json.Marshal(inheritedParts) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: parent.ID, - CreatedBy: []uuid.UUID{user.ID}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, - Content: []string{string(content)}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: parent.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: content, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, }) - require.NoError(t, err) parentChat, err := db.GetChatByID(ctx, parent.ID) require.NoError(t, err) @@ -329,7 +318,7 @@ func createParentChatWithRotatedInheritedContext( ) database.Chat { t.Helper() - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, @@ -379,26 +368,22 @@ func createParentChatWithRotatedInheritedContext( }) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: parent.ID, - CreatedBy: []uuid.UUID{user.ID, user.ID}, - ModelConfigID: []uuid.UUID{model.ID, model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleUser}, - Content: []string{string(oldContent), string(newContent)}, - ContentVersion: []int16{chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth, database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0, 0}, - OutputTokens: []int64{0, 0}, - TotalTokens: []int64{0, 0}, - ReasoningTokens: []int64{0, 0}, - CacheCreationTokens: []int64{0, 0}, - CacheReadTokens: []int64{0, 0}, - ContextLimit: []int64{0, 0}, - Compressed: []bool{false, false}, - TotalCostMicros: []int64{0, 0}, - RuntimeMs: []int64{0, 0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: parent.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: oldContent, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + }) + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: parent.ID, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: newContent, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, }) - require.NoError(t, err) parentChat, err := db.GetChatByID(ctx, parent.ID) require.NoError(t, err) @@ -476,7 +461,7 @@ func TestSpawnComputerUseAgentInheritsContext(t *testing.T) { ctx := chatdTestContext(t) parentChat := createParentChatWithInheritedContext(ctx, t, db, server) - insertEnabledAnthropicProvider(ctx, t, db, parentChat.OwnerID) + insertEnabledAnthropicProvider(t, db, parentChat.OwnerID) // The direct DB insert above bypasses the pubsub event that // production uses to invalidate the provider cache. Explicitly // invalidate here so the background processing goroutine does diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 708233ee84..de4c6c60fc 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -10,6 +10,7 @@ import ( "charm.land/fantasy" "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -172,7 +173,6 @@ func (s *subagentTestLogSink) entriesAtLevelWithMessage( // and model. This deliberately does NOT create an Anthropic // provider. func seedInternalChatDeps( - ctx context.Context, t *testing.T, db database.Store, ) (database.User, database.Organization, database.ChatModelConfig) { @@ -184,31 +184,14 @@ func seedInternalChatDeps( UserID: user.ID, OrganizationID: org.ID, }) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - APIKey: "test-key", - BaseUrl: "", - ApiKeyKeyID: sql.NullString{}, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + IsDefault: true, }) - require.NoError(t, err) return user, org, model } @@ -217,24 +200,18 @@ func seedInternalChatDeps( // the current test user so computer_use flows keep Anthropic credentials // after provider-key pruning. func insertEnabledAnthropicProvider( - ctx context.Context, t *testing.T, db database.Store, userID uuid.UUID, ) { t.Helper() - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "anthropic", - DisplayName: "Anthropic", - APIKey: "test-anthropic-key", - BaseUrl: "", - ApiKeyKeyID: sql.NullString{}, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "anthropic", + DisplayName: "Anthropic", + APIKey: "test-anthropic-key", + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, }) - require.NoError(t, err) } func TestResolveUserProviderAPIKeys_PreservesAnthropicKeyFromDBProvider(t *testing.T) { @@ -247,8 +224,8 @@ func TestResolveUserProviderAPIKeys_PreservesAnthropicKeyFromDBProvider(t *testi server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, _, _ := seedInternalChatDeps(ctx, t, db) - insertEnabledAnthropicProvider(ctx, t, db, user.ID) + user, _, _ := seedInternalChatDeps(t, db) + insertEnabledAnthropicProvider(t, db, user.ID) keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID) require.NoError(t, err) @@ -266,7 +243,7 @@ func TestResolveUserProviderAPIKeys_PreservesAnthropicKeyFromDBProvider(t *testi }) ctx := chatdTestContext(t) - user, _, _ := seedInternalChatDeps(ctx, t, db) + user, _, _ := seedInternalChatDeps(t, db) keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID) require.NoError(t, err) @@ -278,18 +255,14 @@ func TestResolveUserProviderAPIKeys_PreservesAnthropicKeyFromDBProvider(t *testi } func insertInternalChatModelConfig( - ctx context.Context, t *testing.T, db database.Store, - userID uuid.UUID, model string, enabled bool, ) database.ChatModelConfig { return insertInternalChatModelConfigForProvider( - ctx, t, db, - userID, "openai", model, enabled, @@ -297,7 +270,6 @@ func insertInternalChatModelConfig( } func insertInternalChatProvider( - ctx context.Context, t *testing.T, db database.Store, userID uuid.UUID, @@ -309,36 +281,31 @@ func insertInternalChatProvider( ) database.ChatProvider { t.Helper() - providerConfig, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: provider, - DisplayName: provider, - APIKey: apiKey, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: centralAPIKeyEnabled, - AllowUserApiKey: allowUserAPIKey, - AllowCentralApiKeyFallback: allowCentralAPIKeyFallback, + providerConfig := dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: provider, + DisplayName: provider, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + }, func(p *database.InsertChatProviderParams) { + p.APIKey = apiKey + p.CentralApiKeyEnabled = centralAPIKeyEnabled + p.AllowUserApiKey = allowUserAPIKey + p.AllowCentralApiKeyFallback = allowCentralAPIKeyFallback }) - require.NoError(t, err) return providerConfig } func insertInternalChatModelConfigForProvider( - ctx context.Context, t *testing.T, db database.Store, - userID uuid.UUID, provider string, model string, enabled bool, ) database.ChatModelConfig { t.Helper() return insertInternalChatModelConfigWithOptions( - ctx, t, db, - userID, provider, model, enabled, @@ -347,10 +314,8 @@ func insertInternalChatModelConfigForProvider( } func insertInternalChatModelConfigWithOptions( - ctx context.Context, t *testing.T, db database.Store, - userID uuid.UUID, provider string, model string, enabled bool, @@ -358,25 +323,19 @@ func insertInternalChatModelConfigWithOptions( ) database.ChatModelConfig { t.Helper() - modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: provider, - Model: model, - DisplayName: model, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - Enabled: enabled, - IsDefault: false, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: options, + modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: provider, + Model: model, + DisplayName: model, + Options: options, + }, func(p *database.InsertChatModelConfigParams) { + p.Enabled = enabled }) - require.NoError(t, err) return modelConfig } func insertInternalMCPServerConfig( - ctx context.Context, t *testing.T, db database.Store, userID uuid.UUID, @@ -385,23 +344,14 @@ func insertInternalMCPServerConfig( ) database.MCPServerConfig { t.Helper() - cfg, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ + return dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ DisplayName: slug, Slug: slug, Url: "https://" + slug + ".example.com", - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, AllowInPlanMode: allowInPlanMode, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: userID, - UpdatedBy: userID, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, }) - require.NoError(t, err) - - return cfg } func seedWorkspaceBinding( @@ -466,7 +416,7 @@ func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, build, agent := seedWorkspaceBinding(t, db, user.ID) parent, err := server.CreateChat(ctx, CreateOptions{ @@ -634,7 +584,7 @@ func TestCreateChildSubagentChatCopiesPlanMode(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) planMode := database.NullChatPlanMode{ ChatPlanMode: database.ChatPlanModePlan, Valid: true, @@ -671,7 +621,7 @@ func TestSpawnAgent_GeneralInheritsParentModelWhenOmitted(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-inherited-model", ) @@ -697,9 +647,9 @@ func TestSpawnAgent_GeneralUsesConfiguredModelOverride(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) overrideModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "general-override-"+uuid.NewString(), true, + t, db, "general-override-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatGeneralModelOverride(ctx, overrideModel.ID.String())) parentChat := createInternalParentChat( @@ -727,9 +677,8 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenCredentialsUnavailable(t server := newInternalTestServerWithLogger(t, db, ps, chatprovider.ProviderAPIKeys{}, logger) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) insertInternalChatProvider( - ctx, t, db, user.ID, @@ -739,11 +688,10 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenCredentialsUnavailable(t true, false, ) + overrideModel := insertInternalChatModelConfigForProvider( - ctx, t, db, - user.ID, "openai-compat", "gpt-4o-mini", true, @@ -797,23 +745,22 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenProviderDisabled(t *testi ) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai-compat", - DisplayName: "openai-compat", - APIKey: "", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: false, - CentralApiKeyEnabled: false, - AllowUserApiKey: true, - AllowCentralApiKeyFallback: false, + user, org, model := seedInternalChatDeps(t, db) + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai-compat", + DisplayName: "openai-compat", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }, func(p *database.InsertChatProviderParams) { + p.APIKey = "" + p.Enabled = false + p.CentralApiKeyEnabled = false + p.AllowUserApiKey = true + p.AllowCentralApiKeyFallback = false }) - require.NoError(t, err) + overrideModel := insertInternalChatModelConfigForProvider( - ctx, t, db, - user.ID, "openai-compat", "gpt-4o-mini", true, @@ -904,9 +851,9 @@ func TestCreateChildSubagentChat_OverrideWorksWhenParentHasNoModel(t *testing.T) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) overrideModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "override-no-parent-model-"+uuid.NewString(), true, + t, db, "override-no-parent-model-"+uuid.NewString(), true, ) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-no-model", @@ -936,9 +883,9 @@ func TestSpawnAgent_ExploreUsesConfiguredModelOverride(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) overrideModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "explore-override-"+uuid.NewString(), true, + t, db, "explore-override-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatExploreModelOverride(ctx, overrideModel.ID.String())) parentChat := createInternalParentChat( @@ -974,9 +921,9 @@ func TestSpawnAgent_ExploreFallsBackToCurrentTurnModel(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, parentModel := seedInternalChatDeps(ctx, t, db) + user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "explore-current-turn-"+uuid.NewString(), true, + t, db, "explore-current-turn-"+uuid.NewString(), true, ) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explore-fallback", @@ -1006,7 +953,7 @@ func TestCreateChat_ExploreRootStartsWithoutMCPSnapshot(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) root, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, @@ -1033,12 +980,12 @@ func TestResolveExploreToolSnapshot(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) approvedMCP := insertInternalMCPServerConfig( - ctx, t, db, user.ID, "approved-"+uuid.NewString(), true, + t, db, user.ID, "approved-"+uuid.NewString(), true, ) blockedMCP := insertInternalMCPServerConfig( - ctx, t, db, user.ID, "blocked-"+uuid.NewString(), false, + t, db, user.ID, "blocked-"+uuid.NewString(), false, ) askParentRef, err := server.CreateChat(ctx, CreateOptions{ @@ -1130,12 +1077,12 @@ func TestCreateChildSubagentChatWithOptions_ExplorePersistsMCPSnapshot(t *testin server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-explore-snapshot", ) mcpCfg := insertInternalMCPServerConfig( - ctx, t, db, user.ID, "snapshot-"+uuid.NewString(), false, + t, db, user.ID, "snapshot-"+uuid.NewString(), false, ) child, err := server.createChildSubagentChatWithOptions( @@ -1165,12 +1112,12 @@ func TestSpawnAgent_ExploreSnapshotsTurnStateParentState(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) turnStartConfig := insertInternalMCPServerConfig( - ctx, t, db, user.ID, "turn-start-"+uuid.NewString(), false, + t, db, user.ID, "turn-start-"+uuid.NewString(), false, ) mutatedConfig := insertInternalMCPServerConfig( - ctx, t, db, user.ID, "mutated-"+uuid.NewString(), true, + t, db, user.ID, "mutated-"+uuid.NewString(), true, ) parent, err := server.CreateChat(ctx, CreateOptions{ @@ -1246,9 +1193,9 @@ func TestSpawnAgent_ExploreFallsBackOnInvalidUUID(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, parentModel := seedInternalChatDeps(ctx, t, db) + user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "explore-invalid-override-"+uuid.NewString(), true, + t, db, "explore-invalid-override-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatExploreModelOverride(ctx, "not-a-uuid")) parentChat := createInternalParentChat( @@ -1278,12 +1225,12 @@ func TestSpawnAgent_ExploreFallsBackWhenOverrideIsUnavailable(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, parentModel := seedInternalChatDeps(ctx, t, db) + user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "explore-fallback-current-"+uuid.NewString(), true, + t, db, "explore-fallback-current-"+uuid.NewString(), true, ) disabledModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "explore-disabled-"+uuid.NewString(), false, + t, db, "explore-disabled-"+uuid.NewString(), false, ) require.NoError(t, db.UpsertChatExploreModelOverride(ctx, disabledModel.ID.String())) parentChat := createInternalParentChat( @@ -1313,35 +1260,25 @@ func TestSpawnAgent_ExploreFallsBackWhenOverrideCredentialsAreUnavailable(t *tes server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, parentModel := seedInternalChatDeps(ctx, t, db) + user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - ctx, t, db, user.ID, "explore-missing-user-key-current-"+uuid.NewString(), true, + t, db, "explore-missing-user-key-current-"+uuid.NewString(), true, ) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai-compat", - DisplayName: "OpenAI Compat", - APIKey: "", - BaseUrl: "", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: false, - AllowUserApiKey: true, - AllowCentralApiKeyFallback: false, + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai-compat", + DisplayName: "OpenAI Compat", + }, func(p *database.InsertChatProviderParams) { + p.APIKey = "" + p.CentralApiKeyEnabled = false + p.AllowUserApiKey = true + p.AllowCentralApiKeyFallback = false }) - require.NoError(t, err) - overrideModel, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai-compat", - Model: "gpt-4o-mini", - DisplayName: "Explore Override Missing User Key", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: false, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + + overrideModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai-compat", + Model: "gpt-4o-mini", + DisplayName: "Explore Override Missing User Key", }) - require.NoError(t, err) require.NoError(t, db.UpsertChatExploreModelOverride(ctx, overrideModel.ID.String())) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explore-missing-user-key", @@ -1373,7 +1310,7 @@ func TestSpawnAgent_DescriptionListsAllAvailableTypes(t *testing.T) { }) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-description-all", ) @@ -1395,7 +1332,7 @@ func TestSpawnAgent_DescriptionOmitsComputerUseWhenUnavailable(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-description-unavailable", ) @@ -1419,7 +1356,7 @@ func TestSpawnAgent_PlanModeDescriptionOmitsComputerUse(t *testing.T) { }) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, @@ -1455,7 +1392,7 @@ func TestSpawnAgent_PlanModeRejectsComputerUse(t *testing.T) { }) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, @@ -1499,7 +1436,7 @@ func TestSpawnAgent_InvalidTypeAndUnavailableTypeAreDistinct(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-invalid-type", ) @@ -1539,7 +1476,7 @@ func TestSpawnAgent_BlankTypeReturnsValidOptions(t *testing.T) { }) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-blank-type", ) @@ -1580,7 +1517,7 @@ func TestSpawnAgent_NotAvailableForChildChats(t *testing.T) { }) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) _, child := createParentChildChats(ctx, t, server, user, org, model) childChat, err := db.GetChatByID(ctx, child.ID) @@ -1608,7 +1545,7 @@ func TestSpawnAgent_NotAvailableForExploreChats(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) exploreChat, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, @@ -1662,9 +1599,9 @@ func TestSubagentLifecycleToolsIncludePersistedSubagentTypeAcrossVariants(t *tes server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) if tt.variant == subagentTypeComputerUse { - insertEnabledAnthropicProvider(ctx, t, db, user.ID) + insertEnabledAnthropicProvider(t, db, user.ID) } parentChat := createInternalParentChat( ctx, @@ -1687,7 +1624,7 @@ func TestSubagentLifecycleToolsIncludePersistedSubagentTypeAcrossVariants(t *tes require.NoError(t, err) setChatStatus(ctx, t, db, childID, database.ChatStatusWaiting, "") - insertAssistantMessage(ctx, t, db, childID, model.ID, "task complete") + insertAssistantMessage(t, db, childID, model.ID, "task complete") waitResult := requireToolResponseMap(t, runSubagentTool( ctx, t, @@ -1732,7 +1669,7 @@ func TestSubagentLifecycleToolErrorsIncludePersistedSubagentType(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) _, child := createParentChildChats(ctx, t, server, user, org, model) unrelated, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, @@ -1798,8 +1735,8 @@ func TestSpawnAgent_ComputerUseUsesComputerUseModelNotParent(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) - insertEnabledAnthropicProvider(ctx, t, db, user.ID) + user, org, model := seedInternalChatDeps(t, db) + insertEnabledAnthropicProvider(t, db, user.ID) workspace, build, agent := seedWorkspaceBinding(t, db, user.ID) require.Equal(t, "openai", model.Provider, "seed helper must create an OpenAI model") @@ -1855,23 +1792,16 @@ func TestSpawnAgent_ComputerUseInheritsMCPServerIDs(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) - insertEnabledAnthropicProvider(ctx, t, db, user.ID) + user, org, model := seedInternalChatDeps(t, db) + insertEnabledAnthropicProvider(t, db, user.ID) - mcpCfg, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "MCP Test", - Slug: "mcp-test", - Url: "https://mcp.example.com", - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + mcpCfg := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "MCP Test", + Slug: "mcp-test", + Url: "https://mcp.example.com", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) parentMCPIDs := []uuid.UUID{mcpCfg.ID} @@ -1912,39 +1842,25 @@ func TestCreateChildSubagentChat_InheritsMCPServerIDs(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) // Insert two MCP server configs so we can verify both are // inherited by the child chat. - mcpA, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "MCP A", - Slug: "mcp-a", - Url: "https://mcp-a.example.com", - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + mcpA := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "MCP A", + Slug: "mcp-a", + Url: "https://mcp-a.example.com", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - mcpB, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "MCP B", - Slug: "mcp-b", - Url: "https://mcp-b.example.com", - Transport: "streamable_http", - AuthType: "none", - Availability: "default_off", - Enabled: true, - ToolAllowList: []string{}, - ToolDenyList: []string{}, - CreatedBy: user.ID, - UpdatedBy: user.ID, + mcpB := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ + DisplayName: "MCP B", + Slug: "mcp-b", + Url: "https://mcp-b.example.com", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) parentMCPIDs := []uuid.UUID{mcpA.ID, mcpB.ID} @@ -1988,7 +1904,7 @@ func TestCreateChildSubagentChat_NoMCPServersStaysEmpty(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) // Create a parent chat without any MCP servers. parent, err := server.CreateChat(ctx, CreateOptions{ @@ -2025,7 +1941,7 @@ func TestIsSubagentDescendant(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) // Build a chain: root -> child -> grandchild. root, err := server.CreateChat(ctx, CreateOptions{ @@ -2225,7 +2141,6 @@ func setChatStatus( // insertAssistantMessage inserts an assistant message with v1 content // into a chat. func insertAssistantMessage( - ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, @@ -2238,26 +2153,14 @@ func insertAssistantMessage( data, err := json.Marshal(parts) require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chatID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - Content: []string{string(data)}, - ContentVersion: []int16{chatprompt.ContentVersionV1}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{}, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: data, Valid: true}, + ContentVersion: chatprompt.ContentVersionV1, }) - require.NoError(t, err) } func insertLinkedChatFile( @@ -2298,12 +2201,12 @@ func TestWaitAgentDoesNotRelayComputerUseSubagentAttachments(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, agent := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) parent, child := createComputerUseParentChild( - ctx, t, server, user, org, model, workspace, agent, + t, server, user, org, model, workspace, agent, "parent-relay", "child-relay", ) @@ -2318,7 +2221,7 @@ func TestWaitAgentDoesNotRelayComputerUseSubagentAttachments(t *testing.T) { "image/png", []byte("fake-png"), ) - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Shared the screenshot.") + insertAssistantMessage(t, db, child.ID, model.ID, "Shared the screenshot.") setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") resp, err := invokeWaitAgentTool(ctx, t, server, db, parent.ID, child.ID, 5) @@ -2366,7 +2269,7 @@ func TestWaitAgentDoesNotRelayRegularSubagentAttachments(t *testing.T) { db, ps := dbtestutil.NewDB(t) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) workspace, _, _ := seedWorkspaceBinding(t, db, user.ID) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) @@ -2384,7 +2287,7 @@ func TestWaitAgentDoesNotRelayRegularSubagentAttachments(t *testing.T) { "text/plain", []byte("release notes"), ) - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Shared the release notes.") + insertAssistantMessage(t, db, child.ID, model.ID, "Shared the release notes.") setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") resp, err := invokeWaitAgentTool(ctx, t, server, db, parent.ID, child.ID, 5) @@ -2422,8 +2325,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { // also use the mock clock. db, ps := dbtestutil.NewDB(t) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) - ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) t.Run("NotDescendant", func(t *testing.T) { t.Parallel() @@ -2453,7 +2355,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { parent, child := createParentChildChats(ctx, t, server, user, org, model) setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "task complete") + insertAssistantMessage(t, db, child.ID, model.ID, "task complete") gotChat, report, err := server.awaitSubagentCompletion( ctx, parent.ID, child.ID, time.Second, @@ -2471,7 +2373,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { parent, child := createParentChildChats(ctx, t, server, user, org, model) setChatStatus(ctx, t, db, child.ID, database.ChatStatusError, "something broke") - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "partial work done") + insertAssistantMessage(t, db, child.ID, model.ID, "partial work done") _, _, err := server.awaitSubagentCompletion( ctx, parent.ID, child.ID, time.Second, @@ -2504,7 +2406,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { mClock := quartz.NewMock(t) server := newInternalTestServerWithClock(t, db, nil, chatprovider.ProviderAPIKeys{}, mClock) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parent, child := createParentChildChats(ctx, t, server, user, org, model) @@ -2534,7 +2436,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { // Now set the state and advance the clock to the next // tick so the poll detects the transition. setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "poll result") + insertAssistantMessage(t, db, child.ID, model.ID, "poll result") mClock.Advance(subagentAwaitPollInterval).MustWait(ctx) result := testutil.RequireReceive(ctx, t, resultCh) @@ -2550,7 +2452,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { mClock := quartz.NewMock(t) server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parent, child := createParentChildChats(ctx, t, server, user, org, model) @@ -2612,7 +2514,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { // see done=true (Waiting) with an empty report. By // inserting the message first, the report is guaranteed // to be committed before the status makes it visible. - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "pubsub result") + insertAssistantMessage(t, db, child.ID, model.ID, "pubsub result") setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") require.EventuallyWithT(t, func(c *assert.CollectT) { chat, report, done, err := server.checkSubagentCompletion(ctx, child.ID) @@ -2661,7 +2563,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { mClock := quartz.NewMock(t) server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock) ctx := chatdTestContext(t) - user, org, model := seedInternalChatDeps(ctx, t, db) + user, org, model := seedInternalChatDeps(t, db) parent, child := createParentChildChats(ctx, t, server, user, org, model) @@ -2733,7 +2635,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { // Pre-complete the child so it returns immediately. setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "") - insertAssistantMessage(ctx, t, db, child.ID, model.ID, "zero timeout ok") + insertAssistantMessage(t, db, child.ID, model.ID, "zero timeout ok") gotChat, report, err := server.awaitSubagentCompletion( ctx, parent.ID, child.ID, 0, diff --git a/coderd/x/chatd/subagent_test.go b/coderd/x/chatd/subagent_test.go index 157b872262..a768f3487e 100644 --- a/coderd/x/chatd/subagent_test.go +++ b/coderd/x/chatd/subagent_test.go @@ -21,7 +21,7 @@ func TestSpawnComputerUseAgent_CreatesChildWithChatMode(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Create a parent chat. parent, err := server.CreateChat(ctx, chatd.CreateOptions{ @@ -77,7 +77,7 @@ func TestSpawnComputerUseAgent_SystemPromptFormat(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) parent, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -136,7 +136,7 @@ func TestSpawnComputerUseAgent_ChildIsListedUnderParent(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) parent, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -181,7 +181,7 @@ func TestSpawnComputerUseAgent_RootChatIDPropagation(t *testing.T) { db, ps := dbtestutil.NewDB(t) server := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Create a root parent chat (no parent of its own). parent, err := server.CreateChat(ctx, chatd.CreateOptions{ diff --git a/coderd/x/gitsync/worker_test.go b/coderd/x/gitsync/worker_test.go index 97ee0720c5..833ad5fae9 100644 --- a/coderd/x/gitsync/worker_test.go +++ b/coderd/x/gitsync/worker_test.go @@ -3,7 +3,6 @@ package gitsync_test import ( "context" "database/sql" - "encoding/json" "fmt" "sync" "sync/atomic" @@ -946,37 +945,22 @@ func TestWorker(t *testing.T) { org := dbgen.Organization(t, db, database.Organization{}) // 3. Set up FK chain: chat_providers -> chat_model_configs -> chats. - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - Enabled: true, - CentralApiKeyEnabled: true, - }) - require.NoError(t, err) + _ = dbgen.ChatProvider(t, db, database.ChatProvider{}) - modelCfg, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "test-model", - DisplayName: "Test Model", - Enabled: true, - ContextLimit: 100000, - CompressionThreshold: 70, - Options: json.RawMessage("{}"), + modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "test-model", + ContextLimit: 100000, }) - require.NoError(t, err) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: modelCfg.ID, Title: "integration-test", }) - require.NoError(t, err) // 4. Seed a stale diff status row so the worker picks it up. - _, err = db.UpsertChatDiffStatusReference(ctx, database.UpsertChatDiffStatusReferenceParams{ + _, err := db.UpsertChatDiffStatusReference(ctx, database.UpsertChatDiffStatusReferenceParams{ ChatID: chat.ID, GitBranch: "feature", GitRemoteOrigin: "https://github.com/o/r", diff --git a/enterprise/coderd/x/chatd/chatd_retry_test.go b/enterprise/coderd/x/chatd/chatd_retry_test.go index 3135796116..d21a15b9ba 100644 --- a/enterprise/coderd/x/chatd/chatd_retry_test.go +++ b/enterprise/coderd/x/chatd/chatd_retry_test.go @@ -126,8 +126,8 @@ func TestRelayReconnectUsesExponentialBackoff(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-backoff") + user, org, model := seedChatDependencies(t, db) + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-backoff") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -214,8 +214,8 @@ func TestRelayRepeatedDropsHitCap(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-drops") + user, org, model := seedChatDependencies(t, db) + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-drops") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -305,8 +305,8 @@ func TestRelayStopsAfterIntermittentCap(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-cap") + user, org, model := seedChatDependencies(t, db) + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-cap") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -416,8 +416,8 @@ func TestRelayReconnectStopsAfterDBErrorCap(t *testing.T) { failingDB.okRemain.Store(1) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, realDB) - chat := seedWaitingChat(ctx, t, realDB, org.ID, user, model, "relay-db-error") + user, org, model := seedChatDependencies(t, realDB) + chat := seedWaitingChat(t, realDB, org.ID, user, model, "relay-db-error") subscriber := newTestServer(t, failingDB, ps, subscriberID, dialer, mclk) _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) @@ -509,8 +509,8 @@ func TestRelayStopsImmediatelyOnUnauthorized(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, + user, org, model := seedChatDependencies(t, db) + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-unrec-"+tc.name) _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) @@ -584,8 +584,8 @@ func TestRelayBackoffResetsOnStatusChange(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-reset-on-status") + user, org, model := seedChatDependencies(t, db) + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-reset-on-status") _, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -661,8 +661,8 @@ func TestRelayBackoffRespectsContextCancel(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-cancel") + user, org, model := seedChatDependencies(t, db) + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-cancel") subCtx, subCancel := context.WithCancel(ctx) _, events, cancel, ok := subscriber.Subscribe(subCtx, chat.ID, nil, 0) @@ -740,11 +740,11 @@ func TestDialRelayReal401(t *testing.T) { subscribeFn := entchatd.NewMultiReplicaSubscribeFn(cfg) ctx := testutil.Context(t, testutil.WaitMedium) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Seed a waiting chat - no sync dial - then push a running // status notification to trigger the async dial via the real // dialRelay path. - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-real-401") + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-real-401") statusCh := make(chan osschatd.StatusNotification, 1) evs := subscribeFn(ctx, osschatd.SubscribeFnParams{ diff --git a/enterprise/coderd/x/chatd/chatd_test.go b/enterprise/coderd/x/chatd/chatd_test.go index 86a10e2ad0..37d30e23c2 100644 --- a/enterprise/coderd/x/chatd/chatd_test.go +++ b/enterprise/coderd/x/chatd/chatd_test.go @@ -91,7 +91,6 @@ func newActiveWorkerServer( // seedChatDependencies creates a user, organization, and chat model // config in the database for use in relay tests. func seedChatDependencies( - ctx context.Context, t *testing.T, db database.Store, ) (database.User, database.Organization, database.ChatModelConfig) { @@ -110,35 +109,19 @@ func seedChatDependencies( UserID: user.ID, OrganizationID: org.ID, }) - _, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - APIKey: "test-key", - BaseUrl: safetyNet.URL, - CentralApiKeyEnabled: true, - ApiKeyKeyID: sql.NullString{}, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + BaseUrl: safetyNet.URL, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + IsDefault: true, }) - require.NoError(t, err) return user, org, model } func seedWaitingChat( - ctx context.Context, t *testing.T, db database.Store, orgID uuid.UUID, @@ -148,16 +131,12 @@ func seedWaitingChat( ) database.Chat { t.Helper() - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: orgID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: user.ID, LastModelConfigID: model.ID, Title: title, - MCPServerIDs: []uuid.UUID{}, }) - require.NoError(t, err) return chat } @@ -173,7 +152,7 @@ func seedRemoteRunningChat( ) database.Chat { t.Helper() - chat := seedWaitingChat(ctx, t, db, orgID, user, model, title) + chat := seedWaitingChat(t, db, orgID, user, model, title) now := time.Now() chat, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ ID: chat.ID, @@ -258,7 +237,7 @@ func TestSubscribeRelayReconnectsOnDrop(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-reconnect") @@ -336,11 +315,11 @@ func TestSubscribeRelayAsyncDoesNotBlock(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Seed a waiting chat so Subscribe does not trigger a synchronous // relay. - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-async-nonblock") + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-async-nonblock") // Subscribe before the chat is marked running so the relay opens // via pubsub notification (openRelayAsync path). @@ -438,7 +417,7 @@ func TestSubscribeRelaySnapshotDelivered(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-snapshot") @@ -526,7 +505,7 @@ func TestSubscribeRetryEventAcrossInstances(t *testing.T) { }, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) chat, err := worker.CreateChat(ctx, osschatd.CreateOptions{ @@ -663,11 +642,11 @@ func TestSubscribeRelayStaleDialDiscardedAfterInterrupt(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Seed the chat in waiting state so Subscribe does not try an initial // relay. - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "stale-dial-test") + chat := seedWaitingChat(t, db, org.ID, user, model, "stale-dial-test") // Subscribe while chat is in "waiting" state — no relay opened. _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) @@ -815,11 +794,11 @@ func TestSubscribeCancelDuringInFlightDial(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Seed the chat in waiting state so Subscribe does not open a // synchronous relay. - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "cancel-inflight-dial") + chat := seedWaitingChat(t, db, org.ID, user, model, "cancel-inflight-dial") _, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -901,10 +880,10 @@ func TestSubscribeRelayRunningToRunningSwitch(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Seed the chat in waiting state so Subscribe does not open a relay. - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "running-to-running") + chat := seedWaitingChat(t, db, org.ID, user, model, "running-to-running") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -1009,11 +988,11 @@ func TestSubscribeRelayFailedDialRetries(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) // Seed the chat in waiting state so Subscribe does not open a // synchronous relay dial. - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "failed-dial-retry") + chat := seedWaitingChat(t, db, org.ID, user, model, "failed-dial-retry") _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) @@ -1105,7 +1084,7 @@ func TestSubscribeRunningLocalWorkerClosesRelay(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat := seedRemoteRunningChat( ctx, @@ -1205,7 +1184,7 @@ func TestSubscribeRelayMultipleReconnects(t *testing.T) { subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) chat := seedRemoteRunningChat( ctx, @@ -1349,13 +1328,13 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) { }, nil) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // Create the chat in waiting state so the subscriber sees it // before the worker picks it up (avoids the synchronous relay // path in Subscribe). - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "fast-completion-relay-race") + chat := seedWaitingChat(t, db, org.ID, user, model, "fast-completion-relay-race") // Subscribe from the subscriber replica while the chat is idle. // No relay is opened because the chat is in waiting state. @@ -1505,10 +1484,10 @@ func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) { }, subscriberClock) ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-drain-characterization") + chat := seedWaitingChat(t, db, org.ID, user, model, "relay-drain-characterization") // Attach before processing so the relay opens as soon as // status=running arrives. @@ -1699,11 +1678,11 @@ func TestSubscribeRelayEstablishedMidStream(t *testing.T) { // call) involves multiple DB round-trips that can be slow under // load. ctx := testutil.Context(t, testutil.WaitSuperLong) - user, org, model := seedChatDependencies(ctx, t, db) + user, org, model := seedChatDependencies(t, db) setOpenAIProviderBaseURL(ctx, t, db, openAIURL) // Create the chat in waiting state. - chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "mid-stream-relay") + chat := seedWaitingChat(t, db, org.ID, user, model, "mid-stream-relay") // Subscribe from the subscriber replica while the chat is idle. _, events, subCancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) diff --git a/enterprise/coderd/x/chatd/usagelimit_test.go b/enterprise/coderd/x/chatd/usagelimit_test.go index aeea060725..9f44bfa07c 100644 --- a/enterprise/coderd/x/chatd/usagelimit_test.go +++ b/enterprise/coderd/x/chatd/usagelimit_test.go @@ -1,11 +1,13 @@ package chatd_test import ( + "database/sql" "encoding/json" "testing" "time" "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" "github.com/coder/coder/v2/coderd/database" @@ -86,28 +88,14 @@ func TestResolveUsageLimitStatus_OrgScoped(t *testing.T) { require.NoError(t, err) // We need a chat provider + model config for inserting chats. - _, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "openai", - APIKey: "test-key", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - CentralApiKeyEnabled: true, + _ = dbgen.ChatProvider(t, db, database.ChatProvider{ + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) - modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ - Provider: "openai", - Model: "gpt-4o-mini", - DisplayName: "Test Model", - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - Enabled: true, - IsDefault: true, - ContextLimit: 128000, - CompressionThreshold: 70, - Options: json.RawMessage(`{}`), + modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + IsDefault: true, }) - require.NoError(t, err) now := time.Now().UTC() @@ -115,38 +103,25 @@ func TestResolveUsageLimitStatus_OrgScoped(t *testing.T) { // given org and inserts a single message with the specified cost. insertChatWithSpend := func(t *testing.T, ownerID, orgID, modelCfgID uuid.UUID, costMicros int64) { t.Helper() - tctx := testutil.Context(t, testutil.WaitLong) - c, err := db.InsertChat(tctx, database.InsertChatParams{ + c := dbgen.Chat(t, db, database.Chat{ OrganizationID: orgID, OwnerID: ownerID, LastModelConfigID: modelCfgID, Title: "test chat", - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, - MCPServerIDs: []uuid.UUID{}, }) - require.NoError(t, err) - _, err = db.InsertChatMessages(tctx, database.InsertChatMessagesParams{ - ChatID: c.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{modelCfgID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - Content: []string{`[{"type":"text","text":"hello"}]`}, - ContentVersion: []int16{1}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{100}, - OutputTokens: []int64{50}, - TotalTokens: []int64{150}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{128000}, - Compressed: []bool{false}, - TotalCostMicros: []int64{costMicros}, - RuntimeMs: []int64{500}, - ProviderResponseID: []string{uuid.NewString()}, + _ = dbgen.ChatMessage(t, db, database.ChatMessage{ + ChatID: c.ID, + ModelConfigID: uuid.NullUUID{UUID: modelCfgID, Valid: true}, + Role: database.ChatMessageRoleAssistant, + Content: pqtype.NullRawMessage{RawMessage: json.RawMessage(`[{"type":"text","text":"hello"}]`), Valid: true}, + InputTokens: sql.NullInt64{Int64: 100, Valid: true}, + OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, + TotalTokens: sql.NullInt64{Int64: 150, Valid: true}, + ContextLimit: sql.NullInt64{Int64: 128000, Valid: true}, + TotalCostMicros: sql.NullInt64{Int64: costMicros, Valid: true}, + RuntimeMs: sql.NullInt64{Int64: 500, Valid: true}, + ProviderResponseID: sql.NullString{String: uuid.NewString(), Valid: true}, }) - require.NoError(t, err) } t.Run("OrgA_gets_orgA_limit", func(t *testing.T) { diff --git a/enterprise/dbcrypt/dbcrypt_internal_test.go b/enterprise/dbcrypt/dbcrypt_internal_test.go index abece5e7e7..f6d24270d7 100644 --- a/enterprise/dbcrypt/dbcrypt_internal_test.go +++ b/enterprise/dbcrypt/dbcrypt_internal_test.go @@ -926,30 +926,19 @@ func TestMCPServerConfigs(t *testing.T) { apiKeyValue = "my-api-key" customHeaders = `{"X-Custom":"header-value"}` ) - // insertConfig is a small helper that creates a user and an MCP - // server config through the encrypted store, returning both. + // insertConfig is a small helper that creates an MCP server + // config through the encrypted store with secret fields set. insertConfig := func(t *testing.T, crypt *dbCrypt, ciphers []Cipher) database.MCPServerConfig { t.Helper() - user := dbgen.User(t, crypt, database.User{}) - cfg, err := crypt.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Test MCP Server", - Slug: "test-mcp-" + uuid.New().String()[:8], + cfg := dbgen.MCPServerConfig(t, crypt, database.MCPServerConfig{ Description: "test description", - Url: "https://mcp.example.com", - Transport: "streamable_http", AuthType: "oauth2", OAuth2ClientID: "client-id", OAuth2ClientSecret: oauthSecret, APIKeyValue: apiKeyValue, CustomHeaders: customHeaders, - ToolAllowList: []string{}, - ToolDenyList: []string{}, Availability: "force_on", - Enabled: true, - CreatedBy: user.ID, - UpdatedBy: user.ID, }) - require.NoError(t, err) requireMCPServerConfigDecrypted(t, cfg, ciphers, oauthSecret, apiKeyValue, customHeaders) return cfg } @@ -1084,20 +1073,12 @@ func TestMCPServerUserTokens(t *testing.T) { ) (database.MCPServerConfig, database.MCPServerUserToken) { t.Helper() user := dbgen.User(t, crypt, database.User{}) - cfg, err := crypt.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{ - DisplayName: "Token Test MCP", - Slug: "tok-mcp-" + uuid.New().String()[:8], - Url: "https://mcp.example.com", - Transport: "streamable_http", - AuthType: "oauth2", - ToolAllowList: []string{}, - ToolDenyList: []string{}, - Availability: "default_off", - Enabled: true, - CreatedBy: user.ID, - UpdatedBy: user.ID, + cfg := dbgen.MCPServerConfig(t, crypt, database.MCPServerConfig{ + DisplayName: "Token Test MCP", + AuthType: "oauth2", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) - require.NoError(t, err) tok, err := crypt.UpsertMCPServerUserToken(ctx, database.UpsertMCPServerUserTokenParams{ MCPServerConfigID: cfg.ID, @@ -1196,14 +1177,11 @@ func TestUserChatProviderKeys(t *testing.T) { ) (database.ChatProvider, database.UserChatProviderKey) { t.Helper() user := dbgen.User(t, crypt, database.User{}) - provider, err := crypt.InsertChatProvider(ctx, database.InsertChatProviderParams{ - Provider: "openai", - DisplayName: "OpenAI", - APIKey: "", - Enabled: true, + provider := dbgen.ChatProvider(t, crypt, database.ChatProvider{ AllowUserApiKey: true, + }, func(params *database.InsertChatProviderParams) { + params.APIKey = "" }) - require.NoError(t, err) key, err := crypt.UpsertUserChatProviderKey(ctx, database.UpsertUserChatProviderKeyParams{ UserID: user.ID,