diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index c36aea4efb..49f18e82af 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -12493,6 +12493,162 @@ func TestAdvisorChainMode_SnapshotKeepsFullHistory(t *testing.T) { "advisor snapshot must retain the turn 1 assistant message even when chain mode is active") } +// TestProviderSwitchSanitizesAndRestoresPEToolHistory verifies the A→B→A +// provider-switch contract: +// +// 1. A turn using model MA (backed by provider A) produces a +// provider-executed (PE) tool call in the DB. +// 2. A subsequent turn using model MB (backed by provider B) does NOT +// send that PE tool call to provider B. +// 3. A further turn back to MA sends the PE tool call to provider A again. +// 4. The DB row is never mutated; the filter is read-time only. +func TestProviderSwitchSanitizesAndRestoresPEToolHistory(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + 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, + }) + + // Given: two AI providers A and B + const peToolCallID = "pe_switch_test_id" + + chanA := make(chan string, 4) + chanB := make(chan string, 4) + + serverAURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse(`{"title":"switch-test"}`) + } + chanA <- string(req.RawBody) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("answer from A")...) + }) + serverBURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse(`{"title":"switch-test"}`) + } + chanB <- string(req.RawBody) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("answer from B")...) + }) + cpA := dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai-compat", + BaseUrl: serverAURL, + }) + cpB := dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai-compat", + BaseUrl: serverBURL, + }) + + mA := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai-compat", + Model: "gpt-4o-mini", + DisplayName: "Model A", + Enabled: true, + AIProviderID: uuid.NullUUID{UUID: cpA.ID, Valid: true}, + }) + mB := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai-compat", + Model: "gpt-4o-mini", + DisplayName: "Model B", + Enabled: true, + AIProviderID: uuid.NullUUID{UUID: cpB.ID, Valid: true}, + }) + + server := newActiveTestServer(t, db, ps) + + // Given: an initial conversation turn with model A that produces provider-executed + // tool call results + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "provider-switch-test", + ModelConfigID: mA.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertChatMessageParts(ctx, t, db, chat.ID, database.ChatMessageRoleAssistant, mA.ID, uuid.Nil, + []codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: peToolCallID, + ToolName: "web_search", + Args: json.RawMessage(`{"query":"coder"}`), + ProviderExecuted: true, + }, + { + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: peToolCallID, + ToolName: "web_search", + Result: json.RawMessage(`"search results"`), + ProviderExecuted: true, + }, + codersdk.ChatMessageText("here is the answer"), + }, + ) + + // When: a conversation turn is executed with model B + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: mB.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue with B")}, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + // When: a further conversation turn is executed with model A again + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: mA.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("back to A")}, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + // Then: the provider-executed tool call results should still be in the database + allMessages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + }) + require.NoError(t, err) + var peRowFound bool + for _, msg := range allMessages { + if msg.Role != database.ChatMessageRoleAssistant || msg.ModelConfigID.UUID != mA.ID { + continue + } + parts, parseErr := chatprompt.ParseContent(msg) + require.NoError(t, parseErr) + for _, p := range parts { + if p.ProviderExecuted && p.ToolCallID == peToolCallID { + peRowFound = true + } + } + } + require.True(t, peRowFound, "PE tool call must still be in the DB after provider switches") + + // Skip initial generation request + _ = testutil.TryReceive(ctx, t, chanA) + + // Then: PE tool call ID from A must not appear in the request to provider B + turn2Body := testutil.TryReceive(ctx, t, chanB) + require.NotContains(t, turn2Body, peToolCallID, + "provider B must not receive the PE tool call from provider A") + + // Then: PE tool call ID must appear in the second request to provider A + turn3Body := testutil.TryReceive(ctx, t, chanA) + require.Contains(t, turn3Body, peToolCallID, + "provider A must receive its own PE tool call when switching back") +} + func seedAdvisorConfig( ctx context.Context, t *testing.T, diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index dee28e1e66..b4e107dcfc 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -216,6 +216,13 @@ func (server *Server) prepareGeneration( planPathBlock string ) + // Drop provider-executed tool history produced by a different provider + // before building the prompt. A provider that shares another's wire format + // (e.g. Bedrock and Anthropic) can still reject the other's + // provider-executed blocks, so a mid-chat provider switch must not replay + // them. + promptRows = server.sanitizeForeignProviderExecutedToolRows(ctx, logger, promptRows, modelConfig.ID) + if chat.WorkspaceID.Valid { // Resolve the workspace agent so the chat row's AgentID and // BuildID bindings are up to date before the chatworker diff --git a/coderd/x/chatd/provider_switch_sanitize.go b/coderd/x/chatd/provider_switch_sanitize.go new file mode 100644 index 0000000000..3b300003f1 --- /dev/null +++ b/coderd/x/chatd/provider_switch_sanitize.go @@ -0,0 +1,153 @@ +package chatd + +import ( + "context" + + "github.com/google/uuid" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/codersdk" +) + +// providerSwitchStripStats counts provider-executed tool history removed +// during a provider switch. +type providerSwitchStripStats struct { + RemovedToolCalls int + RemovedToolResults int + DroppedMessages int +} + +// modelConfigProviderIdentity returns a stable identity for the upstream provider +// behind a model config. When the config has an AIProviderID (the modern path), +// the identity is the provider instance UUID, so two providers of the same type +// (e.g. two openai-compat providers at different base URLs) are distinguished. +// When AIProviderID is invalid (legacy configs with no provider row), the +// identity falls back to the normalized provider type name. +func modelConfigProviderIdentity(modelConfig database.ChatModelConfig, normalizedProvider string) string { + if modelConfig.AIProviderID.Valid { + return modelConfig.AIProviderID.UUID.String() + } + return normalizedProvider +} + +// stripForeignProviderExecutedToolRows drops provider-executed tool blocks +// (calls and results) from assistant rows whose producing provider differs +// from targetIdentity. Rows with an unknown origin are treated as foreign +// (fail closed). Rows emptied by stripping are dropped; rows that fail to parse +// or re-marshal are kept unchanged. +// +// See modelConfigProviderIdentity for how identity is derived. +func stripForeignProviderExecutedToolRows( + rows []database.ChatMessage, + targetIdentity string, + originProvider func(uuid.NullUUID) (string, bool), +) ([]database.ChatMessage, providerSwitchStripStats) { + var stats providerSwitchStripStats + if targetIdentity == "" || len(rows) == 0 { + return rows, stats + } + + out := make([]database.ChatMessage, 0, len(rows)) + for _, row := range rows { + if row.Role != database.ChatMessageRoleAssistant { + out = append(out, row) + continue + } + if origin, ok := originProvider(row.ModelConfigID); ok && origin == targetIdentity { + out = append(out, row) + continue + } + + parts, err := chatprompt.ParseContent(row) + if err != nil { + out = append(out, row) + continue + } + + kept := make([]codersdk.ChatMessagePart, 0, len(parts)) + var removedCalls, removedResults int + for _, part := range parts { + switch { + case part.Type == codersdk.ChatMessagePartTypeToolCall && part.ProviderExecuted: + removedCalls++ + case part.Type == codersdk.ChatMessagePartTypeToolResult && part.ProviderExecuted: + removedResults++ + default: + kept = append(kept, part) + } + } + if removedCalls == 0 && removedResults == 0 { + out = append(out, row) + continue + } + stats.RemovedToolCalls += removedCalls + stats.RemovedToolResults += removedResults + if len(kept) == 0 { + stats.DroppedMessages++ + continue + } + + content, err := chatprompt.MarshalParts(kept) + if err != nil { + out = append(out, row) + continue + } + row.Content = content + row.ContentVersion = chatprompt.CurrentContentVersion + out = append(out, row) + } + return out, stats +} + +func (server *Server) sanitizeForeignProviderExecutedToolRows( + ctx context.Context, + logger slog.Logger, + rows []database.ChatMessage, + modelConfigID uuid.UUID, +) []database.ChatMessage { + targetCfg, targetProvider, err := server.resolveModelConfigAndNormalizedProvider(ctx, modelConfigID) + if err != nil || targetProvider == "" { + logger.Debug(ctx, "skipping provider-switch sanitization: target provider unresolved", + slog.F("model_config_id", modelConfigID), + slog.Error(err), + ) + return rows + } + targetIdentity := modelConfigProviderIdentity(targetCfg, targetProvider) + + cache := make(map[uuid.UUID]string) + originProvider := func(id uuid.NullUUID) (string, bool) { + if !id.Valid { + return "", false + } + if identity, seen := cache[id.UUID]; seen { + return identity, identity != "" + } + originCfg, provider, rErr := server.resolveModelConfigAndNormalizedProvider(ctx, id.UUID) + if rErr != nil { + logger.Debug(ctx, "provider-switch sanitization: origin provider unresolved, treating as foreign", + slog.F("model_config_id", id.UUID), + slog.Error(rErr), + ) + cache[id.UUID] = "" + return "", false + } + identity := modelConfigProviderIdentity(originCfg, provider) + cache[id.UUID] = identity + return identity, identity != "" + } + + sanitized, stats := stripForeignProviderExecutedToolRows(rows, targetIdentity, originProvider) + if stats != (providerSwitchStripStats{}) { + logger.Debug(ctx, "stripped foreign provider-executed tool history", + slog.F("phase", "provider_switch"), + slog.F("target_provider_identity", targetIdentity), + slog.F("removed_tool_calls", stats.RemovedToolCalls), + slog.F("removed_tool_results", stats.RemovedToolResults), + slog.F("dropped_messages", stats.DroppedMessages), + ) + } + return sanitized +} diff --git a/coderd/x/chatd/provider_switch_sanitize_internal_test.go b/coderd/x/chatd/provider_switch_sanitize_internal_test.go new file mode 100644 index 0000000000..5d9adc74e4 --- /dev/null +++ b/coderd/x/chatd/provider_switch_sanitize_internal_test.go @@ -0,0 +1,247 @@ +package chatd + +import ( + "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/x/chatd/chatprompt" + "github.com/coder/coder/v2/codersdk" +) + +func TestStripForeignProviderExecutedToolRows(t *testing.T) { + t.Parallel() + + const ( + anthropic = "anthropic" + bedrock = "bedrock" + openai = "openai" + ) + + anthropicCfg := uuid.New() + openAICfg := uuid.New() + unknownCfg := uuid.New() + + vllmProviderID := uuid.New() + togetherProviderID := uuid.New() + vllmCfg := uuid.New() + + peCall := func(id string) codersdk.ChatMessagePart { + p := codersdk.ChatMessageToolCall(id, "web_search", json.RawMessage(`{"query":"x"}`)) + p.ProviderExecuted = true + return p + } + peResult := func(id string) codersdk.ChatMessagePart { + p := codersdk.ChatMessageToolResult(id, "web_search", json.RawMessage(`{"ok":true}`), false, false) + p.ProviderExecuted = true + return p + } + localCall := func(id string) codersdk.ChatMessagePart { + return codersdk.ChatMessageToolCall(id, "read_file", json.RawMessage(`{}`)) + } + text := codersdk.ChatMessageText + + assistantRow := func(t *testing.T, cfg uuid.UUID, parts ...codersdk.ChatMessagePart) database.ChatMessage { + t.Helper() + content, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return database.ChatMessage{ + Role: database.ChatMessageRoleAssistant, + ModelConfigID: uuid.NullUUID{UUID: cfg, Valid: cfg != uuid.Nil}, + Content: content, + ContentVersion: chatprompt.ContentVersionV1, + } + } + userRow := func(t *testing.T, s string) database.ChatMessage { + t.Helper() + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{text(s)}) + require.NoError(t, err) + return database.ChatMessage{ + Role: database.ChatMessageRoleUser, + Content: content, + ContentVersion: chatprompt.ContentVersionV1, + } + } + + origin := func(providerByConfig map[uuid.UUID]string) func(uuid.NullUUID) (string, bool) { + return func(id uuid.NullUUID) (string, bool) { + if !id.Valid { + return "", false + } + provider, ok := providerByConfig[id.UUID] + return provider, ok + } + } + resolver := origin(map[uuid.UUID]string{ + anthropicCfg: anthropic, + openAICfg: openai, + vllmCfg: vllmProviderID.String(), + }) + + partsOf := func(t *testing.T, row database.ChatMessage) []codersdk.ChatMessagePart { + t.Helper() + parts, err := chatprompt.ParseContent(row) + require.NoError(t, err) + return parts + } + + t.Run("same provider kept", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + userRow(t, "hi"), + assistantRow(t, anthropicCfg, peCall("ws"), peResult("ws"), text("done")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, anthropic, resolver) + require.Equal(t, rows, got) + require.Zero(t, stats) + }) + + t.Run("anthropic to bedrock drops provider blocks", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + userRow(t, "hi"), + assistantRow(t, anthropicCfg, peCall("ws"), peResult("ws"), text("done")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, bedrock, resolver) + require.Len(t, got, 2) + require.Equal(t, []codersdk.ChatMessagePart{text("done")}, partsOf(t, got[1])) + require.Equal(t, providerSwitchStripStats{RemovedToolCalls: 1, RemovedToolResults: 1}, stats) + }) + + t.Run("foreign-only row dropped", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + userRow(t, "hi"), + assistantRow(t, anthropicCfg, peCall("ws")), + userRow(t, "again"), + } + got, stats := stripForeignProviderExecutedToolRows(rows, bedrock, resolver) + require.Len(t, got, 2) + require.Equal(t, database.ChatMessageRoleUser, got[0].Role) + require.Equal(t, database.ChatMessageRoleUser, got[1].Role) + require.Equal(t, providerSwitchStripStats{RemovedToolCalls: 1, DroppedMessages: 1}, stats) + }) + + t.Run("multi-provider keeps native strips foreign", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + assistantRow(t, openAICfg, peCall("os"), peResult("os"), text("openai")), + assistantRow(t, anthropicCfg, peCall("as"), peResult("as"), text("anthropic")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, anthropic, resolver) + require.Len(t, got, 2) + require.Equal(t, []codersdk.ChatMessagePart{text("openai")}, partsOf(t, got[0])) + require.Equal(t, rows[1], got[1]) + require.Equal(t, providerSwitchStripStats{RemovedToolCalls: 1, RemovedToolResults: 1}, stats) + }) + + t.Run("non-provider-executed parts untouched", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + assistantRow(t, anthropicCfg, text("hello"), localCall("local")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, bedrock, resolver) + require.Equal(t, rows, got) + require.Zero(t, stats) + }) + + t.Run("empty target is a no-op", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + assistantRow(t, anthropicCfg, peCall("ws"), peResult("ws")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, "", resolver) + require.Equal(t, rows, got) + require.Zero(t, stats) + }) + + t.Run("unknown origin fails closed", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + assistantRow(t, unknownCfg, peResult("ws"), text("done")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, bedrock, resolver) + require.Len(t, got, 1) + require.Equal(t, []codersdk.ChatMessagePart{text("done")}, partsOf(t, got[0])) + require.Equal(t, providerSwitchStripStats{RemovedToolResults: 1}, stats) + }) + + t.Run("unparsable foreign row kept unchanged", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{{ + Role: database.ChatMessageRoleAssistant, + ModelConfigID: uuid.NullUUID{UUID: anthropicCfg, Valid: true}, + Content: pqtype.NullRawMessage{RawMessage: []byte("{not json"), Valid: true}, + ContentVersion: chatprompt.ContentVersionV1, + }} + got, stats := stripForeignProviderExecutedToolRows(rows, bedrock, resolver) + require.Equal(t, rows, got) + require.Zero(t, stats) + }) + + t.Run("same type different instance drops provider blocks", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + userRow(t, "hi"), + assistantRow(t, vllmCfg, peCall("ws"), peResult("ws"), text("done")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, togetherProviderID.String(), resolver) + require.Len(t, got, 2) + require.Equal(t, []codersdk.ChatMessagePart{text("done")}, partsOf(t, got[1])) + require.Equal(t, providerSwitchStripStats{RemovedToolCalls: 1, RemovedToolResults: 1}, stats) + }) + + t.Run("same instance keeps provider blocks", func(t *testing.T) { + t.Parallel() + rows := []database.ChatMessage{ + userRow(t, "hi"), + assistantRow(t, vllmCfg, peCall("ws"), peResult("ws"), text("done")), + } + got, stats := stripForeignProviderExecutedToolRows(rows, vllmProviderID.String(), resolver) + require.Equal(t, rows, got) + require.Zero(t, stats) + }) +} + +func TestModelConfigProviderIdentity(t *testing.T) { + t.Parallel() + + providerID := uuid.New() + + t.Run("AIProviderID valid returns provider UUID", func(t *testing.T) { + t.Parallel() + cfg := database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + } + got := modelConfigProviderIdentity(cfg, "openai-compat") + require.Equal(t, providerID.String(), got) + }) + + t.Run("AIProviderID invalid falls back to normalized type", func(t *testing.T) { + t.Parallel() + cfg := database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{}, + } + got := modelConfigProviderIdentity(cfg, "anthropic") + require.Equal(t, "anthropic", got) + }) + + t.Run("same type different provider IDs are distinguished", func(t *testing.T) { + t.Parallel() + otherProviderID := uuid.New() + cfgA := database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + } + cfgB := database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: otherProviderID, Valid: true}, + } + require.NotEqual(t, + modelConfigProviderIdentity(cfgA, "openai-compat"), + modelConfigProviderIdentity(cfgB, "openai-compat"), + ) + }) +}