mirror of
https://github.com/coder/coder.git
synced 2026-10-29 17:44:33 +08:00
feat: snapshot explore subagent tool entitlements (#24638)
Explore sub-agents previously could not use `web_search` or external MCP tools. `runChat` hard-skipped both for Explore. Lifting those guards naively would over-grant tools, because a child chat could outlive the spawning turn's plan-mode filter. This change persists the spawning parent turn's filtered external MCP server IDs onto the child Explore chat, and simplifies the Explore provider-tool filter in `runChat`: - New `resolveExploreToolSnapshot` helper: computes the child's inherited external MCP subset by running the parent's configs through `filterExternalMCPConfigsForTurn` (plan-mode policy) and, if the parent is itself an Explore child, further narrowing to the parent's own persisted `MCPServerIDs`. The result is written to the child's `MCPServerIDs` column at spawn time. - The existing `mcp_server_ids` column is the sole durable snapshot. No new chat column is added. - `runChat` for Explore children: loads MCP tools from the persisted snapshot, and keeps only `web_search` from provider-native tools (to block computer-use and other write-style tools, since Explore is read-only). Whether `web_search` is actually available is a per-model decision, determined by the current model config, just like a main chat. - Built-in Explore allowlist is unchanged. Workspace-local MCP remains excluded for Explore. Verification: `go build ./...`, `go test ./coderd/x/chatd/... -count=1`, `make gen` (clean tree), `make lint/emdash`, `go vet`. Deep-review ran 12 reviewers on the feature and 5 on the clarity refactor; CAR reviewed and approved; a subsequent scope reduction dropped a temporary `allow_web_search` column in favor of per-model handling. > Mux is acting on Mike's behalf.
This commit is contained in:
+65
-17
@@ -1223,13 +1223,21 @@ func (p *Server) SendMessage(
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Update MCP server IDs on the chat when explicitly provided.
|
// Update MCP server IDs on the chat when explicitly provided.
|
||||||
|
// Explore child chats keep the spawn-time snapshot immutable.
|
||||||
if opts.MCPServerIDs != nil {
|
if opts.MCPServerIDs != nil {
|
||||||
lockedChat, err = tx.UpdateChatMCPServerIDs(ctx, database.UpdateChatMCPServerIDsParams{
|
if isExploreSubagentMode(lockedChat.Mode) {
|
||||||
ID: opts.ChatID,
|
p.logger.Warn(ctx,
|
||||||
MCPServerIDs: *opts.MCPServerIDs,
|
"ignoring explore subagent mcp server ids update, snapshot is immutable after spawn",
|
||||||
})
|
slog.F("chat_id", opts.ChatID),
|
||||||
if err != nil {
|
)
|
||||||
return xerrors.Errorf("update chat mcp server ids: %w", err)
|
} else {
|
||||||
|
lockedChat, err = tx.UpdateChatMCPServerIDs(ctx, database.UpdateChatMCPServerIDsParams{
|
||||||
|
ID: opts.ChatID,
|
||||||
|
MCPServerIDs: *opts.MCPServerIDs,
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return xerrors.Errorf("update chat mcp server ids: %w", err)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -5228,6 +5236,9 @@ func isExploreSubagentMode(mode database.NullChatMode) bool {
|
|||||||
return mode.Valid && mode.ChatMode == database.ChatModeExplore
|
return mode.Valid && mode.ChatMode == database.ChatModeExplore
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// filterExternalMCPConfigsForTurn returns the external MCP server configs
|
||||||
|
// visible on the current turn. Explore children snapshot this filtered set at
|
||||||
|
// spawn time so later model overrides cannot widen the external-tool boundary.
|
||||||
func filterExternalMCPConfigsForTurn(
|
func filterExternalMCPConfigsForTurn(
|
||||||
configs []database.MCPServerConfig,
|
configs []database.MCPServerConfig,
|
||||||
mode database.NullChatPlanMode,
|
mode database.NullChatPlanMode,
|
||||||
@@ -5354,6 +5365,15 @@ func allowedExploreToolNames(allTools []fantasy.AgentTool) []string {
|
|||||||
name := tool.Info().Name
|
name := tool.Info().Name
|
||||||
if builtinExplorePolicy[name] {
|
if builtinExplorePolicy[name] {
|
||||||
toolNames = append(toolNames, name)
|
toolNames = append(toolNames, name)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
// External MCP tools pass through here. They were snapshot-filtered
|
||||||
|
// at spawn time on chat.MCPServerIDs. WorkspaceMCPTool does not
|
||||||
|
// implement MCPToolIdentifier, so workspace tools are excluded
|
||||||
|
// here too, in addition to the structural exclusion in runChat
|
||||||
|
// tool assembly.
|
||||||
|
if _, ok := tool.(mcpclient.MCPToolIdentifier); ok {
|
||||||
|
toolNames = append(toolNames, name)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return toolNames
|
return toolNames
|
||||||
@@ -5711,11 +5731,25 @@ func (p *Server) runChat(
|
|||||||
isPlanModeTurn := currentPlanMode.Valid && currentPlanMode.ChatPlanMode == database.ChatPlanModePlan
|
isPlanModeTurn := currentPlanMode.Valid && currentPlanMode.ChatPlanMode == database.ChatPlanModePlan
|
||||||
isExploreSubagent := isExploreSubagentMode(chat.Mode)
|
isExploreSubagent := isExploreSubagentMode(chat.Mode)
|
||||||
isRootChat := !chat.ParentChatID.Valid
|
isRootChat := !chat.ParentChatID.Valid
|
||||||
mcpConnectConfigs, approvedPlanMCPConfigIDs := filterExternalMCPConfigsForTurn(
|
var mcpConnectConfigs []database.MCPServerConfig
|
||||||
|
var approvedPlanMCPConfigIDs map[uuid.UUID]struct{}
|
||||||
|
// Explore subagents rely on the immutable spawn-time snapshot
|
||||||
|
// persisted in chat.MCPServerIDs. SendMessage cannot mutate that
|
||||||
|
// snapshot, so no runtime re-filter against parent state is needed.
|
||||||
|
// The child's persisted set is authoritative.
|
||||||
|
mcpConnectConfigs, approvedPlanMCPConfigIDs = filterExternalMCPConfigsForTurn(
|
||||||
mcpConfigs,
|
mcpConfigs,
|
||||||
currentPlanMode,
|
currentPlanMode,
|
||||||
chat.ParentChatID,
|
chat.ParentChatID,
|
||||||
)
|
)
|
||||||
|
if isExploreSubagent && isRootChat {
|
||||||
|
// Root Explore chats stay builtin-only per the accepted plan, so
|
||||||
|
// strip any persisted external MCP configs at runtime regardless of
|
||||||
|
// what's on the chat row. Explore children get their snapshot via
|
||||||
|
// the spawn-time inheritance path and are handled below.
|
||||||
|
mcpConnectConfigs = nil
|
||||||
|
approvedPlanMCPConfigIDs = map[uuid.UUID]struct{}{}
|
||||||
|
}
|
||||||
planModeInstructions := p.loadPlanModeInstructions(ctx, currentPlanMode, logger)
|
planModeInstructions := p.loadPlanModeInstructions(ctx, currentPlanMode, logger)
|
||||||
|
|
||||||
chainInfo := resolveChainMode(messages)
|
chainInfo := resolveChainMode(messages)
|
||||||
@@ -6429,14 +6463,12 @@ func (p *Server) runChat(
|
|||||||
builtinToolNames[t.Info().Name] = true
|
builtinToolNames[t.Info().Name] = true
|
||||||
}
|
}
|
||||||
|
|
||||||
// Append external and workspace MCP tools after the built-ins so the
|
// Append external MCP tools from the chat's persisted snapshot after the
|
||||||
// LLM sees them as additional capabilities. Explore subagents keep
|
// built-ins so the LLM sees them as additional capabilities. Explore chats
|
||||||
// the narrower built-in-only boundary from main. Root plan mode gets
|
// trust only the persisted MCPServerIDs snapshot, and workspace-local MCP
|
||||||
// only approved external MCP tools because mcpConnectConfigs was
|
// tools stay unavailable to Explore chats.
|
||||||
// pre-filtered above, and filterToolsForTurn removes any remaining
|
tools = append(tools, mcpTools...)
|
||||||
// plan-mode ineligible tools from the assembled set.
|
|
||||||
if !isExploreSubagent {
|
if !isExploreSubagent {
|
||||||
tools = append(tools, mcpTools...)
|
|
||||||
tools = append(tools, workspaceMCPTools...)
|
tools = append(tools, workspaceMCPTools...)
|
||||||
}
|
}
|
||||||
tools = filterToolsForTurn(
|
tools = filterToolsForTurn(
|
||||||
@@ -6462,11 +6494,23 @@ func (p *Server) runChat(
|
|||||||
return result, err
|
return result, err
|
||||||
}
|
}
|
||||||
|
|
||||||
// Build provider-native tools (e.g., web search) based on
|
// Build provider-native tools (e.g. web search) based on the
|
||||||
// the model configuration.
|
// current model configuration. Root Explore chats stay builtin-only per
|
||||||
|
// the accepted plan, so delegated Explore children are the only Explore
|
||||||
|
// chats that can inherit web_search. Write-style provider tools stay
|
||||||
|
// blocked for all Explore chats.
|
||||||
var providerTools []chatloop.ProviderTool
|
var providerTools []chatloop.ProviderTool
|
||||||
if !isPlanModeTurn && !isExploreSubagent && callConfig.ProviderOptions != nil {
|
if !isPlanModeTurn && callConfig.ProviderOptions != nil {
|
||||||
providerTools = buildProviderTools(model.Provider(), callConfig.ProviderOptions)
|
providerTools = buildProviderTools(model.Provider(), callConfig.ProviderOptions)
|
||||||
|
if isExploreSubagent {
|
||||||
|
if !chat.ParentChatID.Valid {
|
||||||
|
providerTools = nil
|
||||||
|
} else {
|
||||||
|
providerTools = slices.DeleteFunc(providerTools, func(tool chatloop.ProviderTool) bool {
|
||||||
|
return tool.Definition.GetName() != "web_search"
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if !isPlanModeTurn && !isExploreSubagent && isComputerUse {
|
if !isPlanModeTurn && !isExploreSubagent && isComputerUse {
|
||||||
@@ -6719,6 +6763,10 @@ func (p *Server) runChat(
|
|||||||
func buildProviderTools(_ string, options *codersdk.ChatModelProviderOptions) []chatloop.ProviderTool {
|
func buildProviderTools(_ string, options *codersdk.ChatModelProviderOptions) []chatloop.ProviderTool {
|
||||||
var tools []chatloop.ProviderTool
|
var tools []chatloop.ProviderTool
|
||||||
|
|
||||||
|
if options == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
if options.Anthropic != nil && options.Anthropic.WebSearchEnabled != nil && *options.Anthropic.WebSearchEnabled {
|
if options.Anthropic != nil && options.Anthropic.WebSearchEnabled != nil && *options.Anthropic.WebSearchEnabled {
|
||||||
tools = append(tools, chatloop.ProviderTool{
|
tools = append(tools, chatloop.ProviderTool{
|
||||||
Definition: anthropic.WebSearchTool(&anthropic.WebSearchToolOptions{
|
Definition: anthropic.WebSearchTool(&anthropic.WebSearchToolOptions{
|
||||||
|
|||||||
@@ -300,36 +300,32 @@ func TestActiveToolNamesForTurn(t *testing.T) {
|
|||||||
func TestAllowedExploreToolNames(t *testing.T) {
|
func TestAllowedExploreToolNames(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
makeTools := func(names ...string) []fantasy.AgentTool {
|
externalConfigID := uuid.New()
|
||||||
tools := make([]fantasy.AgentTool, 0, len(names))
|
got := allowedExploreToolNames([]fantasy.AgentTool{
|
||||||
for _, name := range names {
|
newTestAgentTool("read_file"),
|
||||||
tools = append(tools, newTestAgentTool(name))
|
newTestAgentTool("write_file"),
|
||||||
}
|
newTestMCPAgentTool("external-mcp__echo", externalConfigID),
|
||||||
return tools
|
newTestAgentTool("workspace-mcp__echo"),
|
||||||
}
|
newTestAgentTool("execute"),
|
||||||
|
newTestAgentTool("process_output"),
|
||||||
got := allowedExploreToolNames(makeTools(
|
newTestAgentTool("process_list"),
|
||||||
"read_file",
|
newTestAgentTool("process_signal"),
|
||||||
"write_file",
|
newTestAgentTool("spawn_agent"),
|
||||||
"edit_files",
|
newTestAgentTool("wait_agent"),
|
||||||
"execute",
|
newTestAgentTool("read_skill"),
|
||||||
"process_output",
|
newTestAgentTool("read_skill_file"),
|
||||||
"process_list",
|
newTestAgentTool("ask_user_question"),
|
||||||
"process_signal",
|
})
|
||||||
"spawn_agent",
|
|
||||||
"wait_agent",
|
|
||||||
"read_skill",
|
|
||||||
"read_skill_file",
|
|
||||||
"ask_user_question",
|
|
||||||
))
|
|
||||||
|
|
||||||
require.Equal(t, []string{
|
require.Equal(t, []string{
|
||||||
"read_file",
|
"read_file",
|
||||||
|
"external-mcp__echo",
|
||||||
"execute",
|
"execute",
|
||||||
"process_output",
|
"process_output",
|
||||||
"read_skill",
|
"read_skill",
|
||||||
"read_skill_file",
|
"read_skill_file",
|
||||||
}, got)
|
}, got)
|
||||||
|
require.NotContains(t, got, "workspace-mcp__echo")
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestAllowedBehaviorToolNames(t *testing.T) {
|
func TestAllowedBehaviorToolNames(t *testing.T) {
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package chatd_test
|
package chatd_test
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"cmp"
|
||||||
"context"
|
"context"
|
||||||
"database/sql"
|
"database/sql"
|
||||||
"encoding/base64"
|
"encoding/base64"
|
||||||
@@ -62,11 +63,15 @@ type recordedOpenAIRequest struct {
|
|||||||
ContentLength int64
|
ContentLength int64
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func openAIToolName(tool chattest.OpenAITool) string {
|
||||||
|
return cmp.Or(tool.Function.Name, tool.Name, tool.Type)
|
||||||
|
}
|
||||||
|
|
||||||
func recordOpenAIRequest(req *chattest.OpenAIRequest) recordedOpenAIRequest {
|
func recordOpenAIRequest(req *chattest.OpenAIRequest) recordedOpenAIRequest {
|
||||||
messages := append([]chattest.OpenAIMessage(nil), req.Messages...)
|
messages := append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||||
tools := make([]string, 0, len(req.Tools))
|
tools := make([]string, 0, len(req.Tools))
|
||||||
for _, tool := range req.Tools {
|
for _, tool := range req.Tools {
|
||||||
tools = append(tools, tool.Function.Name)
|
tools = append(tools, openAIToolName(tool))
|
||||||
}
|
}
|
||||||
|
|
||||||
var store *bool
|
var store *bool
|
||||||
@@ -697,6 +702,554 @@ func TestExploreSubagentIsReadOnly(t *testing.T) {
|
|||||||
require.Len(t, exploreChildren, 1)
|
require.Len(t, exploreChildren, 1)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
db, ps := dbtestutil.NewDB(t)
|
||||||
|
ctx := testutil.Context(t, testutil.WaitLong)
|
||||||
|
|
||||||
|
externalMCP := mcpserver.NewMCPServer("external-snapshot-mcp", "1.0.0")
|
||||||
|
externalMCP.AddTools(mcpserver.ServerTool{
|
||||||
|
Tool: mcpgo.NewTool("echo",
|
||||||
|
mcpgo.WithDescription("Echoes the input"),
|
||||||
|
mcpgo.WithString("input",
|
||||||
|
mcpgo.Description("The input string"),
|
||||||
|
mcpgo.Required(),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Handler: func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
|
||||||
|
input, _ := req.GetArguments()["input"].(string)
|
||||||
|
return mcpgo.NewToolResultText("echo: " + input), nil
|
||||||
|
},
|
||||||
|
})
|
||||||
|
externalMCPServer := httptest.NewServer(mcpserver.NewStreamableHTTPServer(externalMCP))
|
||||||
|
defer externalMCPServer.Close()
|
||||||
|
|
||||||
|
secondMCP := mcpserver.NewMCPServer("second-mcp", "1.0.0")
|
||||||
|
secondMCP.AddTools(mcpserver.ServerTool{
|
||||||
|
Tool: mcpgo.NewTool("echo",
|
||||||
|
mcpgo.WithDescription("Echoes the input"),
|
||||||
|
mcpgo.WithString("input",
|
||||||
|
mcpgo.Description("The input string"),
|
||||||
|
mcpgo.Required(),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Handler: func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
|
||||||
|
input, _ := req.GetArguments()["input"].(string)
|
||||||
|
return mcpgo.NewToolResultText("echo: " + input), nil
|
||||||
|
},
|
||||||
|
})
|
||||||
|
secondMCPServer := httptest.NewServer(mcpserver.NewStreamableHTTPServer(secondMCP))
|
||||||
|
defer secondMCPServer.Close()
|
||||||
|
|
||||||
|
var (
|
||||||
|
requestsMu sync.Mutex
|
||||||
|
requests []recordedOpenAIRequest
|
||||||
|
)
|
||||||
|
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||||
|
if !req.Stream {
|
||||||
|
return chattest.OpenAINonStreamingResponse("ok")
|
||||||
|
}
|
||||||
|
|
||||||
|
requestsMu.Lock()
|
||||||
|
requests = append(requests, recordOpenAIRequest(req))
|
||||||
|
requestsMu.Unlock()
|
||||||
|
|
||||||
|
return chattest.OpenAIStreamingResponse(
|
||||||
|
chattest.OpenAITextChunks("done")...,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
|
||||||
|
user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL)
|
||||||
|
webSearchEnabled := true
|
||||||
|
storeEnabled := true
|
||||||
|
// OpenAI only serializes web_search through the Responses API.
|
||||||
|
// Store=true routes there only for supported Responses models.
|
||||||
|
webSearchModel := insertChatModelConfigWithCallConfig(
|
||||||
|
ctx,
|
||||||
|
t,
|
||||||
|
db,
|
||||||
|
user.ID,
|
||||||
|
"openai",
|
||||||
|
"gpt-4o",
|
||||||
|
codersdk.ChatModelCallConfig{
|
||||||
|
ProviderOptions: &codersdk.ChatModelProviderOptions{
|
||||||
|
OpenAI: &codersdk.ChatModelOpenAIProviderOptions{
|
||||||
|
Store: &storeEnabled,
|
||||||
|
WebSearchEnabled: &webSearchEnabled,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||||
|
DisplayName: "External Snapshot MCP",
|
||||||
|
Slug: "external-snapshot-mcp",
|
||||||
|
Url: externalMCPServer.URL,
|
||||||
|
Transport: "streamable_http",
|
||||||
|
AuthType: "none",
|
||||||
|
Availability: "default_off",
|
||||||
|
Enabled: true,
|
||||||
|
ToolAllowList: []string{},
|
||||||
|
ToolDenyList: []string{},
|
||||||
|
CreatedBy: user.ID,
|
||||||
|
UpdatedBy: user.ID,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
_, err = db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||||
|
DisplayName: "Second MCP",
|
||||||
|
Slug: "second-mcp",
|
||||||
|
Url: secondMCPServer.URL,
|
||||||
|
Transport: "streamable_http",
|
||||||
|
AuthType: "none",
|
||||||
|
Availability: "default_off",
|
||||||
|
Enabled: true,
|
||||||
|
ToolAllowList: []string{},
|
||||||
|
ToolDenyList: []string{},
|
||||||
|
CreatedBy: user.ID,
|
||||||
|
UpdatedBy: user.ID,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID)
|
||||||
|
rootChat, err := db.InsertChat(ctx, database.InsertChatParams{
|
||||||
|
OrganizationID: org.ID,
|
||||||
|
OwnerID: user.ID,
|
||||||
|
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
|
||||||
|
AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true},
|
||||||
|
LastModelConfigID: webSearchModel.ID,
|
||||||
|
Title: "root",
|
||||||
|
Status: database.ChatStatusWaiting,
|
||||||
|
ClientType: database.ChatClientTypeApi,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
exploreChat, err := db.InsertChat(ctx, database.InsertChatParams{
|
||||||
|
OrganizationID: org.ID,
|
||||||
|
OwnerID: user.ID,
|
||||||
|
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
|
||||||
|
AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true},
|
||||||
|
ParentChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true},
|
||||||
|
RootChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true},
|
||||||
|
LastModelConfigID: webSearchModel.ID,
|
||||||
|
Title: "explore",
|
||||||
|
Mode: database.NullChatMode{
|
||||||
|
ChatMode: database.ChatModeExplore,
|
||||||
|
Valid: true,
|
||||||
|
},
|
||||||
|
Status: database.ChatStatusPending,
|
||||||
|
MCPServerIDs: []uuid.UUID{mcpConfig.ID},
|
||||||
|
ClientType: database.ChatClientTypeApi,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
insertUserTextMessage(ctx, t, db, exploreChat.ID, user.ID, webSearchModel.ID, "inspect the codebase")
|
||||||
|
|
||||||
|
ctrl := gomock.NewController(t)
|
||||||
|
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||||
|
mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes()
|
||||||
|
mockConn.EXPECT().ContextConfig(gomock.Any()).
|
||||||
|
Return(workspacesdk.ContextConfigResponse{}, xerrors.New("not supported")).AnyTimes()
|
||||||
|
workspaceToolName := "workspace-snapshot-mcp__echo"
|
||||||
|
mockConn.EXPECT().ListMCPTools(gomock.Any()).
|
||||||
|
Return(workspacesdk.ListMCPToolsResponse{Tools: []workspacesdk.MCPToolInfo{{
|
||||||
|
ServerName: "workspace-snapshot-mcp",
|
||||||
|
Name: workspaceToolName,
|
||||||
|
Description: "Workspace echo tool",
|
||||||
|
Schema: map[string]any{
|
||||||
|
"input": map[string]any{"type": "string"},
|
||||||
|
},
|
||||||
|
Required: []string{"input"},
|
||||||
|
}}}, nil).
|
||||||
|
AnyTimes()
|
||||||
|
mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()).
|
||||||
|
Return(workspacesdk.LSResponse{AbsolutePathString: "/home/coder"}, nil).AnyTimes()
|
||||||
|
mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
|
||||||
|
Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes()
|
||||||
|
|
||||||
|
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
|
||||||
|
cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||||
|
require.Equal(t, dbAgent.ID, agentID)
|
||||||
|
return mockConn, func() {}, nil
|
||||||
|
}
|
||||||
|
})
|
||||||
|
_ = server
|
||||||
|
|
||||||
|
chatResult := waitForTerminalChat(ctx, t, db, exploreChat.ID)
|
||||||
|
if chatResult.Status == database.ChatStatusError {
|
||||||
|
require.FailNowf(t, "explore chat failed", "last_error=%q", chatResult.LastError.String)
|
||||||
|
}
|
||||||
|
|
||||||
|
requestsMu.Lock()
|
||||||
|
recorded := append([]recordedOpenAIRequest(nil), requests...)
|
||||||
|
requestsMu.Unlock()
|
||||||
|
require.Len(t, recorded, 1)
|
||||||
|
|
||||||
|
tools := recorded[0].Tools
|
||||||
|
require.Contains(t, tools, "read_file")
|
||||||
|
require.Contains(t, tools, "execute")
|
||||||
|
require.Contains(t, tools, "process_output")
|
||||||
|
require.Contains(t, tools, "external-snapshot-mcp__echo")
|
||||||
|
require.Contains(t, tools, "web_search", "Explore provider tool filter should let web_search through when the current model supports it")
|
||||||
|
require.NotContains(t, tools, "second-mcp__echo")
|
||||||
|
require.NotContains(t, tools, workspaceToolName)
|
||||||
|
require.NotContains(t, tools, "write_file")
|
||||||
|
require.NotContains(t, tools, "edit_files")
|
||||||
|
require.NotContains(t, tools, "spawn_agent")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRootExploreChatStaysBuiltinOnlyAtRuntime(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
db, ps := dbtestutil.NewDB(t)
|
||||||
|
ctx := testutil.Context(t, testutil.WaitLong)
|
||||||
|
|
||||||
|
externalMCP := mcpserver.NewMCPServer("root-explore-runtime-mcp", "1.0.0")
|
||||||
|
externalMCP.AddTools(mcpserver.ServerTool{
|
||||||
|
Tool: mcpgo.NewTool("echo",
|
||||||
|
mcpgo.WithDescription("Echoes the input"),
|
||||||
|
mcpgo.WithString("input",
|
||||||
|
mcpgo.Description("The input string"),
|
||||||
|
mcpgo.Required(),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Handler: func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
|
||||||
|
input, _ := req.GetArguments()["input"].(string)
|
||||||
|
return mcpgo.NewToolResultText("echo: " + input), nil
|
||||||
|
},
|
||||||
|
})
|
||||||
|
externalMCPServer := httptest.NewServer(mcpserver.NewStreamableHTTPServer(externalMCP))
|
||||||
|
defer externalMCPServer.Close()
|
||||||
|
|
||||||
|
var (
|
||||||
|
requestsMu sync.Mutex
|
||||||
|
requests []recordedOpenAIRequest
|
||||||
|
)
|
||||||
|
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||||
|
if !req.Stream {
|
||||||
|
return chattest.OpenAINonStreamingResponse("ok")
|
||||||
|
}
|
||||||
|
|
||||||
|
requestsMu.Lock()
|
||||||
|
requests = append(requests, recordOpenAIRequest(req))
|
||||||
|
requestsMu.Unlock()
|
||||||
|
|
||||||
|
return chattest.OpenAIStreamingResponse(
|
||||||
|
chattest.OpenAITextChunks("done")...,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
|
||||||
|
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||||
|
mcpConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||||
|
DisplayName: "Root Explore Runtime MCP",
|
||||||
|
Slug: "root-explore-runtime-mcp",
|
||||||
|
Url: externalMCPServer.URL,
|
||||||
|
Transport: "streamable_http",
|
||||||
|
AuthType: "none",
|
||||||
|
Availability: "default_off",
|
||||||
|
Enabled: true,
|
||||||
|
ToolAllowList: []string{},
|
||||||
|
ToolDenyList: []string{},
|
||||||
|
CreatedBy: user.ID,
|
||||||
|
UpdatedBy: user.ID,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
server := newActiveTestServer(t, db, ps)
|
||||||
|
|
||||||
|
exploreChat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||||
|
OrganizationID: org.ID,
|
||||||
|
OwnerID: user.ID,
|
||||||
|
Title: "root-explore-builtin-only",
|
||||||
|
ModelConfigID: model.ID,
|
||||||
|
ChatMode: database.NullChatMode{
|
||||||
|
ChatMode: database.ChatModeExplore,
|
||||||
|
Valid: true,
|
||||||
|
},
|
||||||
|
MCPServerIDs: []uuid.UUID{mcpConfig.ID},
|
||||||
|
InitialUserContent: []codersdk.ChatMessagePart{
|
||||||
|
codersdk.ChatMessageText("Inspect the codebase."),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
waitForChatProcessed(ctx, t, db, exploreChat.ID, server)
|
||||||
|
|
||||||
|
storedChat, err := db.GetChatByID(ctx, exploreChat.ID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
if storedChat.Status == database.ChatStatusError {
|
||||||
|
require.FailNowf(t, "explore chat failed", "last_error=%q", storedChat.LastError.String)
|
||||||
|
}
|
||||||
|
require.Equal(t, database.ChatStatusWaiting, storedChat.Status)
|
||||||
|
require.ElementsMatch(t, []uuid.UUID{mcpConfig.ID}, storedChat.MCPServerIDs)
|
||||||
|
|
||||||
|
requestsMu.Lock()
|
||||||
|
recorded := append([]recordedOpenAIRequest(nil), requests...)
|
||||||
|
requestsMu.Unlock()
|
||||||
|
require.Len(t, recorded, 1)
|
||||||
|
|
||||||
|
tools := recorded[0].Tools
|
||||||
|
require.Contains(t, tools, "read_file")
|
||||||
|
require.Contains(t, tools, "execute")
|
||||||
|
require.NotContains(t, tools, "write_file")
|
||||||
|
require.NotContains(t, tools, "root-explore-runtime-mcp__echo",
|
||||||
|
"root Explore chats should strip persisted external MCP tools at runtime")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRootExploreChatExcludesWebSearchProviderToolAtRuntime(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
db, ps := dbtestutil.NewDB(t)
|
||||||
|
ctx := testutil.Context(t, testutil.WaitLong)
|
||||||
|
|
||||||
|
var (
|
||||||
|
requestsMu sync.Mutex
|
||||||
|
requests []recordedOpenAIRequest
|
||||||
|
)
|
||||||
|
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||||
|
if !req.Stream {
|
||||||
|
return chattest.OpenAINonStreamingResponse("ok")
|
||||||
|
}
|
||||||
|
|
||||||
|
requestsMu.Lock()
|
||||||
|
requests = append(requests, recordOpenAIRequest(req))
|
||||||
|
requestsMu.Unlock()
|
||||||
|
|
||||||
|
return chattest.OpenAIStreamingResponse(
|
||||||
|
chattest.OpenAITextChunks("done")...,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
|
||||||
|
user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL)
|
||||||
|
webSearchEnabled := true
|
||||||
|
storeEnabled := true
|
||||||
|
// OpenAI only serializes web_search through the Responses API.
|
||||||
|
// Store=true routes there only for supported Responses models.
|
||||||
|
webSearchModel := insertChatModelConfigWithCallConfig(
|
||||||
|
ctx,
|
||||||
|
t,
|
||||||
|
db,
|
||||||
|
user.ID,
|
||||||
|
"openai",
|
||||||
|
"gpt-4o",
|
||||||
|
codersdk.ChatModelCallConfig{
|
||||||
|
ProviderOptions: &codersdk.ChatModelProviderOptions{
|
||||||
|
OpenAI: &codersdk.ChatModelOpenAIProviderOptions{
|
||||||
|
Store: &storeEnabled,
|
||||||
|
WebSearchEnabled: &webSearchEnabled,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
server := newActiveTestServer(t, db, ps)
|
||||||
|
|
||||||
|
exploreChat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||||
|
OrganizationID: org.ID,
|
||||||
|
OwnerID: user.ID,
|
||||||
|
Title: "root-explore-no-provider-web-search",
|
||||||
|
ModelConfigID: webSearchModel.ID,
|
||||||
|
ChatMode: database.NullChatMode{
|
||||||
|
ChatMode: database.ChatModeExplore,
|
||||||
|
Valid: true,
|
||||||
|
},
|
||||||
|
InitialUserContent: []codersdk.ChatMessagePart{
|
||||||
|
codersdk.ChatMessageText("Inspect the codebase."),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
waitForChatProcessed(ctx, t, db, exploreChat.ID, server)
|
||||||
|
|
||||||
|
storedChat, err := db.GetChatByID(ctx, exploreChat.ID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
if storedChat.Status == database.ChatStatusError {
|
||||||
|
require.FailNowf(t, "explore chat failed", "last_error=%q", storedChat.LastError.String)
|
||||||
|
}
|
||||||
|
require.Equal(t, database.ChatStatusWaiting, storedChat.Status)
|
||||||
|
|
||||||
|
requestsMu.Lock()
|
||||||
|
recorded := append([]recordedOpenAIRequest(nil), requests...)
|
||||||
|
requestsMu.Unlock()
|
||||||
|
require.Len(t, recorded, 1)
|
||||||
|
|
||||||
|
tools := recorded[0].Tools
|
||||||
|
require.Contains(t, tools, "read_file")
|
||||||
|
require.Contains(t, tools, "execute")
|
||||||
|
require.NotContains(t, tools, "web_search",
|
||||||
|
"root Explore chats should stay builtin-only and must not inherit provider-native web_search at runtime")
|
||||||
|
require.NotContains(t, tools, "write_file")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestExploreChatSendMessageCannotMutateMCPSnapshot(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
|
||||||
|
db, ps := dbtestutil.NewDB(t)
|
||||||
|
ctx := testutil.Context(t, testutil.WaitLong)
|
||||||
|
|
||||||
|
newEchoMCPServer := func(name string) *httptest.Server {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
mcpSrv := mcpserver.NewMCPServer(name, "1.0.0")
|
||||||
|
mcpSrv.AddTools(mcpserver.ServerTool{
|
||||||
|
Tool: mcpgo.NewTool("echo",
|
||||||
|
mcpgo.WithDescription("Echoes the input"),
|
||||||
|
mcpgo.WithString("input",
|
||||||
|
mcpgo.Description("The input string"),
|
||||||
|
mcpgo.Required(),
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Handler: func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) {
|
||||||
|
input, _ := req.GetArguments()["input"].(string)
|
||||||
|
return mcpgo.NewToolResultText("echo: " + input), nil
|
||||||
|
},
|
||||||
|
})
|
||||||
|
mcpTS := httptest.NewServer(mcpserver.NewStreamableHTTPServer(mcpSrv))
|
||||||
|
t.Cleanup(mcpTS.Close)
|
||||||
|
return mcpTS
|
||||||
|
}
|
||||||
|
|
||||||
|
parentTS := newEchoMCPServer("runtime-parent-mcp")
|
||||||
|
injectedTS := newEchoMCPServer("runtime-injected-mcp")
|
||||||
|
|
||||||
|
var (
|
||||||
|
requestsMu sync.Mutex
|
||||||
|
requests []recordedOpenAIRequest
|
||||||
|
)
|
||||||
|
childRequests := func() []recordedOpenAIRequest {
|
||||||
|
requestsMu.Lock()
|
||||||
|
defer requestsMu.Unlock()
|
||||||
|
|
||||||
|
filtered := make([]recordedOpenAIRequest, 0, len(requests))
|
||||||
|
for _, req := range requests {
|
||||||
|
if requestHasSystemSubstring(req, "You are in Explore Mode as a delegated sub-agent.") {
|
||||||
|
filtered = append(filtered, req)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return filtered
|
||||||
|
}
|
||||||
|
|
||||||
|
var streamCallCount atomic.Int32
|
||||||
|
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||||
|
if !req.Stream {
|
||||||
|
return chattest.OpenAINonStreamingResponse("ok")
|
||||||
|
}
|
||||||
|
|
||||||
|
requestsMu.Lock()
|
||||||
|
requests = append(requests, recordOpenAIRequest(req))
|
||||||
|
requestsMu.Unlock()
|
||||||
|
|
||||||
|
if streamCallCount.Add(1) == 1 {
|
||||||
|
return chattest.OpenAIStreamingResponse(
|
||||||
|
chattest.OpenAIToolCallChunk("spawn_agent", `{"type":"explore","prompt":"inspect the codebase","title":"sub"}`),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
return chattest.OpenAIStreamingResponse(
|
||||||
|
chattest.OpenAITextChunks("done")...,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
|
||||||
|
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||||
|
parentConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||||
|
DisplayName: "Runtime Parent MCP",
|
||||||
|
Slug: "runtime-parent-mcp",
|
||||||
|
Url: parentTS.URL,
|
||||||
|
Transport: "streamable_http",
|
||||||
|
AuthType: "none",
|
||||||
|
Availability: "default_off",
|
||||||
|
Enabled: true,
|
||||||
|
ToolAllowList: []string{},
|
||||||
|
ToolDenyList: []string{},
|
||||||
|
CreatedBy: user.ID,
|
||||||
|
UpdatedBy: user.ID,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
injectedConfig, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||||
|
DisplayName: "Runtime Injected MCP",
|
||||||
|
Slug: "runtime-injected-mcp",
|
||||||
|
Url: injectedTS.URL,
|
||||||
|
Transport: "streamable_http",
|
||||||
|
AuthType: "none",
|
||||||
|
Availability: "default_off",
|
||||||
|
Enabled: true,
|
||||||
|
ToolAllowList: []string{},
|
||||||
|
ToolDenyList: []string{},
|
||||||
|
CreatedBy: user.ID,
|
||||||
|
UpdatedBy: user.ID,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
server := newActiveTestServer(t, db, ps)
|
||||||
|
|
||||||
|
rootChat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||||
|
OrganizationID: org.ID,
|
||||||
|
OwnerID: user.ID,
|
||||||
|
Title: "runtime-parent",
|
||||||
|
ModelConfigID: model.ID,
|
||||||
|
MCPServerIDs: []uuid.UUID{parentConfig.ID},
|
||||||
|
InitialUserContent: []codersdk.ChatMessagePart{
|
||||||
|
codersdk.ChatMessageText("Spawn an Explore subagent to inspect the codebase."),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
var exploreChat database.Chat
|
||||||
|
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||||||
|
childRows, err := db.GetChildChatsByParentIDs(dbauthz.AsChatd(ctx), database.GetChildChatsByParentIDsParams{
|
||||||
|
ParentIds: []uuid.UUID{rootChat.ID},
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
for _, candidate := range childRows {
|
||||||
|
if candidate.Chat.Mode.Valid && candidate.Chat.Mode.ChatMode == database.ChatModeExplore {
|
||||||
|
exploreChat = candidate.Chat
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}, testutil.IntervalFast)
|
||||||
|
|
||||||
|
chatResult := waitForTerminalChat(ctx, t, db, exploreChat.ID)
|
||||||
|
if chatResult.Status == database.ChatStatusError {
|
||||||
|
require.FailNowf(t, "explore chat failed", "last_error=%q", chatResult.LastError.String)
|
||||||
|
}
|
||||||
|
|
||||||
|
exploreChat, err = db.GetChatByID(ctx, exploreChat.ID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.ElementsMatch(t, []uuid.UUID{parentConfig.ID}, exploreChat.MCPServerIDs)
|
||||||
|
|
||||||
|
initialChildRequestCount := len(childRequests())
|
||||||
|
require.GreaterOrEqual(t, initialChildRequestCount, 1)
|
||||||
|
|
||||||
|
updatedMCPServerIDs := []uuid.UUID{injectedConfig.ID}
|
||||||
|
_, err = server.SendMessage(ctx, chatd.SendMessageOptions{
|
||||||
|
ChatID: exploreChat.ID,
|
||||||
|
CreatedBy: user.ID,
|
||||||
|
Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("inspect the codebase again")},
|
||||||
|
MCPServerIDs: &updatedMCPServerIDs,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
storedExploreChat, err := db.GetChatByID(ctx, exploreChat.ID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.ElementsMatch(t, []uuid.UUID{parentConfig.ID}, storedExploreChat.MCPServerIDs)
|
||||||
|
|
||||||
|
testutil.Eventually(ctx, t, func(ctx context.Context) bool {
|
||||||
|
return len(childRequests()) > initialChildRequestCount
|
||||||
|
}, testutil.IntervalFast)
|
||||||
|
|
||||||
|
chatResult = waitForTerminalChat(ctx, t, db, exploreChat.ID)
|
||||||
|
if chatResult.Status == database.ChatStatusError {
|
||||||
|
require.FailNowf(t, "explore chat failed", "last_error=%q", chatResult.LastError.String)
|
||||||
|
}
|
||||||
|
|
||||||
|
recordedChildRequests := childRequests()
|
||||||
|
require.GreaterOrEqual(t, len(recordedChildRequests), initialChildRequestCount+1)
|
||||||
|
|
||||||
|
tools := recordedChildRequests[len(recordedChildRequests)-1].Tools
|
||||||
|
require.Contains(t, tools, "runtime-parent-mcp__echo")
|
||||||
|
require.NotContains(t, tools, "runtime-injected-mcp__echo",
|
||||||
|
"Explore child runtime should keep the spawn-time MCP snapshot after SendMessage")
|
||||||
|
}
|
||||||
|
|
||||||
func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) {
|
func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
@@ -4884,6 +5437,73 @@ func waitForTerminalChat(
|
|||||||
return chatResult
|
return chatResult
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func insertChatModelConfigWithCallConfig(
|
||||||
|
ctx context.Context,
|
||||||
|
t *testing.T,
|
||||||
|
db database.Store,
|
||||||
|
userID uuid.UUID,
|
||||||
|
provider string,
|
||||||
|
model string,
|
||||||
|
callConfig codersdk.ChatModelCallConfig,
|
||||||
|
) database.ChatModelConfig {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
options, err := json.Marshal(callConfig)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||||
|
Provider: provider,
|
||||||
|
Model: model,
|
||||||
|
DisplayName: model,
|
||||||
|
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||||
|
UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||||
|
Enabled: true,
|
||||||
|
IsDefault: false,
|
||||||
|
ContextLimit: 128000,
|
||||||
|
CompressionThreshold: 70,
|
||||||
|
Options: options,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
return modelConfig
|
||||||
|
}
|
||||||
|
|
||||||
|
func insertUserTextMessage(
|
||||||
|
ctx context.Context,
|
||||||
|
t *testing.T,
|
||||||
|
db database.Store,
|
||||||
|
chatID uuid.UUID,
|
||||||
|
userID uuid.UUID,
|
||||||
|
modelConfigID uuid.UUID,
|
||||||
|
text string,
|
||||||
|
) {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
_, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{
|
||||||
|
ChatID: chatID,
|
||||||
|
CreatedBy: []uuid.UUID{userID},
|
||||||
|
ModelConfigID: []uuid.UUID{modelConfigID},
|
||||||
|
Role: []database.ChatMessageRole{database.ChatMessageRoleUser},
|
||||||
|
Content: []string{string(content.RawMessage)},
|
||||||
|
ContentVersion: []int16{chatprompt.CurrentContentVersion},
|
||||||
|
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth},
|
||||||
|
InputTokens: []int64{0},
|
||||||
|
OutputTokens: []int64{0},
|
||||||
|
TotalTokens: []int64{0},
|
||||||
|
ReasoningTokens: []int64{0},
|
||||||
|
CacheCreationTokens: []int64{0},
|
||||||
|
CacheReadTokens: []int64{0},
|
||||||
|
ContextLimit: []int64{0},
|
||||||
|
Compressed: []bool{false},
|
||||||
|
TotalCostMicros: []int64{0},
|
||||||
|
RuntimeMs: []int64{0},
|
||||||
|
ProviderResponseID: []string{""},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
}
|
||||||
|
|
||||||
// seedWorkspaceWithAgent creates a full workspace chain with a connected
|
// seedWorkspaceWithAgent creates a full workspace chain with a connected
|
||||||
// agent. This is the common setup needed by tests that exercise tool
|
// agent. This is the common setup needed by tests that exercise tool
|
||||||
// execution against a workspace.
|
// execution against a workspace.
|
||||||
|
|||||||
@@ -71,6 +71,7 @@ type OpenAIToolFunction struct {
|
|||||||
// OpenAITool represents a tool definition in an OpenAI request.
|
// OpenAITool represents a tool definition in an OpenAI request.
|
||||||
type OpenAITool struct {
|
type OpenAITool struct {
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
|
Name string `json:"name,omitempty"`
|
||||||
Function OpenAIToolFunction `json:"function"`
|
Function OpenAIToolFunction `json:"function"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ import (
|
|||||||
"database/sql"
|
"database/sql"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"slices"
|
||||||
"sort"
|
"sort"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
@@ -189,10 +190,16 @@ func (p *Server) subagentTools(
|
|||||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
turnParent := currentChatSnapshot
|
||||||
|
if turnParent.ID == uuid.Nil {
|
||||||
|
turnParent = parent
|
||||||
|
}
|
||||||
|
|
||||||
options, err := definition.buildOptions(
|
options, err := definition.buildOptions(
|
||||||
ctx,
|
ctx,
|
||||||
p,
|
p,
|
||||||
parent,
|
parent,
|
||||||
|
turnParent,
|
||||||
currentModelConfigID,
|
currentModelConfigID,
|
||||||
args.Prompt,
|
args.Prompt,
|
||||||
)
|
)
|
||||||
@@ -462,11 +469,66 @@ func parseSubagentToolChatID(raw string) (uuid.UUID, error) {
|
|||||||
return chatID, nil
|
return chatID, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// childSubagentChatOptions carries per-child overrides for subagent chat
|
||||||
|
// creation. modelConfigIDOverride and planModeOverride apply to any
|
||||||
|
// subagent. inheritedMCPServerIDs is an Explore-only snapshot of the
|
||||||
|
// spawning parent turn's effective external MCP entitlement.
|
||||||
|
// resolveExploreToolSnapshot computes and persists it on the child chat.
|
||||||
|
// Non-Explore children ignore this field.
|
||||||
type childSubagentChatOptions struct {
|
type childSubagentChatOptions struct {
|
||||||
chatMode database.NullChatMode
|
chatMode database.NullChatMode
|
||||||
systemPrompt string
|
systemPrompt string
|
||||||
modelConfigIDOverride *uuid.UUID
|
modelConfigIDOverride *uuid.UUID
|
||||||
planModeOverride *database.NullChatPlanMode
|
planModeOverride *database.NullChatPlanMode
|
||||||
|
inheritedMCPServerIDs []uuid.UUID
|
||||||
|
}
|
||||||
|
|
||||||
|
// resolveExploreToolSnapshot computes the child chat's inherited MCP
|
||||||
|
// server snapshot from the spawning parent turn.
|
||||||
|
//
|
||||||
|
// The MCP set is filtered in two stages. First,
|
||||||
|
// filterExternalMCPConfigsForTurn applies the parent turn's plan-mode
|
||||||
|
// policy to the parent's MCP configs, producing visibleConfigs. Second,
|
||||||
|
// if the parent is itself an Explore child, the visible set is narrowed to
|
||||||
|
// the parent's persisted MCPServerIDs so an Explore chain cannot
|
||||||
|
// re-escalate beyond the original grant. Non-Explore parents pass
|
||||||
|
// through the second stage unchanged.
|
||||||
|
func (p *Server) resolveExploreToolSnapshot(
|
||||||
|
ctx context.Context,
|
||||||
|
parent database.Chat,
|
||||||
|
) ([]uuid.UUID, error) {
|
||||||
|
inheritedMCPServerIDs := []uuid.UUID{}
|
||||||
|
if len(parent.MCPServerIDs) > 0 {
|
||||||
|
configs, err := p.db.GetMCPServerConfigsByIDs(ctx, parent.MCPServerIDs)
|
||||||
|
if err != nil {
|
||||||
|
return nil, xerrors.Errorf("get parent MCP server configs for chat %s: %w", parent.ID, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
visibleConfigs, _ := filterExternalMCPConfigsForTurn(
|
||||||
|
configs,
|
||||||
|
parent.PlanMode,
|
||||||
|
parent.ParentChatID,
|
||||||
|
)
|
||||||
|
// Empty means the parent is not Explore, so all plan-filtered
|
||||||
|
// configs remain eligible. Populated means the parent is
|
||||||
|
// Explore, so only its persisted snapshot can pass.
|
||||||
|
allowedParentIDs := map[uuid.UUID]struct{}{}
|
||||||
|
if isExploreSubagentMode(parent.Mode) {
|
||||||
|
for _, id := range parent.MCPServerIDs {
|
||||||
|
allowedParentIDs[id] = struct{}{}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, cfg := range visibleConfigs {
|
||||||
|
if len(allowedParentIDs) > 0 {
|
||||||
|
if _, ok := allowedParentIDs[cfg.ID]; !ok {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
}
|
||||||
|
inheritedMCPServerIDs = append(inheritedMCPServerIDs, cfg.ID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return inheritedMCPServerIDs, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (p *Server) createChildSubagentChat(
|
func (p *Server) createChildSubagentChat(
|
||||||
@@ -518,6 +580,9 @@ func (p *Server) createChildSubagentChatWithOptions(
|
|||||||
}
|
}
|
||||||
|
|
||||||
mcpServerIDs := parent.MCPServerIDs
|
mcpServerIDs := parent.MCPServerIDs
|
||||||
|
if isExploreSubagentMode(opts.chatMode) {
|
||||||
|
mcpServerIDs = slices.Clone(opts.inheritedMCPServerIDs)
|
||||||
|
}
|
||||||
if mcpServerIDs == nil {
|
if mcpServerIDs == nil {
|
||||||
mcpServerIDs = []uuid.UUID{}
|
mcpServerIDs = []uuid.UUID{}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -35,7 +35,7 @@ type subagentDefinition struct {
|
|||||||
id string
|
id string
|
||||||
description string
|
description string
|
||||||
unavailableReason func(context.Context, *Server, database.Chat) string
|
unavailableReason func(context.Context, *Server, database.Chat) string
|
||||||
buildOptions func(context.Context, *Server, database.Chat, uuid.UUID, string) (childSubagentChatOptions, error)
|
buildOptions func(context.Context, *Server, database.Chat, database.Chat, uuid.UUID, string) (childSubagentChatOptions, error)
|
||||||
}
|
}
|
||||||
|
|
||||||
func allSubagentDefinitions() []subagentDefinition {
|
func allSubagentDefinitions() []subagentDefinition {
|
||||||
@@ -43,22 +43,31 @@ func allSubagentDefinitions() []subagentDefinition {
|
|||||||
{
|
{
|
||||||
id: subagentTypeGeneral,
|
id: subagentTypeGeneral,
|
||||||
description: "delegated work that may inspect or modify workspace files",
|
description: "delegated work that may inspect or modify workspace files",
|
||||||
buildOptions: func(_ context.Context, _ *Server, _ database.Chat, _ uuid.UUID, _ string) (childSubagentChatOptions, error) {
|
buildOptions: func(_ context.Context, _ *Server, _ database.Chat, _ database.Chat, _ uuid.UUID, _ string) (childSubagentChatOptions, error) {
|
||||||
return childSubagentChatOptions{}, nil
|
return childSubagentChatOptions{}, nil
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
id: subagentTypeExplore,
|
id: subagentTypeExplore,
|
||||||
description: "read-only discovery, code tracing, and system understanding",
|
description: "read-only discovery, code tracing, and system understanding",
|
||||||
buildOptions: func(ctx context.Context, p *Server, parent database.Chat, currentModelConfigID uuid.UUID, _ string) (childSubagentChatOptions, error) {
|
buildOptions: func(ctx context.Context, p *Server, _ database.Chat, turnParent database.Chat, currentModelConfigID uuid.UUID, _ string) (childSubagentChatOptions, error) {
|
||||||
modelConfigID, err := p.resolveExploreSubagentModelConfigID(
|
modelConfigID, err := p.resolveExploreSubagentModelConfigID(
|
||||||
ctx,
|
ctx,
|
||||||
parent.OwnerID,
|
turnParent.OwnerID,
|
||||||
currentModelConfigID,
|
currentModelConfigID,
|
||||||
)
|
)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return childSubagentChatOptions{}, err
|
return childSubagentChatOptions{}, err
|
||||||
}
|
}
|
||||||
|
inheritedMCPServerIDs, err := p.resolveExploreToolSnapshot(
|
||||||
|
ctx,
|
||||||
|
turnParent,
|
||||||
|
)
|
||||||
|
if err != nil {
|
||||||
|
return childSubagentChatOptions{}, err
|
||||||
|
}
|
||||||
|
// Clearing plan mode changes only the Explore model behavior.
|
||||||
|
// The inherited tool snapshot still comes from the parent turn.
|
||||||
clearPlanMode := database.NullChatPlanMode{}
|
clearPlanMode := database.NullChatPlanMode{}
|
||||||
return childSubagentChatOptions{
|
return childSubagentChatOptions{
|
||||||
chatMode: database.NullChatMode{
|
chatMode: database.NullChatMode{
|
||||||
@@ -67,6 +76,7 @@ func allSubagentDefinitions() []subagentDefinition {
|
|||||||
},
|
},
|
||||||
modelConfigIDOverride: &modelConfigID,
|
modelConfigIDOverride: &modelConfigID,
|
||||||
planModeOverride: &clearPlanMode,
|
planModeOverride: &clearPlanMode,
|
||||||
|
inheritedMCPServerIDs: inheritedMCPServerIDs,
|
||||||
}, nil
|
}, nil
|
||||||
},
|
},
|
||||||
},
|
},
|
||||||
@@ -82,7 +92,7 @@ func allSubagentDefinitions() []subagentDefinition {
|
|||||||
}
|
}
|
||||||
return ""
|
return ""
|
||||||
},
|
},
|
||||||
buildOptions: func(_ context.Context, _ *Server, _ database.Chat, _ uuid.UUID, prompt string) (childSubagentChatOptions, error) {
|
buildOptions: func(_ context.Context, _ *Server, _ database.Chat, _ database.Chat, _ uuid.UUID, prompt string) (childSubagentChatOptions, error) {
|
||||||
return childSubagentChatOptions{
|
return childSubagentChatOptions{
|
||||||
chatMode: database.NullChatMode{
|
chatMode: database.NullChatMode{
|
||||||
ChatMode: database.ChatModeComputerUse,
|
ChatMode: database.ChatModeComputerUse,
|
||||||
|
|||||||
@@ -220,6 +220,27 @@ func insertInternalChatModelConfig(
|
|||||||
enabled bool,
|
enabled bool,
|
||||||
) database.ChatModelConfig {
|
) database.ChatModelConfig {
|
||||||
t.Helper()
|
t.Helper()
|
||||||
|
return insertInternalChatModelConfigWithOptions(
|
||||||
|
ctx,
|
||||||
|
t,
|
||||||
|
db,
|
||||||
|
userID,
|
||||||
|
model,
|
||||||
|
enabled,
|
||||||
|
json.RawMessage(`{}`),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
func insertInternalChatModelConfigWithOptions(
|
||||||
|
ctx context.Context,
|
||||||
|
t *testing.T,
|
||||||
|
db database.Store,
|
||||||
|
userID uuid.UUID,
|
||||||
|
model string,
|
||||||
|
enabled bool,
|
||||||
|
options json.RawMessage,
|
||||||
|
) database.ChatModelConfig {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
modelConfig, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||||
Provider: "openai",
|
Provider: "openai",
|
||||||
@@ -231,13 +252,42 @@ func insertInternalChatModelConfig(
|
|||||||
IsDefault: false,
|
IsDefault: false,
|
||||||
ContextLimit: 128000,
|
ContextLimit: 128000,
|
||||||
CompressionThreshold: 70,
|
CompressionThreshold: 70,
|
||||||
Options: json.RawMessage(`{}`),
|
Options: options,
|
||||||
})
|
})
|
||||||
require.NoError(t, err)
|
require.NoError(t, err)
|
||||||
|
|
||||||
return modelConfig
|
return modelConfig
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func insertInternalMCPServerConfig(
|
||||||
|
ctx context.Context,
|
||||||
|
t *testing.T,
|
||||||
|
db database.Store,
|
||||||
|
userID uuid.UUID,
|
||||||
|
slug string,
|
||||||
|
allowInPlanMode bool,
|
||||||
|
) database.MCPServerConfig {
|
||||||
|
t.Helper()
|
||||||
|
|
||||||
|
cfg, err := db.InsertMCPServerConfig(ctx, database.InsertMCPServerConfigParams{
|
||||||
|
DisplayName: slug,
|
||||||
|
Slug: slug,
|
||||||
|
Url: "https://" + slug + ".example.com",
|
||||||
|
Transport: "streamable_http",
|
||||||
|
AuthType: "none",
|
||||||
|
Availability: "default_off",
|
||||||
|
Enabled: true,
|
||||||
|
AllowInPlanMode: allowInPlanMode,
|
||||||
|
ToolAllowList: []string{},
|
||||||
|
ToolDenyList: []string{},
|
||||||
|
CreatedBy: userID,
|
||||||
|
UpdatedBy: userID,
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
return cfg
|
||||||
|
}
|
||||||
|
|
||||||
func seedWorkspaceBinding(
|
func seedWorkspaceBinding(
|
||||||
t *testing.T,
|
t *testing.T,
|
||||||
db database.Store,
|
db database.Store,
|
||||||
@@ -626,6 +676,246 @@ func TestSpawnAgent_ExploreFallsBackToCurrentTurnModel(t *testing.T) {
|
|||||||
require.Equal(t, parentModel.ID, parentChat.LastModelConfigID)
|
require.Equal(t, parentModel.ID, parentChat.LastModelConfigID)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
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(ctx, 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{})
|
||||||
|
|
||||||
|
ctx := chatdTestContext(t)
|
||||||
|
user, org, model := seedInternalChatDeps(ctx, t, db)
|
||||||
|
approvedMCP := insertInternalMCPServerConfig(
|
||||||
|
ctx, t, db, user.ID, "approved-"+uuid.NewString(), true,
|
||||||
|
)
|
||||||
|
blockedMCP := insertInternalMCPServerConfig(
|
||||||
|
ctx, t, db, user.ID, "blocked-"+uuid.NewString(), false,
|
||||||
|
)
|
||||||
|
|
||||||
|
askParentRef, err := server.CreateChat(ctx, CreateOptions{
|
||||||
|
OrganizationID: org.ID,
|
||||||
|
OwnerID: user.ID,
|
||||||
|
Title: "ask-parent",
|
||||||
|
ModelConfigID: model.ID,
|
||||||
|
MCPServerIDs: []uuid.UUID{approvedMCP.ID, blockedMCP.ID},
|
||||||
|
InitialUserContent: []codersdk.ChatMessagePart{
|
||||||
|
codersdk.ChatMessageText("hello"),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
askParent, err := db.GetChatByID(ctx, askParentRef.ID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
planParentRef, err := server.CreateChat(ctx, CreateOptions{
|
||||||
|
OrganizationID: org.ID,
|
||||||
|
OwnerID: user.ID,
|
||||||
|
Title: "plan-parent",
|
||||||
|
ModelConfigID: model.ID,
|
||||||
|
PlanMode: database.NullChatPlanMode{
|
||||||
|
ChatPlanMode: database.ChatPlanModePlan,
|
||||||
|
Valid: true,
|
||||||
|
},
|
||||||
|
MCPServerIDs: []uuid.UUID{approvedMCP.ID, blockedMCP.ID},
|
||||||
|
InitialUserContent: []codersdk.ChatMessagePart{
|
||||||
|
codersdk.ChatMessageText("hello"),
|
||||||
|
},
|
||||||
|
})
|
||||||
|
require.NoError(t, err)
|
||||||
|
planParent, err := db.GetChatByID(ctx, planParentRef.ID)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
subagentPlanParent := planParent
|
||||||
|
subagentPlanParent.ParentChatID = uuid.NullUUID{UUID: uuid.New(), Valid: true}
|
||||||
|
|
||||||
|
exploreParent := askParent
|
||||||
|
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()
|
||||||
|
|
||||||
|
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(ctx, t, db)
|
||||||
|
parentChat := createInternalParentChat(
|
||||||
|
ctx, t, server, db, org.ID, user.ID, model.ID, "parent-explore-snapshot",
|
||||||
|
)
|
||||||
|
mcpCfg := insertInternalMCPServerConfig(
|
||||||
|
ctx, t, db, user.ID, "snapshot-"+uuid.NewString(), false,
|
||||||
|
)
|
||||||
|
|
||||||
|
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(ctx, t, db)
|
||||||
|
turnStartConfig := insertInternalMCPServerConfig(
|
||||||
|
ctx, t, db, user.ID, "turn-start-"+uuid.NewString(), false,
|
||||||
|
)
|
||||||
|
mutatedConfig := insertInternalMCPServerConfig(
|
||||||
|
ctx, 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) {
|
func TestSpawnAgent_ExploreFallsBackOnInvalidUUID(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user