From 5b7e54f5da7a369077243fddfae7d7b4e7117f8d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Almeida?= <57680069+zeval@users.noreply.github.com> Date: Sun, 22 Mar 2026 23:19:37 +0000 Subject: [PATCH] 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) --- .../internal/agent/lifecycle/session.go | 5 ++-- .../backend/internal/agentctl/client/agent.go | 9 ++++--- .../internal/agentctl/client/agent_test.go | 4 +-- .../agentctl/server/adapter/adapter.go | 5 +++- .../agentctl/server/adapter/amp_adapter.go | 2 +- .../server/adapter/copilot_adapter.go | 5 ++-- .../server/adapter/e2e/mock_agent_test.go | 2 +- .../server/adapter/transport/acp/adapter.go | 13 +++++++-- .../adapter/transport/codex/codex_session.go | 2 +- .../adapter/transport/opencode/adapter.go | 2 +- .../streamjson/streamjson_session.go | 2 +- .../internal/agentctl/server/api/agent.go | 27 +++++++++++++++++-- .../agentctl/server/process/manager_test.go | 2 +- 13 files changed, 60 insertions(+), 20 deletions(-) diff --git a/apps/backend/internal/agent/lifecycle/session.go b/apps/backend/internal/agent/lifecycle/session.go index 11727efa97..9c662088ff 100644 --- a/apps/backend/internal/agent/lifecycle/session.go +++ b/apps/backend/internal/agent/lifecycle/session.go @@ -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), diff --git a/apps/backend/internal/agentctl/client/agent.go b/apps/backend/internal/agentctl/client/agent.go index fb81533823..b516dcdc79 100644 --- a/apps/backend/internal/agentctl/client/agent.go +++ b/apps/backend/internal/agentctl/client/agent.go @@ -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 { diff --git a/apps/backend/internal/agentctl/client/agent_test.go b/apps/backend/internal/agentctl/client/agent_test.go index 6cb515dc3f..e51506976c 100644 --- a/apps/backend/internal/agentctl/client/agent_test.go +++ b/apps/backend/internal/agentctl/client/agent_test.go @@ -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") } diff --git a/apps/backend/internal/agentctl/server/adapter/adapter.go b/apps/backend/internal/agentctl/server/adapter/adapter.go index 0b34cdf466..ed9ae2bdc4 100644 --- a/apps/backend/internal/agentctl/server/adapter/adapter.go +++ b/apps/backend/internal/agentctl/server/adapter/adapter.go @@ -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. diff --git a/apps/backend/internal/agentctl/server/adapter/amp_adapter.go b/apps/backend/internal/agentctl/server/adapter/amp_adapter.go index 8faa7884b4..24eac2150f 100644 --- a/apps/backend/internal/agentctl/server/adapter/amp_adapter.go +++ b/apps/backend/internal/agentctl/server/adapter/amp_adapter.go @@ -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 diff --git a/apps/backend/internal/agentctl/server/adapter/copilot_adapter.go b/apps/backend/internal/agentctl/server/adapter/copilot_adapter.go index 1a25b8c9ca..8998ab18ff 100644 --- a/apps/backend/internal/agentctl/server/adapter/copilot_adapter.go +++ b/apps/backend/internal/agentctl/server/adapter/copilot_adapter.go @@ -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)) diff --git a/apps/backend/internal/agentctl/server/adapter/e2e/mock_agent_test.go b/apps/backend/internal/agentctl/server/adapter/e2e/mock_agent_test.go index 9a65fb83d3..c3ccf4433e 100644 --- a/apps/backend/internal/agentctl/server/adapter/e2e/mock_agent_test.go +++ b/apps/backend/internal/agentctl/server/adapter/e2e/mock_agent_test.go @@ -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) } diff --git a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go index f9c59dd5b4..76e2ad26f9 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/acp/adapter.go @@ -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 { diff --git a/apps/backend/internal/agentctl/server/adapter/transport/codex/codex_session.go b/apps/backend/internal/agentctl/server/adapter/transport/codex/codex_session.go index 2a0aa46817..b25e82b707 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/codex/codex_session.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/codex/codex_session.go @@ -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() diff --git a/apps/backend/internal/agentctl/server/adapter/transport/opencode/adapter.go b/apps/backend/internal/agentctl/server/adapter/transport/opencode/adapter.go index 375b92d649..12bcc07a91 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/opencode/adapter.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/opencode/adapter.go @@ -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() diff --git a/apps/backend/internal/agentctl/server/adapter/transport/streamjson/streamjson_session.go b/apps/backend/internal/agentctl/server/adapter/transport/streamjson/streamjson_session.go index 265437d3cd..b521946149 100644 --- a/apps/backend/internal/agentctl/server/adapter/transport/streamjson/streamjson_session.go +++ b/apps/backend/internal/agentctl/server/adapter/transport/streamjson/streamjson_session.go @@ -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() diff --git a/apps/backend/internal/agentctl/server/api/agent.go b/apps/backend/internal/agentctl/server/api/agent.go index 24f42e3a7b..0caba608c7 100644 --- a/apps/backend/internal/agentctl/server/api/agent.go +++ b/apps/backend/internal/agentctl/server/api/agent.go @@ -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 diff --git a/apps/backend/internal/agentctl/server/process/manager_test.go b/apps/backend/internal/agentctl/server/process/manager_test.go index e8cf3fd56a..3ce1e83c71 100644 --- a/apps/backend/internal/agentctl/server/process/manager_test.go +++ b/apps/backend/internal/agentctl/server/process/manager_test.go @@ -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 }