diff --git a/cmd/gateway_http_wiring.go b/cmd/gateway_http_wiring.go index 35eb9dfd..e55a955f 100644 --- a/cmd/gateway_http_wiring.go +++ b/cmd/gateway_http_wiring.go @@ -64,6 +64,9 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer( d.server.SetMCPHandler(h.mcp) } if h.mcpUserCreds != nil { + if mcpPool != nil { + h.mcpUserCreds.SetPoolEvictor(mcpPool) + } d.server.SetMCPUserCredentialsHandler(h.mcpUserCreds) } if h.channelInstances != nil { diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index 1a0c991e..39b84398 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -229,6 +229,22 @@ func (l *Loop) makeBuildFilteredTools(req *RunRequest) func(state *pipeline.RunS allMsgs := state.Messages.All() toolDefs, _, returnedMsgs := l.buildFilteredTools(req, state.Context.HadBootstrap, state.Iteration, maxIter, allMsgs) + if state.Iteration != maxIter && len(userTools) > 0 { + existing := make(map[string]struct{}, len(toolDefs)+len(userTools)) + for _, td := range toolDefs { + if td.Function != nil { + existing[td.Function.Name] = struct{}{} + } + } + for _, t := range userTools { + name := t.Name() + if _, ok := existing[name]; ok { + continue + } + toolDefs = append(toolDefs, tools.ToProviderDef(t)) + existing[name] = struct{}{} + } + } // buildFilteredTools returns the full messages slice; only messages appended // beyond the original length are injections (e.g. final-iteration hint). // Appending the entire slice would duplicate system+history into pending. diff --git a/internal/http/mcp_user_credentials.go b/internal/http/mcp_user_credentials.go index 71c23edf..c67e3f95 100644 --- a/internal/http/mcp_user_credentials.go +++ b/internal/http/mcp_user_credentials.go @@ -1,6 +1,7 @@ package http import ( + "context" "encoding/json" "log/slog" "net/http" @@ -15,6 +16,12 @@ import ( type MCPUserCredentialsHandler struct { store store.MCPServerStore tenantStore store.TenantStore + poolEvictor MCPUserPoolEvictor +} + +// MCPUserPoolEvictor evicts pooled user connections after credential rotation. +type MCPUserPoolEvictor interface { + EvictUser(tenantID uuid.UUID, serverName, userID string) } // NewMCPUserCredentialsHandler creates a handler for MCP user credential endpoints. @@ -22,6 +29,9 @@ func NewMCPUserCredentialsHandler(s store.MCPServerStore, ts store.TenantStore) return &MCPUserCredentialsHandler{store: s, tenantStore: ts} } +// SetPoolEvictor wires the MCP connection pool used by runtime execution. +func (h *MCPUserCredentialsHandler) SetPoolEvictor(e MCPUserPoolEvictor) { h.poolEvictor = e } + // RegisterRoutes registers MCP user credential routes. func (h *MCPUserCredentialsHandler) RegisterRoutes(mux *http.ServeMux) { mux.HandleFunc("PUT /v1/mcp/servers/{id}/user-credentials", h.auth(h.handleSet)) @@ -104,6 +114,7 @@ func (h *MCPUserCredentialsHandler) handleSet(w http.ResponseWriter, r *http.Req return } + h.evictUserConnection(r.Context(), serverID, userID) writeJSON(w, http.StatusOK, map[string]string{"status": "updated"}) } @@ -127,12 +138,19 @@ func (h *MCPUserCredentialsHandler) handleGet(w http.ResponseWriter, r *http.Req } creds, err := h.store.GetUserCredentials(r.Context(), serverID, userID) - if err != nil { - writeJSON(w, http.StatusOK, map[string]any{"has_credentials": false}) + if err != nil || creds == nil { + writeJSON(w, http.StatusOK, map[string]any{ + "user_id": userID, + "has_credentials": false, + "has_api_key": false, + "has_headers": false, + "has_env": false, + }) return } writeJSON(w, http.StatusOK, map[string]any{ + "user_id": userID, "has_credentials": true, "has_api_key": creds.APIKey != "", "has_headers": len(creds.Headers) > 0, @@ -164,9 +182,23 @@ func (h *MCPUserCredentialsHandler) handleDelete(w http.ResponseWriter, r *http. return } + h.evictUserConnection(r.Context(), serverID, userID) writeJSON(w, http.StatusOK, map[string]string{"status": "deleted"}) } +func (h *MCPUserCredentialsHandler) evictUserConnection(ctx context.Context, serverID uuid.UUID, userID string) { + if h.poolEvictor == nil || userID == "" { + return + } + srv, err := h.store.GetServer(ctx, serverID) + if err != nil || srv == nil { + slog.Warn("mcp.user_credentials.evict_lookup_failed", "server_id", serverID, "user", userID, "error", err) + return + } + tid := store.TenantIDFromContext(ctx) + h.poolEvictor.EvictUser(tid, srv.Name, userID) +} + // httpStatusText returns a short error message for common HTTP status codes. func httpStatusText(code int) string { switch code { diff --git a/internal/mcp/pool.go b/internal/mcp/pool.go index a6414bc9..614a77bf 100644 --- a/internal/mcp/pool.go +++ b/internal/mcp/pool.go @@ -51,9 +51,9 @@ type poolEntry struct { // Per-user connections are keyed by tenantID/serverName/user:userID. type Pool struct { mu sync.Mutex - servers map[string]*poolEntry // shared connections: tenantID/serverName - userServers map[string]*poolEntry // user connections: tenantID/serverName/user:userID - userSlots map[string]chan struct{} // per-server semaphores: tenantID/serverName → capacity MaxUserConns + servers map[string]*poolEntry // shared connections: tenantID/serverName + userServers map[string]*poolEntry // user connections: tenantID/serverName/user:userID + userSlots map[string]chan struct{} // per-server semaphores: tenantID/serverName → capacity MaxUserConns cfg PoolConfig slot chan struct{} // semaphore for MaxSize stopCh chan struct{} @@ -434,6 +434,36 @@ func (p *Pool) Evict(tenantID uuid.UUID, serverName string) { slog.Info("mcp.pool.evicted_on_rotation", "key", key) } +// EvictUser closes a specific user-scoped pooled connection by tenant, server, +// and user. Called when per-user credentials are rotated so the next request +// reconnects with fresh headers/env. +func (p *Pool) EvictUser(tenantID uuid.UUID, serverName, userID string) { + key := UserPoolKey(tenantID, serverName, userID) + slotKey := userSlotKey(tenantID, serverName) + p.mu.Lock() + defer p.mu.Unlock() + + entry, ok := p.userServers[key] + if !ok { + return + } + entry.state.connected.Store(false) + if entry.state.cancel != nil { + entry.state.cancel() + } + if client := entry.state.clientPtr.Load(); client != nil { + _ = client.Close() + } + delete(p.userServers, key) + if sem, ok := p.userSlots[slotKey]; ok { + select { + case <-sem: + default: + } + } + slog.Info("mcp.pool.user.evicted_on_rotation", "key", key) +} + // evictLoop runs periodically to close idle connections over MaxIdle count. func (p *Pool) evictLoop() { ticker := time.NewTicker(60 * time.Second) diff --git a/internal/store/mcp_store.go b/internal/store/mcp_store.go index 34ceace7..fc7334db 100644 --- a/internal/store/mcp_store.go +++ b/internal/store/mcp_store.go @@ -78,9 +78,9 @@ type MCPAccessInfo struct { // MCPUserCredentials holds per-user credential overrides for an MCP server. type MCPUserCredentials struct { - APIKey string `json:"api_key,omitempty" db:"-"` // decrypted - Headers map[string]string `json:"headers,omitempty" db:"-"` // decrypted - Env map[string]string `json:"env,omitempty" db:"-"` // decrypted + APIKey string `json:"api_key,omitempty" db:"-"` // decrypted + Headers map[string]string `json:"headers,omitempty" db:"-"` // decrypted + Env map[string]string `json:"env,omitempty" db:"-"` // decrypted } // MCPServerStore manages MCP server configs and access grants. diff --git a/ui/web/src/pages/mcp/mcp-user-credentials-dialog.tsx b/ui/web/src/pages/mcp/mcp-user-credentials-dialog.tsx index c40a46ad..c61a2f80 100644 --- a/ui/web/src/pages/mcp/mcp-user-credentials-dialog.tsx +++ b/ui/web/src/pages/mcp/mcp-user-credentials-dialog.tsx @@ -111,8 +111,10 @@ export function MCPUserCredentialsDialog({ if (Object.keys(data.env).length > 0) creds.env = data.env as Record; const targetUser = canManageUsers ? selectedUserId : undefined; await onSetCredentials(server.id, creds, targetUser); + const nextStatus = await onGetCredentials(server.id, targetUser); + setStatus(nextStatus); + reset({ apiKey: "", headers: {}, env: {} }); toast.success(i18next.t("mcp:userCredentials.saved")); - onOpenChange(false); } catch (err) { toast.error(i18next.t("mcp:userCredentials.saveFailed"), err instanceof Error ? err.message : ""); } finally { @@ -125,8 +127,9 @@ export function MCPUserCredentialsDialog({ try { const targetUser = canManageUsers ? selectedUserId : undefined; await onDeleteCredentials(server.id, targetUser); + const nextStatus = await onGetCredentials(server.id, targetUser); + setStatus(nextStatus); toast.success(i18next.t("mcp:userCredentials.deleted")); - onOpenChange(false); } catch (err) { toast.error(i18next.t("mcp:userCredentials.deleteFailed"), err instanceof Error ? err.message : ""); } finally { diff --git a/ui/web/src/types/mcp.ts b/ui/web/src/types/mcp.ts index 10f7a5ce..8e471bfd 100644 --- a/ui/web/src/types/mcp.ts +++ b/ui/web/src/types/mcp.ts @@ -62,6 +62,7 @@ export interface MCPAgentGrant { } export interface MCPUserCredentialStatus { + user_id?: string; has_credentials: boolean; has_api_key: boolean; has_headers: boolean;