From 9e4e40c0fa43b5f10d3489bfb7c24313a9250271 Mon Sep 17 00:00:00 2001 From: Vitalii Pedchenko Date: Wed, 8 Jul 2026 23:49:34 +0300 Subject: [PATCH] feat(mcp): session-scoped notification delivery and wire-format negotiation --- cmd/gortex/daemon_mcp.go | 36 +++++++- internal/daemon/mcp_roundtrip_test.go | 72 +++++++++++++++ internal/daemon/server.go | 31 ++++++- internal/mcp/conn_session.go | 114 +++++++++++++++++++++++ internal/mcp/conn_session_test.go | 115 ++++++++++++++++++++++++ internal/mcp/format_negotiation_test.go | 52 ++++++++++- internal/mcp/server.go | 83 +++++++++++++++-- 7 files changed, 494 insertions(+), 9 deletions(-) create mode 100644 internal/mcp/conn_session.go create mode 100644 internal/mcp/conn_session_test.go diff --git a/cmd/gortex/daemon_mcp.go b/cmd/gortex/daemon_mcp.go index a5bc8333..e5e8e7fd 100644 --- a/cmd/gortex/daemon_mcp.go +++ b/cmd/gortex/daemon_mcp.go @@ -97,6 +97,11 @@ func (d *mcpDispatcher) Dispatch(ctx context.Context, sess *daemon.Session, fram // every query runs against the whole multi-workspace graph and a // session in workspace A can see workspace B's nodes. ctx = gortexmcp.WithSessionCWD(ctx, sess.CWD) + // Run session-aware: attaching the connected client session (wired + // by SessionStarted) lets mcp-go's initialize handler mark it + // initialized — the gate SendNotificationToAllClients applies + // before delivering tools/list_changed and other pushes. + ctx = d.srv.WithClientSession(ctx, sess.ID) // Relay the client-forwarded tool-surface preference (GORTEX_TOOLS / // --tools of the proxy) so the MCP server resolves THIS session's @@ -224,11 +229,29 @@ func rewriteUntrackedResponse(method string, out []byte, cwd string) []byte { return out } +// SessionStarted implements daemon.SessionStartedHook: register the +// connection as a live mcp-go client session so server-initiated +// notifications (tools/list_changed after a tools_search promotion, +// graph_invalidated, diagnostics pushes, ...) reach this client over +// the socket. Best-effort — on failure the session still works +// request/response, it just misses pushes. +func (d *mcpDispatcher) SessionStarted(sess *daemon.Session, write func([]byte) error) { + if d.srv == nil || sess == nil { + return + } + if err := d.srv.ConnectSession(sess.ID, write); err != nil { + d.logger.Debug("daemon: notification session unavailable", + zap.String("session_id", sess.ID), zap.Error(err)) + } +} + // SessionEnded implements daemon.SessionEndedHook. When a proxy -// disconnects, drop its entry from the MCP server's session map so idle -// per-session state doesn't accumulate for the daemon's lifetime. +// disconnects, unregister its notification session and drop its entry +// from the MCP server's session map so idle per-session state doesn't +// accumulate for the daemon's lifetime. func (d *mcpDispatcher) SessionEnded(sess *daemon.Session) { if d.srv != nil && sess != nil { + d.srv.DisconnectSession(sess.ID) d.srv.ReleaseSession(sess.ID) } } @@ -389,6 +412,9 @@ func (d *mcpDispatcher) maybeSnoopInitialize(sess *daemon.Session, frame []byte) Name string `json:"name"` Version string `json:"version"` } `json:"clientInfo"` + Capabilities struct { + WireFormats []string `json:"gortex/wire"` + } `json:"capabilities"` } `json:"params"` } if err := json.Unmarshal(frame, &peek); err != nil { @@ -397,6 +423,12 @@ func (d *mcpDispatcher) maybeSnoopInitialize(sess *daemon.Session, frame []byte) if peek.Method != "initialize" { return } + // A gortex/wire capability self-declares the client's compact-format + // decoders; session format resolution prefers it over the built-in + // client-name allowlist, so new clients need no server-side entry. + if formats := peek.Params.Capabilities.WireFormats; len(formats) > 0 && d.srv != nil { + d.srv.NoteSessionWireFormats(sess.ID, formats) + } if peek.Params.ClientInfo.Name == "" && peek.Params.ClientInfo.Version == "" { return } diff --git a/internal/daemon/mcp_roundtrip_test.go b/internal/daemon/mcp_roundtrip_test.go index e02ffc3c..3cab4dc6 100644 --- a/internal/daemon/mcp_roundtrip_test.go +++ b/internal/daemon/mcp_roundtrip_test.go @@ -186,3 +186,75 @@ func TestDaemon_SessionEndedHook_FiresOnDisconnect(t *testing.T) { t.Fatal("SessionEnded hook never fired after client close") } } + +// pushDispatcher satisfies MCPDispatcher and SessionStartedHook so we +// can prove a dispatcher-initiated frame (an MCP notification) reaches +// the client over the socket, interleaved with the reply path. +type pushDispatcher struct { + echoDispatcher + started chan func([]byte) error +} + +func (p *pushDispatcher) SessionStarted(_ *Session, write func([]byte) error) { + p.started <- write +} + +// TestDaemon_SessionStartedHook_PushesFrames pins the server-initiated +// notification path: the daemon hands dispatchers a write callback on +// connect, and frames written through it arrive at the client without a +// pending request — the transport leg that lets tools/list_changed +// (and every other MCP push) reach unix-socket proxy clients. +func TestDaemon_SessionStartedHook_PushesFrames(t *testing.T) { + dir, err := os.MkdirTemp("/tmp", "gx") + require.NoError(t, err) + defer func() { _ = os.RemoveAll(dir) }() + socket := filepath.Join(dir, "s") + t.Setenv("GORTEX_DAEMON_SOCKET", socket) + t.Setenv("GORTEX_DAEMON_PIDFILE", filepath.Join(dir, "p")) + + disp := &pushDispatcher{ + echoDispatcher: echoDispatcher{seen: make(chan echoFrame, 4)}, + started: make(chan func([]byte) error, 1), + } + srv := New(socket, "test", zap.NewNop()) + srv.MCPDispatcher = disp + srv.Controller = &fakeController{} + require.NoError(t, srv.Listen()) + go func() { _ = srv.Serve() }() + defer func() { _ = srv.Shutdown() }() + + require.Eventually(t, func() bool { return IsRunningAt(socket) }, + 2*time.Second, 10*time.Millisecond) + + client, err := DialTo(socket, Handshake{Mode: ModeMCP, CWD: "/tmp/fake-repo"}) + require.NoError(t, err) + defer client.Close() + + var write func([]byte) error + select { + case write = <-disp.started: + case <-time.After(2 * time.Second): + t.Fatal("SessionStarted hook never fired after connect") + } + + // Push an unprompted notification; the client must read it even + // though it never sent a request. + notif := `{"jsonrpc":"2.0","method":"notifications/tools/list_changed"}` + require.NoError(t, write([]byte(notif))) + + frame, err := client.ReadMCPFrame() + require.NoError(t, err) + var got map[string]any + require.NoError(t, json.Unmarshal(frame, &got)) + assert.Equal(t, "notifications/tools/list_changed", got["method"]) + + // The request/reply path still works alongside the push path. + rpc := `{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}` + require.NoError(t, client.WriteMCPFrame([]byte(rpc))) + reply, err := client.ReadMCPFrame() + require.NoError(t, err) + var replyObj map[string]any + require.NoError(t, json.Unmarshal(reply, &replyObj)) + _, hasResult := replyObj["result"] + assert.True(t, hasResult, "reply must carry a result: %s", reply) +} diff --git a/internal/daemon/server.go b/internal/daemon/server.go index ce608398..8e3b1775 100644 --- a/internal/daemon/server.go +++ b/internal/daemon/server.go @@ -94,6 +94,20 @@ type SessionEndedHook interface { SessionEnded(sess *Session) } +// SessionStartedHook is an optional extension that MCPDispatcher +// implementations can satisfy to receive a connect callback for each +// ModeMCP session, fired after the handshake ack and before the first +// frame is dispatched. write delivers one server-initiated JSON-RPC +// frame (without trailing newline) to the client; it is safe for +// concurrent use with the reply path — the daemon serialises all +// writes to the connection — and returns an error once the connection +// is gone. This is how server-initiated MCP notifications +// (tools/list_changed, graph_invalidated, ...) reach socket clients; +// the matching SessionEnded fires on teardown. +type SessionStartedHook interface { + SessionStarted(sess *Session, write func([]byte) error) +} + // Controller implements the daemon's control surface. Separated from // MCPDispatcher so the two can evolve independently and so control-only // tests don't need a full MCP stack. @@ -477,6 +491,21 @@ func (s *Server) serveMCP(conn net.Conn, reader *bufio.Reader, sess *Session) { }) return } + + // Server-initiated frames (notifications pushed by the dispatcher) + // share the connection with request replies; serialise every write + // under one lock so frames never interleave mid-line. + var writeMu sync.Mutex + writeFrame := func(frame []byte) error { + writeMu.Lock() + defer writeMu.Unlock() + _, err := conn.Write(append(frame, '\n')) + return err + } + if hook, ok := s.MCPDispatcher.(SessionStartedHook); ok && hook != nil { + hook.SessionStarted(sess, writeFrame) + } + for { line, err := reader.ReadBytes('\n') if err != nil { @@ -506,7 +535,7 @@ func (s *Server) serveMCP(conn net.Conn, reader *bufio.Reader, sess *Session) { continue } // The dispatcher returns a full JSON-RPC frame; re-append newline. - if _, werr := conn.Write(append(reply, '\n')); werr != nil { + if werr := writeFrame(reply); werr != nil { s.Logger.Debug("daemon: mcp write failed", zap.String("session_id", sess.ID), zap.Error(werr)) return diff --git a/internal/mcp/conn_session.go b/internal/mcp/conn_session.go new file mode 100644 index 00000000..502e9657 --- /dev/null +++ b/internal/mcp/conn_session.go @@ -0,0 +1,114 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "sync/atomic" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" +) + +// connClientSession adapts one daemon-socket connection to mcp-go's +// server.ClientSession so server-initiated notifications +// (notifications/tools/list_changed after a tools_search promotion, +// graph_invalidated, diagnostics pushes, ...) reach stdio-proxy +// clients. Without this registration, HandleMessage-dispatched +// sessions are invisible to SendNotificationToAllClients: the +// broadcast iterates registered sessions only, so unix-socket clients +// silently miss every push — the streamable HTTP transport and the +// embedded stdio server were the only paths that ever delivered one. +// +// The pump goroutine serialises each notification and hands it to the +// transport-owned write callback; the daemon interleaves those frames +// with request replies under its own connection write lock. +type connClientSession struct { + id string + notifCh chan mcp.JSONRPCNotification + done chan struct{} + initialized atomic.Bool +} + +func (c *connClientSession) SessionID() string { return c.id } +func (c *connClientSession) NotificationChannel() chan<- mcp.JSONRPCNotification { + return c.notifCh +} +func (c *connClientSession) Initialize() { c.initialized.Store(true) } +func (c *connClientSession) Initialized() bool { return c.initialized.Load() } + +// connNotifBuffer sizes the per-session notification channel. mcp-go +// delivers non-blocking and drops on overflow, so the buffer only has +// to absorb short bursts (a promotion sweep fires a single frame). +const connNotifBuffer = 16 + +// ConnectSession registers a daemon-socket connection as a live MCP +// client session on the underlying mcp-go server. write must deliver +// one complete JSON-RPC frame to the client (no trailing newline) and +// be safe to call from the pump goroutine concurrently with the +// request/reply path; once it errors the pump stops (the daemon's +// teardown calls DisconnectSession shortly after). +// +// Pair every ConnectSession with a DisconnectSession — the mcp-go +// session registry and the pump goroutine live until then. +func (s *Server) ConnectSession(sessionID string, write func([]byte) error) error { + if s.mcpServer == nil { + return fmt.Errorf("connect session %s: mcp server not initialised", sessionID) + } + cs := &connClientSession{ + id: sessionID, + notifCh: make(chan mcp.JSONRPCNotification, connNotifBuffer), + done: make(chan struct{}), + } + if _, loaded := s.connSessions.LoadOrStore(sessionID, cs); loaded { + return fmt.Errorf("connect session %s: already connected", sessionID) + } + if err := s.mcpServer.RegisterSession(context.Background(), cs); err != nil { + s.connSessions.Delete(sessionID) + return fmt.Errorf("connect session %s: %w", sessionID, err) + } + go func() { + for { + select { + case n := <-cs.notifCh: + frame, err := json.Marshal(n) + if err != nil { + continue + } + if write(frame) != nil { + return + } + case <-cs.done: + return + } + } + }() + return nil +} + +// DisconnectSession unregisters a session connected via ConnectSession +// and stops its notification pump. Unknown ids are a no-op, so callers +// can invoke it unconditionally on connection teardown. +func (s *Server) DisconnectSession(sessionID string) { + v, ok := s.connSessions.LoadAndDelete(sessionID) + if !ok { + return + } + if s.mcpServer != nil { + s.mcpServer.UnregisterSession(context.Background(), sessionID) + } + close(v.(*connClientSession).done) +} + +// WithClientSession attaches the connected session for sessionID to +// ctx so HandleMessage runs session-aware. This is what lets mcp-go's +// initialize handler mark the session initialized — the gate +// sendNotificationToAllClients applies before delivering a broadcast — +// and what enables session-targeted APIs generally. A sessionID +// without a connected session returns ctx unchanged. +func (s *Server) WithClientSession(ctx context.Context, sessionID string) context.Context { + if v, ok := s.connSessions.Load(sessionID); ok { + return s.mcpServer.WithContext(ctx, v.(server.ClientSession)) + } + return ctx +} diff --git a/internal/mcp/conn_session_test.go b/internal/mcp/conn_session_test.go new file mode 100644 index 00000000..293798ff --- /dev/null +++ b/internal/mcp/conn_session_test.go @@ -0,0 +1,115 @@ +package mcp + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/mark3labs/mcp-go/mcp" + "go.uber.org/zap" + + "github.com/zzet/gortex/internal/graph" + "github.com/zzet/gortex/internal/query" +) + +func newConnSessionTestServer(t *testing.T) *Server { + t.Helper() + g := graph.New() + eng := query.NewEngine(g) + return NewServer(eng, g, nil, nil, zap.NewNop(), nil) +} + +// initializeSession runs the MCP initialize handshake through the +// session-aware dispatch path — the same shape the daemon dispatcher +// uses — so mcp-go marks the connected session initialized (the gate +// broadcasts apply before delivering). +func initializeSession(t *testing.T, s *Server, sessionID string) { + t.Helper() + ctx := s.WithClientSession(context.Background(), sessionID) + init := json.RawMessage(`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"0"}}}`) + if reply := s.MCPServer().HandleMessage(ctx, init); reply == nil { + t.Fatal("initialize returned no reply") + } +} + +// TestConnectSessionDeliversToolsListChanged guards the daemon-socket +// notification path end to end: a session connected via ConnectSession +// and initialized through the session-aware dispatch must receive +// notifications/tools/list_changed when a tool is added (the exact +// broadcast a tools_search promotion fires). Before ConnectSession +// existed, HandleMessage-dispatched socket sessions were invisible to +// SendNotificationToAllClients and every push was silently dropped. +func TestConnectSessionDeliversToolsListChanged(t *testing.T) { + s := newConnSessionTestServer(t) + frames := make(chan []byte, 4) + if err := s.ConnectSession("sess-notify", func(b []byte) error { + frames <- append([]byte(nil), b...) + return nil + }); err != nil { + t.Fatalf("ConnectSession: %v", err) + } + defer s.DisconnectSession("sess-notify") + + initializeSession(t, s, "sess-notify") + + s.MCPServer().AddTool( + mcp.NewTool("late_promoted_tool"), + func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return &mcp.CallToolResult{}, nil + }, + ) + + select { + case frame := <-frames: + var n struct { + Method string `json:"method"` + } + if err := json.Unmarshal(frame, &n); err != nil { + t.Fatalf("notification frame is not JSON: %v (%s)", err, frame) + } + if n.Method != "notifications/tools/list_changed" { + t.Fatalf("method = %q, want notifications/tools/list_changed", n.Method) + } + case <-time.After(2 * time.Second): + t.Fatal("no notification delivered to the connected session") + } +} + +// TestConnectSessionLifecycle covers the bookkeeping edges: a duplicate +// connect is refused, an uninitialized session receives no broadcast, +// and after DisconnectSession nothing is delivered. +func TestConnectSessionLifecycle(t *testing.T) { + s := newConnSessionTestServer(t) + frames := make(chan []byte, 4) + write := func(b []byte) error { + frames <- append([]byte(nil), b...) + return nil + } + if err := s.ConnectSession("sess-life", write); err != nil { + t.Fatalf("ConnectSession: %v", err) + } + if err := s.ConnectSession("sess-life", write); err == nil { + t.Fatal("duplicate ConnectSession succeeded, want error") + } + + // Not initialized yet — a broadcast must not reach the session. + s.MCPServer().SendNotificationToAllClients("notifications/tools/list_changed", nil) + select { + case f := <-frames: + t.Fatalf("uninitialized session received %s", f) + case <-time.After(50 * time.Millisecond): + } + + initializeSession(t, s, "sess-life") + s.DisconnectSession("sess-life") + // Unknown / already-disconnected ids are a no-op. + s.DisconnectSession("sess-life") + + s.MCPServer().SendNotificationToAllClients("notifications/tools/list_changed", nil) + select { + case f := <-frames: + t.Fatalf("disconnected session received %s", f) + case <-time.After(50 * time.Millisecond): + } +} diff --git a/internal/mcp/format_negotiation_test.go b/internal/mcp/format_negotiation_test.go index 1c166c24..93a989b4 100644 --- a/internal/mcp/format_negotiation_test.go +++ b/internal/mcp/format_negotiation_test.go @@ -33,8 +33,11 @@ func TestDefaultFormatForClient(t *testing.T) { {"codex", "gcx"}, {"omp-coding-agent", "gcx"}, - // Unknown / unset → JSON fallback. + // Unknown / unset → JSON fallback. "pi" is deliberately NOT in + // the allowlist: its bridge declares the gortex/wire capability + // instead (TestResolveSessionFormat_WireCapability). {"", ""}, + {"pi", ""}, {"some-other-client", ""}, {"unknown", ""}, } @@ -61,6 +64,53 @@ func TestResolveSessionFormat_KnownClient(t *testing.T) { assert.Equal(t, "gcx", srv.resolveSessionFormat(ctx)) } +// TestFormatFromWireCapability pins the capability → format mapping: the +// first declared format the server can emit wins, unknown names are +// skipped, and an empty declaration yields the JSON fallback. +func TestFormatFromWireCapability(t *testing.T) { + cases := []struct { + formats []string + want string + }{ + {[]string{"gcx"}, "gcx"}, + {[]string{"toon"}, "toon"}, + {[]string{" GCX "}, "gcx"}, // trimmed + case-insensitive + {[]string{"gcx2-draft", "gcx"}, "gcx"}, // unknown skipped, fallback honoured + {[]string{"toon", "gcx"}, "toon"}, // client preference order wins + {[]string{"bogus"}, ""}, // nothing usable → JSON + {nil, ""}, // no declaration → JSON + } + for _, tc := range cases { + assert.Equal(t, tc.want, formatFromWireCapability(tc.formats), + "formatFromWireCapability(%q)", tc.formats) + } +} + +// TestResolveSessionFormat_WireCapability verifies the full per-session +// capability path: NoteSessionWireFormats stores the declaration (the +// daemon dispatcher snoops it from initialize.capabilities +// ["gortex/wire"]), and resolveSessionFormat prefers it over the +// client-name allowlist — this is how the Pi bridge gets gcx without an +// allowlist entry. +func TestResolveSessionFormat_WireCapability(t *testing.T) { + srv, _ := setupTestServer(t) + srv.NoteSessionClient("session_pi", "pi", "1.0.0") // NOT in the allowlist + srv.NoteSessionWireFormats("session_pi", []string{"gcx"}) + ctx := WithSessionID(context.Background(), "session_pi") + assert.Equal(t, "gcx", srv.resolveSessionFormat(ctx)) +} + +// TestNoteSessionWireFormats_NilSafe mirrors TestNoteSessionClient_NilSafe +// for the capability path. +func TestNoteSessionWireFormats_NilSafe(t *testing.T) { + var srv *Server + srv.NoteSessionWireFormats("sess", []string{"gcx"}) + + srv2, _ := setupTestServer(t) + srv2.NoteSessionWireFormats("", []string{"gcx"}) // empty session id → no-op + srv2.NoteSessionWireFormats("sess", nil) // empty declaration → no-op +} + func TestResolveSessionFormat_UnknownClient(t *testing.T) { srv, _ := setupTestServer(t) srv.NoteSessionClient("session_X", "some-bespoke-client", "0.1") diff --git a/internal/mcp/server.go b/internal/mcp/server.go index 40007e4d..4c98f733 100644 --- a/internal/mcp/server.go +++ b/internal/mcp/server.go @@ -87,7 +87,12 @@ func (sh *symbolHistory) All() map[string][]SymbolModification { // Server wraps the MCP server with Gortex-specific tools. type Server struct { - mcpServer *server.MCPServer + mcpServer *server.MCPServer + // connSessions tracks daemon-socket connections registered as live + // mcp-go client sessions so server-initiated notifications reach + // them (ConnectSession / DisconnectSession / WithClientSession in + // conn_session.go). Keyed by daemon session id. + connSessions sync.Map engine *query.Engine graph graph.Store indexer *indexer.Indexer @@ -241,9 +246,9 @@ type Server struct { // instead of re-spawning the language server. An entry is recorded after // the first attempt (success or empty) to bound per-query LSP work. refsConfirmed sync.Map // symbolID → struct{} - feedback *feedbackManager - notes *notesManager - memories *memoryManager + feedback *feedbackManager + notes *notesManager + memories *memoryManager // promotedTools is the per-workspace learned tool surface: deferred // tools promoted into the eager surface after use, persisted so the // learning survives daemon restarts (with demotion after disuse). Nil @@ -524,6 +529,13 @@ type sessionState struct { // toolSpec / clientName is (re)recorded. resolvedToolPolicy *toolPolicy toolPolicyResolved bool + // wireFormats is the client's self-declared gortex/wire capability + // (initialize.capabilities["gortex/wire"]) — the + // compact wire formats its decoder supports, in preference order. + // When set it wins over the clientName allowlist in + // resolveSessionFormat, so new clients don't need a server-side + // allowlist entry. + wireFormats []string // lastSearch captures the most recent search_symbols call so that a // subsequent get_symbol_source / get_editing_context on one of its // results can be attributed back to the query — this is the raw input @@ -795,7 +807,6 @@ func (ss *sessionState) recordToolPolicy(spec, mode string) { ss.toolPolicyResolved = false } - // snapshotClientName returns the captured client name under the // session lock. Returns empty when the `initialize` frame hasn't // arrived yet (very early tool calls — rare but possible during @@ -806,6 +817,26 @@ func (ss *sessionState) snapshotClientName() string { return ss.clientName } +// recordWireFormats captures the client's gortex/wire capability from +// the `initialize` frame. Empty input is ignored so a re-init without +// the capability can't clobber a prior declaration. +func (ss *sessionState) recordWireFormats(formats []string) { + if len(formats) == 0 { + return + } + ss.mu.Lock() + defer ss.mu.Unlock() + ss.wireFormats = append([]string(nil), formats...) +} + +// snapshotWireFormats returns the declared gortex/wire capability under +// the session lock; nil when the client declared none. +func (ss *sessionState) snapshotWireFormats() []string { + ss.mu.Lock() + defer ss.mu.Unlock() + return ss.wireFormats +} + // NoteSessionClient is called by the daemon dispatcher after it // snoops the MCP `initialize.clientInfo.name` value, so the per- // session sessionState can default tool wire-format based on the @@ -862,6 +893,24 @@ func (s *Server) NoteSessionToolPolicy(sessionID, spec, mode string) { s.sessions.get(sessionID).session.recordToolPolicy(spec, mode) } +// NoteSessionWireFormats is called by the daemon dispatcher after it +// snoops the MCP initialize capabilities["gortex/wire"] +// declaration, so format resolution can honour a client's self-declared +// decoder capability without a client-name allowlist entry. +func (s *Server) NoteSessionWireFormats(sessionID string, formats []string) { + if s == nil || sessionID == "" || len(formats) == 0 { + return + } + if s.sessions == nil { + // Embedded mode — single shared session. + if s.session != nil { + s.session.recordWireFormats(formats) + } + return + } + s.sessions.get(sessionID).session.recordWireFormats(formats) +} + // defaultFormatForClient returns the most-compressed wire format the // named MCP client is known to decode. Resolution order is gcx > // toon > json: @@ -878,6 +927,10 @@ func (s *Server) NoteSessionToolPolicy(sessionID, spec, mode string) { // // Lower-cased client name is matched. Unknown clients are not a // failure — they just keep the JSON default until they're added. +// +// This allowlist is the legacy fallback: a client that declares the +// gortex/wire capability in its initialize request (see +// formatFromWireCapability) never needs an entry here. func defaultFormatForClient(name string) string { if isKnownAgentClient(name) { return "gcx" @@ -909,6 +962,23 @@ func isKnownAgentClient(name string) bool { return knownAgentClients[strings.ToLower(strings.TrimSpace(name))] } +// formatFromWireCapability maps a client's self-declared gortex/wire +// capability (initialize.capabilities["gortex/wire"], in +// the client's preference order) to the first format this server can +// emit. Unknown names are skipped so a future client can declare a +// newer format ahead of an older fallback. +func formatFromWireCapability(formats []string) string { + for _, f := range formats { + switch strings.ToLower(strings.TrimSpace(f)) { + case "gcx": + return "gcx" + case "toon": + return "toon" + } + } + return "" +} + // resolveSessionFormat returns the format the current session prefers // when a tool's `format` arg is absent. Pure read — used by isGCX / // isTOON when the caller didn't pin a format explicitly. @@ -920,6 +990,9 @@ func (s *Server) resolveSessionFormat(ctx context.Context) string { if sess == nil { return "" } + if f := formatFromWireCapability(sess.snapshotWireFormats()); f != "" { + return f + } return defaultFormatForClient(sess.snapshotClientName()) }