fix(nomos): auto-reconnect stale MCP session; enlarge SSE scan buffer
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>
This commit is contained in:
@@ -5,6 +5,7 @@ import (
|
|||||||
"bytes"
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"net/http"
|
"net/http"
|
||||||
@@ -351,9 +352,51 @@ type mcpJSONRPCResponse struct {
|
|||||||
Error json.RawMessage `json:"error"`
|
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) {
|
func (c *mcpClient) doRequest(method string, params map[string]any) (*mcpJSONRPCResponse, error) {
|
||||||
c.mu.Lock()
|
c.mu.Lock()
|
||||||
defer c.mu.Unlock()
|
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++
|
c.nextID++
|
||||||
body, _ := json.Marshal(map[string]any{
|
body, _ := json.Marshal(map[string]any{
|
||||||
"jsonrpc": "2.0",
|
"jsonrpc": "2.0",
|
||||||
@@ -378,13 +421,21 @@ func (c *mcpClient) doRequest(method string, params map[string]any) (*mcpJSONRPC
|
|||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
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 := &mcpJSONRPCResponse{}
|
||||||
result.sessionID = resp.Header.Get("Mcp-Session-Id")
|
result.sessionID = resp.Header.Get("Mcp-Session-Id")
|
||||||
|
|
||||||
scanner := bufio.NewScanner(resp.Body)
|
scanner := bufio.NewScanner(resp.Body)
|
||||||
|
scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
|
||||||
|
gotData := false
|
||||||
for scanner.Scan() {
|
for scanner.Scan() {
|
||||||
line := scanner.Text()
|
line := scanner.Text()
|
||||||
if strings.HasPrefix(line, "data: ") {
|
if strings.HasPrefix(line, "data: ") {
|
||||||
|
gotData = true
|
||||||
data := line[6:]
|
data := line[6:]
|
||||||
if err := json.Unmarshal([]byte(data), result); err != nil {
|
if err := json.Unmarshal([]byte(data), result); err != nil {
|
||||||
return nil, fmt.Errorf("parse response: %w", err)
|
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))
|
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 != "" {
|
if result.sessionID != "" {
|
||||||
c.sessionID = result.sessionID
|
c.sessionID = result.sessionID
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user