mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
fix(coderd/x/chatd): drop foreign provider-executed tools on model switch (#26555)
Drops provider-executed tool history (calls and results) from assistant rows whose producing provider ID differs from the target turn's provider ID, before the prompt is built. Same-provider history is left untouched, so normal `web_search` replay is unaffected. - When a model config has an `AIProviderID`, use this as identity so two providers of the same type (e.g. two `openai-compat` providers at different base URLs) are correctly distinguished. Falls back to the normalized provider type name. - Sanitization runs at the `database.ChatMessage` row level in `prepareGeneration`. The `chatloop` pre-request and reload paths are untouched. - Foreign provider-executed results are dropped and not converted to text. - Unknown origin (unresolvable `ModelConfigID`) fails closed (strip). - Adds tests for the pure `stripForeignProviderExecutedToolRows`. - Adds unit tests for `modelConfigProviderIdentity`. _This pull request was created by Coder Agents on behalf of @johnstcn._
This commit is contained in:
@@ -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")
|
"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(
|
func seedAdvisorConfig(
|
||||||
ctx context.Context,
|
ctx context.Context,
|
||||||
t *testing.T,
|
t *testing.T,
|
||||||
|
|||||||
@@ -216,6 +216,13 @@ func (server *Server) prepareGeneration(
|
|||||||
planPathBlock string
|
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 {
|
if chat.WorkspaceID.Valid {
|
||||||
// Resolve the workspace agent so the chat row's AgentID and
|
// Resolve the workspace agent so the chat row's AgentID and
|
||||||
// BuildID bindings are up to date before the chatworker
|
// BuildID bindings are up to date before the chatworker
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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"),
|
||||||
|
)
|
||||||
|
})
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user