mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
refactor: add dbgen chat generators and migrate test boilerplate (#24497)
- Adds chat-related dbgen generators covering defaults, overrides, and message field mapping. - Replaces raw single-row chat, message, provider, and model-config setup in tests with dbgen helpers. - Simplifies chat seed helpers after moving fixture setup into dbgen. > Generated with [Coder Agents](https://coder.com/agents).
This commit is contained in:
@@ -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()),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
+206
-522
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
}
|
||||
|
||||
+137
-130
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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:
|
||||
|
||||
+278
-566
File diff suppressed because it is too large
Load Diff
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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{})
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user