mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +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.
|
||||
// Explore child chats keep the spawn-time snapshot immutable.
|
||||
if opts.MCPServerIDs != nil {
|
||||
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)
|
||||
if isExploreSubagentMode(lockedChat.Mode) {
|
||||
p.logger.Warn(ctx,
|
||||
"ignoring explore subagent mcp server ids update, snapshot is immutable after spawn",
|
||||
slog.F("chat_id", opts.ChatID),
|
||||
)
|
||||
} 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
|
||||
}
|
||||
|
||||
// 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(
|
||||
configs []database.MCPServerConfig,
|
||||
mode database.NullChatPlanMode,
|
||||
@@ -5354,6 +5365,15 @@ func allowedExploreToolNames(allTools []fantasy.AgentTool) []string {
|
||||
name := tool.Info().Name
|
||||
if builtinExplorePolicy[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
|
||||
@@ -5711,11 +5731,25 @@ func (p *Server) runChat(
|
||||
isPlanModeTurn := currentPlanMode.Valid && currentPlanMode.ChatPlanMode == database.ChatPlanModePlan
|
||||
isExploreSubagent := isExploreSubagentMode(chat.Mode)
|
||||
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,
|
||||
currentPlanMode,
|
||||
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)
|
||||
|
||||
chainInfo := resolveChainMode(messages)
|
||||
@@ -6429,14 +6463,12 @@ func (p *Server) runChat(
|
||||
builtinToolNames[t.Info().Name] = true
|
||||
}
|
||||
|
||||
// Append external and workspace MCP tools after the built-ins so the
|
||||
// LLM sees them as additional capabilities. Explore subagents keep
|
||||
// the narrower built-in-only boundary from main. Root plan mode gets
|
||||
// only approved external MCP tools because mcpConnectConfigs was
|
||||
// pre-filtered above, and filterToolsForTurn removes any remaining
|
||||
// plan-mode ineligible tools from the assembled set.
|
||||
// Append external MCP tools from the chat's persisted snapshot after the
|
||||
// built-ins so the LLM sees them as additional capabilities. Explore chats
|
||||
// trust only the persisted MCPServerIDs snapshot, and workspace-local MCP
|
||||
// tools stay unavailable to Explore chats.
|
||||
tools = append(tools, mcpTools...)
|
||||
if !isExploreSubagent {
|
||||
tools = append(tools, mcpTools...)
|
||||
tools = append(tools, workspaceMCPTools...)
|
||||
}
|
||||
tools = filterToolsForTurn(
|
||||
@@ -6462,11 +6494,23 @@ func (p *Server) runChat(
|
||||
return result, err
|
||||
}
|
||||
|
||||
// Build provider-native tools (e.g., web search) based on
|
||||
// the model configuration.
|
||||
// Build provider-native tools (e.g. web search) based on the
|
||||
// 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
|
||||
if !isPlanModeTurn && !isExploreSubagent && callConfig.ProviderOptions != nil {
|
||||
if !isPlanModeTurn && callConfig.ProviderOptions != nil {
|
||||
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 {
|
||||
@@ -6719,6 +6763,10 @@ func (p *Server) runChat(
|
||||
func buildProviderTools(_ string, options *codersdk.ChatModelProviderOptions) []chatloop.ProviderTool {
|
||||
var tools []chatloop.ProviderTool
|
||||
|
||||
if options == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
if options.Anthropic != nil && options.Anthropic.WebSearchEnabled != nil && *options.Anthropic.WebSearchEnabled {
|
||||
tools = append(tools, chatloop.ProviderTool{
|
||||
Definition: anthropic.WebSearchTool(&anthropic.WebSearchToolOptions{
|
||||
|
||||
@@ -300,36 +300,32 @@ func TestActiveToolNamesForTurn(t *testing.T) {
|
||||
func TestAllowedExploreToolNames(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
makeTools := func(names ...string) []fantasy.AgentTool {
|
||||
tools := make([]fantasy.AgentTool, 0, len(names))
|
||||
for _, name := range names {
|
||||
tools = append(tools, newTestAgentTool(name))
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
got := allowedExploreToolNames(makeTools(
|
||||
"read_file",
|
||||
"write_file",
|
||||
"edit_files",
|
||||
"execute",
|
||||
"process_output",
|
||||
"process_list",
|
||||
"process_signal",
|
||||
"spawn_agent",
|
||||
"wait_agent",
|
||||
"read_skill",
|
||||
"read_skill_file",
|
||||
"ask_user_question",
|
||||
))
|
||||
externalConfigID := uuid.New()
|
||||
got := allowedExploreToolNames([]fantasy.AgentTool{
|
||||
newTestAgentTool("read_file"),
|
||||
newTestAgentTool("write_file"),
|
||||
newTestMCPAgentTool("external-mcp__echo", externalConfigID),
|
||||
newTestAgentTool("workspace-mcp__echo"),
|
||||
newTestAgentTool("execute"),
|
||||
newTestAgentTool("process_output"),
|
||||
newTestAgentTool("process_list"),
|
||||
newTestAgentTool("process_signal"),
|
||||
newTestAgentTool("spawn_agent"),
|
||||
newTestAgentTool("wait_agent"),
|
||||
newTestAgentTool("read_skill"),
|
||||
newTestAgentTool("read_skill_file"),
|
||||
newTestAgentTool("ask_user_question"),
|
||||
})
|
||||
|
||||
require.Equal(t, []string{
|
||||
"read_file",
|
||||
"external-mcp__echo",
|
||||
"execute",
|
||||
"process_output",
|
||||
"read_skill",
|
||||
"read_skill_file",
|
||||
}, got)
|
||||
require.NotContains(t, got, "workspace-mcp__echo")
|
||||
}
|
||||
|
||||
func TestAllowedBehaviorToolNames(t *testing.T) {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package chatd_test
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/base64"
|
||||
@@ -62,11 +63,15 @@ type recordedOpenAIRequest struct {
|
||||
ContentLength int64
|
||||
}
|
||||
|
||||
func openAIToolName(tool chattest.OpenAITool) string {
|
||||
return cmp.Or(tool.Function.Name, tool.Name, tool.Type)
|
||||
}
|
||||
|
||||
func recordOpenAIRequest(req *chattest.OpenAIRequest) recordedOpenAIRequest {
|
||||
messages := append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
tools := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
tools = append(tools, tool.Function.Name)
|
||||
tools = append(tools, openAIToolName(tool))
|
||||
}
|
||||
|
||||
var store *bool
|
||||
@@ -697,6 +702,554 @@ func TestExploreSubagentIsReadOnly(t *testing.T) {
|
||||
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) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -4884,6 +5437,73 @@ func waitForTerminalChat(
|
||||
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
|
||||
// agent. This is the common setup needed by tests that exercise tool
|
||||
// execution against a workspace.
|
||||
|
||||
@@ -71,6 +71,7 @@ type OpenAIToolFunction struct {
|
||||
// OpenAITool represents a tool definition in an OpenAI request.
|
||||
type OpenAITool struct {
|
||||
Type string `json:"type"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Function OpenAIToolFunction `json:"function"`
|
||||
}
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"slices"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -189,10 +190,16 @@ func (p *Server) subagentTools(
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
}
|
||||
|
||||
turnParent := currentChatSnapshot
|
||||
if turnParent.ID == uuid.Nil {
|
||||
turnParent = parent
|
||||
}
|
||||
|
||||
options, err := definition.buildOptions(
|
||||
ctx,
|
||||
p,
|
||||
parent,
|
||||
turnParent,
|
||||
currentModelConfigID,
|
||||
args.Prompt,
|
||||
)
|
||||
@@ -462,11 +469,66 @@ func parseSubagentToolChatID(raw string) (uuid.UUID, error) {
|
||||
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 {
|
||||
chatMode database.NullChatMode
|
||||
systemPrompt string
|
||||
modelConfigIDOverride *uuid.UUID
|
||||
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(
|
||||
@@ -518,6 +580,9 @@ func (p *Server) createChildSubagentChatWithOptions(
|
||||
}
|
||||
|
||||
mcpServerIDs := parent.MCPServerIDs
|
||||
if isExploreSubagentMode(opts.chatMode) {
|
||||
mcpServerIDs = slices.Clone(opts.inheritedMCPServerIDs)
|
||||
}
|
||||
if mcpServerIDs == nil {
|
||||
mcpServerIDs = []uuid.UUID{}
|
||||
}
|
||||
|
||||
@@ -35,7 +35,7 @@ type subagentDefinition struct {
|
||||
id string
|
||||
description 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 {
|
||||
@@ -43,22 +43,31 @@ func allSubagentDefinitions() []subagentDefinition {
|
||||
{
|
||||
id: subagentTypeGeneral,
|
||||
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
|
||||
},
|
||||
},
|
||||
{
|
||||
id: subagentTypeExplore,
|
||||
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(
|
||||
ctx,
|
||||
parent.OwnerID,
|
||||
turnParent.OwnerID,
|
||||
currentModelConfigID,
|
||||
)
|
||||
if err != nil {
|
||||
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{}
|
||||
return childSubagentChatOptions{
|
||||
chatMode: database.NullChatMode{
|
||||
@@ -67,6 +76,7 @@ func allSubagentDefinitions() []subagentDefinition {
|
||||
},
|
||||
modelConfigIDOverride: &modelConfigID,
|
||||
planModeOverride: &clearPlanMode,
|
||||
inheritedMCPServerIDs: inheritedMCPServerIDs,
|
||||
}, nil
|
||||
},
|
||||
},
|
||||
@@ -82,7 +92,7 @@ func allSubagentDefinitions() []subagentDefinition {
|
||||
}
|
||||
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{
|
||||
chatMode: database.NullChatMode{
|
||||
ChatMode: database.ChatModeComputerUse,
|
||||
|
||||
@@ -220,6 +220,27 @@ func insertInternalChatModelConfig(
|
||||
enabled bool,
|
||||
) database.ChatModelConfig {
|
||||
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{
|
||||
Provider: "openai",
|
||||
@@ -231,13 +252,42 @@ func insertInternalChatModelConfig(
|
||||
IsDefault: false,
|
||||
ContextLimit: 128000,
|
||||
CompressionThreshold: 70,
|
||||
Options: json.RawMessage(`{}`),
|
||||
Options: options,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
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(
|
||||
t *testing.T,
|
||||
db database.Store,
|
||||
@@ -626,6 +676,246 @@ func TestSpawnAgent_ExploreFallsBackToCurrentTurnModel(t *testing.T) {
|
||||
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) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user