From cff05c07682f01cd20e1c723a5fda62b9ecf5f2a Mon Sep 17 00:00:00 2001
From: dtoro <d.toro.v@pm.me>
Date: Wed, 8 Jul 2026 15:55:10 +0200
Subject: [PATCH] fix(nomos): auto-reconnect stale MCP session; enlarge SSE
 scan buffer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

After an api (MCP server) restart, nomos held a dead session id and every
tool call failed with "unexpected end of JSON input" until nomos was manually
restarted — which happens on every deploy. The MCP client now detects a
rejected session (4xx or empty body) and transparently re-initializes and
retries once. Also raise the SSE scanner buffer to 4MB so large tool results
don't exceed the 64KB default token limit.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 cmd/nomos/main.go | 57 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/cmd/nomos/main.go b/cmd/nomos/main.go
index 3267a2a..eeec8b3 100644
--- a/cmd/nomos/main.go
+++ b/cmd/nomos/main.go
@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log/slog"
 	"net/http"
@@ -351,9 +352,51 @@ type mcpJSONRPCResponse struct {
 	Error     json.RawMessage `json:"error"`
 }
 
+// errStaleSession signals that the MCP server rejected our session id (e.g.
+// after an api/MCP restart), so the client should re-initialize and retry.
+var errStaleSession = fmt.Errorf("mcp session stale")
+
+// doRequest serializes MCP calls and transparently re-initializes the session
+// if the server has forgotten it (common after an api redeploy), retrying the
+// original call once. Without this, an api restart permanently breaks nomos
+// until it is itself restarted.
 func (c *mcpClient) doRequest(method string, params map[string]any) (*mcpJSONRPCResponse, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+
+	resp, err := c.send(method, params)
+	if err != nil && method != "initialize" && errors.Is(err, errStaleSession) {
+		slog.Warn("nomos: mcp session stale, reconnecting")
+		if rerr := c.reconnectLocked(); rerr != nil {
+			return nil, fmt.Errorf("mcp reconnect: %w (original: %v)", rerr, err)
+		}
+		return c.send(method, params)
+	}
+	return resp, err
+}
+
+// reconnectLocked re-initializes the MCP session. The caller must hold c.mu.
+func (c *mcpClient) reconnectLocked() error {
+	c.sessionID = ""
+	resp, err := c.send("initialize", map[string]any{
+		"protocolVersion": "2024-11-05",
+		"capabilities":    map[string]any{},
+		"clientInfo":      map[string]any{"name": "nomos", "version": "2.0"},
+	})
+	if err != nil {
+		return err
+	}
+	if resp.sessionID == "" {
+		return fmt.Errorf("no session ID on re-initialize")
+	}
+	c.sessionID = resp.sessionID
+	_, _ = c.send("notifications/initialized", map[string]any{})
+	slog.Info("nomos: mcp reconnected", "session", c.sessionID[:16]+"...")
+	return nil
+}
+
+// send performs one MCP round-trip. It does not lock; callers hold c.mu.
+func (c *mcpClient) send(method string, params map[string]any) (*mcpJSONRPCResponse, error) {
 	c.nextID++
 	body, _ := json.Marshal(map[string]any{
 		"jsonrpc": "2.0",
@@ -378,13 +421,21 @@ func (c *mcpClient) doRequest(method string, params map[string]any) (*mcpJSONRPC
 	}
 	defer resp.Body.Close()
 
+	// A rejected/unknown session comes back as 4xx (commonly 400/404).
+	if resp.StatusCode == http.StatusNotFound || resp.StatusCode == http.StatusBadRequest {
+		return nil, errStaleSession
+	}
+
 	result := &mcpJSONRPCResponse{}
 	result.sessionID = resp.Header.Get("Mcp-Session-Id")
 
 	scanner := bufio.NewScanner(resp.Body)
+	scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
+	gotData := false
 	for scanner.Scan() {
 		line := scanner.Text()
 		if strings.HasPrefix(line, "data: ") {
+			gotData = true
 			data := line[6:]
 			if err := json.Unmarshal([]byte(data), result); err != nil {
 				return nil, fmt.Errorf("parse response: %w", err)
@@ -396,6 +447,12 @@ func (c *mcpClient) doRequest(method string, params map[string]any) (*mcpJSONRPC
 		return nil, fmt.Errorf("rpc error: %s", string(result.Error))
 	}
 
+	// Empty body with no result and a session set: the server likely dropped
+	// our session. Notifications legitimately return no data, so exempt them.
+	if !gotData && result.Result == nil && method != "notifications/initialized" {
+		return nil, errStaleSession
+	}
+
 	if result.sessionID != "" {
 		c.sessionID = result.sessionID
 	}