server: add declarative auto-enable for MCP servers
Add configuration-driven auto-enable system allowing MCP servers to declare when they should activate based on context. MCPDefinitions (server/mcp_definitions.go): - Static server configuration from config files - AutoEnableMode: never, always, with_path, if_match - EnableCondition: file_exists, env_set conditions - Config loading: ~/.ollama/, /etc/ollama/, OLLAMA_MCP_SERVERS MCPSessionManager (server/mcp_sessions.go): - Runtime session tracking with TTL - Session pooling for API requests CommandResolver (server/mcp_command_resolver.go): - CommandResolverInterface for dependency injection - DefaultCommandResolver singleton for production use Public API (server/mcp.go): - GetMCPServersForTools(): CLI tool resolution - GetMCPManager(): Session-based manager access - ListMCPServers(): Server discovery Relates to #7865
This commit is contained in:
code4me2
parent
abcb81bb07
commit
c747bedbfa
|
|
@ -0,0 +1,59 @@
|
|||
{
|
||||
"servers": [
|
||||
{
|
||||
"name": "filesystem",
|
||||
"description": "File system operations with path-based access control",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-filesystem"],
|
||||
"requires_path": true,
|
||||
"path_arg_index": -1,
|
||||
"capabilities": ["read", "write", "list", "search"],
|
||||
"auto_enable": "with_path"
|
||||
},
|
||||
{
|
||||
"name": "git",
|
||||
"description": "Git repository operations",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@cyanheads/git-mcp-server"],
|
||||
"requires_path": true,
|
||||
"path_arg_index": -1,
|
||||
"capabilities": ["diff", "log", "status", "branch"],
|
||||
"auto_enable": "if_match",
|
||||
"enable_if": {
|
||||
"file_exists": ".git"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "postgres",
|
||||
"description": "PostgreSQL database operations",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-postgres"],
|
||||
"env": {
|
||||
"POSTGRES_CONNECTION": ""
|
||||
},
|
||||
"capabilities": ["query", "schema"],
|
||||
"auto_enable": "if_match",
|
||||
"enable_if": {
|
||||
"env_set": "POSTGRES_CONNECTION"
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "python",
|
||||
"description": "Python code execution in sandboxed environment",
|
||||
"command": "python",
|
||||
"args": ["-m", "mcp_server_python"],
|
||||
"requires_path": false,
|
||||
"capabilities": ["execute", "eval"],
|
||||
"auto_enable": "never"
|
||||
},
|
||||
{
|
||||
"name": "web",
|
||||
"description": "Web browsing and content extraction",
|
||||
"command": "npx",
|
||||
"args": ["-y", "@modelcontextprotocol/server-puppeteer"],
|
||||
"requires_path": false,
|
||||
"capabilities": ["browse", "screenshot", "extract"],
|
||||
"auto_enable": "never"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,114 @@
|
|||
// Package server provides MCP (Model Context Protocol) integration for Ollama.
|
||||
//
|
||||
// MCP Architecture:
|
||||
//
|
||||
// ┌─────────────────────────────────────────────────────────────────┐
|
||||
// │ Public API (this file) │
|
||||
// │ GetMCPServersForTools() - Get servers for --tools flag │
|
||||
// │ GetMCPManager() - Get manager for explicit configs │
|
||||
// │ GetMCPManagerForPath() - Get manager for tools path │
|
||||
// │ ListMCPServers() - List available server definitions │
|
||||
// └─────────────────────────────────────────────────────────────────┘
|
||||
// │
|
||||
// ┌───────────────────┴───────────────────┐
|
||||
// ▼ ▼
|
||||
// ┌─────────────────────┐ ┌─────────────────────┐
|
||||
// │ MCPDefinitions │ │ MCPSessionManager │
|
||||
// │ (mcp_definitions) │ │ (mcp_sessions) │
|
||||
// │ │ │ │
|
||||
// │ Static config of │ │ Runtime sessions │
|
||||
// │ available servers │ │ with connections │
|
||||
// └─────────────────────┘ └─────────────────────┘
|
||||
// │
|
||||
// ▼
|
||||
// ┌─────────────────────┐
|
||||
// │ MCPManager │
|
||||
// │ (mcp_manager) │
|
||||
// │ │
|
||||
// │ Multi-client mgmt │
|
||||
// │ Tool execution │
|
||||
// └─────────────────────┘
|
||||
// │
|
||||
// ▼
|
||||
// ┌─────────────────────┐
|
||||
// │ MCPClient │
|
||||
// │ (mcp_client) │
|
||||
// │ │
|
||||
// │ Single JSON-RPC │
|
||||
// │ connection │
|
||||
// └─────────────────────┘
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
// ============================================================================
|
||||
// Public API - Clean interface for external code
|
||||
// ============================================================================
|
||||
|
||||
// GetMCPServersForTools returns the MCP server configs that should be enabled
|
||||
// for the given tools spec. It handles path normalization:
|
||||
// - "." or "true" → current working directory
|
||||
// - "~/path" → expands to home directory
|
||||
// - relative paths → resolved to absolute paths
|
||||
//
|
||||
// Returns the server configs and the resolved absolute path.
|
||||
// On error, still returns the resolved path so callers can implement fallback.
|
||||
// This is used by the --tools CLI flag.
|
||||
func GetMCPServersForTools(toolsSpec string) ([]api.MCPServerConfig, string, error) {
|
||||
// Normalize the tools path first (needed even for fallback on error)
|
||||
toolsPath := toolsSpec
|
||||
if toolsSpec == "." || toolsSpec == "true" {
|
||||
if cwd, err := os.Getwd(); err == nil {
|
||||
toolsPath = cwd
|
||||
}
|
||||
}
|
||||
|
||||
// Expand tilde to home directory
|
||||
if strings.HasPrefix(toolsPath, "~") {
|
||||
if home := os.Getenv("HOME"); home != "" {
|
||||
toolsPath = filepath.Join(home, toolsPath[1:])
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve to absolute path
|
||||
if absPath, err := filepath.Abs(toolsPath); err == nil {
|
||||
toolsPath = absPath
|
||||
}
|
||||
|
||||
// Load definitions
|
||||
defs, err := LoadMCPDefinitions()
|
||||
if err != nil {
|
||||
return nil, toolsPath, err
|
||||
}
|
||||
|
||||
ctx := AutoEnableContext{ToolsPath: toolsPath}
|
||||
return defs.GetAutoEnableServers(ctx), toolsPath, nil
|
||||
}
|
||||
|
||||
// GetMCPManager returns an MCP manager for the given session and configs.
|
||||
// If a session with matching configs already exists, it will be reused.
|
||||
func GetMCPManager(sessionID string, configs []api.MCPServerConfig) (*MCPManager, error) {
|
||||
return GetMCPSessionManager().GetOrCreateManager(sessionID, configs)
|
||||
}
|
||||
|
||||
// GetMCPManagerForPath returns an MCP manager for servers that auto-enable
|
||||
// for the given tools path. Used by CLI: `ollama run model --tools /path`
|
||||
func GetMCPManagerForPath(model string, toolsPath string) (*MCPManager, error) {
|
||||
return GetMCPSessionManager().GetManagerForToolsPath(model, toolsPath)
|
||||
}
|
||||
|
||||
// ListMCPServers returns information about all available MCP server definitions.
|
||||
func ListMCPServers() ([]MCPServerInfo, error) {
|
||||
defs, err := LoadMCPDefinitions()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return defs.ListServers(), nil
|
||||
}
|
||||
|
|
@ -8,13 +8,40 @@ import (
|
|||
"sync"
|
||||
)
|
||||
|
||||
// =============================================================================
|
||||
// Command Resolver Interface & Default Implementation
|
||||
// =============================================================================
|
||||
//
|
||||
// SECURITY REVIEW: This component determines which executables are launched
|
||||
// for MCP servers. Changes here should be reviewed carefully.
|
||||
|
||||
// CommandResolverInterface defines the contract for command resolution.
|
||||
// Implementations resolve command names (like "npx", "python") to actual
|
||||
// executable paths, with support for fallbacks and environment overrides.
|
||||
//
|
||||
// This interface enables dependency injection for testing MCPClient without
|
||||
// requiring actual executables to be present on the system.
|
||||
type CommandResolverInterface interface {
|
||||
// ResolveCommand finds the actual executable for a command name.
|
||||
// Returns the resolved path/command or an error if not found.
|
||||
ResolveCommand(command string) (string, error)
|
||||
|
||||
// ResolveForEnvironment resolves a command, checking environment
|
||||
// variable overrides first (e.g., OLLAMA_NPX_COMMAND for "npx").
|
||||
// Returns the original command as fallback if resolution fails.
|
||||
ResolveForEnvironment(command string) string
|
||||
}
|
||||
|
||||
// CommandResolver handles resolving commands to their actual executables
|
||||
// with fallback detection for different system configurations
|
||||
// with fallback detection for different system configurations.
|
||||
type CommandResolver struct {
|
||||
mu sync.RWMutex
|
||||
resolved map[string]string
|
||||
}
|
||||
|
||||
// Ensure CommandResolver implements CommandResolverInterface
|
||||
var _ CommandResolverInterface = (*CommandResolver)(nil)
|
||||
|
||||
// NewCommandResolver creates a new command resolver
|
||||
func NewCommandResolver() *CommandResolver {
|
||||
return &CommandResolver{
|
||||
|
|
|
|||
|
|
@ -0,0 +1,280 @@
|
|||
package server
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
// AutoEnableMode determines when a server auto-enables with --tools
|
||||
type AutoEnableMode string
|
||||
|
||||
const (
|
||||
// AutoEnableNever means the server must be explicitly configured (default)
|
||||
AutoEnableNever AutoEnableMode = "never"
|
||||
// AutoEnableAlways means the server enables whenever --tools is used
|
||||
AutoEnableAlways AutoEnableMode = "always"
|
||||
// AutoEnableWithPath means the server enables when --tools has a path
|
||||
AutoEnableWithPath AutoEnableMode = "with_path"
|
||||
// AutoEnableIfMatch means the server enables if EnableIf condition matches
|
||||
AutoEnableIfMatch AutoEnableMode = "if_match"
|
||||
)
|
||||
|
||||
// EnableCondition specifies conditions for AutoEnableIfMatch mode
|
||||
type EnableCondition struct {
|
||||
// FileExists checks if a specific file exists in the tools path
|
||||
FileExists string `json:"file_exists,omitempty"`
|
||||
// EnvSet checks if an environment variable is set (non-empty)
|
||||
EnvSet string `json:"env_set,omitempty"`
|
||||
}
|
||||
|
||||
// AutoEnableContext provides context for auto-enable decisions
|
||||
type AutoEnableContext struct {
|
||||
// ToolsPath is the path from --tools flag (may be empty)
|
||||
ToolsPath string
|
||||
// Env contains environment variables (optional, falls back to os.Getenv)
|
||||
Env map[string]string
|
||||
}
|
||||
|
||||
// MCPDefinitions holds available MCP server definitions loaded from configuration.
|
||||
// This is the static configuration of what servers CAN be used.
|
||||
type MCPDefinitions struct {
|
||||
Servers map[string]MCPServerDefinition `json:"servers"`
|
||||
}
|
||||
|
||||
// MCPServerDefinition defines an available MCP server type
|
||||
type MCPServerDefinition struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Command string `json:"command"`
|
||||
Args []string `json:"args,omitempty"`
|
||||
RequiresPath bool `json:"requires_path,omitempty"`
|
||||
PathArgIndex int `json:"path_arg_index,omitempty"`
|
||||
Env map[string]string `json:"env,omitempty"`
|
||||
Capabilities []string `json:"capabilities,omitempty"`
|
||||
|
||||
// AutoEnable determines when this server auto-enables with --tools
|
||||
// Default is "never" (must be explicitly configured via API)
|
||||
AutoEnable AutoEnableMode `json:"auto_enable,omitempty"`
|
||||
|
||||
// EnableIf specifies conditions for AutoEnableIfMatch mode
|
||||
EnableIf EnableCondition `json:"enable_if,omitempty"`
|
||||
}
|
||||
|
||||
// MCPServerInfo provides information about an available MCP server
|
||||
type MCPServerInfo struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
RequiresPath bool `json:"requires_path"`
|
||||
Capabilities []string `json:"capabilities,omitempty"`
|
||||
AutoEnable AutoEnableMode `json:"auto_enable,omitempty"`
|
||||
}
|
||||
|
||||
// DefaultMCPServers returns minimal built-in MCP server definitions
|
||||
// Full examples are provided in examples/mcp-servers.json
|
||||
func DefaultMCPServers() map[string]MCPServerDefinition {
|
||||
// Only include filesystem by default - it requires only npx which is commonly available
|
||||
// Users can add more servers via ~/.ollama/mcp-servers.json
|
||||
return map[string]MCPServerDefinition{
|
||||
"filesystem": {
|
||||
Name: "filesystem",
|
||||
Description: "File system operations with path-based access control",
|
||||
Command: "npx",
|
||||
Args: []string{"-y", "@modelcontextprotocol/server-filesystem"},
|
||||
RequiresPath: true,
|
||||
PathArgIndex: -1,
|
||||
Capabilities: []string{"read", "write", "list", "search"},
|
||||
AutoEnable: AutoEnableWithPath, // Enable when --tools has a path
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// LoadMCPDefinitions loads MCP server definitions from configuration files.
|
||||
// Priority order: user config (~/.ollama) > system config (/etc/ollama) > defaults
|
||||
func LoadMCPDefinitions() (*MCPDefinitions, error) {
|
||||
defs := &MCPDefinitions{
|
||||
Servers: DefaultMCPServers(),
|
||||
}
|
||||
|
||||
// Load from user config if exists
|
||||
configPaths := []string{
|
||||
filepath.Join(os.Getenv("HOME"), ".ollama", "mcp-servers.json"),
|
||||
"/etc/ollama/mcp-servers.json",
|
||||
"./mcp-servers.json",
|
||||
}
|
||||
|
||||
for _, path := range configPaths {
|
||||
if err := defs.LoadFromFile(path); err == nil {
|
||||
slog.Debug("Loaded MCP definitions", "path", path)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Load from environment variable if set
|
||||
if mcpConfig := os.Getenv("OLLAMA_MCP_SERVERS"); mcpConfig != "" {
|
||||
if err := defs.LoadFromJSON([]byte(mcpConfig)); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse OLLAMA_MCP_SERVERS: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return defs, nil
|
||||
}
|
||||
|
||||
// LoadFromFile loads MCP server definitions from a JSON file
|
||||
func (d *MCPDefinitions) LoadFromFile(path string) error {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return d.LoadFromJSON(data)
|
||||
}
|
||||
|
||||
// LoadFromJSON loads MCP server definitions from JSON data
|
||||
func (d *MCPDefinitions) LoadFromJSON(data []byte) error {
|
||||
var config struct {
|
||||
Servers []MCPServerDefinition `json:"servers"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(data, &config); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for _, server := range config.Servers {
|
||||
d.Servers[server.Name] = server
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListServers returns information about all available MCP servers
|
||||
func (d *MCPDefinitions) ListServers() []MCPServerInfo {
|
||||
var servers []MCPServerInfo
|
||||
for _, def := range d.Servers {
|
||||
servers = append(servers, MCPServerInfo{
|
||||
Name: def.Name,
|
||||
Description: def.Description,
|
||||
RequiresPath: def.RequiresPath,
|
||||
Capabilities: def.Capabilities,
|
||||
AutoEnable: def.AutoEnable,
|
||||
})
|
||||
}
|
||||
return servers
|
||||
}
|
||||
|
||||
// GetAutoEnableServers returns servers that should auto-enable for the given context.
|
||||
// This method checks each server's AutoEnable mode and EnableIf conditions.
|
||||
func (d *MCPDefinitions) GetAutoEnableServers(ctx AutoEnableContext) []api.MCPServerConfig {
|
||||
var configs []api.MCPServerConfig
|
||||
|
||||
for _, def := range d.Servers {
|
||||
if !d.shouldAutoEnable(def, ctx) {
|
||||
continue
|
||||
}
|
||||
|
||||
config, err := d.buildConfigForAutoEnable(def, ctx)
|
||||
if err != nil {
|
||||
slog.Warn("Failed to build config for auto-enable server",
|
||||
"name", def.Name, "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
configs = append(configs, config)
|
||||
}
|
||||
|
||||
return configs
|
||||
}
|
||||
|
||||
// shouldAutoEnable checks if a server should auto-enable for the given context
|
||||
func (d *MCPDefinitions) shouldAutoEnable(def MCPServerDefinition, ctx AutoEnableContext) bool {
|
||||
switch def.AutoEnable {
|
||||
case AutoEnableNever, "":
|
||||
return false
|
||||
|
||||
case AutoEnableAlways:
|
||||
return true
|
||||
|
||||
case AutoEnableWithPath:
|
||||
return ctx.ToolsPath != ""
|
||||
|
||||
case AutoEnableIfMatch:
|
||||
return d.checkEnableCondition(def.EnableIf, ctx)
|
||||
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// checkEnableCondition evaluates an EnableCondition against the context
|
||||
func (d *MCPDefinitions) checkEnableCondition(cond EnableCondition, ctx AutoEnableContext) bool {
|
||||
// All specified conditions must match (AND logic)
|
||||
|
||||
if cond.FileExists != "" {
|
||||
checkPath := filepath.Join(ctx.ToolsPath, cond.FileExists)
|
||||
if _, err := os.Stat(checkPath); err != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
if cond.EnvSet != "" {
|
||||
// Check context env first, fall back to os.Getenv
|
||||
val := ""
|
||||
if ctx.Env != nil {
|
||||
val = ctx.Env[cond.EnvSet]
|
||||
}
|
||||
if val == "" {
|
||||
val = os.Getenv(cond.EnvSet)
|
||||
}
|
||||
if val == "" {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// buildConfigForAutoEnable creates an MCPServerConfig for auto-enabled servers
|
||||
func (d *MCPDefinitions) buildConfigForAutoEnable(def MCPServerDefinition, ctx AutoEnableContext) (api.MCPServerConfig, error) {
|
||||
// Resolve the command using the command resolver
|
||||
resolvedCommand := DefaultCommandResolver.ResolveForEnvironment(def.Command)
|
||||
|
||||
config := api.MCPServerConfig{
|
||||
Name: def.Name,
|
||||
Command: resolvedCommand,
|
||||
Args: append([]string{}, def.Args...), // Copy args
|
||||
Env: make(map[string]string),
|
||||
}
|
||||
|
||||
// Copy environment variables
|
||||
for k, v := range def.Env {
|
||||
config.Env[k] = v
|
||||
}
|
||||
|
||||
// Add path if required
|
||||
if def.RequiresPath {
|
||||
if ctx.ToolsPath == "" {
|
||||
return config, fmt.Errorf("server '%s' requires a path but none provided", def.Name)
|
||||
}
|
||||
|
||||
// Validate path exists
|
||||
if _, err := os.Stat(ctx.ToolsPath); err != nil {
|
||||
return config, fmt.Errorf("invalid path for server '%s': %w", def.Name, err)
|
||||
}
|
||||
|
||||
// Add path to args at specified position
|
||||
if def.PathArgIndex < 0 {
|
||||
config.Args = append(config.Args, ctx.ToolsPath)
|
||||
} else if def.PathArgIndex <= len(config.Args) {
|
||||
config.Args = append(config.Args[:def.PathArgIndex],
|
||||
append([]string{ctx.ToolsPath}, config.Args[def.PathArgIndex:]...)...)
|
||||
} else {
|
||||
// PathArgIndex out of bounds, append to end
|
||||
config.Args = append(config.Args, ctx.ToolsPath)
|
||||
}
|
||||
}
|
||||
|
||||
return config, nil
|
||||
}
|
||||
|
|
@ -0,0 +1,178 @@
|
|||
package server
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"log/slog"
|
||||
"reflect"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/ollama/ollama/api"
|
||||
)
|
||||
|
||||
// MCPSessionManager manages active MCP sessions with automatic cleanup.
|
||||
// This is the runtime component that tracks active connections.
|
||||
type MCPSessionManager struct {
|
||||
mu sync.RWMutex
|
||||
sessions map[string]*MCPSession // session ID -> session
|
||||
ttl time.Duration // session timeout
|
||||
stopCleanup chan struct{} // signals cleanup goroutine to stop
|
||||
}
|
||||
|
||||
// MCPSession wraps an MCPManager with session metadata
|
||||
type MCPSession struct {
|
||||
*MCPManager
|
||||
lastAccess time.Time
|
||||
sessionID string
|
||||
configs []api.MCPServerConfig
|
||||
}
|
||||
|
||||
var (
|
||||
globalSessionManager *MCPSessionManager
|
||||
sessionManagerOnce sync.Once
|
||||
)
|
||||
|
||||
// GetMCPSessionManager returns the singleton MCP session manager
|
||||
func GetMCPSessionManager() *MCPSessionManager {
|
||||
sessionManagerOnce.Do(func() {
|
||||
globalSessionManager = &MCPSessionManager{
|
||||
sessions: make(map[string]*MCPSession),
|
||||
ttl: 30 * time.Minute, // Sessions expire after 30 min
|
||||
stopCleanup: make(chan struct{}),
|
||||
}
|
||||
// Start cleanup goroutine
|
||||
go globalSessionManager.cleanupExpired()
|
||||
})
|
||||
return globalSessionManager
|
||||
}
|
||||
|
||||
// GetOrCreateManager gets existing or creates new MCP manager for session
|
||||
func (sm *MCPSessionManager) GetOrCreateManager(sessionID string, configs []api.MCPServerConfig) (*MCPManager, error) {
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
|
||||
// Check if session exists and configs match
|
||||
if session, exists := sm.sessions[sessionID]; exists {
|
||||
if configsMatch(session.configs, configs) {
|
||||
session.lastAccess = time.Now()
|
||||
slog.Debug("Reusing existing MCP session", "session", sessionID, "clients", len(session.clients))
|
||||
return session.MCPManager, nil
|
||||
}
|
||||
// Configs changed, shutdown old session
|
||||
slog.Info("MCP configs changed, recreating session", "session", sessionID)
|
||||
session.Shutdown()
|
||||
delete(sm.sessions, sessionID)
|
||||
}
|
||||
|
||||
// Create new session
|
||||
slog.Info("Creating new MCP session", "session", sessionID, "configs", len(configs))
|
||||
manager := NewMCPManager(10)
|
||||
for _, config := range configs {
|
||||
if err := manager.AddServer(config); err != nil {
|
||||
slog.Warn("Failed to add MCP server", "name", config.Name, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
sm.sessions[sessionID] = &MCPSession{
|
||||
MCPManager: manager,
|
||||
lastAccess: time.Now(),
|
||||
sessionID: sessionID,
|
||||
configs: configs,
|
||||
}
|
||||
|
||||
return manager, nil
|
||||
}
|
||||
|
||||
// GetManagerForToolsPath creates a manager for a tools directory path.
|
||||
// It uses the definitions system to get auto-enabled servers for the path.
|
||||
func (sm *MCPSessionManager) GetManagerForToolsPath(model string, toolsPath string) (*MCPManager, error) {
|
||||
// Generate consistent session ID for model + tools path
|
||||
sessionID := generateToolsSessionID(model, toolsPath)
|
||||
|
||||
// Use definitions to get auto-enabled servers (single source of truth)
|
||||
defs, err := LoadMCPDefinitions()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx := AutoEnableContext{ToolsPath: toolsPath}
|
||||
configs := defs.GetAutoEnableServers(ctx)
|
||||
|
||||
return sm.GetOrCreateManager(sessionID, configs)
|
||||
}
|
||||
|
||||
// cleanupExpired removes expired sessions
|
||||
func (sm *MCPSessionManager) cleanupExpired() {
|
||||
ticker := time.NewTicker(5 * time.Minute)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-sm.stopCleanup:
|
||||
return
|
||||
case <-ticker.C:
|
||||
sm.mu.Lock()
|
||||
now := time.Now()
|
||||
for sessionID, session := range sm.sessions {
|
||||
if now.Sub(session.lastAccess) > sm.ttl {
|
||||
slog.Info("Cleaning up expired MCP session", "session", sessionID)
|
||||
session.Shutdown()
|
||||
delete(sm.sessions, sessionID)
|
||||
}
|
||||
}
|
||||
sm.mu.Unlock()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown closes all sessions and stops the cleanup goroutine
|
||||
func (sm *MCPSessionManager) Shutdown() {
|
||||
// Signal cleanup goroutine to stop
|
||||
close(sm.stopCleanup)
|
||||
|
||||
sm.mu.Lock()
|
||||
defer sm.mu.Unlock()
|
||||
|
||||
slog.Info("Shutting down MCP session manager", "sessions", len(sm.sessions))
|
||||
for sessionID, session := range sm.sessions {
|
||||
slog.Debug("Shutting down session", "session", sessionID)
|
||||
session.Shutdown()
|
||||
}
|
||||
sm.sessions = make(map[string]*MCPSession)
|
||||
}
|
||||
|
||||
// configsMatch checks if two sets of MCP configs are equivalent
|
||||
func configsMatch(a, b []api.MCPServerConfig) bool {
|
||||
if len(a) != len(b) {
|
||||
return false
|
||||
}
|
||||
return reflect.DeepEqual(a, b)
|
||||
}
|
||||
|
||||
// generateToolsSessionID creates a consistent session ID for model + tools path
|
||||
func generateToolsSessionID(model, toolsPath string) string {
|
||||
h := sha256.New()
|
||||
h.Write([]byte(model))
|
||||
h.Write([]byte(toolsPath))
|
||||
return "tools-" + hex.EncodeToString(h.Sum(nil))[:16]
|
||||
}
|
||||
|
||||
// GenerateSessionID creates a session ID based on the request
|
||||
func GenerateSessionID(req api.ChatRequest) string {
|
||||
// If explicit session ID provided
|
||||
if req.SessionID != "" {
|
||||
return req.SessionID
|
||||
}
|
||||
|
||||
// For interactive mode with tools path
|
||||
if req.ToolsPath != "" {
|
||||
return generateToolsSessionID(req.Model, req.ToolsPath)
|
||||
}
|
||||
|
||||
// Default: use request-specific ID (no persistence)
|
||||
h := sha256.New()
|
||||
h.Write([]byte(time.Now().Format(time.RFC3339Nano)))
|
||||
h.Write([]byte(req.Model))
|
||||
return "req-" + hex.EncodeToString(h.Sum(nil))[:16]
|
||||
}
|
||||
Loading…
Reference in New Issue