mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
Michael Suchacz
c17bed25e0
Wires chat lifecycle hooks into chatd, gated by the `agent-lifecycle-hooks` experiment. Part of the lifecycle hooks stack (#27401, #27428, #27430). See `docs/admin/setup/chat-lifecycle-hooks.md` for the consumer-facing contract. ## Summary When a hook URL is configured, chatd dispatches `session_start`, `user_prompt_submit`, `pre_tool_use`, `post_tool_use`, `pre_compact`, `post_compact`, and `stop` events to the consumer and applies its responses. ## Design - **Stateless**: Coder stores no hook dispatch or decision state. Delivery is at least once; consumers deduplicate on stable payload identifiers (chat ID, event type, tool-use ID) and answer duplicates with the same decision. - **Admission-time prompt effects**: `user_prompt_submit` dispatches exactly once per submission (create, send, queue, edit, subagent spawn) and folds its effects into the stored prompt as typed message parts: original-or-overridden user parts, then model-only `hook-context`, then a user-visible `hook-notice`. Hook context is stripped from every client-facing conversion; hook notices are excluded from model prompts. The server rejects hook parts in client-submitted content. - **Tool gating**: `pre_tool_use` allow can override tool input; deny becomes a synthetic denied tool result, with any returned model context persisted as a model-only transcript row so it never reaches clients. The denial text identifies an external policy (the deployment's lifecycle hook) as the source and marks the decision as persistent, so the model explains the denial instead of retrying it or misreporting it as an infrastructure failure. - **Fail closed**: a dispatch failure rejects the triggering request or moves the chat to the error state in the same transaction as the affected step, so a runnable state is never published with unapproved content. - **Admission before persistence**: `pre_tool_use` is dispatched for the calls the model produced, before the assistant message is stored. See "Staged tool admission" below. - **Fresh dispatch per tool call**: every non-provider-executed tool call is decided by its own `pre_tool_use` dispatch; Coder never reuses an earlier decision on the consumer's behalf. Retries re-dispatch the same logical event. ## Structure All hook dispatch flows through one seam: entry points build a `chathooks.Chat` (chat identity) and a `chathooks.Message` (event details) and call `Trigger.Trigger`, the only component that talks to the dispatcher. The integration lives in the `coderd/x/chatd/chathooks` subpackage, split by responsibility: - `trigger.go`: the trigger seam; builds the wire envelope per event, normalizes deny into a typed error, and holds the package's single enabled-check. - `effects.go`: pure conversion of hook results into transcript rows and prompt parts. - `errors.go`: failure classification (dispatch error messages, denial mapping, tool-result dispatch-failure scanning). - `tooluse.go`: the tool-call gate (`pre_tool_use` preflight, `post_tool_use` payloads, applying admitted input to the step). Server-bound glue stays in `coderd/x/chatd/hook_server.go`: the chat-parking dispatch error handlers, the step-commit row insertion wrappers, and the dynamic post-tool-use state loader, which depends on chatd validation types. This PR adopts the `codersdk/x/agenthooks` and `coderd/x/agenthooks/dispatch` import paths introduced at the tip of #27401; intermediate commits still reference the pre-move paths and are not individually buildable. ## Staged tool admission `pre_tool_use` originally ran at tool execution time, which is after the assistant message carrying the tool call was already committed. An `input_override` therefore had to rewrite stored message content in place. @hugodutka pointed out that chatd treats message content as immutable, and that the rewrite was a shortcut rather than a requirement. It was also a correctness problem in its own right: the rewrite only updated the database, so the transcript could show one input while a different one had executed. The hook now runs before the step is persisted: ```text provider stream ends (tool calls complete, in memory) -> pre_tool_use dispatch per call -> ONE transaction: assistant row with admitted inputs, synthetic denials, hook rows -> execute ``` The step is inserted once, carrying the input the tool runs with. `UpdateChatMessageContentByID` and `Tx.UpdateMessageContent` are deleted from #27428, so message content stays immutable. Two consequences, both intentional: - **Clients converge rather than wait.** Tool-call parts still stream live, so a rewritten call briefly shows the model's proposed input before the committed message replaces it. The chat store already clears stream state when an assistant message arrives, so the stored input wins with no frontend change and no added latency before tool cards appear. - **A call already in history was already admitted.** Execution consumes the stored input instead of dispatching a second decision, which keeps one dispatch and one set of hook effects per call. A consumer policy change between admission and execution applies to later calls, not to calls already admitted. The per-chat debug endpoint still records the provider's original tool input. Its purpose is to report provider behavior, and it requires an explicit per-chat debug flag; the invariant here covers the transcript. ## Configuration Adds `chat-hook-url`, `chat-hook-secret`, `chat-hook-timeout`, and `chat-hook-enabled` deployment options with startup validation. The flags are hidden from `coder server --help` while the feature is experimental; the setup guide documents them. ## Tool input validation Built-in tool arguments reach a consumer as raw JSON with key spelling preserved, but the tools decode those bytes with Go, which matches struct fields case-insensitively and keeps the last match. A policy reading `path` could therefore authorize one value while the tool executed another, and a lone case variant such as `{"PATH":"/secret"}` was invisible to a policy checking for `path`. Coder now rejects a built-in tool call whose input repeats a key or spells a schema property with different capitalization, before the `pre_tool_use` dispatch, so a consumer is never asked to authorize bytes whose meaning depends on the reader. Rejected calls produce an error result the model can retry; unambiguous calls in the same batch still run. A consumer-authored `input_override` is rechecked after the dispatch and fails the turn closed, because the model cannot correct it. Dynamic and MCP inputs are excluded because the client and the workspace agent execute those calls rather than coderd. Two paths needed more than a schema check. Execution resolves a deprecated tool name to its canonical tool, so validation resolves aliases first. The `edit_files` decoder also reads `search` and `replace`, which its schema does not advertise, so those aliases are now matched exactly and their case variants ignored. A hook denial now returns a structured 403 carrying `kind: "hook_denied"`, mirroring the dispatch-failure response that already carries its own kind. Without it a client cannot tell a policy decision apart from a generic failure, and the chat UI titled a denial "Request failed". Adding a kind needs no migration: `ChatErrorKind` is persisted only inside the JSONB `chats.last_error` column, whose decoder accepts unknown kinds. The hook docs also correct the tool-input convergence window. A batch dispatches sequentially before the assistant row commits, so the original input stays visible for a span that scales with the number of tool calls in the step rather than a single hook timeout. > This PR was written by Mux, an AI coding agent, on Mike's behalf.
4835 lines
150 KiB
Go
4835 lines
150 KiB
Go
package chatd
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"slices"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"charm.land/fantasy"
|
|
"github.com/google/uuid"
|
|
"github.com/prometheus/client_golang/prometheus"
|
|
"github.com/sqlc-dev/pqtype"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/xerrors"
|
|
|
|
"cdr.dev/slog/v3"
|
|
"cdr.dev/slog/v3/sloggers/slogtest"
|
|
"github.com/coder/coder/v2/coderd/aibridge"
|
|
"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/dbtestutil"
|
|
"github.com/coder/coder/v2/coderd/database/pubsub"
|
|
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
|
"github.com/coder/coder/v2/coderd/util/ptr"
|
|
"github.com/coder/coder/v2/coderd/x/agenthooks/dispatch"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chathooks"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatloop"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
|
|
"github.com/coder/coder/v2/codersdk"
|
|
"github.com/coder/coder/v2/testutil"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
func TestSubagentFallbackChatTitle(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
input string
|
|
want string
|
|
}{
|
|
{
|
|
name: "EmptyPrompt",
|
|
input: "",
|
|
want: "New Chat",
|
|
},
|
|
{
|
|
name: "ShortPrompt",
|
|
input: "Open Firefox",
|
|
want: "Open Firefox",
|
|
},
|
|
{
|
|
name: "LongPrompt",
|
|
input: "Please open the Firefox browser and navigate to the settings page",
|
|
want: "Please open the Firefox browser and...",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
got := subagentFallbackChatTitle(tt.input)
|
|
assert.Equal(t, tt.want, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
type internalTestServerConfig struct {
|
|
logger slog.Logger
|
|
clock quartz.Clock
|
|
startWorker bool
|
|
experiments codersdk.Experiments
|
|
transportFactory *atomic.Pointer[aibridge.TransportFactory]
|
|
}
|
|
|
|
type internalTestServerOpt func(*internalTestServerConfig)
|
|
|
|
func withInternalTestServerClock(clk quartz.Clock) internalTestServerOpt {
|
|
return func(cfg *internalTestServerConfig) {
|
|
cfg.clock = clk
|
|
}
|
|
}
|
|
|
|
func withInternalTestServerLogger(logger slog.Logger) internalTestServerOpt {
|
|
return func(cfg *internalTestServerConfig) {
|
|
cfg.logger = logger
|
|
}
|
|
}
|
|
|
|
func withInternalTestServerWorker() internalTestServerOpt {
|
|
return func(cfg *internalTestServerConfig) {
|
|
cfg.startWorker = true
|
|
}
|
|
}
|
|
|
|
func withInternalTestServerExperiments(experiments codersdk.Experiments) internalTestServerOpt {
|
|
return func(cfg *internalTestServerConfig) {
|
|
cfg.experiments = experiments
|
|
}
|
|
}
|
|
|
|
// withInternalTestServerTransportFactory wires an [aibridge.TransportFactory]
|
|
// into the server's Config so tests that drive real model generation through
|
|
// runSubagentTool or processChat can control the HTTP transport AI Gateway
|
|
// routing uses.
|
|
func withInternalTestServerTransportFactory(factory aibridge.TransportFactory) internalTestServerOpt {
|
|
return func(cfg *internalTestServerConfig) {
|
|
cfg.transportFactory = aibridgeTestFactoryPointer(factory)
|
|
}
|
|
}
|
|
|
|
func experimentsOrDefault(experiments codersdk.Experiments) codersdk.Experiments {
|
|
if experiments == nil {
|
|
return codersdk.ExperimentsKnown
|
|
}
|
|
return experiments
|
|
}
|
|
|
|
// newInternalTestServer creates a passive Server for internal tests with
|
|
// custom provider API keys. Pass withInternalTestServerWorker to start the
|
|
// background chat worker for tests that need real execution.
|
|
func newInternalTestServer(
|
|
t *testing.T,
|
|
db database.Store,
|
|
ps pubsub.Pubsub,
|
|
keys chatprovider.ProviderAPIKeys,
|
|
opts ...internalTestServerOpt,
|
|
) *Server {
|
|
t.Helper()
|
|
|
|
cfg := internalTestServerConfig{
|
|
logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
}
|
|
for _, opt := range opts {
|
|
opt(&cfg)
|
|
}
|
|
|
|
server := New(ps, Config{
|
|
Logger: cfg.logger,
|
|
Database: db,
|
|
ReplicaID: uuid.New(),
|
|
Clock: cfg.clock,
|
|
// Use a very long interval so the background loop
|
|
// does not interfere with test assertions.
|
|
PendingChatAcquireInterval: testutil.WaitLong,
|
|
ProviderAPIKeys: keys,
|
|
Experiments: experimentsOrDefault(cfg.experiments),
|
|
AIBridgeTransportFactory: cfg.transportFactory,
|
|
})
|
|
if cfg.startWorker {
|
|
server.Start()
|
|
}
|
|
t.Cleanup(func() {
|
|
require.NoError(t, server.Close())
|
|
})
|
|
return server
|
|
}
|
|
|
|
type subscribeFailingPubsub struct {
|
|
pubsub.Pubsub
|
|
}
|
|
|
|
func (subscribeFailingPubsub) Subscribe(_ string, _ pubsub.Listener) (func(), error) {
|
|
return nil, xerrors.New("subscribe disabled")
|
|
}
|
|
|
|
func (subscribeFailingPubsub) SubscribeWithErr(_ string, _ pubsub.ListenerWithErr) (func(), error) {
|
|
return nil, xerrors.New("subscribe disabled")
|
|
}
|
|
|
|
type subagentTestLogSink struct {
|
|
mu sync.Mutex
|
|
entries []slog.SinkEntry
|
|
}
|
|
|
|
func (s *subagentTestLogSink) LogEntry(_ context.Context, entry slog.SinkEntry) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
s.entries = append(s.entries, entry)
|
|
}
|
|
|
|
func (*subagentTestLogSink) Sync() {}
|
|
|
|
func (s *subagentTestLogSink) entriesAtLevelWithMessage(
|
|
level slog.Level,
|
|
message string,
|
|
) []slog.SinkEntry {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
entries := make([]slog.SinkEntry, 0, len(s.entries))
|
|
for _, entry := range s.entries {
|
|
if entry.Level == level && entry.Message == message {
|
|
entries = append(entries, entry)
|
|
}
|
|
}
|
|
return entries
|
|
}
|
|
|
|
// seedInternalChatDeps inserts an OpenAI provider and model config
|
|
// into the database and returns the created user, organization,
|
|
// and model. This deliberately does NOT create an Anthropic
|
|
// provider.
|
|
func seedInternalChatDeps(
|
|
t *testing.T,
|
|
db database.Store,
|
|
) (database.User, database.Organization, database.ChatModelConfig) {
|
|
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,
|
|
})
|
|
provider := dbgen.ChatProvider(t, db, database.ChatProvider{
|
|
Provider: "openai",
|
|
DisplayName: "OpenAI",
|
|
})
|
|
|
|
model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
|
AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true},
|
|
IsDefault: true,
|
|
})
|
|
|
|
return user, org, model
|
|
}
|
|
|
|
// insertEnabledAnthropicProvider inserts an enabled Anthropic provider for
|
|
// the current test user so computer_use flows keep Anthropic credentials
|
|
// after provider-key pruning.
|
|
func insertEnabledAnthropicProvider(
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
) {
|
|
t.Helper()
|
|
|
|
dbgen.ChatProvider(t, db, database.ChatProvider{
|
|
Provider: "anthropic",
|
|
DisplayName: "Anthropic",
|
|
APIKey: "test-anthropic-key",
|
|
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
|
})
|
|
}
|
|
|
|
func insertInternalAIProvider(
|
|
t *testing.T,
|
|
db database.Store,
|
|
providerType database.AIProviderType,
|
|
apiKey string,
|
|
enabled bool,
|
|
) database.AIProvider {
|
|
t.Helper()
|
|
return dbgen.AIProviderWithOptionalKey(t, db, database.AIProvider{
|
|
Type: providerType,
|
|
}, apiKey, func(params *database.InsertAIProviderParams) {
|
|
params.Enabled = enabled
|
|
})
|
|
}
|
|
|
|
func TestCreateChildSubagentChatDispatchesUserPromptSubmit(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
newFixture := func(t *testing.T, handler http.HandlerFunc) (context.Context, database.Store, database.Chat, *Server) {
|
|
t.Helper()
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
db, _ := dbtestutil.NewDB(t)
|
|
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})
|
|
model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{})
|
|
parent := dbgen.Chat(t, db, database.Chat{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
LastModelConfigID: model.ID,
|
|
})
|
|
apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID})
|
|
ctx = aibridge.WithDelegatedAPIKeyID(ctx, apiKey.ID)
|
|
|
|
consumer := httptest.NewServer(handler)
|
|
t.Cleanup(consumer.Close)
|
|
server := &Server{
|
|
db: db,
|
|
logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
hooks: chathooks.NewTrigger(dispatch.New(
|
|
slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
|
|
consumer.Client(),
|
|
consumer.URL,
|
|
"test-hook-secret-32-bytes-minimum!!",
|
|
time.Second,
|
|
"test-deployment",
|
|
"test-version",
|
|
prometheus.NewRegistry(),
|
|
)),
|
|
}
|
|
return ctx, db, parent, server
|
|
}
|
|
|
|
t.Run("Rewrite", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var meta struct {
|
|
sync.Mutex
|
|
parentChatID string
|
|
prompt string
|
|
}
|
|
ctx, db, parent, server := newFixture(t, func(rw http.ResponseWriter, r *http.Request) {
|
|
var request struct {
|
|
Type string `json:"type"`
|
|
Meta struct {
|
|
ParentChatID *uuid.UUID `json:"parent_chat_id"`
|
|
} `json:"meta"`
|
|
Data struct {
|
|
Prompt string `json:"prompt"`
|
|
} `json:"data"`
|
|
}
|
|
require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
|
|
require.Equal(t, "user_prompt_submit", request.Type)
|
|
meta.Lock()
|
|
if request.Meta.ParentChatID != nil {
|
|
meta.parentChatID = request.Meta.ParentChatID.String()
|
|
}
|
|
meta.prompt = request.Data.Prompt
|
|
meta.Unlock()
|
|
rw.Header().Set("Content-Type", "application/json")
|
|
_, _ = rw.Write([]byte(`{
|
|
"permission": {"decision": "allow", "input_override": {"prompt": "REVIEWED: inspect"}}
|
|
}`))
|
|
})
|
|
|
|
child, err := server.createChildSubagentChatWithOptions(ctx, parent, "inspect the workspace", "", childSubagentChatOptions{})
|
|
require.NoError(t, err)
|
|
|
|
meta.Lock()
|
|
require.Equal(t, parent.ID.String(), meta.parentChatID, "spawn dispatch must identify the parent chat")
|
|
require.Equal(t, "inspect the workspace", meta.prompt)
|
|
meta.Unlock()
|
|
|
|
messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: child.ID})
|
|
require.NoError(t, err)
|
|
var childUserMessage database.ChatMessage
|
|
for _, message := range messages {
|
|
if message.Role == database.ChatMessageRoleUser {
|
|
childUserMessage = message
|
|
break
|
|
}
|
|
}
|
|
require.NotZero(t, childUserMessage.ID)
|
|
require.True(t, childUserMessage.Content.Valid)
|
|
require.Contains(t, string(childUserMessage.Content.RawMessage), "REVIEWED: inspect",
|
|
"the hook rewrite must land as the child's initial prompt")
|
|
require.NotContains(t, string(childUserMessage.Content.RawMessage), "inspect the workspace")
|
|
})
|
|
|
|
t.Run("DenyRefusesSpawn", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx, db, parent, server := newFixture(t, func(rw http.ResponseWriter, _ *http.Request) {
|
|
rw.Header().Set("Content-Type", "application/json")
|
|
_, _ = rw.Write([]byte(`{"permission": {"decision": "deny", "reason": "spawn blocked"}, "user_message": "not allowed"}`))
|
|
})
|
|
|
|
_, err := server.createChildSubagentChatWithOptions(ctx, parent, "exfiltrate secrets", "", childSubagentChatOptions{})
|
|
var denied *chathooks.UserPromptDeniedError
|
|
require.ErrorAs(t, err, &denied)
|
|
require.Equal(t, "not allowed", denied.UserMessage)
|
|
|
|
chats, err := db.GetChildChatsByParentIDs(ctx, database.GetChildChatsByParentIDsParams{
|
|
ParentIds: []uuid.UUID{parent.ID},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Empty(t, chats, "a denied spawn must not create a child chat")
|
|
})
|
|
|
|
t.Run("DispatchFailurePropagatesFromSpawnTool", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx, db, parent, server := newFixture(t, func(rw http.ResponseWriter, _ *http.Request) {
|
|
http.Error(rw, "hook consumer down", http.StatusInternalServerError)
|
|
})
|
|
|
|
tools := server.subagentTools(ctx, func() database.Chat { return parent }, parent.LastModelConfigID)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool)
|
|
input, err := json.Marshal(spawnAgentArgs{
|
|
Type: subagentTypeExplore,
|
|
Prompt: "inspect the workspace",
|
|
Title: "sub",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
_, runErr := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: uuid.NewString(),
|
|
Name: spawnAgentToolName,
|
|
Input: string(input),
|
|
})
|
|
var hookErr *dispatch.Error
|
|
require.ErrorAs(t, runErr, &hookErr,
|
|
"dispatch failures must fail closed, not degrade to a tool error the model can ignore")
|
|
|
|
chats, err := db.GetChildChatsByParentIDs(ctx, database.GetChildChatsByParentIDsParams{
|
|
ParentIds: []uuid.UUID{parent.ID},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Empty(t, chats)
|
|
})
|
|
}
|
|
|
|
func TestResolveUserProviderAPIKeys_AIProvider(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("UserKeyWinsWhenBYOKEnabled", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
provider := insertInternalAIProvider(t, db, database.AIProviderTypeOpenai, "provider-api-key", true)
|
|
now := time.Now()
|
|
_, err := db.UpsertUserAIProviderKey(ctx, database.UpsertUserAIProviderKeyParams{
|
|
ID: uuid.New(),
|
|
UserID: user.ID,
|
|
AIProviderID: provider.ID,
|
|
APIKey: "user-api-key",
|
|
CreatedAt: now,
|
|
UpdatedAt: now,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID, provider.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "user-api-key", keys.APIKey("openai"))
|
|
// The expected URL is dbgen's default AIProvider BaseUrl.
|
|
require.Equal(t, "invalid://test.invalid/", keys.BaseURL("openai"))
|
|
})
|
|
|
|
t.Run("ProviderKeyUsedWhenBYOKDisabled", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
server.allowBYOK = false
|
|
ctx := chatdTestContext(t)
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
provider := insertInternalAIProvider(t, db, database.AIProviderTypeOpenai, "provider-api-key", true)
|
|
now := time.Now()
|
|
_, err := db.UpsertUserAIProviderKey(ctx, database.UpsertUserAIProviderKeyParams{
|
|
ID: uuid.New(),
|
|
UserID: user.ID,
|
|
AIProviderID: provider.ID,
|
|
APIKey: "user-api-key",
|
|
CreatedAt: now,
|
|
UpdatedAt: now,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID, provider.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "provider-api-key", keys.APIKey("openai"))
|
|
})
|
|
|
|
t.Run("ProviderTypeUsesAIProvider", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
insertInternalAIProvider(t, db, database.AIProviderTypeAzure, "provider-api-key", true)
|
|
|
|
keys, err := server.resolveUserProviderAPIKeysForProviderType(ctx, user.ID, "azure")
|
|
require.NoError(t, err)
|
|
require.Equal(t, "provider-api-key", keys.APIKey("azure"))
|
|
})
|
|
|
|
t.Run("BedrockUsesAmbientAuth", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
provider := insertInternalAIProvider(t, db, database.AIProviderTypeBedrock, "", true)
|
|
|
|
keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID, provider.ID)
|
|
require.NoError(t, err)
|
|
require.True(t, keys.HasProvider("bedrock"))
|
|
require.Empty(t, keys.APIKey("bedrock"))
|
|
})
|
|
|
|
t.Run("RejectsAmbiguousProviderTypeWithoutSelectedProvider", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
insertInternalAIProvider(t, db, database.AIProviderTypeOpenai, "first-provider-api-key", true)
|
|
insertInternalAIProvider(t, db, database.AIProviderTypeOpenai, "second-provider-api-key", true)
|
|
|
|
keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID, uuid.Nil)
|
|
require.ErrorContains(t, err, "multiple enabled AI providers use provider type")
|
|
require.Equal(t, chatprovider.ProviderAPIKeys{}, keys)
|
|
})
|
|
}
|
|
|
|
func TestResolveChatModel_AIProviderDisabled(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := chatdTestContext(t)
|
|
db, ps := dbtestutil.NewDB(t)
|
|
user, org, _ := seedInternalChatDeps(t, db)
|
|
provider := insertInternalAIProvider(t, db, database.AIProviderTypeOpenai, "provider-api-key", false)
|
|
modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
|
Model: "gpt-4o-mini",
|
|
AIProviderID: uuid.NullUUID{
|
|
UUID: provider.ID,
|
|
Valid: true,
|
|
},
|
|
})
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
chat := dbgen.Chat(t, db, database.Chat{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
LastModelConfigID: modelConfig.ID,
|
|
})
|
|
|
|
model, config, _, debugEnabled, resolvedProvider, resolvedModel, err := server.resolveChatModel(ctx, chat, modelBuildOptions{})
|
|
require.ErrorContains(t, err, "is disabled")
|
|
require.Nil(t, model)
|
|
require.Equal(t, database.ChatModelConfig{}, config)
|
|
require.False(t, debugEnabled)
|
|
require.Empty(t, resolvedProvider)
|
|
require.Empty(t, resolvedModel)
|
|
}
|
|
|
|
func TestResolveUserProviderAPIKeys_PreservesAnthropicKeyFromDBProvider(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("PreservesDBProviderKeyWithoutFallback", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
insertEnabledAnthropicProvider(t, db, user.ID)
|
|
|
|
keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID, uuid.Nil)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "test-anthropic-key", keys.Anthropic)
|
|
require.Equal(t, "test-anthropic-key", keys.APIKey("anthropic"))
|
|
require.Equal(t, "test-anthropic-key", keys.ByProvider["anthropic"])
|
|
})
|
|
|
|
t.Run("PrunesFallbackKeyWithoutEnabledProvider", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
|
Anthropic: "test-anthropic-key",
|
|
})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
|
|
keys, err := server.resolveUserProviderAPIKeys(ctx, user.ID, uuid.Nil)
|
|
require.NoError(t, err)
|
|
require.Empty(t, keys.Anthropic)
|
|
require.Empty(t, keys.APIKey("anthropic"))
|
|
_, ok := keys.ByProvider["anthropic"]
|
|
require.False(t, ok)
|
|
})
|
|
}
|
|
|
|
func insertInternalChatModelConfig(
|
|
t *testing.T,
|
|
db database.Store,
|
|
model string,
|
|
enabled bool,
|
|
) database.ChatModelConfig {
|
|
return insertInternalChatModelConfigForProvider(
|
|
t,
|
|
db,
|
|
"openai",
|
|
model,
|
|
enabled,
|
|
)
|
|
}
|
|
|
|
func insertInternalChatProvider(
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
provider string,
|
|
apiKey string,
|
|
centralAPIKeyEnabled bool,
|
|
allowUserAPIKey bool,
|
|
allowCentralAPIKeyFallback bool,
|
|
) database.AIProvider {
|
|
t.Helper()
|
|
|
|
providerConfig := dbgen.AIProvider(t, db, database.AIProvider{
|
|
Type: database.AIProviderType(provider),
|
|
Name: "test-" + uuid.NewString(),
|
|
DisplayName: sql.NullString{String: provider, Valid: true},
|
|
})
|
|
if apiKey != "" {
|
|
dbgen.AIProviderKey(t, db, database.AIProviderKey{
|
|
ProviderID: providerConfig.ID,
|
|
APIKey: apiKey,
|
|
})
|
|
}
|
|
|
|
return providerConfig
|
|
}
|
|
|
|
func insertInternalChatModelConfigForProvider(
|
|
t *testing.T,
|
|
db database.Store,
|
|
provider string,
|
|
model string,
|
|
enabled bool,
|
|
) database.ChatModelConfig {
|
|
t.Helper()
|
|
return insertInternalChatModelConfigWithOptions(
|
|
t,
|
|
db,
|
|
provider,
|
|
model,
|
|
enabled,
|
|
json.RawMessage(`{}`),
|
|
)
|
|
}
|
|
|
|
func insertInternalChatModelConfigWithOptions(
|
|
t *testing.T,
|
|
db database.Store,
|
|
provider string,
|
|
model string,
|
|
enabled bool,
|
|
options json.RawMessage,
|
|
) database.ChatModelConfig {
|
|
t.Helper()
|
|
|
|
// Reuse the newest AI provider of this type (creating a bare credential-less
|
|
// one only when none exists) so the config links the provider already
|
|
// carrying the test's credentials, or lack thereof, rather than a fresh one.
|
|
providers, err := db.GetAIProviders(context.Background(), database.GetAIProvidersParams{IncludeDisabled: true})
|
|
require.NoError(t, err)
|
|
var aiProvider database.AIProvider
|
|
for _, candidate := range providers {
|
|
if candidate.Type != database.AIProviderType(provider) {
|
|
continue
|
|
}
|
|
if aiProvider.ID == uuid.Nil || candidate.CreatedAt.After(aiProvider.CreatedAt) {
|
|
aiProvider = candidate
|
|
}
|
|
}
|
|
if aiProvider.ID == uuid.Nil {
|
|
aiProvider = dbgen.AIProvider(t, db, database.AIProvider{
|
|
Type: database.AIProviderType(provider),
|
|
})
|
|
}
|
|
modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
|
AIProviderID: uuid.NullUUID{UUID: aiProvider.ID, Valid: true},
|
|
Model: model,
|
|
DisplayName: model,
|
|
Options: options,
|
|
}, func(p *database.InsertChatModelConfigParams) {
|
|
p.Enabled = enabled
|
|
})
|
|
|
|
return modelConfig
|
|
}
|
|
|
|
func insertInternalMCPServerConfig(
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
slug string,
|
|
allowInPlanMode bool,
|
|
) database.MCPServerConfig {
|
|
t.Helper()
|
|
|
|
return dbgen.MCPServerConfig(t, db, database.MCPServerConfig{
|
|
DisplayName: slug,
|
|
Slug: slug,
|
|
Url: "https://" + slug + ".example.com",
|
|
AllowInPlanMode: allowInPlanMode,
|
|
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
|
UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
|
})
|
|
}
|
|
|
|
func seedWorkspaceBinding(
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
) (database.WorkspaceTable, database.WorkspaceBuild, database.WorkspaceAgent) {
|
|
t.Helper()
|
|
|
|
org := dbgen.Organization(t, db, database.Organization{})
|
|
tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{
|
|
OrganizationID: org.ID,
|
|
CreatedBy: userID,
|
|
})
|
|
tpl := dbgen.Template(t, db, database.Template{
|
|
CreatedBy: userID,
|
|
OrganizationID: org.ID,
|
|
ActiveVersionID: tv.ID,
|
|
})
|
|
workspace := dbgen.Workspace(t, db, database.WorkspaceTable{
|
|
TemplateID: tpl.ID,
|
|
OwnerID: userID,
|
|
OrganizationID: org.ID,
|
|
})
|
|
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
|
|
InitiatorID: userID,
|
|
OrganizationID: org.ID,
|
|
})
|
|
build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
|
|
TemplateVersionID: tv.ID,
|
|
WorkspaceID: workspace.ID,
|
|
JobID: job.ID,
|
|
})
|
|
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
|
|
Transition: database.WorkspaceTransitionStart,
|
|
JobID: job.ID,
|
|
})
|
|
agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: resource.ID})
|
|
return workspace, build, agent
|
|
}
|
|
|
|
// findToolByName returns the tool with the given name from the
|
|
// slice, or nil if no match is found.
|
|
func findToolByName(tools []fantasy.AgentTool, name string) fantasy.AgentTool {
|
|
for _, tool := range tools {
|
|
if tool.Info().Name == name {
|
|
return tool
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func chatdTestContext(t *testing.T) context.Context {
|
|
t.Helper()
|
|
return dbauthz.AsChatd(testutil.Context(t, testutil.WaitLong))
|
|
}
|
|
|
|
func systemRestrictedTestContext(t *testing.T) context.Context {
|
|
t.Helper()
|
|
return dbauthz.AsSystemRestricted(testutil.Context(t, testutil.WaitLong))
|
|
}
|
|
|
|
func enableInternalChatPersonalModelOverrides(
|
|
t *testing.T,
|
|
db database.Store,
|
|
) {
|
|
t.Helper()
|
|
require.NoError(
|
|
t,
|
|
db.UpsertChatPersonalModelOverridesEnabled(
|
|
systemRestrictedTestContext(t),
|
|
true,
|
|
),
|
|
)
|
|
}
|
|
|
|
func upsertInternalUserChatPersonalModelOverride(
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
overrideContext codersdk.ChatPersonalModelOverrideContext,
|
|
raw string,
|
|
) {
|
|
t.Helper()
|
|
require.NoError(
|
|
t,
|
|
db.UpsertUserChatPersonalModelOverride(
|
|
systemRestrictedTestContext(t),
|
|
database.UpsertUserChatPersonalModelOverrideParams{
|
|
UserID: userID,
|
|
Key: ChatPersonalModelOverrideKey(overrideContext),
|
|
Value: raw,
|
|
},
|
|
),
|
|
)
|
|
}
|
|
|
|
func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
workspace, build, agent := seedWorkspaceBinding(t, db, user.ID)
|
|
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
WorkspaceID: uuid.NullUUID{
|
|
UUID: workspace.ID,
|
|
Valid: true,
|
|
},
|
|
BuildID: uuid.NullUUID{
|
|
UUID: build.ID,
|
|
Valid: true,
|
|
},
|
|
AgentID: uuid.NullUUID{
|
|
UUID: agent.ID,
|
|
Valid: true,
|
|
},
|
|
Title: "bound-parent",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
child, err := server.createChildSubagentChatWithOptions(ctx, parentChat, "inspect bindings", "", childSubagentChatOptions{})
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, parentChat.OrganizationID, childChat.OrganizationID)
|
|
require.Equal(t, parentChat.WorkspaceID, childChat.WorkspaceID)
|
|
require.Equal(t, parentChat.BuildID, childChat.BuildID)
|
|
require.Equal(t, parentChat.AgentID, childChat.AgentID)
|
|
}
|
|
|
|
func createInternalParentChat(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
server *Server,
|
|
db database.Store,
|
|
orgID uuid.UUID,
|
|
userID uuid.UUID,
|
|
modelConfigID uuid.UUID,
|
|
title string,
|
|
) database.Chat {
|
|
t.Helper()
|
|
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: orgID,
|
|
OwnerID: userID,
|
|
Title: title,
|
|
ModelConfigID: modelConfigID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
return parentChat
|
|
}
|
|
|
|
func runSubagentTool(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
server *Server,
|
|
parentChat database.Chat,
|
|
currentModelConfigID uuid.UUID,
|
|
toolName string,
|
|
args any,
|
|
) fantasy.ToolResponse {
|
|
t.Helper()
|
|
|
|
tools := server.subagentTools(
|
|
ctx,
|
|
func() database.Chat { return parentChat },
|
|
currentModelConfigID,
|
|
)
|
|
tool := findToolByName(tools, toolName)
|
|
require.NotNil(t, tool, "%s tool must be present", toolName)
|
|
|
|
input, err := json.Marshal(args)
|
|
require.NoError(t, err)
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: uuid.NewString(),
|
|
Name: toolName,
|
|
Input: string(input),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
return resp
|
|
}
|
|
|
|
func runSpawnAgentTool(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
server *Server,
|
|
parentChat database.Chat,
|
|
args spawnAgentArgs,
|
|
) fantasy.ToolResponse {
|
|
t.Helper()
|
|
return runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
spawnAgentToolName,
|
|
args,
|
|
)
|
|
}
|
|
|
|
func requireSpawnAgentResponse(t *testing.T, resp fantasy.ToolResponse) struct {
|
|
ChatID string `json:"chat_id"`
|
|
SubagentType string `json:"type"`
|
|
} {
|
|
t.Helper()
|
|
require.False(t, resp.IsError, "expected success but got: %s", resp.Content)
|
|
|
|
var result struct {
|
|
ChatID string `json:"chat_id"`
|
|
SubagentType string `json:"type"`
|
|
}
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
require.NotEmpty(t, result.ChatID, "response must contain chat_id")
|
|
require.NotEmpty(t, result.SubagentType, "response must contain type")
|
|
return result
|
|
}
|
|
|
|
func requireSpawnAgentChildChatID(t *testing.T, resp fantasy.ToolResponse) uuid.UUID {
|
|
t.Helper()
|
|
require.False(t, resp.IsError, "expected success but got: %s", resp.Content)
|
|
|
|
var result struct {
|
|
ChatID string `json:"chat_id"`
|
|
}
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
require.NotEmpty(t, result.ChatID, "response must contain chat_id")
|
|
|
|
childID, err := uuid.Parse(result.ChatID)
|
|
require.NoError(t, err)
|
|
return childID
|
|
}
|
|
|
|
func requireToolResponseMap(
|
|
t *testing.T,
|
|
resp fantasy.ToolResponse,
|
|
wantError bool,
|
|
) map[string]any {
|
|
t.Helper()
|
|
require.Equal(t, wantError, resp.IsError, "unexpected tool error state: %s", resp.Content)
|
|
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
return result
|
|
}
|
|
|
|
func TestCreateChildSubagentChatCopiesPlanMode(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
planMode := database.NullChatPlanMode{
|
|
ChatPlanMode: database.ChatPlanModePlan,
|
|
Valid: true,
|
|
}
|
|
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "plan-parent",
|
|
ModelConfigID: model.ID,
|
|
PlanMode: planMode,
|
|
InitialUserContent: []codersdk.ChatMessagePart{
|
|
codersdk.ChatMessageText("plan this change"),
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, planMode, parentChat.PlanMode)
|
|
|
|
child, err := server.createChildSubagentChatWithOptions(ctx, parentChat, "inspect bindings", "", childSubagentChatOptions{})
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, planMode, childChat.PlanMode)
|
|
}
|
|
|
|
func TestSpawnAgent_GeneralInheritsParentModelWhenOmitted(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-inherited-model",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate work",
|
|
})
|
|
result := requireSpawnAgentResponse(t, resp)
|
|
require.Equal(t, subagentTypeGeneral, result.SubagentType)
|
|
childID, err := uuid.Parse(result.ChatID)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, parentChat.LastModelConfigID, childChat.LastModelConfigID)
|
|
}
|
|
|
|
func TestSpawnAgent_GeneralUsesConfiguredModelOverride(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
overrideModel := insertInternalChatModelConfig(
|
|
t, db, "general-override-"+uuid.NewString(), true,
|
|
)
|
|
require.NoError(t, db.UpsertChatGeneralModelOverride(ctx, overrideModel.ID.String()))
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-general-override",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate general work",
|
|
})
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, overrideModel.ID, childChat.LastModelConfigID)
|
|
require.False(t, childChat.PlanMode.Valid)
|
|
}
|
|
|
|
func TestSpawnAgent_GeneralHonorsPersonalModelOverrides(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
enablePersonalOverride bool
|
|
personalRaw func(database.ChatModelConfig) string
|
|
personalModel func(context.Context, *testing.T, database.Store, uuid.UUID) database.ChatModelConfig
|
|
wantModelID func(
|
|
database.ChatModelConfig,
|
|
database.ChatModelConfig,
|
|
database.ChatModelConfig,
|
|
) uuid.UUID
|
|
}{
|
|
{
|
|
name: "UnsetUsesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
wantModelID: func(_, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "DeploymentDefaultUsesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeDeploymentDefault)
|
|
},
|
|
wantModelID: func(_, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "ChatDefaultBypassesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeChatDefault)
|
|
},
|
|
wantModelID: func(parentModel, _, _ database.ChatModelConfig) uuid.UUID {
|
|
return parentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "ModelUsesPersonalOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(personalModel database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeModel) + ":" +
|
|
personalModel.ID.String()
|
|
},
|
|
wantModelID: func(_, _, personalModel database.ChatModelConfig) uuid.UUID {
|
|
return personalModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "AdminFlagOffIgnoresPersonalOverride",
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeChatDefault)
|
|
},
|
|
wantModelID: func(_, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "DisabledPersonalModelFallsBackToDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalModel: func(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
) database.ChatModelConfig {
|
|
return insertInternalChatModelConfig(
|
|
t,
|
|
db,
|
|
"general-personal-disabled-"+uuid.NewString(),
|
|
false,
|
|
)
|
|
},
|
|
personalRaw: func(personalModel database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeModel) + ":" +
|
|
personalModel.ID.String()
|
|
},
|
|
wantModelID: func(_, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "MissingCredentialsFallsBackToDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalModel: func(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
) database.ChatModelConfig {
|
|
insertInternalChatProvider(
|
|
t,
|
|
db,
|
|
userID,
|
|
"openai-compat",
|
|
"",
|
|
false,
|
|
true,
|
|
false,
|
|
)
|
|
return insertInternalChatModelConfigForProvider(
|
|
t,
|
|
db,
|
|
"openai-compat",
|
|
"gpt-4o-mini",
|
|
true,
|
|
)
|
|
},
|
|
personalRaw: func(personalModel database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeModel) + ":" +
|
|
personalModel.ID.String()
|
|
},
|
|
wantModelID: func(_, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "MalformedValueUsesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return "model:not-a-uuid"
|
|
},
|
|
wantModelID: func(_, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
deploymentModel := insertInternalChatModelConfig(
|
|
t,
|
|
db,
|
|
"general-deployment-"+uuid.NewString(),
|
|
true,
|
|
)
|
|
require.NoError(t, db.UpsertChatGeneralModelOverride(ctx, deploymentModel.ID.String()))
|
|
personalModel := insertInternalChatModelConfig(
|
|
t,
|
|
db,
|
|
"general-personal-"+uuid.NewString(),
|
|
true,
|
|
)
|
|
if tt.personalModel != nil {
|
|
personalModel = tt.personalModel(ctx, t, db, user.ID)
|
|
}
|
|
if tt.enablePersonalOverride {
|
|
enableInternalChatPersonalModelOverrides(t, db)
|
|
}
|
|
if tt.personalRaw != nil {
|
|
upsertInternalUserChatPersonalModelOverride(
|
|
t,
|
|
db,
|
|
user.ID,
|
|
codersdk.ChatPersonalModelOverrideContextGeneral,
|
|
tt.personalRaw(personalModel),
|
|
)
|
|
}
|
|
parentChat := createInternalParentChat(
|
|
ctx,
|
|
t,
|
|
server,
|
|
db,
|
|
org.ID,
|
|
user.ID,
|
|
parentModel.ID,
|
|
"parent-general-personal-override",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate general work",
|
|
})
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(
|
|
t,
|
|
tt.wantModelID(parentModel, deploymentModel, personalModel),
|
|
childChat.LastModelConfigID,
|
|
)
|
|
require.False(t, childChat.PlanMode.Valid)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenCredentialsUnavailable(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
logSink := &subagentTestLogSink{}
|
|
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).AppendSinks(logSink)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerLogger(logger))
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
insertInternalChatProvider(
|
|
t,
|
|
db,
|
|
user.ID,
|
|
"openai-compat",
|
|
"",
|
|
false,
|
|
true,
|
|
false,
|
|
)
|
|
|
|
overrideModel := insertInternalChatModelConfigForProvider(
|
|
t,
|
|
db,
|
|
"openai-compat",
|
|
"gpt-4o-mini",
|
|
true,
|
|
)
|
|
require.NoError(t, db.UpsertChatGeneralModelOverride(ctx, overrideModel.ID.String()))
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "parent-general-credentials-fallback",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{
|
|
codersdk.ChatMessageText("delegate work"),
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "inspect provider credentials",
|
|
})
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, model.ID, childChat.LastModelConfigID)
|
|
require.False(t, childChat.PlanMode.Valid)
|
|
require.Len(t, logSink.entriesAtLevelWithMessage(
|
|
slog.LevelInfo,
|
|
"model override credentials are unavailable, ignoring",
|
|
), 1)
|
|
}
|
|
|
|
func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenProviderDisabled(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
logSink := &subagentTestLogSink{}
|
|
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).AppendSinks(logSink)
|
|
server := newInternalTestServer(
|
|
t,
|
|
db,
|
|
ps,
|
|
chatprovider.ProviderAPIKeys{
|
|
ByProvider: map[string]string{
|
|
"openai-compat": "fallback-key",
|
|
},
|
|
},
|
|
withInternalTestServerLogger(logger),
|
|
)
|
|
|
|
ctx := chatdTestContext(t)
|
|
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
|
|
})
|
|
|
|
overrideModel := insertInternalChatModelConfigForProvider(
|
|
t,
|
|
db,
|
|
"openai-compat",
|
|
"gpt-4o-mini",
|
|
true,
|
|
)
|
|
require.NoError(t, db.UpsertChatGeneralModelOverride(ctx, overrideModel.ID.String()))
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "parent-general-disabled-provider-fallback",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{
|
|
codersdk.ChatMessageText("delegate work"),
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "inspect disabled providers",
|
|
})
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, model.ID, childChat.LastModelConfigID)
|
|
require.False(t, childChat.PlanMode.Valid)
|
|
require.Len(t, logSink.entriesAtLevelWithMessage(
|
|
slog.LevelInfo,
|
|
"model override is unavailable, ignoring",
|
|
), 1)
|
|
}
|
|
|
|
func TestResolveConfiguredModelOverride_AcceptsAmbientCredentialsProvider(
|
|
t *testing.T,
|
|
) {
|
|
t.Parallel()
|
|
|
|
logSink := &subagentTestLogSink{}
|
|
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).AppendSinks(logSink)
|
|
server := &Server{logger: logger}
|
|
ctx := chatdTestContext(t)
|
|
ownerID := uuid.New()
|
|
modelConfig := database.ChatModelConfig{
|
|
ID: uuid.New(),
|
|
Model: "anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
DisplayName: "Ambient Bedrock Override",
|
|
Enabled: true,
|
|
}
|
|
|
|
resolvedModelConfig, _, reasoningEffort, ok, err := server.resolveConfiguredModelOverride(
|
|
ctx,
|
|
"plan",
|
|
modelConfig.ID.String(),
|
|
ownerID,
|
|
func(
|
|
_ context.Context,
|
|
configuredModelConfigID uuid.UUID,
|
|
) (database.ChatModelConfig, string, error) {
|
|
require.Equal(t, modelConfig.ID, configuredModelConfigID)
|
|
return modelConfig, "bedrock", nil
|
|
},
|
|
func(
|
|
_ context.Context,
|
|
resolvedOwnerID uuid.UUID,
|
|
_ uuid.UUID,
|
|
) (chatprovider.ProviderAPIKeys, error) {
|
|
require.Equal(t, ownerID, resolvedOwnerID)
|
|
return chatprovider.ProviderAPIKeys{
|
|
ByProvider: map[string]string{"bedrock": ""},
|
|
}, nil
|
|
},
|
|
modelOverrideFailureModeSoft,
|
|
)
|
|
require.NoError(t, err)
|
|
require.Nil(t, reasoningEffort)
|
|
require.True(t, ok)
|
|
require.Equal(t, modelConfig, resolvedModelConfig)
|
|
require.Empty(t, logSink.entriesAtLevelWithMessage(
|
|
slog.LevelInfo,
|
|
"model override credentials are unavailable, ignoring",
|
|
))
|
|
}
|
|
|
|
func TestWithResolvedReasoningEffort(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
baseOptions, err := json.Marshal(codersdk.ChatModelCallConfig{
|
|
MaxOutputTokens: ptr.Ref(int64(123)),
|
|
ReasoningEffort: &codersdk.ChatModelReasoningEffortConfig{
|
|
Default: ptr.Ref(codersdk.ChatModelReasoningEffortLow),
|
|
Max: ptr.Ref(codersdk.ChatModelReasoningEffortHigh),
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
t.Run("NilEffortReturnsOriginalConfig", func(t *testing.T) {
|
|
t.Parallel()
|
|
modelConfig := database.ChatModelConfig{Options: baseOptions}
|
|
require.Equal(t, modelConfig, withResolvedReasoningEffort(modelConfig, nil))
|
|
})
|
|
|
|
t.Run("InvalidJSONReturnsOriginalConfig", func(t *testing.T) {
|
|
t.Parallel()
|
|
modelConfig := database.ChatModelConfig{Options: []byte(`{`)}
|
|
require.Equal(t, modelConfig, withResolvedReasoningEffort(modelConfig, ptr.Ref(codersdk.ChatModelReasoningEffortHigh)))
|
|
})
|
|
|
|
t.Run("NoReasoningConfigReturnsOriginalConfig", func(t *testing.T) {
|
|
t.Parallel()
|
|
modelConfig := database.ChatModelConfig{Options: []byte(`{"max_output_tokens":123}`)}
|
|
require.Equal(t, modelConfig, withResolvedReasoningEffort(modelConfig, ptr.Ref(codersdk.ChatModelReasoningEffortHigh)))
|
|
})
|
|
|
|
t.Run("ClampsRequestedEffortAndPreservesOtherOptions", func(t *testing.T) {
|
|
t.Parallel()
|
|
modelConfig := database.ChatModelConfig{Options: baseOptions}
|
|
|
|
got := withResolvedReasoningEffort(modelConfig, ptr.Ref(codersdk.ChatModelReasoningEffortXHigh))
|
|
|
|
var callConfig codersdk.ChatModelCallConfig
|
|
require.NoError(t, json.Unmarshal(got.Options, &callConfig))
|
|
require.Equal(t, ptr.Ref(int64(123)), callConfig.MaxOutputTokens)
|
|
require.Equal(t, ptr.Ref(codersdk.ChatModelReasoningEffortHigh), callConfig.ReasoningEffort.Default)
|
|
require.Equal(t, ptr.Ref(codersdk.ChatModelReasoningEffortHigh), callConfig.ReasoningEffort.Max)
|
|
})
|
|
}
|
|
|
|
func TestCreateChildSubagentChat_StoresReasoningEffortOverride(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-effort-override",
|
|
)
|
|
child, err := server.createChildSubagentChatWithOptions(
|
|
ctx,
|
|
parentChat,
|
|
"delegate work",
|
|
"",
|
|
childSubagentChatOptions{reasoningEffortOverride: ptr.Ref("high")},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortHigh, Valid: true}, childChat.LastReasoningEffort)
|
|
}
|
|
|
|
func TestCreateChildSubagentChat_OverrideWorksWhenParentHasNoModel(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
overrideModel := insertInternalChatModelConfig(
|
|
t, db, "override-no-parent-model-"+uuid.NewString(), true,
|
|
)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-no-model",
|
|
)
|
|
|
|
// The chats table enforces a foreign key for last_model_config_id, so
|
|
// use a synthetic parent value here to exercise the override path.
|
|
parentChat.LastModelConfigID = uuid.Nil
|
|
child, err := server.createChildSubagentChatWithOptions(
|
|
ctx,
|
|
parentChat,
|
|
"delegate work",
|
|
"",
|
|
childSubagentChatOptions{modelConfigIDOverride: &overrideModel.ID},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, overrideModel.ID, childChat.LastModelConfigID)
|
|
}
|
|
|
|
func TestSpawnAgent_ExplicitModelConfigID(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
for _, subagentType := range []string{subagentTypeGeneral, subagentTypeExplore} {
|
|
t.Run(subagentType, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
selectedModel := insertInternalChatModelConfig(
|
|
t, db, "spawn-explicit-"+uuid.NewString(), true,
|
|
)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explicit-model",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentType,
|
|
Prompt: "delegate work",
|
|
ModelConfigID: selectedModel.ID.String(),
|
|
})
|
|
result := requireSpawnAgentResponse(t, resp)
|
|
require.Equal(t, subagentType, result.SubagentType)
|
|
childID, err := uuid.Parse(result.ChatID)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, selectedModel.ID, childChat.LastModelConfigID)
|
|
require.False(t, childChat.LastReasoningEffort.Valid)
|
|
|
|
parentAfter, err := db.GetChatByID(ctx, parentChat.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, parentModel.ID, parentAfter.LastModelConfigID)
|
|
})
|
|
}
|
|
}
|
|
|
|
type personalOverrideErrorStore struct {
|
|
database.Store
|
|
}
|
|
|
|
func (*personalOverrideErrorStore) GetUserChatPersonalModelOverride(
|
|
context.Context,
|
|
database.GetUserChatPersonalModelOverrideParams,
|
|
) (string, error) {
|
|
return "", xerrors.New("unexpected personal override lookup")
|
|
}
|
|
|
|
func TestSpawnAgent_ExplicitModelSkipsConfiguredOverrides(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
for _, subagentType := range []string{subagentTypeGeneral, subagentTypeExplore} {
|
|
t.Run(subagentType, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
store := &personalOverrideErrorStore{Store: db}
|
|
server := newInternalTestServer(t, store, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
enableInternalChatPersonalModelOverrides(t, db)
|
|
selectedModel := insertInternalChatModelConfig(
|
|
t, db, "selected-skip-overrides-"+uuid.NewString(), true,
|
|
)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, store, org.ID, user.ID, parentModel.ID, "parent-skip-overrides",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentType,
|
|
Prompt: "delegate work",
|
|
ModelConfigID: selectedModel.ID.String(),
|
|
})
|
|
childChat, err := db.GetChatByID(ctx, requireSpawnAgentChildChatID(t, resp))
|
|
require.NoError(t, err)
|
|
require.Equal(t, selectedModel.ID, childChat.LastModelConfigID)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSpawnAgent_ExplicitModelBeatsConfiguredOverrides(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
deploymentModel := insertInternalChatModelConfig(
|
|
t, db, "general-deployment-"+uuid.NewString(), true,
|
|
)
|
|
personalModel := insertInternalChatModelConfig(
|
|
t, db, "general-personal-"+uuid.NewString(), true,
|
|
)
|
|
selectedModel := insertInternalChatModelConfig(
|
|
t, db, "general-selected-"+uuid.NewString(), true,
|
|
)
|
|
require.NoError(t, db.UpsertChatGeneralModelOverride(
|
|
ctx, deploymentModel.ID.String()+":medium",
|
|
))
|
|
enableInternalChatPersonalModelOverrides(t, db)
|
|
upsertInternalUserChatPersonalModelOverride(
|
|
t,
|
|
db,
|
|
user.ID,
|
|
codersdk.ChatPersonalModelOverrideContextGeneral,
|
|
string(codersdk.ChatPersonalModelOverrideModeModel)+":"+personalModel.ID.String()+":low",
|
|
)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explicit-beats-overrides",
|
|
)
|
|
|
|
baselineResp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate baseline work",
|
|
})
|
|
baselineChild, err := db.GetChatByID(ctx, requireSpawnAgentChildChatID(t, baselineResp))
|
|
require.NoError(t, err)
|
|
require.Equal(t, personalModel.ID, baselineChild.LastModelConfigID)
|
|
require.Equal(
|
|
t,
|
|
database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortLow, Valid: true},
|
|
baselineChild.LastReasoningEffort,
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate explicit work",
|
|
ModelConfigID: selectedModel.ID.String(),
|
|
})
|
|
childChat, err := db.GetChatByID(ctx, requireSpawnAgentChildChatID(t, resp))
|
|
require.NoError(t, err)
|
|
require.Equal(t, selectedModel.ID, childChat.LastModelConfigID)
|
|
require.False(t, childChat.LastReasoningEffort.Valid)
|
|
}
|
|
|
|
func TestSpawnAgent_ExplicitReasoningEffort(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("EffortOnlyKeepsModelChain", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-effort-only",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate work",
|
|
ReasoningEffort: "high",
|
|
})
|
|
childChat, err := db.GetChatByID(ctx, requireSpawnAgentChildChatID(t, resp))
|
|
require.NoError(t, err)
|
|
require.Equal(t, parentModel.ID, childChat.LastModelConfigID)
|
|
require.Equal(
|
|
t,
|
|
database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortHigh, Valid: true},
|
|
childChat.LastReasoningEffort,
|
|
)
|
|
})
|
|
|
|
t.Run("CombinedModelAndEffort", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
selectedModel := insertInternalChatModelConfig(
|
|
t, db, "combined-"+uuid.NewString(), true,
|
|
)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-combined",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate work",
|
|
ModelConfigID: selectedModel.ID.String(),
|
|
ReasoningEffort: "low",
|
|
})
|
|
childChat, err := db.GetChatByID(ctx, requireSpawnAgentChildChatID(t, resp))
|
|
require.NoError(t, err)
|
|
require.Equal(t, selectedModel.ID, childChat.LastModelConfigID)
|
|
require.Equal(
|
|
t,
|
|
database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortLow, Valid: true},
|
|
childChat.LastReasoningEffort,
|
|
)
|
|
})
|
|
|
|
t.Run("BeatsOverrideCarriedEffort", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
deploymentModel := insertInternalChatModelConfig(
|
|
t, db, "general-deployment-"+uuid.NewString(), true,
|
|
)
|
|
require.NoError(t, db.UpsertChatGeneralModelOverride(
|
|
ctx, deploymentModel.ID.String()+":low",
|
|
))
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-effort-beats-override",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeGeneral,
|
|
Prompt: "delegate work",
|
|
ReasoningEffort: "high",
|
|
})
|
|
childChat, err := db.GetChatByID(ctx, requireSpawnAgentChildChatID(t, resp))
|
|
require.NoError(t, err)
|
|
require.Equal(t, deploymentModel.ID, childChat.LastModelConfigID)
|
|
require.Equal(
|
|
t,
|
|
database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffortHigh, Valid: true},
|
|
childChat.LastReasoningEffort,
|
|
)
|
|
})
|
|
}
|
|
|
|
func TestSpawnAgent_ModelArgsRejections(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
args func(ctx context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs
|
|
wantError string
|
|
}{
|
|
{
|
|
name: "MalformedModelConfigID",
|
|
args: func(context.Context, *testing.T, database.Store, database.User) spawnAgentArgs {
|
|
return spawnAgentArgs{ModelConfigID: "not-a-uuid"}
|
|
},
|
|
wantError: "invalid model_config_id: must be a valid UUID",
|
|
},
|
|
{
|
|
name: "UnknownModelConfigID",
|
|
args: func(context.Context, *testing.T, database.Store, database.User) spawnAgentArgs {
|
|
return spawnAgentArgs{ModelConfigID: uuid.NewString()}
|
|
},
|
|
wantError: "model_config_id not found or is disabled",
|
|
},
|
|
{
|
|
name: "DisabledModelConfig",
|
|
args: func(_ context.Context, t *testing.T, db database.Store, _ database.User) spawnAgentArgs {
|
|
disabled := insertInternalChatModelConfig(
|
|
t, db, "reject-disabled-"+uuid.NewString(), false,
|
|
)
|
|
return spawnAgentArgs{ModelConfigID: disabled.ID.String()}
|
|
},
|
|
wantError: "model_config_id not found or is disabled",
|
|
},
|
|
{
|
|
name: "DisabledProvider",
|
|
args: func(_ context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs {
|
|
dbgen.ChatProvider(t, db, database.ChatProvider{
|
|
Provider: "openai-compat",
|
|
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
|
}, func(p *database.InsertChatProviderParams) {
|
|
p.Enabled = false
|
|
})
|
|
config := insertInternalChatModelConfigForProvider(
|
|
t, db, "openai-compat", "gpt-4o-mini", true,
|
|
)
|
|
return spawnAgentArgs{ModelConfigID: config.ID.String()}
|
|
},
|
|
wantError: "model_config_id not found or is disabled",
|
|
},
|
|
{
|
|
name: "DeletedProvider",
|
|
args: func(_ context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs {
|
|
provider := dbgen.ChatProvider(t, db, database.ChatProvider{
|
|
Provider: "openai-compat",
|
|
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
|
})
|
|
config := insertInternalChatModelConfigForProvider(
|
|
t, db, "openai-compat", "gpt-4o-mini", true,
|
|
)
|
|
require.NoError(t, db.DeleteAIProviderByID(
|
|
systemRestrictedTestContext(t), provider.ID,
|
|
))
|
|
return spawnAgentArgs{ModelConfigID: config.ID.String()}
|
|
},
|
|
wantError: "model_config_id not found or is disabled",
|
|
},
|
|
{
|
|
name: "MissingCredentials",
|
|
args: func(_ context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs {
|
|
insertInternalChatProvider(
|
|
t, db, user.ID, "openai-compat", "", false, true, false,
|
|
)
|
|
config := insertInternalChatModelConfigForProvider(
|
|
t, db, "openai-compat", "gpt-4o-mini", true,
|
|
)
|
|
return spawnAgentArgs{ModelConfigID: config.ID.String()}
|
|
},
|
|
wantError: `credentials are unavailable for provider "openai-compat"`,
|
|
},
|
|
{
|
|
name: "InvalidReasoningEffort",
|
|
args: func(context.Context, *testing.T, database.Store, database.User) spawnAgentArgs {
|
|
return spawnAgentArgs{ReasoningEffort: "ultra"}
|
|
},
|
|
wantError: "invalid reasoning_effort: must be one of",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-rejection",
|
|
)
|
|
|
|
args := tt.args(ctx, t, db, user)
|
|
args.Type = subagentTypeGeneral
|
|
args.Prompt = "delegate work"
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, args)
|
|
require.True(t, resp.IsError, "expected tool error but got: %s", resp.Content)
|
|
require.Contains(t, resp.Content, tt.wantError)
|
|
|
|
children, err := db.GetChildChatsByParentIDs(ctx, database.GetChildChatsByParentIDsParams{
|
|
ParentIds: []uuid.UUID{parentChat.ID},
|
|
Archived: sql.NullBool{},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Empty(t, children, "no child chat may be created on rejection")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSpawnAgent_ComputerUseRejectsModelArgs(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
muter func(*spawnAgentArgs, database.ChatModelConfig)
|
|
}{
|
|
{
|
|
name: "ModelConfigID",
|
|
muter: func(args *spawnAgentArgs, model database.ChatModelConfig) {
|
|
args.ModelConfigID = model.ID.String()
|
|
},
|
|
},
|
|
{
|
|
name: "ReasoningEffort",
|
|
muter: func(args *spawnAgentArgs, _ database.ChatModelConfig) {
|
|
args.ReasoningEffort = "high"
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
insertEnabledAnthropicProvider(t, db, user.ID)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-computer-use-args",
|
|
)
|
|
|
|
args := spawnAgentArgs{
|
|
Type: subagentTypeComputerUse,
|
|
Prompt: "take a screenshot",
|
|
}
|
|
tt.muter(&args, model)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, args)
|
|
require.True(t, resp.IsError, "expected tool error but got: %s", resp.Content)
|
|
require.Contains(t, resp.Content, `not supported for type "computer_use"`)
|
|
|
|
children, err := db.GetChildChatsByParentIDs(ctx, database.GetChildChatsByParentIDsParams{
|
|
ParentIds: []uuid.UUID{parentChat.ID},
|
|
Archived: sql.NullBool{},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Empty(t, children, "no child chat may be created on rejection")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestListSubagentModels(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, defaultModel := seedInternalChatDeps(t, db)
|
|
effortModel := insertInternalChatModelConfigWithOptions(
|
|
t,
|
|
db,
|
|
"openai",
|
|
"gpt-effort-"+uuid.NewString(),
|
|
true,
|
|
json.RawMessage(`{"reasoning_effort":{"default":"medium","max":"high"}}`),
|
|
)
|
|
disabledModel := insertInternalChatModelConfig(
|
|
t, db, "list-disabled-"+uuid.NewString(), false,
|
|
)
|
|
insertInternalChatProvider(t, db, user.ID, "anthropic", "", false, true, false)
|
|
noCredsModel := insertInternalChatModelConfigForProvider(
|
|
t, db, "anthropic", "claude-sonnet-4-5", true,
|
|
)
|
|
dbgen.ChatProvider(t, db, database.ChatProvider{
|
|
Provider: "openai-compat",
|
|
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
|
}, func(p *database.InsertChatProviderParams) {
|
|
p.Enabled = false
|
|
})
|
|
disabledProviderModel := insertInternalChatModelConfigForProvider(
|
|
t, db, "openai-compat", "gpt-4o-mini", true,
|
|
)
|
|
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, defaultModel.ID, "parent-list-models",
|
|
)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
listSubagentModelsToolName,
|
|
struct{}{},
|
|
)
|
|
require.False(t, resp.IsError, "expected success but got: %s", resp.Content)
|
|
|
|
type listedModel struct {
|
|
ModelConfigID string `json:"model_config_id"`
|
|
DisplayName string `json:"display_name"`
|
|
Model string `json:"model"`
|
|
Provider string `json:"provider"`
|
|
ContextLimit int64 `json:"context_limit"`
|
|
IsDefault bool `json:"is_default"`
|
|
ReasoningEfforts []string `json:"reasoning_efforts"`
|
|
}
|
|
var result struct {
|
|
Models []listedModel `json:"models"`
|
|
}
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
|
|
byID := make(map[string]listedModel, len(result.Models))
|
|
for _, m := range result.Models {
|
|
byID[m.ModelConfigID] = m
|
|
}
|
|
require.Contains(t, byID, defaultModel.ID.String())
|
|
require.Contains(t, byID, effortModel.ID.String())
|
|
require.NotContains(t, byID, disabledModel.ID.String())
|
|
require.NotContains(t, byID, noCredsModel.ID.String())
|
|
require.NotContains(t, byID, disabledProviderModel.ID.String())
|
|
|
|
defaultEntry := byID[defaultModel.ID.String()]
|
|
require.Equal(t, defaultModel.DisplayName, defaultEntry.DisplayName)
|
|
require.Equal(t, defaultModel.Model, defaultEntry.Model)
|
|
require.Equal(t, "openai", defaultEntry.Provider)
|
|
require.Equal(t, defaultModel.ContextLimit, defaultEntry.ContextLimit)
|
|
require.True(t, defaultEntry.IsDefault)
|
|
|
|
effortEntry := byID[effortModel.ID.String()]
|
|
require.False(t, effortEntry.IsDefault)
|
|
require.Equal(
|
|
t,
|
|
[]string{"none", "minimal", "low", "medium", "high"},
|
|
effortEntry.ReasoningEfforts,
|
|
)
|
|
}
|
|
|
|
func TestListSubagentModels_NotAvailableForChildChats(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
_, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.True(t, childChat.ParentChatID.Valid, "child chat must have a parent")
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
childChat,
|
|
childChat.LastModelConfigID,
|
|
listSubagentModelsToolName,
|
|
struct{}{},
|
|
)
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "delegated chats cannot create child subagents")
|
|
}
|
|
|
|
func TestSpawnAgent_ExploreUsesConfiguredModelOverride(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
overrideModel := insertInternalChatModelConfig(
|
|
t, db, "explore-override-"+uuid.NewString(), true,
|
|
)
|
|
require.NoError(t, db.UpsertChatExploreModelOverride(ctx, overrideModel.ID.String()))
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-explore-override",
|
|
)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeExplore, Prompt: "investigate the codebase"},
|
|
)
|
|
result := requireSpawnAgentResponse(t, resp)
|
|
require.Equal(t, subagentTypeExplore, result.SubagentType)
|
|
childID, err := uuid.Parse(result.ChatID)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, overrideModel.ID, childChat.LastModelConfigID)
|
|
require.True(t, childChat.Mode.Valid)
|
|
require.Equal(t, database.ChatModeExplore, childChat.Mode.ChatMode)
|
|
require.False(t, childChat.PlanMode.Valid)
|
|
}
|
|
|
|
func TestSpawnAgent_ExploreFallsBackToCurrentTurnModel(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
currentTurnModel := insertInternalChatModelConfig(
|
|
t, db, "explore-current-turn-"+uuid.NewString(), true,
|
|
)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explore-fallback",
|
|
)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
currentTurnModel.ID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeExplore, Prompt: "trace the request flow"},
|
|
)
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, currentTurnModel.ID, childChat.LastModelConfigID)
|
|
require.Equal(t, parentModel.ID, parentChat.LastModelConfigID)
|
|
}
|
|
|
|
func TestSpawnAgent_ExploreHonorsPersonalModelOverrides(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
enablePersonalOverride bool
|
|
personalRaw func(database.ChatModelConfig) string
|
|
personalModel func(context.Context, *testing.T, database.Store, uuid.UUID) database.ChatModelConfig
|
|
wantModelID func(
|
|
database.ChatModelConfig,
|
|
database.ChatModelConfig,
|
|
database.ChatModelConfig,
|
|
database.ChatModelConfig,
|
|
) uuid.UUID
|
|
}{
|
|
{
|
|
name: "UnsetUsesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
wantModelID: func(_, _, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "DeploymentDefaultUsesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeDeploymentDefault)
|
|
},
|
|
wantModelID: func(_, _, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "ChatDefaultBypassesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeChatDefault)
|
|
},
|
|
wantModelID: func(_, currentTurnModel, _, _ database.ChatModelConfig) uuid.UUID {
|
|
return currentTurnModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "ModelUsesPersonalOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(personalModel database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeModel) + ":" +
|
|
personalModel.ID.String()
|
|
},
|
|
wantModelID: func(_, _, _, personalModel database.ChatModelConfig) uuid.UUID {
|
|
return personalModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "AdminFlagOffIgnoresPersonalOverride",
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeChatDefault)
|
|
},
|
|
wantModelID: func(_, _, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "DisabledPersonalModelFallsBackToDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalModel: func(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
) database.ChatModelConfig {
|
|
return insertInternalChatModelConfig(
|
|
t,
|
|
db,
|
|
"explore-personal-disabled-"+uuid.NewString(),
|
|
false,
|
|
)
|
|
},
|
|
personalRaw: func(personalModel database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeModel) + ":" +
|
|
personalModel.ID.String()
|
|
},
|
|
wantModelID: func(_, _, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "MissingCredentialsFallsBackToDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalModel: func(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
db database.Store,
|
|
userID uuid.UUID,
|
|
) database.ChatModelConfig {
|
|
insertInternalChatProvider(
|
|
t,
|
|
db,
|
|
userID,
|
|
"openai-compat",
|
|
"",
|
|
false,
|
|
true,
|
|
false,
|
|
)
|
|
return insertInternalChatModelConfigForProvider(
|
|
t,
|
|
db,
|
|
"openai-compat",
|
|
"gpt-4o-mini",
|
|
true,
|
|
)
|
|
},
|
|
personalRaw: func(personalModel database.ChatModelConfig) string {
|
|
return string(codersdk.ChatPersonalModelOverrideModeModel) + ":" +
|
|
personalModel.ID.String()
|
|
},
|
|
wantModelID: func(_, _, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
{
|
|
name: "MalformedValueUsesDeploymentOverride",
|
|
enablePersonalOverride: true,
|
|
personalRaw: func(database.ChatModelConfig) string {
|
|
return "not-a-mode"
|
|
},
|
|
wantModelID: func(_, _, deploymentModel, _ database.ChatModelConfig) uuid.UUID {
|
|
return deploymentModel.ID
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
currentTurnModel := insertInternalChatModelConfig(
|
|
t,
|
|
db,
|
|
"explore-current-turn-"+uuid.NewString(),
|
|
true,
|
|
)
|
|
deploymentModel := insertInternalChatModelConfig(
|
|
t,
|
|
db,
|
|
"explore-deployment-"+uuid.NewString(),
|
|
true,
|
|
)
|
|
require.NoError(t, db.UpsertChatExploreModelOverride(ctx, deploymentModel.ID.String()))
|
|
personalModel := insertInternalChatModelConfig(
|
|
t,
|
|
db,
|
|
"explore-personal-"+uuid.NewString(),
|
|
true,
|
|
)
|
|
if tt.personalModel != nil {
|
|
personalModel = tt.personalModel(ctx, t, db, user.ID)
|
|
}
|
|
if tt.enablePersonalOverride {
|
|
enableInternalChatPersonalModelOverrides(t, db)
|
|
}
|
|
if tt.personalRaw != nil {
|
|
upsertInternalUserChatPersonalModelOverride(
|
|
t,
|
|
db,
|
|
user.ID,
|
|
codersdk.ChatPersonalModelOverrideContextExplore,
|
|
tt.personalRaw(personalModel),
|
|
)
|
|
}
|
|
parentChat := createInternalParentChat(
|
|
ctx,
|
|
t,
|
|
server,
|
|
db,
|
|
org.ID,
|
|
user.ID,
|
|
parentModel.ID,
|
|
"parent-explore-personal-override",
|
|
)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
currentTurnModel.ID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeExplore, Prompt: "inspect the codebase"},
|
|
)
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(
|
|
t,
|
|
tt.wantModelID(parentModel, currentTurnModel, deploymentModel, personalModel),
|
|
childChat.LastModelConfigID,
|
|
)
|
|
require.True(t, childChat.Mode.Valid)
|
|
require.Equal(t, database.ChatModeExplore, childChat.Mode.ChatMode)
|
|
require.False(t, childChat.PlanMode.Valid)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCreateChat_ExploreRootStartsWithoutMCPSnapshot(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
root, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "root-explore",
|
|
ModelConfigID: model.ID,
|
|
ChatMode: database.NullChatMode{
|
|
ChatMode: database.ChatModeExplore,
|
|
Valid: true,
|
|
},
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("inspect the codebase")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
rootChat, err := db.GetChatByID(ctx, root.ID)
|
|
require.NoError(t, err)
|
|
require.Empty(t, rootChat.MCPServerIDs)
|
|
}
|
|
|
|
func TestResolveExploreToolSnapshot(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
user, _, _ := seedInternalChatDeps(t, db)
|
|
approvedMCP := insertInternalMCPServerConfig(
|
|
t, db, user.ID, "approved-"+uuid.NewString(), true,
|
|
)
|
|
blockedMCP := insertInternalMCPServerConfig(
|
|
t, db, user.ID, "blocked-"+uuid.NewString(), false,
|
|
)
|
|
|
|
// Build parent chats in memory rather than via server.CreateChat.
|
|
// resolveExploreToolSnapshot only reads ID, MCPServerIDs, PlanMode,
|
|
// ParentChatID, and Mode from its parent argument, so persisting
|
|
// the chats is unnecessary. Skipping CreateChat avoids waking the
|
|
// background acquireLoop, which would otherwise try to dial the
|
|
// fake MCP URLs and call OpenAI with the dbgen test API key. Those
|
|
// side effects were the root cause of the flake tracked in
|
|
// CODAGT-367.
|
|
askParent := database.Chat{
|
|
ID: uuid.New(),
|
|
MCPServerIDs: []uuid.UUID{approvedMCP.ID, blockedMCP.ID},
|
|
}
|
|
planParent := database.Chat{
|
|
ID: uuid.New(),
|
|
PlanMode: database.NullChatPlanMode{
|
|
ChatPlanMode: database.ChatPlanModePlan,
|
|
Valid: true,
|
|
},
|
|
MCPServerIDs: []uuid.UUID{approvedMCP.ID, blockedMCP.ID},
|
|
}
|
|
|
|
subagentPlanParent := planParent
|
|
subagentPlanParent.ID = uuid.New()
|
|
subagentPlanParent.ParentChatID = uuid.NullUUID{UUID: uuid.New(), Valid: true}
|
|
|
|
exploreParent := askParent
|
|
exploreParent.ID = uuid.New()
|
|
exploreParent.Mode = database.NullChatMode{ChatMode: database.ChatModeExplore, Valid: true}
|
|
exploreParent.ParentChatID = uuid.NullUUID{UUID: uuid.New(), Valid: true}
|
|
exploreParent.MCPServerIDs = []uuid.UUID{approvedMCP.ID}
|
|
|
|
tests := []struct {
|
|
name string
|
|
parent database.Chat
|
|
wantMCPServerIDs []uuid.UUID
|
|
}{
|
|
{
|
|
name: "AskModeRootSnapshotsAllExternalTools",
|
|
parent: askParent,
|
|
wantMCPServerIDs: []uuid.UUID{approvedMCP.ID, blockedMCP.ID},
|
|
},
|
|
{
|
|
name: "PlanModeRootKeepsOnlyApprovedExternalTools",
|
|
parent: planParent,
|
|
wantMCPServerIDs: []uuid.UUID{approvedMCP.ID},
|
|
},
|
|
{
|
|
name: "PlanModeSubagentKeepsNoExternalTools",
|
|
parent: subagentPlanParent,
|
|
wantMCPServerIDs: []uuid.UUID{},
|
|
},
|
|
{
|
|
name: "ExploreParentCannotReEscalateSnapshot",
|
|
parent: exploreParent,
|
|
wantMCPServerIDs: []uuid.UUID{approvedMCP.ID},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
tt := tt
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := chatdTestContext(t)
|
|
gotMCPServerIDs, err := server.resolveExploreToolSnapshot(
|
|
ctx,
|
|
tt.parent,
|
|
)
|
|
require.NoError(t, err)
|
|
require.ElementsMatch(t, tt.wantMCPServerIDs, gotMCPServerIDs)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCreateChildSubagentChatWithOptions_ExplorePersistsMCPSnapshot(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-explore-snapshot",
|
|
)
|
|
mcpCfg := insertInternalMCPServerConfig(
|
|
t, db, user.ID, "snapshot-"+uuid.NewString(), false,
|
|
)
|
|
|
|
child, err := server.createChildSubagentChatWithOptions(
|
|
ctx,
|
|
parentChat,
|
|
"inspect the codebase",
|
|
"explore-snapshot",
|
|
childSubagentChatOptions{
|
|
chatMode: database.NullChatMode{
|
|
ChatMode: database.ChatModeExplore,
|
|
Valid: true,
|
|
},
|
|
inheritedMCPServerIDs: []uuid.UUID{mcpCfg.ID},
|
|
},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.ElementsMatch(t, []uuid.UUID{mcpCfg.ID}, childChat.MCPServerIDs)
|
|
}
|
|
|
|
func TestSpawnAgent_ExploreSnapshotsTurnStateParentState(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
turnStartConfig := insertInternalMCPServerConfig(
|
|
t, db, user.ID, "turn-start-"+uuid.NewString(), false,
|
|
)
|
|
mutatedConfig := insertInternalMCPServerConfig(
|
|
t, db, user.ID, "mutated-"+uuid.NewString(), true,
|
|
)
|
|
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "parent-turn-state-snapshot",
|
|
ModelConfigID: model.ID,
|
|
MCPServerIDs: []uuid.UUID{turnStartConfig.ID},
|
|
InitialUserContent: []codersdk.ChatMessagePart{
|
|
codersdk.ChatMessageText("inspect the codebase"),
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
turnParent, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
tools := server.subagentTools(
|
|
ctx,
|
|
func() database.Chat { return turnParent },
|
|
turnParent.LastModelConfigID,
|
|
)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool, "spawn_agent tool must be present")
|
|
|
|
_, err = server.db.UpdateChatPlanModeByID(ctx, database.UpdateChatPlanModeByIDParams{
|
|
ID: turnParent.ID,
|
|
PlanMode: database.NullChatPlanMode{
|
|
ChatPlanMode: database.ChatPlanModePlan,
|
|
Valid: true,
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
_, err = server.db.UpdateChatMCPServerIDs(ctx, database.UpdateChatMCPServerIDsParams{
|
|
ID: turnParent.ID,
|
|
MCPServerIDs: []uuid.UUID{mutatedConfig.ID},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
reloadedParent, err := db.GetChatByID(ctx, turnParent.ID)
|
|
require.NoError(t, err)
|
|
require.True(t, reloadedParent.PlanMode.Valid)
|
|
require.Equal(t, database.ChatPlanModePlan, reloadedParent.PlanMode.ChatPlanMode)
|
|
require.ElementsMatch(t, []uuid.UUID{mutatedConfig.ID}, reloadedParent.MCPServerIDs)
|
|
|
|
input, err := json.Marshal(spawnAgentArgs{
|
|
Type: subagentTypeExplore,
|
|
Prompt: "inspect the codebase",
|
|
Title: "sub",
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: uuid.NewString(),
|
|
Name: spawnAgentToolName,
|
|
Input: string(input),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.True(t, childChat.Mode.Valid)
|
|
require.Equal(t, database.ChatModeExplore, childChat.Mode.ChatMode)
|
|
require.ElementsMatch(t, []uuid.UUID{turnStartConfig.ID}, childChat.MCPServerIDs,
|
|
"Explore child should keep the turn-start MCP snapshot after parent mutations")
|
|
}
|
|
|
|
func TestSpawnAgent_ExploreFallsBackOnInvalidUUID(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
currentTurnModel := insertInternalChatModelConfig(
|
|
t, db, "explore-invalid-override-"+uuid.NewString(), true,
|
|
)
|
|
require.NoError(t, db.UpsertChatExploreModelOverride(ctx, "not-a-uuid"))
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explore-invalid-override",
|
|
)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
currentTurnModel.ID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeExplore, Prompt: "inspect the handler flow"},
|
|
)
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, currentTurnModel.ID, childChat.LastModelConfigID)
|
|
}
|
|
|
|
func TestSpawnAgent_ExploreFallsBackWhenOverrideIsUnavailable(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
currentTurnModel := insertInternalChatModelConfig(
|
|
t, db, "explore-fallback-current-"+uuid.NewString(), true,
|
|
)
|
|
disabledModel := insertInternalChatModelConfig(
|
|
t, db, "explore-disabled-"+uuid.NewString(), false,
|
|
)
|
|
require.NoError(t, db.UpsertChatExploreModelOverride(ctx, disabledModel.ID.String()))
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explore-disabled",
|
|
)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
currentTurnModel.ID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeExplore, Prompt: "inspect the service boundaries"},
|
|
)
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, currentTurnModel.ID, childChat.LastModelConfigID)
|
|
}
|
|
|
|
func TestSpawnAgent_ExploreFallsBackWhenOverrideCredentialsAreUnavailable(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, parentModel := seedInternalChatDeps(t, db)
|
|
currentTurnModel := insertInternalChatModelConfig(
|
|
t, db, "explore-missing-user-key-current-"+uuid.NewString(), true,
|
|
)
|
|
overrideProvider := 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
|
|
})
|
|
|
|
overrideModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
|
AIProviderID: uuid.NullUUID{UUID: overrideProvider.ID, Valid: true},
|
|
Model: "gpt-4o-mini",
|
|
DisplayName: "Explore Override Missing User Key",
|
|
})
|
|
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",
|
|
)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
currentTurnModel.ID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeExplore, Prompt: "inspect provider credential handling"},
|
|
)
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, currentTurnModel.ID, childChat.LastModelConfigID)
|
|
}
|
|
|
|
func TestDefaultSystemPromptPlanningGuidance_SteersSubagentSelection(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
require.Contains(t, defaultSystemPromptPlanningGuidance, `Prefer type="general" for substantial delegated research, analysis, reasoning, review, planning support, or implementation`)
|
|
require.Contains(t, defaultSystemPromptPlanningGuidance, `Use type="general" even for read-only work when the task is open-ended, multi-step, parallel, requires synthesis, or may later need edits`)
|
|
require.Contains(t, defaultSystemPromptPlanningGuidance, `Use type="explore" only for narrow repository-local read-only code discovery or code tracing`)
|
|
require.Contains(t, defaultSystemPromptPlanningGuidance, `Do not use type="explore" for generic research, broad architecture analysis, planning synthesis, external or web research, parallel research, or tasks that may need edits`)
|
|
require.NotContains(t, defaultSystemPromptPlanningGuidance, "research the codebase")
|
|
require.NotContains(t, defaultSystemPromptPlanningGuidance, "Reserve type=\"general\" for writable delegated work")
|
|
}
|
|
|
|
func TestSpawnAgent_DescriptionListsAllAvailableTypes(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
|
Anthropic: "test-anthropic-key",
|
|
})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-description-all",
|
|
)
|
|
|
|
tools := server.subagentTools(ctx, func() database.Chat { return parentChat }, parentChat.LastModelConfigID)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool, "spawn_agent tool must be present")
|
|
description := tool.Info().Description
|
|
require.Contains(t, description, subagentTypeGeneral)
|
|
require.Contains(t, description, subagentTypeExplore)
|
|
require.Contains(t, description, subagentTypeComputerUse)
|
|
}
|
|
|
|
func TestSpawnAgent_DescriptionSteersGeneralForSubstantialResearch(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-description-selection-guidance",
|
|
)
|
|
|
|
tools := server.subagentTools(ctx, func() database.Chat { return parentChat }, parentChat.LastModelConfigID)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool, "spawn_agent tool must be present")
|
|
description := tool.Info().Description
|
|
|
|
require.Contains(t, description, `Prefer type="general" for substantial delegated research, analysis, reasoning, review, planning support, or implementation`)
|
|
require.Contains(t, description, "even when the child should only report findings")
|
|
require.Contains(t, description, `When using type="general" for read-only work, explicitly instruct the child not to modify files and to return findings`)
|
|
require.Contains(t, description, `Use type="explore" only for narrow repository-local read-only code discovery or code tracing`)
|
|
require.Contains(t, description, `Do not use type="explore" for generic research, broad architecture analysis, planning synthesis, external or web research, parallel research, or tasks that may need edits`)
|
|
}
|
|
|
|
func TestSpawnAgent_DescriptionIncludesComputerUseWithMissingProviderKey(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-description-missing-key",
|
|
)
|
|
|
|
tools := server.subagentTools(ctx, func() database.Chat { return parentChat }, parentChat.LastModelConfigID)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool, "spawn_agent tool must be present")
|
|
description := tool.Info().Description
|
|
require.Contains(t, description, subagentTypeGeneral)
|
|
require.Contains(t, description, subagentTypeExplore)
|
|
require.Contains(t, description, subagentTypeComputerUse)
|
|
}
|
|
|
|
func TestSpawnAgent_PlanModeDescriptionOmitsComputerUse(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
|
Anthropic: "test-anthropic-key",
|
|
})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "plan-parent-description",
|
|
ModelConfigID: model.ID,
|
|
PlanMode: database.NullChatPlanMode{
|
|
ChatPlanMode: database.ChatPlanModePlan,
|
|
Valid: true,
|
|
},
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("plan this change")},
|
|
})
|
|
require.NoError(t, err)
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
tools := server.subagentTools(ctx, func() database.Chat { return parentChat }, parentChat.LastModelConfigID)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool, "spawn_agent tool must be present")
|
|
description := tool.Info().Description
|
|
require.Contains(t, description, subagentTypeGeneral)
|
|
require.Contains(t, description, subagentTypeExplore)
|
|
require.NotContains(t, description, subagentTypeComputerUse)
|
|
require.Contains(t, description, `type="general" is for non-mutating substantial investigation and planning support`)
|
|
require.Contains(t, description, `type="explore" is for narrow repository-local lookup or tracing`)
|
|
require.Contains(t, description, `only type="general" should be used for cloning repositories or non-local investigation`)
|
|
require.NotContains(t, description, "Both may use shell commands for exploration, such as cloning repositories")
|
|
require.Contains(t, description, "must not implement changes or intentionally modify workspace files")
|
|
}
|
|
|
|
func TestSpawnAgent_PlanModeRejectsComputerUse(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
|
Anthropic: "test-anthropic-key",
|
|
})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "plan-parent-computer-use-reject",
|
|
ModelConfigID: model.ID,
|
|
PlanMode: database.NullChatPlanMode{
|
|
ChatPlanMode: database.ChatPlanModePlan,
|
|
Valid: true,
|
|
},
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("plan this change")},
|
|
})
|
|
require.NoError(t, err)
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeComputerUse,
|
|
Prompt: "open the browser and click around",
|
|
})
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, `type "computer_use" is unavailable in plan mode`)
|
|
}
|
|
|
|
func TestPlanningOverlaySubagentGuidance_UsesPlanModeSafeDescriptions(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
guidance := planningOverlaySubagentGuidance()
|
|
|
|
require.Contains(t, guidance, subagentTypeGeneral)
|
|
require.Contains(t, guidance, subagentTypeExplore)
|
|
require.Contains(t, guidance, `Use type="general" for substantial investigation, reasoning, and planning support`)
|
|
require.Contains(t, guidance, `Use type="explore" only for narrow repository-local lookup or tracing`)
|
|
require.Contains(t, guidance, "general (non-mutating substantial investigation, analysis, and planning support)")
|
|
require.Contains(t, guidance, "explore (narrow repository-local codebase lookup and code tracing)")
|
|
require.NotContains(t, guidance, subagentTypeComputerUse)
|
|
require.NotContains(t, guidance, "modify")
|
|
require.NotContains(t, guidance, "may inspect or modify workspace files")
|
|
}
|
|
|
|
func TestSpawnAgent_InvalidTypeAndCredentialErrorAreDistinct(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-invalid-type",
|
|
)
|
|
|
|
invalidResp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: "invalid", Prompt: "delegate work"},
|
|
)
|
|
require.True(t, invalidResp.IsError)
|
|
require.Contains(t, invalidResp.Content, "type must be one of: general, explore, computer_use")
|
|
|
|
credentialResp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeComputerUse, Prompt: "open browser"},
|
|
)
|
|
require.True(t, credentialResp.IsError)
|
|
require.Contains(t, credentialResp.Content, "API key")
|
|
require.Contains(t, credentialResp.Content, "computer-use")
|
|
require.Contains(t, credentialResp.Content, "anthropic")
|
|
}
|
|
|
|
func TestSpawnAgent_ComputerUseAvailabilityUsesConfiguredProvider(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
ctx := chatdTestContext(t)
|
|
require.NoError(t, db.UpsertChatComputerUseProvider(
|
|
ctx,
|
|
string(codersdk.ChatComputerUseProviderOpenAI),
|
|
))
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-openai-computer-use",
|
|
)
|
|
|
|
ids := availableSubagentTypeIDs(ctx, server, parentChat)
|
|
require.Contains(t, ids, subagentTypeComputerUse)
|
|
}
|
|
|
|
func TestSpawnAgent_ComputerUseRejectsMissingConfiguredProvider(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
ctx := chatdTestContext(t)
|
|
require.NoError(t, db.UpsertChatComputerUseProvider(
|
|
ctx,
|
|
string(codersdk.ChatComputerUseProviderOpenAI),
|
|
))
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
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,
|
|
})
|
|
model := insertInternalChatModelConfigForProvider(
|
|
t,
|
|
db,
|
|
string(codersdk.ChatComputerUseProviderOpenAI),
|
|
"gpt-4o-mini",
|
|
true,
|
|
)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-openai-missing",
|
|
)
|
|
|
|
ids := availableSubagentTypeIDs(ctx, server, parentChat)
|
|
require.Contains(t, ids, subagentTypeComputerUse)
|
|
beforeChats, err := db.GetChats(ctx, database.GetChatsParams{
|
|
OwnedOnly: true,
|
|
ViewerID: user.ID,
|
|
AfterID: uuid.Nil,
|
|
OffsetOpt: 0,
|
|
LimitOpt: 100,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeComputerUse,
|
|
Prompt: "open the browser",
|
|
})
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "API key")
|
|
require.Contains(t, resp.Content, "computer-use")
|
|
require.Contains(t, resp.Content, "openai")
|
|
afterChats, err := db.GetChats(ctx, database.GetChatsParams{
|
|
OwnedOnly: true,
|
|
ViewerID: user.ID,
|
|
AfterID: uuid.Nil,
|
|
OffsetOpt: 0,
|
|
LimitOpt: 100,
|
|
})
|
|
require.NoError(t, err)
|
|
require.Len(t, afterChats, len(beforeChats))
|
|
}
|
|
|
|
func TestSpawnAgent_ComputerUseRejectsInvalidConfiguredProviderWithStableReason(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
ctx := chatdTestContext(t)
|
|
require.NoError(t, db.UpsertChatComputerUseProvider(ctx, "bogus"))
|
|
logSink := &subagentTestLogSink{}
|
|
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).AppendSinks(logSink)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerLogger(logger))
|
|
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-invalid-computer-use-provider",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeComputerUse,
|
|
Prompt: "open the browser",
|
|
})
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, `type "computer_use" is unavailable because its provider configuration could not be loaded`)
|
|
require.NotContains(t, resp.Content, "bogus")
|
|
require.NotContains(t, resp.Content, "agents_computer_use_provider")
|
|
require.NotEmpty(t, logSink.entriesAtLevelWithMessage(
|
|
slog.LevelWarn,
|
|
"computer-use provider config is unavailable",
|
|
))
|
|
}
|
|
|
|
func TestSpawnAgent_ComputerUseRejectsDesktopDisabled(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
experiments := slices.DeleteFunc(
|
|
slices.Clone(codersdk.ExperimentsKnown),
|
|
func(e codersdk.Experiment) bool { return e == codersdk.ExperimentChatVirtualDesktop },
|
|
)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
|
Anthropic: "test-anthropic-key",
|
|
}, withInternalTestServerExperiments(experiments))
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-desktop-disabled",
|
|
)
|
|
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: subagentTypeComputerUse,
|
|
Prompt: "open the browser",
|
|
})
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, `type "computer_use" is unavailable because the chat-virtual-desktop experiment is not enabled`)
|
|
}
|
|
|
|
func TestSpawnAgent_BlankTypeReturnsValidOptions(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
|
Anthropic: "test-anthropic-key",
|
|
})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parentChat := createInternalParentChat(
|
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-blank-type",
|
|
)
|
|
|
|
tests := []struct {
|
|
name string
|
|
subagentType string
|
|
}{
|
|
{name: "empty", subagentType: ""},
|
|
{name: "space", subagentType: " "},
|
|
{name: "whitespace", subagentType: "\n\t"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := chatdTestContext(t)
|
|
resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: tt.subagentType,
|
|
Prompt: "delegate work",
|
|
})
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "type must be one of:")
|
|
require.Contains(t, resp.Content, subagentTypeGeneral)
|
|
require.Contains(t, resp.Content, subagentTypeExplore)
|
|
require.Contains(t, resp.Content, subagentTypeComputerUse)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSpawnAgent_NotAvailableForChildChats(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
|
Anthropic: "test-anthropic-key",
|
|
})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
_, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.True(t, childChat.ParentChatID.Valid, "child chat must have a parent")
|
|
|
|
tools := server.subagentTools(ctx, func() database.Chat { return childChat }, childChat.LastModelConfigID)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool, "spawn_agent tool must be present")
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: "call-child",
|
|
Name: spawnAgentToolName,
|
|
Input: `{"type":"general","prompt":"open browser"}`,
|
|
})
|
|
require.NoError(t, err)
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "delegated chats cannot create child subagents")
|
|
}
|
|
|
|
func TestSpawnAgent_NotAvailableForExploreChats(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
exploreChat, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "root-explore",
|
|
ModelConfigID: model.ID,
|
|
ChatMode: database.NullChatMode{
|
|
ChatMode: database.ChatModeExplore,
|
|
Valid: true,
|
|
},
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("inspect the codebase")},
|
|
})
|
|
require.NoError(t, err)
|
|
currentChat, err := db.GetChatByID(ctx, exploreChat.ID)
|
|
require.NoError(t, err)
|
|
|
|
tools := server.subagentTools(ctx, func() database.Chat { return currentChat }, currentChat.LastModelConfigID)
|
|
tool := findToolByName(tools, spawnAgentToolName)
|
|
require.NotNil(t, tool, "spawn_agent tool must be present")
|
|
|
|
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
|
ID: "call-explore",
|
|
Name: spawnAgentToolName,
|
|
Input: `{"type":"general","prompt":"delegate work"}`,
|
|
})
|
|
require.NoError(t, err)
|
|
require.True(t, resp.IsError)
|
|
require.Contains(t, resp.Content, "explore chats cannot create child subagents")
|
|
}
|
|
|
|
func TestSubagentLifecycleToolsIncludePersistedSubagentTypeAcrossVariants(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
variant string
|
|
}{
|
|
{name: "General", variant: subagentTypeGeneral},
|
|
{name: "Explore", variant: subagentTypeExplore},
|
|
{name: "ComputerUse", variant: subagentTypeComputerUse},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
if tt.variant == subagentTypeComputerUse {
|
|
insertEnabledAnthropicProvider(t, db, user.ID)
|
|
}
|
|
parentChat := createInternalParentChat(
|
|
ctx,
|
|
t,
|
|
server,
|
|
db,
|
|
org.ID,
|
|
user.ID,
|
|
model.ID,
|
|
"parent-lifecycle-"+tt.variant,
|
|
)
|
|
|
|
spawnResp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{
|
|
Type: tt.variant,
|
|
Prompt: "delegate work",
|
|
})
|
|
spawnResult := requireSpawnAgentResponse(t, spawnResp)
|
|
require.Equal(t, tt.variant, spawnResult.SubagentType)
|
|
childID, err := uuid.Parse(spawnResult.ChatID)
|
|
require.NoError(t, err)
|
|
|
|
setChatStatus(ctx, t, db, childID, database.ChatStatusWaiting, "")
|
|
insertAssistantMessage(t, db, childID, model.ID, "task complete")
|
|
waitResult := requireToolResponseMap(t, runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
"wait_agent",
|
|
waitAgentArgs{ChatID: childID.String()},
|
|
), false)
|
|
require.Equal(t, tt.variant, waitResult["type"])
|
|
|
|
messageResult := requireToolResponseMap(t, runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
"message_agent",
|
|
messageAgentArgs{ChatID: childID.String(), Message: "follow up"},
|
|
), false)
|
|
require.Equal(t, tt.variant, messageResult["type"])
|
|
|
|
setChatStatus(ctx, t, db, childID, database.ChatStatusRunning, "")
|
|
interruptResult := requireToolResponseMap(t, runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
"interrupt_agent",
|
|
interruptAgentArgs{ChatID: childID.String()},
|
|
), false)
|
|
require.Equal(t, tt.variant, interruptResult["type"])
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSubagentLifecycleToolErrorsIncludePersistedSubagentType(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
_, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
unrelated, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "unrelated-lifecycle-parent",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("other")},
|
|
})
|
|
require.NoError(t, err)
|
|
unrelatedChat, err := db.GetChatByID(ctx, unrelated.ID)
|
|
require.NoError(t, err)
|
|
|
|
tests := []struct {
|
|
name string
|
|
toolName string
|
|
args any
|
|
wantError string
|
|
}{
|
|
{
|
|
name: "WaitAgent",
|
|
toolName: "wait_agent",
|
|
args: waitAgentArgs{ChatID: child.ID.String()},
|
|
wantError: ErrSubagentNotDescendant.Error(),
|
|
},
|
|
{
|
|
name: "MessageAgent",
|
|
toolName: "message_agent",
|
|
args: messageAgentArgs{ChatID: child.ID.String(), Message: "follow up"},
|
|
wantError: ErrSubagentNotDescendant.Error(),
|
|
},
|
|
{
|
|
name: "InterruptAgent",
|
|
toolName: "interrupt_agent",
|
|
args: interruptAgentArgs{ChatID: child.ID.String()},
|
|
wantError: ErrSubagentNotDescendant.Error(),
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
ctx := chatdTestContext(t)
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
unrelatedChat,
|
|
unrelatedChat.LastModelConfigID,
|
|
tt.toolName,
|
|
tt.args,
|
|
), true)
|
|
require.Equal(t, subagentTypeGeneral, result["type"])
|
|
require.Equal(t, tt.wantError, result["error"])
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSpawnAgent_ComputerUseUsesComputerUseModelNotParent(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
insertEnabledAnthropicProvider(t, db, user.ID)
|
|
workspace, build, agent := seedWorkspaceBinding(t, db, user.ID)
|
|
|
|
seedProvider, err := db.GetAIProviderByID(ctx, model.AIProviderID.UUID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "openai", string(seedProvider.Type), "seed helper must create an OpenAI model")
|
|
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
|
BuildID: uuid.NullUUID{UUID: build.ID, Valid: true},
|
|
AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true},
|
|
Title: "parent-openai",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeComputerUse, Prompt: "take a screenshot"},
|
|
)
|
|
result := requireSpawnAgentResponse(t, resp)
|
|
require.Equal(t, subagentTypeComputerUse, result.SubagentType)
|
|
childID, err := uuid.Parse(result.ChatID)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
|
|
require.Equal(t, parentChat.WorkspaceID, childChat.WorkspaceID)
|
|
require.Equal(t, parentChat.BuildID, childChat.BuildID)
|
|
require.Equal(t, parentChat.AgentID, childChat.AgentID)
|
|
require.True(t, childChat.Mode.Valid)
|
|
assert.Equal(t, database.ChatModeComputerUse, childChat.Mode.ChatMode)
|
|
computerUseModelProvider, computerUseModelName, ok := chattool.DefaultComputerUseModel(codersdk.ChatComputerUseProviderAnthropic)
|
|
require.True(t, ok)
|
|
assert.NotEqual(t, string(seedProvider.Type), computerUseModelProvider,
|
|
"computer use model provider must differ from parent model provider")
|
|
assert.Equal(t, "anthropic", computerUseModelProvider)
|
|
assert.NotEmpty(t, computerUseModelName)
|
|
}
|
|
|
|
func TestSpawnAgent_ComputerUseInheritsMCPServerIDs(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
insertEnabledAnthropicProvider(t, db, 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},
|
|
})
|
|
|
|
parentMCPIDs := []uuid.UUID{mcpCfg.ID}
|
|
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "parent-cu-mcp",
|
|
ModelConfigID: model.ID,
|
|
MCPServerIDs: parentMCPIDs,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parentChat,
|
|
parentChat.LastModelConfigID,
|
|
spawnAgentToolName,
|
|
spawnAgentArgs{Type: subagentTypeComputerUse, Prompt: "check the UI"},
|
|
)
|
|
childID := requireSpawnAgentChildChatID(t, resp)
|
|
|
|
childChat, err := db.GetChatByID(ctx, childID)
|
|
require.NoError(t, err)
|
|
assert.ElementsMatch(t, parentMCPIDs, childChat.MCPServerIDs,
|
|
"computer use child chat must inherit MCP server IDs from parent")
|
|
}
|
|
|
|
func TestCreateChildSubagentChat_InheritsMCPServerIDs(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
// Insert two MCP server configs so we can verify both are
|
|
// inherited by the child chat.
|
|
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},
|
|
})
|
|
|
|
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},
|
|
})
|
|
|
|
parentMCPIDs := []uuid.UUID{mcpA.ID, mcpB.ID}
|
|
|
|
// Create a parent chat with MCP servers.
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "parent-with-mcp",
|
|
ModelConfigID: model.ID,
|
|
MCPServerIDs: parentMCPIDs,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Refetch the parent to get DB-populated fields.
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
require.ElementsMatch(t, parentMCPIDs, parentChat.MCPServerIDs,
|
|
"parent chat must have the MCP server IDs we set")
|
|
|
|
// Spawn a child subagent chat.
|
|
child, err := server.createChildSubagentChatWithOptions(
|
|
ctx,
|
|
parentChat,
|
|
"do some work",
|
|
"child-task",
|
|
childSubagentChatOptions{},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
// Verify the child inherited the parent's MCP server IDs.
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
assert.ElementsMatch(t, parentMCPIDs, childChat.MCPServerIDs,
|
|
"child chat must inherit MCP server IDs from parent")
|
|
}
|
|
|
|
func TestCreateChildSubagentChat_NoMCPServersStaysEmpty(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
// Create a parent chat without any MCP servers.
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "parent-no-mcp",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
|
|
// Spawn a child.
|
|
child, err := server.createChildSubagentChatWithOptions(
|
|
ctx,
|
|
parentChat,
|
|
"do some work",
|
|
"child-no-mcp",
|
|
childSubagentChatOptions{},
|
|
)
|
|
require.NoError(t, err)
|
|
|
|
childChat, err := db.GetChatByID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, childChat.MCPServerIDs,
|
|
"child chat must have empty MCP server IDs when parent has none")
|
|
}
|
|
|
|
func TestIsSubagentDescendant(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
// Build a chain: root -> child -> grandchild.
|
|
root, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "root",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("root")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
child, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
ParentChatID: uuid.NullUUID{
|
|
UUID: root.ID,
|
|
Valid: true,
|
|
},
|
|
RootChatID: uuid.NullUUID{
|
|
UUID: root.ID,
|
|
Valid: true,
|
|
},
|
|
Title: "child",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("child")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
grandchild, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
ParentChatID: uuid.NullUUID{
|
|
UUID: child.ID,
|
|
Valid: true,
|
|
},
|
|
RootChatID: uuid.NullUUID{
|
|
UUID: root.ID,
|
|
Valid: true,
|
|
},
|
|
Title: "grandchild",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("grandchild")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
// Build a separate, unrelated chain.
|
|
unrelated, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "unrelated-root",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("unrelated")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
unrelatedChild, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
ParentChatID: uuid.NullUUID{
|
|
UUID: unrelated.ID,
|
|
Valid: true,
|
|
},
|
|
RootChatID: uuid.NullUUID{
|
|
UUID: unrelated.ID,
|
|
Valid: true,
|
|
},
|
|
Title: "unrelated-child",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("unrelated-child")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
tests := []struct {
|
|
name string
|
|
ancestor uuid.UUID
|
|
target uuid.UUID
|
|
want bool
|
|
}{
|
|
{
|
|
name: "SameID",
|
|
ancestor: root.ID,
|
|
target: root.ID,
|
|
want: false,
|
|
},
|
|
{
|
|
name: "DirectChild",
|
|
ancestor: root.ID,
|
|
target: child.ID,
|
|
want: true,
|
|
},
|
|
{
|
|
name: "GrandChild",
|
|
ancestor: root.ID,
|
|
target: grandchild.ID,
|
|
want: true,
|
|
},
|
|
{
|
|
name: "Unrelated",
|
|
ancestor: root.ID,
|
|
target: unrelatedChild.ID,
|
|
want: false,
|
|
},
|
|
{
|
|
name: "RootChat",
|
|
ancestor: child.ID,
|
|
target: root.ID,
|
|
want: false,
|
|
},
|
|
{
|
|
name: "BrokenChain",
|
|
ancestor: root.ID,
|
|
target: uuid.New(),
|
|
want: false,
|
|
},
|
|
{
|
|
name: "NotDescendant",
|
|
ancestor: unrelated.ID,
|
|
target: child.ID,
|
|
want: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
got, err := isSubagentDescendant(ctx, db, tt.ancestor, tt.target)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, tt.want, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
// createParentChildChats creates a parent and child chat pair for
|
|
// subagent tests. The child starts in pending status.
|
|
func createParentChildChats(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
server *Server,
|
|
user database.User,
|
|
org database.Organization,
|
|
model database.ChatModelConfig,
|
|
) (parent database.Chat, child database.Chat) {
|
|
t.Helper()
|
|
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "parent-" + t.Name(),
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
child, err = server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
ParentChatID: uuid.NullUUID{
|
|
UUID: parent.ID,
|
|
Valid: true,
|
|
},
|
|
RootChatID: uuid.NullUUID{
|
|
UUID: parent.ID,
|
|
Valid: true,
|
|
},
|
|
Title: "child-" + t.Name(),
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("do work")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
return parent, child
|
|
}
|
|
|
|
// setChatStatus transitions a chat to the given status.
|
|
func setChatStatus(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
db database.Store,
|
|
chatID uuid.UUID,
|
|
status database.ChatStatus,
|
|
lastError string,
|
|
) {
|
|
t.Helper()
|
|
|
|
params := database.UpdateChatStatusParams{
|
|
ID: chatID,
|
|
Status: status,
|
|
}
|
|
if lastError != "" {
|
|
encodedLastError, err := json.Marshal(codersdk.ChatError{
|
|
Message: lastError,
|
|
Kind: codersdk.ChatErrorKindGeneric,
|
|
})
|
|
require.NoError(t, err)
|
|
params.LastError = pqtype.NullRawMessage{RawMessage: encodedLastError, Valid: true}
|
|
}
|
|
_, err := db.UpdateChatStatus(ctx, params)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func setChatStatusWithError(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
db database.Store,
|
|
chatID uuid.UUID,
|
|
status database.ChatStatus,
|
|
chatErr codersdk.ChatError,
|
|
) {
|
|
t.Helper()
|
|
|
|
encoded, err := json.Marshal(chatErr)
|
|
require.NoError(t, err)
|
|
_, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
|
|
ID: chatID,
|
|
Status: status,
|
|
LastError: pqtype.NullRawMessage{RawMessage: encoded, Valid: true},
|
|
})
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
// insertAssistantMessage inserts an assistant message with v1 content
|
|
// into a chat.
|
|
func insertAssistantMessage(
|
|
t *testing.T,
|
|
db database.Store,
|
|
chatID uuid.UUID,
|
|
modelID uuid.UUID,
|
|
text string,
|
|
) {
|
|
t.Helper()
|
|
|
|
parts := []codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}
|
|
data, err := json.Marshal(parts)
|
|
require.NoError(t, err)
|
|
|
|
_ = 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,
|
|
})
|
|
}
|
|
|
|
func insertLinkedChatFile(
|
|
ctx context.Context,
|
|
t *testing.T,
|
|
db database.Store,
|
|
chatID uuid.UUID,
|
|
ownerID uuid.UUID,
|
|
organizationID uuid.UUID,
|
|
name string,
|
|
mediaType string,
|
|
data []byte,
|
|
) uuid.UUID {
|
|
t.Helper()
|
|
|
|
file, err := db.InsertChatFile(ctx, database.InsertChatFileParams{
|
|
OwnerID: ownerID,
|
|
OrganizationID: organizationID,
|
|
Name: name,
|
|
Mimetype: mediaType,
|
|
Data: data,
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
rejected, err := db.LinkChatFiles(ctx, database.LinkChatFilesParams{
|
|
ChatID: chatID,
|
|
MaxFileLinks: int32(codersdk.MaxChatFileIDs),
|
|
FileIds: []uuid.UUID{file.ID},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Zero(t, rejected)
|
|
|
|
return file.ID
|
|
}
|
|
|
|
func TestWaitAgentDoesNotRelayComputerUseSubagentAttachments(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
workspace, _, agent := seedWorkspaceBinding(t, db, user.ID)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
parent, child := createComputerUseParentChild(
|
|
t, server, user, org, model, workspace, agent,
|
|
"parent-relay", "child-relay",
|
|
)
|
|
|
|
insertedFile := insertLinkedChatFile(
|
|
ctx,
|
|
t,
|
|
db,
|
|
child.ID,
|
|
user.ID,
|
|
workspace.OrganizationID,
|
|
"screenshot.png",
|
|
"image/png",
|
|
[]byte("fake-png"),
|
|
)
|
|
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)
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError, "expected successful response, got: %s", resp.Content)
|
|
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
require.Equal(t, "Shared the screenshot.", result["report"])
|
|
require.Equal(t, string(database.ChatStatusWaiting), result["status"])
|
|
assert.NotContains(t, result, "attachment_count")
|
|
assert.NotContains(t, result, "attachment_warning")
|
|
|
|
attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, attachments)
|
|
parts := buildAssistantPartsForPersist(
|
|
context.Background(),
|
|
testutil.Logger(t),
|
|
nil,
|
|
[]fantasy.ToolResultContent{{
|
|
ToolCallID: "call-1",
|
|
ToolName: "wait_agent",
|
|
ClientMetadata: resp.Metadata,
|
|
}},
|
|
chatloop.PersistedStep{},
|
|
nil,
|
|
)
|
|
assert.Empty(t, parts)
|
|
|
|
parentFiles, err := db.GetChatFileMetadataByChatID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, parentFiles)
|
|
|
|
childFiles, err := db.GetChatFileMetadataByChatID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.Len(t, childFiles, 1)
|
|
assert.Equal(t, insertedFile, childFiles[0].ID)
|
|
assert.Equal(t, "screenshot.png", childFiles[0].Name)
|
|
assert.Equal(t, "image/png", childFiles[0].Mimetype)
|
|
}
|
|
|
|
func TestWaitAgentDoesNotRelayRegularSubagentAttachments(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
workspace, _, _ := seedWorkspaceBinding(t, db, user.ID)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
WaitUntilIdleForTest(server)
|
|
|
|
insertedFile := insertLinkedChatFile(
|
|
ctx,
|
|
t,
|
|
db,
|
|
child.ID,
|
|
user.ID,
|
|
workspace.OrganizationID,
|
|
"notes.txt",
|
|
"text/plain",
|
|
[]byte("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)
|
|
require.NoError(t, err)
|
|
require.False(t, resp.IsError, "expected successful response, got: %s", resp.Content)
|
|
|
|
var result map[string]any
|
|
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
|
require.Equal(t, "Shared the release notes.", result["report"])
|
|
assert.NotContains(t, result, "attachment_count")
|
|
assert.NotContains(t, result, "attachment_warning")
|
|
attachments, err := chattool.AttachmentsFromMetadata(resp.Metadata)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, attachments)
|
|
|
|
parentFiles, err := db.GetChatFileMetadataByChatID(ctx, parent.ID)
|
|
require.NoError(t, err)
|
|
assert.Empty(t, parentFiles)
|
|
|
|
childFiles, err := db.GetChatFileMetadataByChatID(ctx, child.ID)
|
|
require.NoError(t, err)
|
|
require.Len(t, childFiles, 1)
|
|
assert.Equal(t, insertedFile, childFiles[0].ID)
|
|
assert.Equal(t, "notes.txt", childFiles[0].Name)
|
|
assert.Equal(t, "text/plain", childFiles[0].Mimetype)
|
|
}
|
|
|
|
func TestAwaitSubagentCompletion(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
// Shared fixtures for subtests that use a real clock. Each
|
|
// subtest creates its own parent+child chats (unique IDs)
|
|
// so they don't collide. Mock-clock subtests need their own
|
|
// DB and server so the wait loop's timers stay isolated.
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
t.Run("NotDescendant", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
|
|
parent, _ := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
unrelated, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: "unrelated",
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("other")},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
_, _, err = server.awaitSubagentCompletion(
|
|
ctx, parent.ID, unrelated.ID, time.Second,
|
|
)
|
|
require.ErrorIs(t, err, ErrSubagentNotDescendant)
|
|
})
|
|
|
|
t.Run("AlreadyWaiting", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
|
insertAssistantMessage(t, db, child.ID, model.ID, "task complete")
|
|
|
|
gotChat, report, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, time.Second,
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, child.ID, gotChat.ID)
|
|
assert.Equal(t, database.ChatStatusWaiting, gotChat.Status)
|
|
assert.Equal(t, "task complete", report)
|
|
})
|
|
|
|
t.Run("AlreadyError", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusError, "something broke")
|
|
insertAssistantMessage(t, db, child.ID, model.ID, "partial work done")
|
|
|
|
_, _, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, time.Second,
|
|
)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "partial work done")
|
|
})
|
|
|
|
t.Run("AlreadyErrorNoReport", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusError, "crash")
|
|
|
|
_, _, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, time.Second,
|
|
)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "agent reached error status")
|
|
})
|
|
|
|
t.Run("CompletesViaPoll", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
// Force subscription failure so awaitSubagentCompletion
|
|
// falls back to the fast 200ms poll interval.
|
|
db, _ := dbtestutil.NewDB(t)
|
|
mClock := quartz.NewMock(t)
|
|
ps := subscribeFailingPubsub{Pubsub: pubsub.NewInMemory()}
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerClock(mClock))
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
setChatStatus(ctx, t, db, parent.ID, database.ChatStatusRunning, "")
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
|
|
|
|
// Set the trap BEFORE starting the goroutine so we
|
|
// deterministically catch the ticker creation.
|
|
tickTrap := mClock.Trap().NewTicker("chatd", "subagent_poll")
|
|
|
|
type awaitResult struct {
|
|
chat database.Chat
|
|
report string
|
|
err error
|
|
}
|
|
resultCh := make(chan awaitResult, 1)
|
|
go func() {
|
|
chat, report, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, 5*time.Second,
|
|
)
|
|
resultCh <- awaitResult{chat, report, err}
|
|
}()
|
|
|
|
// Wait for the poll ticker to be created, confirming
|
|
// the function passed its initial check and entered
|
|
// the loop. Then release the call.
|
|
tickTrap.MustWait(ctx).MustRelease(ctx)
|
|
tickTrap.Close()
|
|
|
|
// 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(t, db, child.ID, model.ID, "poll result")
|
|
mClock.Advance(subagentAwaitPollInterval).MustWait(ctx)
|
|
|
|
result := testutil.RequireReceive(ctx, t, resultCh)
|
|
require.NoError(t, result.err)
|
|
assert.Equal(t, child.ID, result.chat.ID)
|
|
assert.Equal(t, "poll result", result.report)
|
|
})
|
|
|
|
t.Run("CompletesViaPubsub", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
mClock := quartz.NewMock(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerClock(mClock))
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
// signalWake from CreateChat may trigger immediate processing.
|
|
// Wait for it to settle, then reset chats to the state we need.
|
|
WaitUntilIdleForTest(server)
|
|
setChatStatus(ctx, t, db, parent.ID, database.ChatStatusRunning, "")
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
|
|
|
|
// Trap the fallback poll ticker to know when the
|
|
// function has entered the wait setup path. We still
|
|
// need an explicit subscription handshake below because
|
|
// the ticker can be created before SubscribeWithErr has
|
|
// finished registering the listener.
|
|
tickTrap := mClock.Trap().NewTicker("chatd", "subagent_poll")
|
|
|
|
type awaitResult struct {
|
|
chat database.Chat
|
|
report string
|
|
err error
|
|
}
|
|
resultCh := make(chan awaitResult, 1)
|
|
go func() {
|
|
chat, report, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, 5*time.Second,
|
|
)
|
|
resultCh <- awaitResult{chat, report, err}
|
|
}()
|
|
|
|
// Wait for the ticker to be created so the waiter has
|
|
// entered its setup path, then subscribe our own probe on
|
|
// the same channel. Because MemoryPubsub publishes only to
|
|
// listeners already present at Publish time, waiting for
|
|
// our probe to receive a message proves the waiter's
|
|
// subscription is also registered before we assert on the
|
|
// wake-up behavior.
|
|
tickTrap.MustWait(ctx).MustRelease(ctx)
|
|
tickTrap.Close()
|
|
|
|
probeCh := make(chan struct{}, 1)
|
|
cancelProbe, err := ps.SubscribeWithErr(
|
|
coderdpubsub.ChatStateUpdateChannel(child.ID),
|
|
func(_ context.Context, _ []byte, _ error) {
|
|
select {
|
|
case probeCh <- struct{}{}:
|
|
default:
|
|
}
|
|
},
|
|
)
|
|
require.NoError(t, err)
|
|
defer cancelProbe()
|
|
|
|
// Insert the message before transitioning to Waiting so any
|
|
// notification observing the terminal status can also read the
|
|
// committed report.
|
|
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)
|
|
require.NoError(c, err)
|
|
assert.True(c, done)
|
|
assert.Equal(c, child.ID, chat.ID)
|
|
assert.Equal(c, "pubsub result", report)
|
|
}, testutil.WaitMedium, testutil.IntervalFast)
|
|
require.NoError(t, ps.Publish(
|
|
coderdpubsub.ChatStateUpdateChannel(child.ID),
|
|
[]byte("done"),
|
|
))
|
|
testutil.RequireReceive(ctx, t, probeCh)
|
|
|
|
result := testutil.RequireReceive(ctx, t, resultCh)
|
|
require.NoError(t, result.err)
|
|
assert.Equal(t, child.ID, result.chat.ID)
|
|
assert.Equal(t, "pubsub result", result.report)
|
|
})
|
|
|
|
t.Run("AlreadyWaitingNoReport", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
// This case should return immediately, so use the shared
|
|
// real-clock server instead of a mock clock.
|
|
WaitUntilIdleForTest(server)
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
|
|
|
gotChat, report, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, 5*time.Second,
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, child.ID, gotChat.ID)
|
|
assert.Empty(t, report)
|
|
})
|
|
|
|
t.Run("Timeout", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
mClock := quartz.NewMock(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerClock(mClock))
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
// Trap the timeout timer to know when the function
|
|
// has entered its poll loop.
|
|
timerTrap := mClock.Trap().NewTimer("chatd", "subagent_await")
|
|
|
|
type awaitResult struct {
|
|
err error
|
|
}
|
|
resultCh := make(chan awaitResult, 1)
|
|
go func() {
|
|
_, _, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, time.Second,
|
|
)
|
|
resultCh <- awaitResult{err}
|
|
}()
|
|
|
|
// Wait for the timer to be created, release it.
|
|
timerTrap.MustWait(ctx).MustRelease(ctx)
|
|
timerTrap.Close()
|
|
|
|
// Advance to the timeout. With pubsub, the fallback
|
|
// poll is at 5s, so the 1s timer fires first.
|
|
mClock.Advance(time.Second).MustWait(ctx)
|
|
|
|
result := testutil.RequireReceive(ctx, t, resultCh)
|
|
require.Error(t, result.err)
|
|
assert.Contains(t, result.err.Error(), "timed out waiting for delegated subagent completion")
|
|
})
|
|
|
|
t.Run("ContextCanceled", func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
providerCalled := make(chan struct{}, 1)
|
|
providerReleased := make(chan struct{})
|
|
providerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
select {
|
|
case providerCalled <- struct{}{}:
|
|
default:
|
|
}
|
|
|
|
select {
|
|
case <-r.Context().Done():
|
|
case <-providerReleased:
|
|
}
|
|
}))
|
|
t.Cleanup(func() {
|
|
close(providerReleased)
|
|
providerServer.Close()
|
|
})
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
providerServerURL, err := url.Parse(providerServer.URL)
|
|
require.NoError(t, err)
|
|
factory := &aibridgeTestFactory{rt: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
|
cloned := req.Clone(req.Context())
|
|
cloned.URL.Scheme = providerServerURL.Scheme
|
|
cloned.URL.Host = providerServerURL.Host
|
|
cloned.Host = providerServerURL.Host
|
|
return http.DefaultTransport.RoundTrip(cloned)
|
|
})}
|
|
server := newInternalTestServer(
|
|
t, db, ps, chatprovider.ProviderAPIKeys{},
|
|
withInternalTestServerWorker(),
|
|
withInternalTestServerTransportFactory(factory),
|
|
)
|
|
ctx := chatdTestContext(t)
|
|
user, org, _ := seedInternalChatDeps(t, db)
|
|
provider := dbgen.ChatProvider(t, db, database.ChatProvider{
|
|
Provider: "openai",
|
|
DisplayName: "OpenAI",
|
|
BaseUrl: providerServer.URL,
|
|
})
|
|
model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
|
Model: "gpt-4o-mini",
|
|
AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true},
|
|
})
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
testutil.RequireReceive(ctx, t, providerCalled)
|
|
|
|
// Use a short-lived context instead of goroutine + sleep.
|
|
shortCtx, cancel := context.WithTimeout(ctx, testutil.IntervalMedium)
|
|
defer cancel()
|
|
|
|
_, _, err = server.awaitSubagentCompletion(
|
|
shortCtx, parent.ID, child.ID, 5*time.Second,
|
|
)
|
|
require.ErrorIs(t, err, context.DeadlineExceeded)
|
|
})
|
|
|
|
t.Run("ZeroTimeoutUsesDefault", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
// Pre-complete the child so it returns immediately.
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
|
insertAssistantMessage(t, db, child.ID, model.ID, "zero timeout ok")
|
|
|
|
gotChat, report, err := server.awaitSubagentCompletion(
|
|
ctx, parent.ID, child.ID, 0,
|
|
)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, child.ID, gotChat.ID)
|
|
assert.Equal(t, "zero timeout ok", report)
|
|
})
|
|
}
|
|
|
|
func TestWaitAgentToolSchema(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, _ := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
tool := findToolByName(server.subagentTools(ctx, func() database.Chat {
|
|
return parent
|
|
}, parent.LastModelConfigID), "wait_agent")
|
|
require.NotNil(t, tool)
|
|
|
|
timeoutSeconds, ok := tool.Info().Parameters["timeout_seconds"].(map[string]any)
|
|
require.True(t, ok)
|
|
assert.Equal(t, "integer", timeoutSeconds["type"])
|
|
assert.Equal(t, "Defaults to 5 minutes.", timeoutSeconds["description"])
|
|
assert.Contains(t, tool.Info().Description, "Returns immediately when the agent finishes")
|
|
assert.Contains(t, tool.Info().Description, "A timeout does not stop the agent")
|
|
}
|
|
|
|
func TestWaitAgentTimeoutReturnsInformationalPayload(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
mClock := quartz.NewMock(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerClock(mClock))
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
WaitUntilIdleForTest(server)
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
|
|
|
|
timerTrap := mClock.Trap().NewTimer("chatd", "subagent_await")
|
|
|
|
type toolResult struct {
|
|
resp fantasy.ToolResponse
|
|
}
|
|
resultCh := make(chan toolResult, 1)
|
|
oneSecond := 1
|
|
go func() {
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parent,
|
|
parent.LastModelConfigID,
|
|
"wait_agent",
|
|
waitAgentArgs{ChatID: child.ID.String(), TimeoutSeconds: &oneSecond},
|
|
)
|
|
resultCh <- toolResult{resp: resp}
|
|
}()
|
|
|
|
// Wait for the timer to be created, then advance past it.
|
|
timerTrap.MustWait(ctx).MustRelease(ctx)
|
|
timerTrap.Close()
|
|
mClock.Advance(time.Second).MustWait(ctx)
|
|
|
|
result := testutil.RequireReceive(ctx, t, resultCh)
|
|
m := requireToolResponseMap(t, result.resp, false)
|
|
|
|
require.Equal(t, true, m["timed_out"])
|
|
require.Equal(t, child.ID.String(), m["chat_id"])
|
|
require.Equal(t, string(database.ChatStatusRunning), m["status"])
|
|
require.Equal(t, subagentTypeGeneral, m["type"])
|
|
}
|
|
|
|
func TestWaitAgentErrorStatusReturnsStructuredPayload(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
// An errored, non-archived agent is often recoverable. wait_agent
|
|
// must surface a structured payload (status, last_error, report)
|
|
// rather than a bare tool error.
|
|
WaitUntilIdleForTest(server)
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusError, "provider overloaded")
|
|
insertAssistantMessage(t, db, child.ID, model.ID, "partial progress")
|
|
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parent,
|
|
parent.LastModelConfigID,
|
|
"wait_agent",
|
|
waitAgentArgs{ChatID: child.ID.String()},
|
|
), false)
|
|
|
|
require.Equal(t, string(database.ChatStatusError), result["status"])
|
|
require.Equal(t, child.ID.String(), result["chat_id"])
|
|
require.Equal(t, "provider overloaded", result["last_error"])
|
|
require.Equal(t, "partial progress", result["report"])
|
|
require.Equal(t, subagentTypeGeneral, result["type"])
|
|
require.Equal(t, string(codersdk.ChatErrorKindGeneric), result["last_error_kind"])
|
|
require.Equal(t, false, result["last_error_retryable"])
|
|
require.NotContains(t, result, "last_error_detail")
|
|
require.NotContains(t, result, "timed_out")
|
|
}
|
|
|
|
func TestSubagentLastError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
rawJSON := func(s string) pqtype.NullRawMessage {
|
|
return pqtype.NullRawMessage{RawMessage: json.RawMessage(s), Valid: true}
|
|
}
|
|
|
|
cases := []struct {
|
|
name string
|
|
raw pqtype.NullRawMessage
|
|
want string
|
|
wantDecoded bool
|
|
}{
|
|
{name: "Invalid", raw: pqtype.NullRawMessage{}, want: ""},
|
|
// Unrecognized payloads must not leak raw stored bytes into
|
|
// model context; the caller falls back to its status reason.
|
|
{name: "NotChatError", raw: rawJSON(`"oops"`), want: ""},
|
|
{name: "EmptyObject", raw: rawJSON(`{}`), want: ""},
|
|
{
|
|
name: "MessageOnly",
|
|
raw: rawJSON(`{"message":"provider overloaded"}`),
|
|
want: "provider overloaded",
|
|
wantDecoded: true,
|
|
},
|
|
{
|
|
name: "DetailReplacesGenericMessage",
|
|
raw: rawJSON(`{"message":"The chat request failed unexpectedly.","detail":"reasoning model ` + "`max`" + ` not supported"}`),
|
|
want: "reasoning model `max` not supported",
|
|
wantDecoded: true,
|
|
},
|
|
{
|
|
name: "DetailOnly",
|
|
raw: rawJSON(`{"detail":"reasoning model ` + "`max`" + ` not supported"}`),
|
|
want: "reasoning model `max` not supported",
|
|
wantDecoded: true,
|
|
},
|
|
{
|
|
name: "DetailAppendedToMeaningfulMessage",
|
|
raw: rawJSON(`{"kind":"config","message":"Vercel AI Gateway rejected the model configuration.","detail":"unknown model slug"}`),
|
|
want: "Vercel AI Gateway rejected the model configuration. (unknown model slug)",
|
|
wantDecoded: true,
|
|
},
|
|
// Detail passes through exactly as the chat UI renders it,
|
|
// including auth details, provider request IDs, and long
|
|
// opaque diagnostic tokens.
|
|
{
|
|
name: "AuthKindDetailPreserved",
|
|
raw: rawJSON(`{"kind":"auth","message":"Authentication with Anthropic failed.","detail":"401 invalid x-api-key"}`),
|
|
want: "Authentication with Anthropic failed. (401 invalid x-api-key)",
|
|
wantDecoded: true,
|
|
},
|
|
{
|
|
name: "RequestIDPreserved",
|
|
raw: rawJSON(`{"kind":"generic","message":"The chat request failed unexpectedly.","detail":"upstream error (request id req_0a1b2c3d4e5f6a7b8c9d)"}`),
|
|
want: "upstream error (request id req_0a1b2c3d4e5f6a7b8c9d)",
|
|
wantDecoded: true,
|
|
},
|
|
{
|
|
name: "ModelSlugPreserved",
|
|
raw: rawJSON(`{"kind":"generic","message":"The chat request failed unexpectedly.","detail":"model claude-haiku-4-5-20251001 is not available"}`),
|
|
want: "model claude-haiku-4-5-20251001 is not available",
|
|
wantDecoded: true,
|
|
},
|
|
}
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
decoded, message := subagentLastError(tc.raw)
|
|
require.Equal(t, tc.want, message)
|
|
require.Equal(t, tc.wantDecoded, decoded != nil)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestWaitAgentErrorStatusSurfacesLastErrorDetail(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
// A generic fallback message hides the actionable provider detail
|
|
// (e.g. an unsupported reasoning effort); wait_agent must surface
|
|
// the detail so the parent model can self-correct.
|
|
WaitUntilIdleForTest(server)
|
|
setChatStatusWithError(ctx, t, db, child.ID, database.ChatStatusError, codersdk.ChatError{
|
|
Message: "The chat request failed unexpectedly.",
|
|
Detail: "reasoning model `max` not supported",
|
|
Kind: codersdk.ChatErrorKindGeneric,
|
|
Provider: "vercel",
|
|
Retryable: false,
|
|
})
|
|
insertAssistantMessage(t, db, child.ID, model.ID, "partial progress")
|
|
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parent,
|
|
parent.LastModelConfigID,
|
|
"wait_agent",
|
|
waitAgentArgs{ChatID: child.ID.String()},
|
|
), false)
|
|
|
|
require.Equal(t, string(database.ChatStatusError), result["status"])
|
|
require.Equal(t, child.ID.String(), result["chat_id"])
|
|
require.Equal(t, "reasoning model `max` not supported", result["last_error"])
|
|
require.Equal(t, "partial progress", result["report"])
|
|
require.Equal(t, string(codersdk.ChatErrorKindGeneric), result["last_error_kind"])
|
|
require.Equal(t, false, result["last_error_retryable"])
|
|
require.Equal(t, "reasoning model `max` not supported", result["last_error_detail"])
|
|
}
|
|
|
|
func TestWaitAgentTimeoutGapCompletesWithError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
mClock := quartz.NewMock(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerClock(mClock))
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
WaitUntilIdleForTest(server)
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
|
|
|
|
timerTrap := mClock.Trap().NewTimer("chatd", "subagent_await")
|
|
|
|
type toolResult struct {
|
|
resp fantasy.ToolResponse
|
|
}
|
|
resultCh := make(chan toolResult, 1)
|
|
oneSecond := 1
|
|
go func() {
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parent,
|
|
parent.LastModelConfigID,
|
|
"wait_agent",
|
|
waitAgentArgs{ChatID: child.ID.String(), TimeoutSeconds: &oneSecond},
|
|
)
|
|
resultCh <- toolResult{resp: resp}
|
|
}()
|
|
|
|
// Wait for the timer to be created, then advance past it.
|
|
timerTrap.MustWait(ctx).MustRelease(ctx)
|
|
timerTrap.Close()
|
|
|
|
// Flip the child to error before the timer fires so the
|
|
// timeout-gap branch (checkSubagentCompletion after timeout)
|
|
// classifies it through handleSubagentDone.
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusError, "provider overloaded")
|
|
insertAssistantMessage(t, db, child.ID, model.ID, "partial progress")
|
|
|
|
mClock.Advance(time.Second).MustWait(ctx)
|
|
|
|
result := testutil.RequireReceive(ctx, t, resultCh)
|
|
m := requireToolResponseMap(t, result.resp, false)
|
|
|
|
require.Equal(t, string(database.ChatStatusError), m["status"])
|
|
require.Equal(t, "provider overloaded", m["last_error"])
|
|
require.Equal(t, "partial progress", m["report"])
|
|
require.Equal(t, child.ID.String(), m["chat_id"])
|
|
require.Equal(t, subagentTypeGeneral, m["type"])
|
|
require.Equal(t, string(codersdk.ChatErrorKindGeneric), m["last_error_kind"])
|
|
require.Equal(t, false, m["last_error_retryable"])
|
|
require.NotContains(t, m, "last_error_detail")
|
|
require.NotContains(t, m, "timed_out")
|
|
}
|
|
|
|
func TestWaitAgentTimeoutGapSurfacesLastErrorDetail(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
mClock := quartz.NewMock(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}, withInternalTestServerClock(mClock))
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
WaitUntilIdleForTest(server)
|
|
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
|
|
|
|
timerTrap := mClock.Trap().NewTimer("chatd", "subagent_await")
|
|
|
|
type toolResult struct {
|
|
resp fantasy.ToolResponse
|
|
}
|
|
resultCh := make(chan toolResult, 1)
|
|
oneSecond := 1
|
|
go func() {
|
|
resp := runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parent,
|
|
parent.LastModelConfigID,
|
|
"wait_agent",
|
|
waitAgentArgs{ChatID: child.ID.String(), TimeoutSeconds: &oneSecond},
|
|
)
|
|
resultCh <- toolResult{resp: resp}
|
|
}()
|
|
|
|
timerTrap.MustWait(ctx).MustRelease(ctx)
|
|
timerTrap.Close()
|
|
|
|
// The timeout-gap recheck must surface the same detail-aware
|
|
// last_error as the normal poll path.
|
|
setChatStatusWithError(ctx, t, db, child.ID, database.ChatStatusError, codersdk.ChatError{
|
|
Message: "The chat request failed unexpectedly.",
|
|
Detail: "reasoning model `max` not supported",
|
|
Kind: codersdk.ChatErrorKindGeneric,
|
|
Provider: "vercel",
|
|
Retryable: true,
|
|
})
|
|
insertAssistantMessage(t, db, child.ID, model.ID, "partial progress")
|
|
|
|
mClock.Advance(time.Second).MustWait(ctx)
|
|
|
|
result := testutil.RequireReceive(ctx, t, resultCh)
|
|
m := requireToolResponseMap(t, result.resp, false)
|
|
|
|
require.Equal(t, string(database.ChatStatusError), m["status"])
|
|
require.Equal(t, "reasoning model `max` not supported", m["last_error"])
|
|
require.Equal(t, "partial progress", m["report"])
|
|
require.Equal(t, child.ID.String(), m["chat_id"])
|
|
require.Equal(t, string(codersdk.ChatErrorKindGeneric), m["last_error_kind"])
|
|
require.Equal(t, true, m["last_error_retryable"])
|
|
require.Equal(t, "reasoning model `max` not supported", m["last_error_detail"])
|
|
require.NotContains(t, m, "timed_out")
|
|
}
|
|
|
|
func TestWaitAgentErrorStatusUnrecognizedLastError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
ctx := chatdTestContext(t)
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
parent, child := createParentChildChats(ctx, t, server, user, org, model)
|
|
|
|
WaitUntilIdleForTest(server)
|
|
_, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
|
|
ID: child.ID,
|
|
Status: database.ChatStatusError,
|
|
LastError: pqtype.NullRawMessage{RawMessage: json.RawMessage(`"oops"`), Valid: true},
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx,
|
|
t,
|
|
server,
|
|
parent,
|
|
parent.LastModelConfigID,
|
|
"wait_agent",
|
|
waitAgentArgs{ChatID: child.ID.String()},
|
|
), false)
|
|
|
|
// Unrecognized payloads fall back to the status reason instead of
|
|
// leaking raw stored bytes, and omit the structured fields.
|
|
require.Equal(t, string(database.ChatStatusError), result["status"])
|
|
require.Equal(t, "agent reached error status", result["last_error"])
|
|
require.NotContains(t, result, "last_error_kind")
|
|
require.NotContains(t, result, "last_error_retryable")
|
|
require.NotContains(t, result, "last_error_detail")
|
|
}
|
|
|
|
func listAgentsChatIDs(t *testing.T, result map[string]any) []string {
|
|
t.Helper()
|
|
agents, ok := result["agents"].([]any)
|
|
require.True(t, ok, "agents must be an array")
|
|
ids := make([]string, 0, len(agents))
|
|
for _, raw := range agents {
|
|
agent, ok := raw.(map[string]any)
|
|
require.True(t, ok, "each agent must be an object")
|
|
id, ok := agent["chat_id"].(string)
|
|
require.True(t, ok, "each agent must have a chat_id")
|
|
ids = append(ids, id)
|
|
}
|
|
return ids
|
|
}
|
|
|
|
func TestListAgents(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
db, ps := dbtestutil.NewDB(t)
|
|
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
|
user, org, model := seedInternalChatDeps(t, db)
|
|
|
|
// Helpers take the running subtest's t and ctx so a failed require
|
|
// fires on the correct goroutine.
|
|
newParent := func(t *testing.T, ctx context.Context, title string) database.Chat {
|
|
t.Helper()
|
|
parent, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
Title: title,
|
|
ModelConfigID: model.ID,
|
|
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
|
})
|
|
require.NoError(t, err)
|
|
return parent
|
|
}
|
|
newChild := func(t *testing.T, ctx context.Context, parent database.Chat, title string, mode database.NullChatMode) database.Chat {
|
|
t.Helper()
|
|
child, err := server.CreateChat(ctx, CreateOptions{
|
|
OrganizationID: org.ID,
|
|
OwnerID: user.ID,
|
|
ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true},
|
|
RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true},
|
|
Title: title,
|
|
ModelConfigID: model.ID,
|
|
ChatMode: mode,
|
|
InitialUserContent: []codersdk.ChatMessagePart{
|
|
codersdk.ChatMessageText("do work"),
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
return child
|
|
}
|
|
|
|
t.Run("Empty", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
parent := newParent(t, ctx, "list-agents-empty")
|
|
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx, t, server, parent, parent.LastModelConfigID,
|
|
"list_agents", listAgentsArgs{},
|
|
), false)
|
|
|
|
require.Equal(t, float64(0), result["total"])
|
|
require.Equal(t, float64(0), result["returned"])
|
|
require.Equal(t, false, result["has_more"])
|
|
require.Empty(t, listAgentsChatIDs(t, result))
|
|
})
|
|
|
|
t.Run("ReturnsChildren", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
parent := newParent(t, ctx, "list-agents-children")
|
|
generalChild := newChild(t, ctx, parent, "general-child", database.NullChatMode{})
|
|
exploreChild := newChild(t, ctx, parent, "explore-child", database.NullChatMode{
|
|
ChatMode: database.ChatModeExplore,
|
|
Valid: true,
|
|
})
|
|
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx, t, server, parent, parent.LastModelConfigID,
|
|
"list_agents", listAgentsArgs{},
|
|
), false)
|
|
|
|
require.Equal(t, float64(2), result["total"])
|
|
require.Equal(t, float64(2), result["returned"])
|
|
require.Equal(t, false, result["has_more"])
|
|
ids := listAgentsChatIDs(t, result)
|
|
require.Contains(t, ids, generalChild.ID.String())
|
|
require.Contains(t, ids, exploreChild.ID.String())
|
|
|
|
agents, ok := result["agents"].([]any)
|
|
require.True(t, ok)
|
|
typesByID := map[string]string{}
|
|
for _, raw := range agents {
|
|
agent := raw.(map[string]any)
|
|
typesByID[agent["chat_id"].(string)] = agent["type"].(string)
|
|
require.NotEmpty(t, agent["created_at"])
|
|
require.NotEmpty(t, agent["updated_at"])
|
|
}
|
|
require.Equal(t, subagentTypeGeneral, typesByID[generalChild.ID.String()])
|
|
require.Equal(t, subagentTypeExplore, typesByID[exploreChild.ID.String()])
|
|
})
|
|
|
|
t.Run("Pagination", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
parent := newParent(t, ctx, "list-agents-pagination")
|
|
newChild(t, ctx, parent, "child-a", database.NullChatMode{})
|
|
newChild(t, ctx, parent, "child-b", database.NullChatMode{})
|
|
newChild(t, ctx, parent, "child-c", database.NullChatMode{})
|
|
|
|
limit := 2
|
|
first := requireToolResponseMap(t, runSubagentTool(
|
|
ctx, t, server, parent, parent.LastModelConfigID,
|
|
"list_agents", listAgentsArgs{Limit: &limit},
|
|
), false)
|
|
require.Equal(t, float64(3), first["total"])
|
|
require.Equal(t, float64(2), first["returned"])
|
|
require.Equal(t, true, first["has_more"])
|
|
firstIDs := listAgentsChatIDs(t, first)
|
|
require.Len(t, firstIDs, 2)
|
|
|
|
offset := 2
|
|
second := requireToolResponseMap(t, runSubagentTool(
|
|
ctx, t, server, parent, parent.LastModelConfigID,
|
|
"list_agents", listAgentsArgs{Limit: &limit, Offset: &offset},
|
|
), false)
|
|
require.Equal(t, float64(3), second["total"])
|
|
require.Equal(t, float64(1), second["returned"])
|
|
require.Equal(t, false, second["has_more"])
|
|
secondIDs := listAgentsChatIDs(t, second)
|
|
require.Len(t, secondIDs, 1)
|
|
require.NotContains(t, firstIDs, secondIDs[0])
|
|
})
|
|
|
|
t.Run("OrderByUpdatedAtDesc", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
parent := newParent(t, ctx, "list-agents-order")
|
|
older := newChild(t, ctx, parent, "older-child", database.NullChatMode{})
|
|
newChild(t, ctx, parent, "newer-child", database.NullChatMode{})
|
|
|
|
// Touch the older child so its updated_at advances past the
|
|
// newer one; it must then sort first.
|
|
setChatStatus(ctx, t, db, older.ID, database.ChatStatusWaiting, "")
|
|
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx, t, server, parent, parent.LastModelConfigID,
|
|
"list_agents", listAgentsArgs{},
|
|
), false)
|
|
ids := listAgentsChatIDs(t, result)
|
|
require.Len(t, ids, 2)
|
|
require.Equal(t, older.ID.String(), ids[0])
|
|
})
|
|
|
|
t.Run("ExcludesArchived", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
parent := newParent(t, ctx, "list-agents-archived")
|
|
archivedChild := newChild(t, ctx, parent, "archived-child", database.NullChatMode{})
|
|
|
|
WaitUntilIdleForTest(server)
|
|
// SetArchived is only allowed from a waiting/error state, so
|
|
// settle the family into waiting first. Archiving then marks
|
|
// the children archived; they must be excluded from
|
|
// list_agents by default.
|
|
setChatStatus(ctx, t, db, parent.ID, database.ChatStatusWaiting, "")
|
|
setChatStatus(ctx, t, db, archivedChild.ID, database.ChatStatusWaiting, "")
|
|
require.NoError(t, server.ArchiveChat(ctx, parent))
|
|
|
|
result := requireToolResponseMap(t, runSubagentTool(
|
|
ctx, t, server, parent, parent.LastModelConfigID,
|
|
"list_agents", listAgentsArgs{},
|
|
), false)
|
|
require.Equal(t, float64(0), result["total"])
|
|
require.Empty(t, listAgentsChatIDs(t, result))
|
|
})
|
|
|
|
t.Run("DelegatedChatRejected", func(t *testing.T) {
|
|
t.Parallel()
|
|
ctx := chatdTestContext(t)
|
|
parent := newParent(t, ctx, "list-agents-delegated")
|
|
child := newChild(t, ctx, parent, "delegated-caller", database.NullChatMode{})
|
|
|
|
resp := runSubagentTool(
|
|
ctx, t, server, child, child.LastModelConfigID,
|
|
"list_agents", listAgentsArgs{},
|
|
)
|
|
require.True(t, resp.IsError, "list_agents on a delegated chat must return an error")
|
|
msg := resp.Content
|
|
require.Contains(t, msg, "only available on root chats")
|
|
})
|
|
}
|