fix(mcp): refresh per-user credentials without restart (#1375)

Co-authored-by: Trí Đào Nguyễn Minh <tri.dao@buymed.com>

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 {

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.

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 {

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)

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.

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<string, string>;
       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 {

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;