dynamic mcp: harden client catalog discovery

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: nickmisasi

llm/context.go

@@ -4,6 +4,7 @@
 package llm
 
 import (
+	stdcontext "context"
 	"fmt"
 	"strings"
 	"time"
@@ -36,6 +37,13 @@ type Context struct {
 	// User that is making the request
 	RequestingUser *model.User
 
+	// RequestContext carries the caller's request-scoped context for downstream
+	// work such as MCP tool discovery. May be nil in tests.
+	RequestContext stdcontext.Context
+
+	// ConversationID identifies the conversation whose context is being built.
+	ConversationID string
+
 	// Bot Specific
 	BotName            string
 	BotUsername        string
@@ -47,6 +55,42 @@ type Context struct {
 	Tools             *ToolStore
 	DisabledToolsInfo []ToolInfo // Info about tools that are unavailable in the current context (e.g., DM-only tools in a channel)
 	Parameters        map[string]interface{}
+
+	// MCPDynamicToolLoading indicates this context uses strict MCP dynamic loading.
+	MCPDynamicToolLoading bool
+	// MCPDynamicToolTelemetry receives low-cardinality dynamic MCP tool events.
+	MCPDynamicToolTelemetry                 MCPDynamicToolTelemetry
+	MCPDynamicToolSearchUsed                bool
+	MCPDynamicLoadedToolNames               map[string]bool
+	MCPDynamicSearchLoadCallSuccessRecorded map[string]bool
+
+	// DisabledMCPServerOrigins contains per-user disabled MCP server origins that
+	// must be removed before strict registry construction.
+	DisabledMCPServerOrigins []string
+
+	// KeepMCPTool, when non-nil, is applied to MCP tools before strict registry
+	// construction and before flag-off visible MCP insertion.
+	KeepMCPTool func(Tool) bool
+
+	// PreloadedMCPTools contains exact-or-bare MCP tool selectors for internal
+	// predefined flows. They are selected only from the already-authorized MCP
+	// catalog and are request scoped.
+	PreloadedMCPTools []EnabledMCPTool
+
+	// MCPToolRegistry holds the strict MCP tool registry that was built
+	// alongside Tools, when MCP dynamic tool loading is enabled. It is stashed
+	// here so callers can replay loaded-tool restoration after the conversation
+	// row exists without rebuilding the entire tool store.
+	//
+	// Stored as `any` to avoid an llm -> mcp import cycle: the mcp package
+	// already imports llm, and the only consumer that needs the concrete type
+	// is the llmcontext package, which can import both. Type-assert to
+	// *mcp.ToolRegistry there.
+	MCPToolRegistry any
+}
+
+type MCPDynamicToolTelemetry interface {
+	ObserveMCPDynamicToolEvent(botName, event, result string)
 }
 
 // ContextOption defines a function that configures a Context
@@ -99,6 +143,53 @@ func (c *Context) CustomPromptVars() map[string]string {
 	return vars
 }
 
+func (c *Context) ObserveMCPDynamicToolEvent(event, result string) {
+	if c == nil || c.MCPDynamicToolTelemetry == nil {
+		return
+	}
+
+	botName := c.BotUsername
+	if botName == "" {
+		botName = c.BotName
+	}
+	if botName == "" {
+		botName = "unknown"
+	}
+
+	c.MCPDynamicToolTelemetry.ObserveMCPDynamicToolEvent(botName, event, result)
+}
+
+func (c *Context) MarkMCPDynamicToolSearch() {
+	if c == nil {
+		return
+	}
+	c.MCPDynamicToolSearchUsed = true
+}
+
+func (c *Context) MarkMCPDynamicToolLoaded(name string) {
+	if c == nil || name == "" {
+		return
+	}
+	if c.MCPDynamicLoadedToolNames == nil {
+		c.MCPDynamicLoadedToolNames = make(map[string]bool)
+	}
+	c.MCPDynamicLoadedToolNames[name] = true
+}
+
+func (c *Context) ShouldRecordMCPDynamicSearchLoadCallSuccess(name string) bool {
+	if c == nil || name == "" || !c.MCPDynamicToolSearchUsed || !c.MCPDynamicLoadedToolNames[name] {
+		return false
+	}
+	if c.MCPDynamicSearchLoadCallSuccessRecorded == nil {
+		c.MCPDynamicSearchLoadCallSuccessRecorded = make(map[string]bool)
+	}
+	if c.MCPDynamicSearchLoadCallSuccessRecorded[name] {
+		return false
+	}
+	c.MCPDynamicSearchLoadCallSuccessRecorded[name] = true
+	return true
+}
+
 func (c Context) String() string {
 	var result strings.Builder
 	result.WriteString(fmt.Sprintf("Time: %v\nServerName: %v\nCompanyName: %v", c.Time, c.ServerName, c.CompanyName))

llm/context_test.go

@@ -10,6 +10,20 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+type contextTelemetryEvent struct {
+	botName string
+	event   string
+	result  string
+}
+
+type fakeMCPDynamicTelemetry struct {
+	events []contextTelemetryEvent
+}
+
+func (t *fakeMCPDynamicTelemetry) ObserveMCPDynamicToolEvent(botName, event, result string) {
+	t.events = append(t.events, contextTelemetryEvent{botName: botName, event: event, result: result})
+}
+
 func TestContext_SetBotFields(t *testing.T) {
 	c := NewContext()
 	c.SetBotFields("BotDisplay", "botuser", "user-id-123", "gpt-4", "openai", "Be helpful and concise")
@@ -107,3 +121,51 @@ func TestContext_CustomPromptVars(t *testing.T) {
 		})
 	}
 }
+
+func TestContextObserveMCPDynamicToolEventBotLabelFallbacks(t *testing.T) {
+	tests := []struct {
+		name        string
+		context     *Context
+		wantBotName string
+	}{
+		{
+			name:        "username",
+			context:     &Context{BotUsername: "matty", BotName: "Matty"},
+			wantBotName: "matty",
+		},
+		{
+			name:        "display name",
+			context:     &Context{BotName: "Matty"},
+			wantBotName: "Matty",
+		},
+		{
+			name:        "unknown",
+			context:     &Context{},
+			wantBotName: "unknown",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			telemetry := &fakeMCPDynamicTelemetry{}
+			tt.context.MCPDynamicToolTelemetry = telemetry
+
+			tt.context.ObserveMCPDynamicToolEvent("search", "success")
+
+			assert.Equal(t, []contextTelemetryEvent{{botName: tt.wantBotName, event: "search", result: "success"}}, telemetry.events)
+		})
+	}
+}
+
+func TestContextMCPDynamicSearchLoadCallSuccessState(t *testing.T) {
+	c := &Context{}
+
+	assert.False(t, c.ShouldRecordMCPDynamicSearchLoadCallSuccess("jira__get_issue"))
+
+	c.MarkMCPDynamicToolSearch()
+	assert.False(t, c.ShouldRecordMCPDynamicSearchLoadCallSuccess("jira__get_issue"))
+
+	c.MarkMCPDynamicToolLoaded("jira__get_issue")
+	assert.True(t, c.ShouldRecordMCPDynamicSearchLoadCallSuccess("jira__get_issue"))
+	assert.False(t, c.ShouldRecordMCPDynamicSearchLoadCallSuccess("jira__get_issue"))
+}