fix: pass MCP servers through LoadSession so tools survive session resume (#450)

* fix: pass MCP servers through LoadSession so tools survive session resume

* fix: pass MCP servers to Copilot ResumeSession (address CR feedback)
This commit is contained in:
José Almeida
2026-03-22 23:19:37 +00:00
committed by GitHub
parent 564db29abe
commit 5b7e54f5da
13 changed files with 60 additions and 20 deletions
@@ -130,7 +130,7 @@ func (sm *SessionManager) createOrLoadSession(
zap.String("existing_session_id", existingSessionID),
zap.Bool("will_attempt_load", rt.SessionConfig.NativeSessionResume && existingSessionID != ""))
if rt.SessionConfig.NativeSessionResume && existingSessionID != "" {
sessionID, err := sm.loadSession(ctx, client, agentConfig, existingSessionID)
sessionID, err := sm.loadSession(ctx, client, agentConfig, existingSessionID, mcpServers)
if err != nil {
// If session/load fails because the agent doesn't support it, fall back to session/new.
// Check for: JSON-RPC -32601 (Method not found) or our own capability check string.
@@ -199,12 +199,13 @@ func (sm *SessionManager) loadSession(
client *agentctl.Client,
agentConfig agents.Agent,
sessionID string,
mcpServers []agentctltypes.McpServer,
) (string, error) {
sm.logger.Info("sending ACP session/load request",
zap.String("agent_type", agentConfig.ID()),
zap.String("session_id", sessionID))
if err := client.LoadSession(ctx, sessionID); err != nil {
if err := client.LoadSession(ctx, sessionID, mcpServers); err != nil {
sm.logger.Error("ACP session/load failed",
zap.String("agent_type", agentConfig.ID()),
zap.String("session_id", sessionID),
@@ -120,10 +120,13 @@ func (c *Client) ResetSession(ctx context.Context, cwd string, mcpServers []type
}
// LoadSession resumes an existing ACP session via the agent WebSocket stream.
func (c *Client) LoadSession(ctx context.Context, sessionID string) error {
// mcpServers are forwarded to the agentctl handler so agents that receive MCP configs
// via the protocol (e.g. Auggie) can reconnect to MCP servers on the new instance.
func (c *Client) LoadSession(ctx context.Context, sessionID string, mcpServers []types.McpServer) error {
payload := struct {
SessionID string `json:"session_id"`
}{SessionID: sessionID}
SessionID string `json:"session_id"`
McpServers []types.McpServer `json:"mcp_servers,omitempty"`
}{SessionID: sessionID, McpServers: mcpServers}
resp, err := c.sendStreamRequest(ctx, "agent.session.load", payload)
if err != nil {
@@ -417,7 +417,7 @@ func TestLoadSession_Success(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
err := c.LoadSession(ctx, "sess-456")
err := c.LoadSession(ctx, "sess-456", nil)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
@@ -434,7 +434,7 @@ func TestLoadSession_Error(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
err := c.LoadSession(ctx, "sess-456")
err := c.LoadSession(ctx, "sess-456", nil)
if err == nil {
t.Fatal("expected error")
}
@@ -163,7 +163,10 @@ type AgentAdapter interface {
NewSession(ctx context.Context, mcpServers []types.McpServer) (string, error)
// LoadSession resumes an existing session by ID.
LoadSession(ctx context.Context, sessionID string) error
// mcpServers contains the MCP servers to configure for the resumed session.
// Agents that receive MCP configs via the protocol (e.g. ACP with AssumeMcpSse)
// need these to reconnect to MCP servers on a new agentctl instance.
LoadSession(ctx context.Context, sessionID string, mcpServers []types.McpServer) error
// Prompt sends a prompt to the agent.
// The agent's responses are streamed via the Updates channel.
@@ -207,7 +207,7 @@ func (a *AmpAdapter) NewSession(ctx context.Context, _ []types.McpServer) (strin
// LoadSession resumes an existing Amp session (thread).
// For one-shot mode, the session ID is used to build the "threads continue" command.
func (a *AmpAdapter) LoadSession(ctx context.Context, sessionID string) error {
func (a *AmpAdapter) LoadSession(ctx context.Context, sessionID string, _ []types.McpServer) error {
a.mu.Lock()
a.sessionID = sessionID
a.hasAmpThreadID = sessionID != "" // LoadSession always receives a real Amp thread ID
@@ -349,12 +349,13 @@ func (a *CopilotAdapter) NewSession(ctx context.Context, mcpServers []types.McpS
}
// LoadSession resumes an existing Copilot session via the SDK.
func (a *CopilotAdapter) LoadSession(ctx context.Context, sessionID string) error {
func (a *CopilotAdapter) LoadSession(ctx context.Context, sessionID string, mcpServers []types.McpServer) error {
if a.client == nil {
return fmt.Errorf("adapter not initialized")
}
if err := a.client.ResumeSession(ctx, sessionID, nil); err != nil {
mcpConfig := mcpServersToCopilotConfig(mcpServers)
if err := a.client.ResumeSession(ctx, sessionID, mcpConfig); err != nil {
a.logger.Warn("failed to resume SDK session, continuing anyway",
zap.Error(err),
zap.String("session_id", sessionID))
@@ -271,7 +271,7 @@ func TestMockAgent_LoadSession(t *testing.T) {
}
// Load the same session (simulates resume after process restart)
if err := setup.adpt.LoadSession(ctx, sessionID); err != nil {
if err := setup.adpt.LoadSession(ctx, sessionID, nil); err != nil {
t.Fatalf("LoadSession failed: %v", err)
}
@@ -377,7 +377,9 @@ func mapToHTTPHeaders(headers map[string]string) []acp.HttpHeader {
// LoadSession resumes an existing session.
// Returns an error if the agent does not support session loading (LoadSession capability).
func (a *Adapter) LoadSession(ctx context.Context, sessionID string) error {
// mcpServers are passed to the agent so it can reconnect to MCP servers on the new
// agentctl instance (critical for agents that receive MCP configs via the protocol).
func (a *Adapter) LoadSession(ctx context.Context, sessionID string, mcpServers []types.McpServer) error {
a.mu.Lock()
conn := a.acpConn
supportsLoad := a.capabilities.LoadSession
@@ -397,6 +399,13 @@ func (a *Adapter) LoadSession(ctx context.Context, sessionID string) error {
ctx, span := shared.TraceProtocolRequest(ctx, shared.ProtocolACP, a.agentID, "session.load")
defer span.End()
// Filter MCP servers by agent capabilities (same logic as NewSession).
caps := a.capabilities.McpCapabilities
if a.cfg.AssumeMcpSse {
caps.Sse = true
}
filteredServers := filterMcpServersByCapabilities(mcpServers, caps, a.logger)
// Suppress history replay notifications during load.
// ACP session/load replays the entire conversation history asynchronously.
// We set a flag to suppress these notifications to avoid duplicating messages in the database.
@@ -408,7 +417,7 @@ func (a *Adapter) LoadSession(ctx context.Context, sessionID string) error {
resp, err := conn.LoadSession(ctx, acp.LoadSessionRequest{
SessionId: acp.SessionId(sessionID),
Cwd: a.cfg.WorkDir,
McpServers: []acp.McpServer{}, // MCPs configured via CLI args; empty satisfies the required field
McpServers: toACPMcpServers(filteredServers),
})
if err != nil {
@@ -71,7 +71,7 @@ func (a *Adapter) NewSession(ctx context.Context, _ []types.McpServer) (string,
// LoadSession resumes an existing Codex thread.
// It passes the same approval policy and sandbox settings as NewSession to ensure
// permission requirements are preserved across resume (see openai/codex#5322).
func (a *Adapter) LoadSession(ctx context.Context, sessionID string) error {
func (a *Adapter) LoadSession(ctx context.Context, sessionID string, _ []types.McpServer) error {
// Check client under lock, but don't hold lock during Call() to avoid deadlock
// with handleNotification which also needs the lock
a.mu.RLock()
@@ -367,7 +367,7 @@ func (a *Adapter) NewSession(ctx context.Context, _ []types.McpServer) (string,
}
// LoadSession resumes an existing OpenCode session by forking it.
func (a *Adapter) LoadSession(ctx context.Context, sessionID string) error {
func (a *Adapter) LoadSession(ctx context.Context, sessionID string, _ []types.McpServer) error {
// Clear any stale state from previous sessions
a.clearSessionState()
@@ -52,7 +52,7 @@ func (a *Adapter) NewSession(ctx context.Context, _ []types.McpServer) (string,
// LoadSession resumes an existing stream-json session.
// The session ID will be passed to the agent via --resume flag (handled by process manager).
func (a *Adapter) LoadSession(ctx context.Context, sessionID string) error {
func (a *Adapter) LoadSession(ctx context.Context, sessionID string, _ []types.McpServer) error {
a.mu.Lock()
a.sessionID = sessionID
pendingCommands := a.takePendingCommands()
@@ -53,7 +53,8 @@ type NewSessionResponse struct {
// LoadSessionRequest is a request to load an existing ACP session
type LoadSessionRequest struct {
SessionID string `json:"session_id"`
SessionID string `json:"session_id"`
McpServers []types.McpServer `json:"mcp_servers,omitempty"`
}
// LoadSessionResponse is the response to a load session call
@@ -385,7 +386,29 @@ func (s *Server) handleWSLoadSession(ctx context.Context, msg *ws.Message) *ws.M
s.logger.Debug("reset MCP backend client for loaded session")
}
if err := adapter.LoadSession(ctx, req.SessionID); err != nil {
// If MCP server is enabled, prepend the local kandev MCP server to the list.
// This mirrors handleWSNewSession to ensure MCP tools are available after resume.
mcpServers := req.McpServers
if s.mcpServer != nil {
localKandevMcp := types.McpServer{
Name: "kandev",
Type: "sse",
URL: fmt.Sprintf("http://localhost:%d/sse", s.cfg.Port),
}
filtered := make([]types.McpServer, 0, len(mcpServers)+1)
filtered = append(filtered, localKandevMcp)
for _, srv := range mcpServers {
if srv.Name != "kandev" {
filtered = append(filtered, srv)
}
}
mcpServers = filtered
s.logger.Debug("injected local kandev MCP server for loaded session",
zap.String("url", localKandevMcp.URL),
zap.Int("total_servers", len(mcpServers)))
}
if err := adapter.LoadSession(ctx, req.SessionID, mcpServers); err != nil {
s.logger.Error("load session failed", zap.Error(err))
resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, err.Error(), nil)
return resp
@@ -35,7 +35,7 @@ func (s *stubAdapter) GetAgentInfo() *adapter.AgentInfo { return nil }
func (s *stubAdapter) NewSession(_ context.Context, _ []types.McpServer) (string, error) {
return "", nil
}
func (s *stubAdapter) LoadSession(context.Context, string) error { return nil }
func (s *stubAdapter) LoadSession(context.Context, string, []types.McpServer) error { return nil }
func (s *stubAdapter) Prompt(context.Context, string, []v1.MessageAttachment) error { return nil }
func (s *stubAdapter) Cancel(context.Context) error { return nil }
func (s *stubAdapter) Updates() <-chan adapter.AgentEvent { return s.updatesCh }