From 549f979c1ab53da6e708667a617a7bd1dd584036 Mon Sep 17 00:00:00 2001 From: Lorem Dev Date: Fri, 15 May 2026 21:51:53 +0200 Subject: [PATCH 1/7] feat(mcp): add JSON-RPC error inspection helpers --- internal/mcp/jsonrpc_error.go | 56 +++++++++++++ internal/mcp/jsonrpc_error_test.go | 127 +++++++++++++++++++++++++++++ 2 files changed, 183 insertions(+) create mode 100644 internal/mcp/jsonrpc_error.go create mode 100644 internal/mcp/jsonrpc_error_test.go diff --git a/internal/mcp/jsonrpc_error.go b/internal/mcp/jsonrpc_error.go new file mode 100644 index 0000000..4ba6e7b --- /dev/null +++ b/internal/mcp/jsonrpc_error.go @@ -0,0 +1,56 @@ +package mcp + +import ( + "bytes" + "encoding/json" +) + +// rpcEnvelope is the minimal JSON-RPC 2.0 response shape needed to +// classify whether a response carries an error worth one retry. +type rpcEnvelope struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id,omitempty"` + Error json.RawMessage `json:"error,omitempty"` + Result *struct { + IsError bool `json:"isError,omitempty"` + } `json:"result,omitempty"` +} + +// inspectResponse reports whether raw looks like a JSON-RPC response that +// indicates an error worth retrying once with resolved auth headers. +// Returns false for parse failures, notifications, id:null responses, +// non-JSON-RPC payloads, and any response without an error indicator. +func inspectResponse(raw []byte) bool { + var env rpcEnvelope + if err := json.Unmarshal(raw, &env); err != nil { + return false + } + if env.JSONRPC != "2.0" { + return false + } + if len(env.ID) == 0 || bytes.Equal(env.ID, []byte("null")) { + return false + } + if len(env.Error) > 0 { + return true + } + if env.Result != nil && env.Result.IsError { + return true + } + return false +} + +// extractID returns the raw JSON bytes of the response's id field as a +// string, or "" if the id is missing or literal null. The string is +// suitable as a map key: two responses share the same logical id iff +// their extractID values are equal. +func extractID(raw []byte) string { + var env rpcEnvelope + if err := json.Unmarshal(raw, &env); err != nil { + return "" + } + if len(env.ID) == 0 || bytes.Equal(env.ID, []byte("null")) { + return "" + } + return string(env.ID) +} diff --git a/internal/mcp/jsonrpc_error_test.go b/internal/mcp/jsonrpc_error_test.go new file mode 100644 index 0000000..d2364e6 --- /dev/null +++ b/internal/mcp/jsonrpc_error_test.go @@ -0,0 +1,127 @@ +package mcp + +import "testing" + +func TestInspectResponse(t *testing.T) { + cases := []struct { + name string + raw string + want bool + }{ + { + "tool isError true", + `{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[{"type":"text","text":"x"}]}}`, + true, + }, + { + "tool isError false", + `{"jsonrpc":"2.0","id":1,"result":{"isError":false}}`, + false, + }, + { + "tool result no isError", + `{"jsonrpc":"2.0","id":1,"result":{"content":[]}}`, + false, + }, + { + "jsonrpc error object", + `{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"bad"}}`, + true, + }, + { + "notification (no id)", + `{"jsonrpc":"2.0","method":"notify","params":{}}`, + false, + }, + { + "id null", + `{"jsonrpc":"2.0","id":null,"error":{"code":-1}}`, + false, + }, + { + "malformed json", + `garbage{not json`, + false, + }, + { + "non-jsonrpc", + `{"foo":"bar"}`, + false, + }, + { + "wrong jsonrpc version", + `{"jsonrpc":"1.0","id":1,"error":{}}`, + false, + }, + { + "string id with isError", + `{"jsonrpc":"2.0","id":"abc","result":{"isError":true}}`, + true, + }, + { + "empty result object", + `{"jsonrpc":"2.0","id":1,"result":{}}`, + false, + }, + { + "both error and result (error wins)", + `{"jsonrpc":"2.0","id":1,"error":{"code":-1},"result":{}}`, + true, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := inspectResponse([]byte(tc.raw)) + if got != tc.want { + t.Errorf("inspectResponse(%q) = %v, want %v", tc.raw, got, tc.want) + } + }) + } +} + +func TestExtractID(t *testing.T) { + cases := []struct { + name string + raw string + want string + }{ + { + "numeric id", + `{"jsonrpc":"2.0","id":1,"result":{}}`, + "1", + }, + { + "string id", + `{"jsonrpc":"2.0","id":"abc","result":{}}`, + `"abc"`, + }, + { + "id null", + `{"jsonrpc":"2.0","id":null,"result":{}}`, + "", + }, + { + "missing id", + `{"jsonrpc":"2.0","method":"notify"}`, + "", + }, + { + "malformed", + `garbage`, + "", + }, + { + "id with spaces inside", + `{"jsonrpc":"2.0","id":42,"result":{}}`, + "42", + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := extractID([]byte(tc.raw)) + if got != tc.want { + t.Errorf("extractID(%q) = %q, want %q", tc.raw, got, tc.want) + } + }) + } +} From f625e43230697f730ffbe66b0ab3f0b3c7692d26 Mon Sep 17 00:00:00 2001 From: Lorem Dev Date: Fri, 15 May 2026 21:57:05 +0200 Subject: [PATCH 2/7] feat(mcp): add shared authState for one-shot header resolution --- internal/mcp/auth_state.go | 66 ++++++++++++++++++++ internal/mcp/auth_state_test.go | 105 ++++++++++++++++++++++++++++++++ 2 files changed, 171 insertions(+) create mode 100644 internal/mcp/auth_state.go create mode 100644 internal/mcp/auth_state_test.go diff --git a/internal/mcp/auth_state.go b/internal/mcp/auth_state.go new file mode 100644 index 0000000..e65bdbc --- /dev/null +++ b/internal/mcp/auth_state.go @@ -0,0 +1,66 @@ +package mcp + +import ( + "context" + "net/http" + "sync" +) + +// authState carries the shared auth-header resolution lifecycle between +// the proxy run loop (body-level error trigger) and the transports +// (HTTP 401/403 trigger). Resolution happens at most once per session; +// after the first attempt succeeds or fails, the result is cached and +// the resolver is never invoked again. +type authState struct { + mu sync.Mutex + attempted bool + headers http.Header + resolver HeaderResolver +} + +// newAuthState constructs an authState whose resolveOnce invokes +// resolver. A nil resolver makes resolveOnce a no-op that still marks +// the state as attempted (so callers do not loop). +func newAuthState(resolver HeaderResolver) *authState { + return &authState{resolver: resolver} +} + +// resolveOnce invokes the resolver under the mutex if it has not been +// invoked yet. After return, Attempted() reports true regardless of +// outcome. If the resolver returns an error, headers remain nil and +// the error is propagated; subsequent calls short-circuit without +// invoking resolver again. +func (s *authState) resolveOnce(ctx context.Context) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.attempted { + return nil + } + s.attempted = true + if s.resolver == nil { + return nil + } + h, err := s.resolver(ctx) + if err != nil { + return err + } + s.headers = h + return nil +} + +// Headers returns the cached headers (may be nil before resolve or +// after a failed resolve). The returned http.Header must not be +// mutated by callers; it is shared across all subsequent requests. +func (s *authState) Headers() http.Header { + s.mu.Lock() + defer s.mu.Unlock() + return s.headers +} + +// Attempted reports whether resolveOnce has been called (success or +// failure). After it returns true, no further resolves will occur. +func (s *authState) Attempted() bool { + s.mu.Lock() + defer s.mu.Unlock() + return s.attempted +} diff --git a/internal/mcp/auth_state_test.go b/internal/mcp/auth_state_test.go new file mode 100644 index 0000000..aebe7e7 --- /dev/null +++ b/internal/mcp/auth_state_test.go @@ -0,0 +1,105 @@ +package mcp + +import ( + "context" + "errors" + "net/http" + "sync" + "sync/atomic" + "testing" +) + +func TestAuthState_ResolveOnce_Single(t *testing.T) { + var calls atomic.Int32 + resolver := func(ctx context.Context) (http.Header, error) { + calls.Add(1) + h := http.Header{} + h.Set("Authorization", "Bearer x") + return h, nil + } + s := newAuthState(resolver) + + if s.Attempted() { + t.Fatal("Attempted() = true before first call") + } + if err := s.resolveOnce(context.Background()); err != nil { + t.Fatalf("resolveOnce: %v", err) + } + if calls.Load() != 1 { + t.Errorf("resolver calls = %d, want 1", calls.Load()) + } + if !s.Attempted() { + t.Errorf("Attempted() = false after first call") + } + if got := s.Headers().Get("Authorization"); got != "Bearer x" { + t.Errorf("Headers().Get(Authorization) = %q, want %q", got, "Bearer x") + } +} + +func TestAuthState_ResolveOnce_Concurrent(t *testing.T) { + var calls atomic.Int32 + resolver := func(ctx context.Context) (http.Header, error) { + calls.Add(1) + h := http.Header{} + h.Set("X", "y") + return h, nil + } + s := newAuthState(resolver) + + const numGoroutines = 16 + var wg sync.WaitGroup + wg.Add(numGoroutines) + for i := 0; i < numGoroutines; i++ { + go func() { + defer wg.Done() + _ = s.resolveOnce(context.Background()) + }() + } + wg.Wait() + + if got := calls.Load(); got != 1 { + t.Errorf("resolver calls = %d, want 1", got) + } + if got := s.Headers().Get("X"); got != "y" { + t.Errorf("Headers().Get(X) = %q, want %q", got, "y") + } +} + +func TestAuthState_ResolveOnce_ResolverFails(t *testing.T) { + var calls atomic.Int32 + want := errors.New("boom") + resolver := func(ctx context.Context) (http.Header, error) { + calls.Add(1) + return nil, want + } + s := newAuthState(resolver) + + if err := s.resolveOnce(context.Background()); !errors.Is(err, want) { + t.Errorf("resolveOnce err = %v, want errors.Is(_, boom)", err) + } + if !s.Attempted() { + t.Errorf("Attempted() = false after failed call") + } + if s.Headers() != nil { + t.Errorf("Headers() != nil after failed call") + } + if err := s.resolveOnce(context.Background()); err != nil { + t.Errorf("second resolveOnce err = %v, want nil", err) + } + if got := calls.Load(); got != 1 { + t.Errorf("resolver calls = %d, want 1", got) + } +} + +func TestAuthState_NilResolver(t *testing.T) { + s := newAuthState(nil) + if err := s.resolveOnce(context.Background()); err != nil { + t.Errorf("resolveOnce with nil resolver: %v", err) + } + if !s.Attempted() { + t.Errorf("Attempted() = false") + } + if s.Headers() != nil { + t.Errorf("Headers() != nil") + } +} From ba9afea6889545ce0ef197ffe735410235e5271b Mon Sep 17 00:00:00 2001 From: Lorem Dev Date: Fri, 15 May 2026 22:08:17 +0200 Subject: [PATCH 3/7] refactor(mcp): route transport auth retry through shared authState --- internal/mcp/proxy.go | 15 +++---- internal/mcp/transport.go | 27 ++++++------- internal/mcp/transport_http.go | 55 +++++++++++-------------- internal/mcp/transport_http_test.go | 63 +++++++++++++++++++++++------ internal/mcp/transport_sse.go | 36 +++++++---------- internal/mcp/transport_sse_test.go | 12 +++--- internal/mcp/transport_test.go | 14 +++---- 7 files changed, 120 insertions(+), 102 deletions(-) diff --git a/internal/mcp/proxy.go b/internal/mcp/proxy.go index bbaeaee..e103569 100644 --- a/internal/mcp/proxy.go +++ b/internal/mcp/proxy.go @@ -39,7 +39,7 @@ type transportSetup func(ctx context.Context) (Transport, <-chan []byte, error) // connection. func RunProxy(ctx context.Context, fetcher SecretFetcher, cfg ProxyConfig, in io.Reader, out io.Writer) error { static, templates := splitHeaders(cfg.Headers) - resolver := buildAuthResolver(fetcher, templates) + auth := newAuthState(buildAuthResolver(fetcher, templates)) log.Debug(). Str("url", RedactURL(cfg.URL)). Str("transport", cfg.Transport). @@ -48,7 +48,7 @@ func RunProxy(ctx context.Context, fetcher SecretFetcher, cfg ProxyConfig, in io Msg("mcp proxy: starting (lazy auth)") setup := func(ctx context.Context) (Transport, <-chan []byte, error) { - transport, err := NewTransport(cfg.URL, static, resolver, cfg.Transport) + transport, err := NewTransport(cfg.URL, static, auth, cfg.Transport) if err != nil { log.Debug().Err(err).Msg("mcp proxy: creating transport failed") return nil, nil, err @@ -60,7 +60,7 @@ func RunProxy(ctx context.Context, fetcher SecretFetcher, cfg ProxyConfig, in io } return transport, msgCh, nil } - return runLoop(ctx, setup, in, out) + return runLoop(ctx, setup, auth, in, out) } // RunProxyWithTransport runs the proxy loop against a pre-constructed @@ -81,7 +81,7 @@ func RunProxyWithTransport( } return transport, msgCh, nil } - return runLoop(ctx, setup, in, out) + return runLoop(ctx, setup, nil, in, out) } // splitHeaders partitions cfg.Headers by whether the template contains a @@ -103,8 +103,8 @@ func splitHeaders(mappings []HeaderMapping) (http.Header, []HeaderMapping) { // buildAuthResolver returns a HeaderResolver closure that resolves the // supplied templated headers via fetcher on first call. Returns nil if -// templates is empty - signalling to NewTransport that no auth-retry is -// possible for this proxy session. +// templates is empty; wrapping a nil resolver in authState yields a +// no-op resolveOnce so transports skip the retry path. func buildAuthResolver(fetcher SecretFetcher, templates []HeaderMapping) HeaderResolver { if len(templates) == 0 { return nil @@ -125,7 +125,8 @@ func buildAuthResolver(fetcher SecretFetcher, templates []HeaderMapping) HeaderR } } -func runLoop(ctx context.Context, setup transportSetup, in io.Reader, out io.Writer) error { +func runLoop(ctx context.Context, setup transportSetup, auth *authState, in io.Reader, out io.Writer) error { + _ = auth // Reserved for body-level retry (Task 4/5). reader := bufio.NewReader(in) firstLine, err := readFirstNonEmptyLine(reader) diff --git a/internal/mcp/transport.go b/internal/mcp/transport.go index 2c15d23..b17971e 100644 --- a/internal/mcp/transport.go +++ b/internal/mcp/transport.go @@ -40,14 +40,10 @@ type Transport interface { } // HeaderResolver returns the headers that must be attached to a request -// when the remote MCP server demands authentication. Transport -// implementations call it lazily on the first 401 or 403 response and -// cache the result for the lifetime of the connection. The closure may -// touch the vault and may take seconds to return; callers should not -// hold any lock that blocks unrelated work while waiting. -// -// A nil HeaderResolver disables auth-retry: 401/403 responses are -// propagated to the caller as ordinary errors. +// when the remote MCP server demands authentication. authState invokes +// the resolver at most once per proxy session. The closure may touch +// the vault and may take seconds to return; callers should not hold any +// lock that blocks unrelated work while waiting. type HeaderResolver func(ctx context.Context) (http.Header, error) // SSEEvent is a parsed server-sent event. @@ -124,13 +120,14 @@ func readBody(resp *http.Response) ([]byte, error) { // string. transport must be "auto", "sse", or "http". // // staticHeaders are attached to every request from the first send. -// resolveAuth, when non-nil, is invoked on the first 401/403 response; -// its result is cached and merged with staticHeaders on every -// subsequent request. +// auth carries the shared one-shot header-resolution lifecycle; the +// transport calls auth.resolveOnce on 401/403 responses and merges +// auth.Headers() into every request after a successful resolve. A nil +// auth disables 401/403 retry. func NewTransport( baseURL string, staticHeaders http.Header, - resolveAuth HeaderResolver, + auth *authState, transport string, ) (Transport, error) { client := &http.Client{Timeout: httpClientTimout} @@ -139,21 +136,21 @@ func NewTransport( return &StreamableHTTP{ baseURL: baseURL, staticHeaders: staticHeaders, - resolveAuth: resolveAuth, + auth: auth, client: client, }, nil case "sse": return &SSETransport{ baseURL: baseURL, staticHeaders: staticHeaders, - resolveAuth: resolveAuth, + auth: auth, client: client, }, nil case "auto", "": return &AutoTransport{ baseURL: baseURL, staticHeaders: staticHeaders, - resolveAuth: resolveAuth, + auth: auth, client: client, }, nil default: diff --git a/internal/mcp/transport_http.go b/internal/mcp/transport_http.go index 264f8eb..48f17e1 100644 --- a/internal/mcp/transport_http.go +++ b/internal/mcp/transport_http.go @@ -6,7 +6,6 @@ import ( "fmt" "net/http" "strings" - "sync" "github.com/lorem-dev/locksmith/internal/log" ) @@ -20,10 +19,7 @@ type StreamableHTTP struct { cancel context.CancelFunc staticHeaders http.Header - resolveAuth HeaderResolver - - authMu sync.Mutex - cachedAuth http.Header + auth *authState } func (t *StreamableHTTP) Connect(ctx context.Context) (<-chan []byte, error) { @@ -153,16 +149,17 @@ func (t *StreamableHTTP) postOnce(ctx context.Context, msg []byte) (*http.Respon // shouldRetryWithAuth reports whether the given HTTP status is an // auth-related rejection that warrants resolving the auth header and -// retrying once. Only valid before ensureAuth has succeeded; afterwards +// retrying once. Only valid before authState has resolved; afterwards // 401/403 from the server is a genuine failure (the cached auth is // wrong or expired) and is propagated. func (t *StreamableHTTP) shouldRetryWithAuth(status int) bool { + if t.auth == nil { + return false + } if status != http.StatusUnauthorized && status != http.StatusForbidden { return false } - t.authMu.Lock() - defer t.authMu.Unlock() - return t.cachedAuth == nil + return !t.auth.Attempted() } func (t *StreamableHTTP) Close() error { @@ -173,41 +170,35 @@ func (t *StreamableHTTP) Close() error { } // effectiveHeaders returns the headers to attach to the next request: -// staticHeaders alone before authentication, staticHeaders + cachedAuth -// after. The returned header is a fresh copy safe for the caller to -// mutate (e.g. set Content-Type) without affecting future calls. +// staticHeaders alone before authentication, staticHeaders + auth +// headers after. The returned header is a fresh copy safe for the +// caller to mutate (e.g. set Content-Type) without affecting future +// calls. func (t *StreamableHTTP) effectiveHeaders() http.Header { out := make(http.Header, len(t.staticHeaders)) for k, vs := range t.staticHeaders { out[k] = append([]string(nil), vs...) } - t.authMu.Lock() - for k, vs := range t.cachedAuth { + if t.auth == nil { + return out + } + for k, vs := range t.auth.Headers() { out[k] = append([]string(nil), vs...) } - t.authMu.Unlock() return out } -// ensureAuth resolves the auth headers if they have not been resolved -// yet. It returns true if auth is available (already cached or just -// resolved), false if no resolver is configured. A non-nil error means -// the resolver itself failed; the caller should propagate it. +// ensureAuth resolves auth via the shared authState. Returns (true, nil) +// when headers were resolved on this or a prior call, (false, nil) if +// auth is unavailable, or (false, err) if the resolve failed. func (t *StreamableHTTP) ensureAuth(ctx context.Context) (bool, error) { - if t.resolveAuth == nil { + if t.auth == nil { return false, nil } - t.authMu.Lock() - defer t.authMu.Unlock() - if t.cachedAuth != nil { - return true, nil - } - h, err := t.resolveAuth(ctx) - if err != nil { + if err := t.auth.resolveOnce(ctx); err != nil { return false, fmt.Errorf("resolving auth headers: %w", err) } - t.cachedAuth = h - return true, nil + return t.auth.Headers() != nil, nil } // AutoTransport tries Streamable HTTP first; falls back to SSE on 404/405. @@ -218,7 +209,7 @@ type AutoTransport struct { client *http.Client staticHeaders http.Header - resolveAuth HeaderResolver + auth *authState inner Transport outCh chan []byte @@ -234,7 +225,7 @@ func (t *AutoTransport) Connect(ctx context.Context) (<-chan []byte, error) { t.inner = &StreamableHTTP{ baseURL: t.baseURL, staticHeaders: t.staticHeaders, - resolveAuth: t.resolveAuth, + auth: t.auth, client: t.client, } innerCh, err := t.inner.Connect(fwdCtx) @@ -276,7 +267,7 @@ func (t *AutoTransport) Send(ctx context.Context, msg []byte) error { sse := &SSETransport{ baseURL: t.baseURL, staticHeaders: t.staticHeaders, - resolveAuth: t.resolveAuth, + auth: t.auth, client: t.client, } sseCh, connectErr := sse.Connect(ctx) diff --git a/internal/mcp/transport_http_test.go b/internal/mcp/transport_http_test.go index f324d3b..b7fc884 100644 --- a/internal/mcp/transport_http_test.go +++ b/internal/mcp/transport_http_test.go @@ -1,4 +1,4 @@ -package mcp_test +package mcp import ( "context" @@ -6,13 +6,12 @@ import ( "io" "net/http" "net/http/httptest" + "sync/atomic" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - - "github.com/lorem-dev/locksmith/internal/mcp" ) func TestStreamableHTTP_JSONResponse(t *testing.T) { @@ -31,7 +30,7 @@ func TestStreamableHTTP_JSONResponse(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, nil, "http") + transport, err := NewTransport(srv.URL, nil, nil, "http") require.NoError(t, err) defer transport.Close() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) @@ -72,7 +71,7 @@ func TestStreamableHTTP_SSEResponse(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, nil, "http") + transport, err := NewTransport(srv.URL, nil, nil, "http") require.NoError(t, err) defer transport.Close() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) @@ -123,7 +122,7 @@ func TestAutoTransport_FallsBackToSSE(t *testing.T) { srv := httptest.NewServer(mux) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, nil, "auto") + transport, err := NewTransport(srv.URL, nil, nil, "auto") require.NoError(t, err) defer transport.Close() ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) @@ -162,7 +161,7 @@ func TestStreamableHTTP_LazyAuth_200_StaysUnauthenticated(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, resolver, "http") + transport, err := NewTransport(srv.URL, nil, newAuthState(resolver), "http") require.NoError(t, err) defer transport.Close() @@ -189,7 +188,7 @@ func TestStreamableHTTP_LazyAuth_NoResolver_NoRetry(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, nil, "http") + transport, err := NewTransport(srv.URL, nil, nil, "http") require.NoError(t, err) defer transport.Close() @@ -224,7 +223,7 @@ func TestStreamableHTTP_LazyAuth_401_TriggersResolveAndRetry(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, resolver, "http") + transport, err := NewTransport(srv.URL, nil, newAuthState(resolver), "http") require.NoError(t, err) defer transport.Close() @@ -260,7 +259,7 @@ func TestStreamableHTTP_LazyAuth_403_TreatedAsAuth(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, resolver, "http") + transport, err := NewTransport(srv.URL, nil, newAuthState(resolver), "http") require.NoError(t, err) defer transport.Close() @@ -296,7 +295,7 @@ func TestStreamableHTTP_LazyAuth_AuthSticky(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, resolver, "http") + transport, err := NewTransport(srv.URL, nil, newAuthState(resolver), "http") require.NoError(t, err) defer transport.Close() @@ -329,7 +328,7 @@ func TestStreamableHTTP_LazyAuth_AuthFailsAfterRetry(t *testing.T) { })) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, resolver, "http") + transport, err := NewTransport(srv.URL, nil, newAuthState(resolver), "http") require.NoError(t, err) defer transport.Close() @@ -341,3 +340,43 @@ func TestStreamableHTTP_LazyAuth_AuthFailsAfterRetry(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "401") } + +// TestStreamableHTTP_BodyResolveThen401_NoSecondResolve verifies that +// when the shared authState has already been resolved (e.g. by the +// body-level retry path in the proxy run loop), a subsequent 401 from +// the server is treated as a genuine failure: the transport does NOT +// invoke the resolver again, and the error is propagated. +func TestStreamableHTTP_BodyResolveThen401_NoSecondResolve(t *testing.T) { + var calls atomic.Int32 + resolver := func(_ context.Context) (http.Header, error) { + calls.Add(1) + return http.Header{"X-Token": []string{"ok"}}, nil + } + auth := newAuthState(resolver) + if err := auth.resolveOnce(context.Background()); err != nil { + t.Fatalf("seed resolveOnce: %v", err) + } + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer server.Close() + + tr := &StreamableHTTP{ + baseURL: server.URL, + client: server.Client(), + staticHeaders: http.Header{}, + auth: auth, + } + _, err := tr.Connect(context.Background()) + if err != nil { + t.Fatalf("Connect: %v", err) + } + err = tr.Send(context.Background(), []byte(`{"jsonrpc":"2.0","id":1,"method":"x"}`)) + if err == nil { + t.Fatal("Send: expected 401 error, got nil") + } + if got := calls.Load(); got != 1 { + t.Errorf("resolver calls = %d, want 1 (no second resolve)", got) + } +} diff --git a/internal/mcp/transport_sse.go b/internal/mcp/transport_sse.go index d937929..b560732 100644 --- a/internal/mcp/transport_sse.go +++ b/internal/mcp/transport_sse.go @@ -6,7 +6,6 @@ import ( "net/http" "net/url" "strings" - "sync" "time" "github.com/lorem-dev/locksmith/internal/log" @@ -23,10 +22,7 @@ type SSETransport struct { cancel context.CancelFunc staticHeaders http.Header - resolveAuth HeaderResolver - - authMu sync.Mutex - cachedAuth http.Header + auth *authState } func (t *SSETransport) Connect(ctx context.Context) (<-chan []byte, error) { @@ -165,12 +161,13 @@ func (t *SSETransport) postEndpointOnce(ctx context.Context, msg []byte) (*http. // shouldRetryWithAuth reports whether the response status warrants an // auth resolve and retry. Mirrors StreamableHTTP.shouldRetryWithAuth. func (t *SSETransport) shouldRetryWithAuth(status int) bool { + if t.auth == nil { + return false + } if status != http.StatusUnauthorized && status != http.StatusForbidden { return false } - t.authMu.Lock() - defer t.authMu.Unlock() - return t.cachedAuth == nil + return !t.auth.Attempted() } func (t *SSETransport) Close() error { @@ -209,27 +206,24 @@ func (t *SSETransport) effectiveHeaders() http.Header { for k, vs := range t.staticHeaders { out[k] = append([]string(nil), vs...) } - t.authMu.Lock() - for k, vs := range t.cachedAuth { + if t.auth == nil { + return out + } + for k, vs := range t.auth.Headers() { out[k] = append([]string(nil), vs...) } - t.authMu.Unlock() return out } +// ensureAuth resolves auth via the shared authState. Returns (true, nil) +// when headers were resolved on this or a prior call, (false, nil) if +// auth is unavailable, or (false, err) if the resolve failed. func (t *SSETransport) ensureAuth(ctx context.Context) (bool, error) { - if t.resolveAuth == nil { + if t.auth == nil { return false, nil } - t.authMu.Lock() - defer t.authMu.Unlock() - if t.cachedAuth != nil { - return true, nil - } - h, err := t.resolveAuth(ctx) - if err != nil { + if err := t.auth.resolveOnce(ctx); err != nil { return false, fmt.Errorf("resolving auth headers: %w", err) } - t.cachedAuth = h - return true, nil + return t.auth.Headers() != nil, nil } diff --git a/internal/mcp/transport_sse_test.go b/internal/mcp/transport_sse_test.go index 70b30d7..4298b68 100644 --- a/internal/mcp/transport_sse_test.go +++ b/internal/mcp/transport_sse_test.go @@ -1,4 +1,4 @@ -package mcp_test +package mcp import ( "context" @@ -10,8 +10,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - - "github.com/lorem-dev/locksmith/internal/mcp" ) func TestSSETransport_RoundTrip(t *testing.T) { @@ -38,7 +36,7 @@ func TestSSETransport_RoundTrip(t *testing.T) { srv := httptest.NewServer(mux) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, nil, "sse") + transport, err := NewTransport(srv.URL, nil, nil, "sse") require.NoError(t, err) defer transport.Close() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) @@ -76,7 +74,7 @@ func TestSSETransport_AuthHeader(t *testing.T) { defer srv.Close() headers := http.Header{"Authorization": {"Bearer tok-123"}} - transport, err := mcp.NewTransport(srv.URL, headers, nil, "sse") + transport, err := NewTransport(srv.URL, headers, nil, "sse") require.NoError(t, err) defer transport.Close() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) @@ -114,7 +112,7 @@ func TestSSETransport_LazyAuth_Connect_401_ReopensWithAuth(t *testing.T) { srv := httptest.NewServer(mux) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, resolver, "sse") + transport, err := NewTransport(srv.URL, nil, newAuthState(resolver), "sse") require.NoError(t, err) defer transport.Close() @@ -158,7 +156,7 @@ func TestSSETransport_LazyAuth_Send_401_RetriesWithAuth(t *testing.T) { srv := httptest.NewServer(mux) defer srv.Close() - transport, err := mcp.NewTransport(srv.URL, nil, resolver, "sse") + transport, err := NewTransport(srv.URL, nil, newAuthState(resolver), "sse") require.NoError(t, err) defer transport.Close() diff --git a/internal/mcp/transport_test.go b/internal/mcp/transport_test.go index 737e4c9..7b5d9be 100644 --- a/internal/mcp/transport_test.go +++ b/internal/mcp/transport_test.go @@ -1,4 +1,4 @@ -package mcp_test +package mcp import ( "strings" @@ -6,13 +6,11 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - - "github.com/lorem-dev/locksmith/internal/mcp" ) func TestParseSSE(t *testing.T) { input := "event: endpoint\ndata: /messages\n\nevent: message\ndata: {\"jsonrpc\":\"2.0\"}\n\n" - events := mcp.CollectSSE(strings.NewReader(input)) + events := CollectSSE(strings.NewReader(input)) require.Len(t, events, 2) assert.Equal(t, "endpoint", events[0].Type) assert.Equal(t, "/messages", events[0].Data) @@ -22,14 +20,14 @@ func TestParseSSE(t *testing.T) { func TestParseSSE_DataOnly(t *testing.T) { input := "data: hello\n\ndata: world\n\n" - events := mcp.CollectSSE(strings.NewReader(input)) + events := CollectSSE(strings.NewReader(input)) require.Len(t, events, 2) assert.Equal(t, "", events[0].Type) assert.Equal(t, "hello", events[0].Data) } func TestNewTransport_InvalidType(t *testing.T) { - _, err := mcp.NewTransport("https://example.com", nil, nil, "grpc") + _, err := NewTransport("https://example.com", nil, nil, "grpc") require.ErrorContains(t, err, "unknown transport") } @@ -46,8 +44,8 @@ func TestRedactURL(t *testing.T) { } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { - assert.Equal(t, tc.want, mcp.RedactURL(tc.input)) - assert.NotContains(t, mcp.RedactURL(tc.input), "supersecret") + assert.Equal(t, tc.want, RedactURL(tc.input)) + assert.NotContains(t, RedactURL(tc.input), "supersecret") }) } } From 34c9d363a6e4dd94e7dccf4f72469f990b17e754 Mon Sep 17 00:00:00 2001 From: Lorem Dev Date: Fri, 15 May 2026 22:13:12 +0200 Subject: [PATCH 4/7] feat(mcp): track in-flight request ids in proxy run loop --- internal/mcp/proxy.go | 104 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 95 insertions(+), 9 deletions(-) diff --git a/internal/mcp/proxy.go b/internal/mcp/proxy.go index e103569..49ba30f 100644 --- a/internal/mcp/proxy.go +++ b/internal/mcp/proxy.go @@ -126,7 +126,6 @@ func buildAuthResolver(fetcher SecretFetcher, templates []HeaderMapping) HeaderR } func runLoop(ctx context.Context, setup transportSetup, auth *authState, in io.Reader, out io.Writer) error { - _ = auth // Reserved for body-level retry (Task 4/5). reader := bufio.NewReader(in) firstLine, err := readFirstNonEmptyLine(reader) @@ -154,6 +153,7 @@ func runLoop(ctx context.Context, setup transportSetup, auth *authState, in io.R out.Write(append(msg, '\n')) //nolint:errcheck } + state := newProxyState() done := make(chan struct{}) shutdown := func() { select { @@ -167,11 +167,11 @@ func runLoop(ctx context.Context, setup transportSetup, auth *authState, in io.R wg.Add(1) go func() { defer wg.Done() - forwardServerMessages(msgCh, done, writeMsg) + forwardServerMessagesWithTracking(msgCh, done, writeMsg, state, auth) }() var clientMsgCount uint64 - if err := sendClientMessage(ctx, transport, firstLine, &clientMsgCount); err != nil { + if err := sendClientMessageTracked(ctx, transport, firstLine, &clientMsgCount, state, auth); err != nil { shutdown() wg.Wait() return err @@ -182,7 +182,7 @@ func runLoop(ctx context.Context, setup transportSetup, auth *authState, in io.R if len(line) > 0 { trimmed := bytes.TrimRight(line, "\r\n") if len(trimmed) > 0 { - if err := sendClientMessage(ctx, transport, trimmed, &clientMsgCount); err != nil { + if err := sendClientMessageTracked(ctx, transport, trimmed, &clientMsgCount, state, auth); err != nil { shutdown() wg.Wait() return err @@ -225,15 +225,28 @@ func readFirstNonEmptyLine(reader *bufio.Reader) ([]byte, error) { } } -// forwardServerMessages reads server messages from msgCh and dispatches -// them to writeMsg until either msgCh closes or done fires. When done -// fires it first drains any messages already buffered in msgCh so late -// responses still reach the client before shutdown. -func forwardServerMessages(msgCh <-chan []byte, done <-chan struct{}, writeMsg func([]byte)) { +// forwardServerMessagesWithTracking reads server messages from msgCh and +// dispatches them to writeMsg until either msgCh closes or done fires. +// When done fires it first drains any messages already buffered in +// msgCh so late responses still reach the client before shutdown. While +// auth has not been attempted, the response's id is removed from state +// so the in-flight map only holds requests whose response has not been +// forwarded yet. Task 5 swaps this for a variant that also inspects +// responses and triggers resolve+retry. +func forwardServerMessagesWithTracking( + msgCh <-chan []byte, + done <-chan struct{}, + writeMsg func([]byte), + state *proxyState, + auth *authState, +) { var count uint64 forward := func(msg []byte) { count++ log.Debug().Uint64("seq", count).Int("len", len(msg)).Msg("mcp proxy: server -> client") + if auth != nil && !auth.Attempted() { + _ = state.take(extractID(msg)) + } writeMsg(msg) } for { @@ -271,3 +284,76 @@ func sendClientMessage(ctx context.Context, transport Transport, line []byte, co } return nil } + +// sendClientMessageTracked sends to transport and, while auth has not +// been attempted, records the request's id for potential retry. The +// stored bytes are a copy of line so subsequent stdin reads cannot +// mutate them. +func sendClientMessageTracked( + ctx context.Context, + transport Transport, + line []byte, + count *uint64, + state *proxyState, + auth *authState, +) error { + if auth != nil && !auth.Attempted() { + if id := extractID(line); id != "" { + state.record(id, append([]byte(nil), line...)) + } + } + return sendClientMessage(ctx, transport, line, count) +} + +// proxyState tracks in-flight client request bytes by their JSON-RPC id +// while body-level retry is still possible. Once authState.Attempted() +// is true, no further tracking happens (zero overhead via take/record +// short-circuits). +type proxyState struct { + mu sync.Mutex + inFlight map[string][]byte +} + +func newProxyState() *proxyState { + return &proxyState{inFlight: make(map[string][]byte)} +} + +// record stores the request bytes under id. No-op for empty id (e.g. +// notifications) or after clear() has nilled the map. +func (s *proxyState) record(id string, bytes []byte) { + if id == "" { + return + } + s.mu.Lock() + defer s.mu.Unlock() + if s.inFlight == nil { + return + } + s.inFlight[id] = bytes +} + +// take removes and returns the request bytes for id, or nil if the id +// is empty, unknown, or the map has been cleared. +func (s *proxyState) take(id string) []byte { + if id == "" { + return nil + } + s.mu.Lock() + defer s.mu.Unlock() + if s.inFlight == nil { + return nil + } + b := s.inFlight[id] + delete(s.inFlight, id) + return b +} + +// clear releases the inFlight map. Subsequent record/take calls are +// no-ops. Called once authState.Attempted() flips to true. +// +//nolint:unused // Wired up by Task 5 (body-level retry). +func (s *proxyState) clear() { + s.mu.Lock() + defer s.mu.Unlock() + s.inFlight = nil +} From d5416f1f8d8b92fa650dc4bd8f9d0ba5d73c2afe Mon Sep 17 00:00:00 2001 From: Lorem Dev Date: Fri, 15 May 2026 22:32:27 +0200 Subject: [PATCH 5/7] feat(mcp): retry on JSON-RPC body errors with shared authState --- internal/mcp/proxy.go | 151 +++++++++++---- internal/mcp/proxy_retry_test.go | 322 +++++++++++++++++++++++++++++++ 2 files changed, 440 insertions(+), 33 deletions(-) create mode 100644 internal/mcp/proxy_retry_test.go diff --git a/internal/mcp/proxy.go b/internal/mcp/proxy.go index 49ba30f..4163f59 100644 --- a/internal/mcp/proxy.go +++ b/internal/mcp/proxy.go @@ -167,7 +167,7 @@ func runLoop(ctx context.Context, setup transportSetup, auth *authState, in io.R wg.Add(1) go func() { defer wg.Done() - forwardServerMessagesWithTracking(msgCh, done, writeMsg, state, auth) + forwardServerMessagesWithTracking(ctx, msgCh, done, writeMsg, state, auth, transport) }() var clientMsgCount uint64 @@ -227,52 +227,139 @@ func readFirstNonEmptyLine(reader *bufio.Reader) ([]byte, error) { // forwardServerMessagesWithTracking reads server messages from msgCh and // dispatches them to writeMsg until either msgCh closes or done fires. -// When done fires it first drains any messages already buffered in -// msgCh so late responses still reach the client before shutdown. While -// auth has not been attempted, the response's id is removed from state -// so the in-flight map only holds requests whose response has not been -// forwarded yet. Task 5 swaps this for a variant that also inspects -// responses and triggers resolve+retry. +// While auth has not been attempted, each incoming response is inspected +// for a body-level JSON-RPC error. On detection, the original request +// bytes are looked up in state, the authState resolver is invoked once, +// the request is re-sent with the freshly resolved headers, and the +// retry response (matching the same id) is forwarded to the client in +// place of the original error. If resolve or retry-Send fails, the +// original error response is forwarded as-is. func forwardServerMessagesWithTracking( + ctx context.Context, msgCh <-chan []byte, done <-chan struct{}, writeMsg func([]byte), state *proxyState, auth *authState, + transport Transport, ) { var count uint64 forward := func(msg []byte) { count++ log.Debug().Uint64("seq", count).Int("len", len(msg)).Msg("mcp proxy: server -> client") - if auth != nil && !auth.Attempted() { - _ = state.take(extractID(msg)) - } writeMsg(msg) } + + // drained tracks whether done has fired; once it has, subsequent + // recv calls read msgCh non-blockingly so any buffered late + // responses are forwarded without waiting on transports that do + // not close their server channel on shutdown. + drained := false + recv := func() ([]byte, bool) { + if drained { + return tryRecv(msgCh) + } + msg, ok, doneFired := blockingRecv(msgCh, done) + if doneFired { + drained = true + return tryRecv(msgCh) + } + return msg, ok + } + for { - select { - case msg, ok := <-msgCh: - if !ok { - log.Debug().Uint64("count", count).Msg("mcp proxy: server channel closed") + msg, ok := recv() + if !ok { + log.Debug().Uint64("count", count).Msg("mcp proxy: server channel closed") + return + } + if auth != nil && !auth.Attempted() && inspectResponse(msg) { + if exit := handleBodyError(ctx, msg, state, auth, transport, forward, recv); exit { return } - forward(msg) - case <-done: - for { - select { - case msg, ok := <-msgCh: - if !ok { - log.Debug().Uint64("count", count).Msg("mcp proxy: server channel closed during drain") - return - } - forward(msg) - default: - log.Debug().Uint64("count", count).Msg("mcp proxy: server reader stopped") - return - } - } + continue + } + if auth != nil && !auth.Attempted() { + _ = state.take(extractID(msg)) + } + forward(msg) + } +} + +// tryRecv attempts a non-blocking receive on msgCh. The second return +// value reports whether a message was produced; closed channels yield +// (nil, false). +func tryRecv(msgCh <-chan []byte) ([]byte, bool) { + select { + case msg, ok := <-msgCh: + if !ok { + return nil, false + } + return msg, true + default: + return nil, false + } +} + +// blockingRecv waits for a message on msgCh, or for done to fire. The +// third return value distinguishes "done fired" from "msgCh closed", +// so the caller can flip into drain mode for any buffered late +// responses. +func blockingRecv(msgCh <-chan []byte, done <-chan struct{}) ([]byte, bool, bool) { + select { + case m, chOk := <-msgCh: + if !chOk { + return nil, false, false + } + return m, true, false + case <-done: + return nil, false, true + } +} + +// handleBodyError performs the body-level retry handshake for one error +// response. It returns true only when recv reports the server channel +// is gone mid-retry, signalling the caller to exit the forward loop. +func handleBodyError( + ctx context.Context, + msg []byte, + state *proxyState, + auth *authState, + transport Transport, + forward func([]byte), + recv func() ([]byte, bool), +) bool { + id := extractID(msg) + orig := state.take(id) + if orig == nil { + forward(msg) + return false + } + log.Debug().Str("id", id).Msg("mcp proxy: body-level error detected; resolving headers") + if err := auth.resolveOnce(ctx); err != nil { + log.Debug().Err(err).Msg("mcp proxy: resolveOnce failed; forwarding original error") + forward(msg) + state.clear() + return false + } + if err := transport.Send(ctx, orig); err != nil { + log.Debug().Err(err).Msg("mcp proxy: retry Send failed; forwarding original error") + forward(msg) + state.clear() + return false + } + for { + next, ok := recv() + if !ok { + return true + } + forward(next) + if extractID(next) == id { + break } } + state.clear() + return false } func sendClientMessage(ctx context.Context, transport Transport, line []byte, count *uint64) error { @@ -320,7 +407,7 @@ func newProxyState() *proxyState { // record stores the request bytes under id. No-op for empty id (e.g. // notifications) or after clear() has nilled the map. -func (s *proxyState) record(id string, bytes []byte) { +func (s *proxyState) record(id string, raw []byte) { if id == "" { return } @@ -329,7 +416,7 @@ func (s *proxyState) record(id string, bytes []byte) { if s.inFlight == nil { return } - s.inFlight[id] = bytes + s.inFlight[id] = raw } // take removes and returns the request bytes for id, or nil if the id @@ -350,8 +437,6 @@ func (s *proxyState) take(id string) []byte { // clear releases the inFlight map. Subsequent record/take calls are // no-ops. Called once authState.Attempted() flips to true. -// -//nolint:unused // Wired up by Task 5 (body-level retry). func (s *proxyState) clear() { s.mu.Lock() defer s.mu.Unlock() diff --git a/internal/mcp/proxy_retry_test.go b/internal/mcp/proxy_retry_test.go new file mode 100644 index 0000000..f615eb9 --- /dev/null +++ b/internal/mcp/proxy_retry_test.go @@ -0,0 +1,322 @@ +package mcp + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +// scriptedTransport is a Transport that records sent messages and lets +// the test push pre-canned server responses on demand. Suitable for +// driving the proxy run loop without a real HTTP server. +type scriptedTransport struct { + mu sync.Mutex + sent [][]byte + out chan []byte + done chan struct{} + sendErr error + closed bool +} + +func newScriptedTransport() *scriptedTransport { + return &scriptedTransport{out: make(chan []byte, 16), done: make(chan struct{})} +} + +func (t *scriptedTransport) Connect(_ context.Context) (<-chan []byte, error) { + return t.out, nil +} + +func (t *scriptedTransport) Send(_ context.Context, msg []byte) error { + t.mu.Lock() + defer t.mu.Unlock() + if t.sendErr != nil { + return t.sendErr + } + t.sent = append(t.sent, append([]byte(nil), msg...)) + return nil +} + +func (t *scriptedTransport) Close() error { + t.mu.Lock() + defer t.mu.Unlock() + if !t.closed { + t.closed = true + close(t.done) + close(t.out) + } + return nil +} + +func (t *scriptedTransport) push(msg string) { t.out <- []byte(msg) } + +func (t *scriptedTransport) sentCount() int { + t.mu.Lock() + defer t.mu.Unlock() + return len(t.sent) +} + +// delayedEOFReader emits the bytes of `body` and then blocks subsequent +// reads on `wait` so the run loop does not observe stdin EOF until the +// test goroutine signals completion (typically by closing `wait` after +// tr.Close() has been called). This lets scripted tests push server +// responses and trigger retries without racing the shutdown path. +type delayedEOFReader struct { + body *strings.Reader + wait <-chan struct{} +} + +func (r *delayedEOFReader) Read(p []byte) (int, error) { + if r.body.Len() > 0 { + return r.body.Read(p) + } + <-r.wait + return 0, io.EOF +} + +// runProxyForTest drives runLoop with the given scripted transport and +// resolver, then returns the non-empty lines written to stdout. stdin +// is kept open until tr is closed so the scripted server goroutine can +// push responses without racing shutdown. +func runProxyForTest( + t *testing.T, + tr *scriptedTransport, + resolver HeaderResolver, + stdin string, +) ([][]byte, *authState) { + t.Helper() + auth := newAuthState(resolver) + var stdout bytes.Buffer + setup := func(ctx context.Context) (Transport, <-chan []byte, error) { + ch, _ := tr.Connect(ctx) + return tr, ch, nil + } + reader := &delayedEOFReader{body: strings.NewReader(stdin), wait: tr.done} + if err := runLoop(context.Background(), setup, auth, reader, &stdout); err != nil { + t.Fatalf("runLoop: %v", err) + } + var lines [][]byte + for _, l := range bytes.Split(stdout.Bytes(), []byte{'\n'}) { + if len(l) > 0 { + lines = append(lines, l) + } + } + return lines, auth +} + +func TestProxy_BodyError_TriggersResolveAndRetry(t *testing.T) { + tr := newScriptedTransport() + var calls atomic.Int32 + resolver := func(ctx context.Context) (http.Header, error) { + calls.Add(1) + h := http.Header{} + h.Set("X-Token", "ok") + return h, nil + } + go func() { + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[{"type":"text","text":"err"}]}}`) + for tr.sentCount() < 2 { + time.Sleep(5 * time.Millisecond) + } + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":false,"content":[{"type":"text","text":"ok"}]}}`) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + lines, auth := runProxyForTest(t, tr, resolver, stdin) + if calls.Load() != 1 { + t.Errorf("resolver calls = %d, want 1", calls.Load()) + } + if !auth.Attempted() { + t.Errorf("Attempted() = false") + } + if len(lines) != 1 { + t.Fatalf("stdout lines = %d, want 1; got: %q", len(lines), lines) + } + if !bytes.Contains(lines[0], []byte(`"isError":false`)) { + t.Errorf("client did not receive retry success: %q", lines[0]) + } +} + +func TestProxy_BodyError_ResolveFails_ForwardsOriginal(t *testing.T) { + tr := newScriptedTransport() + want := errors.New("vault locked") + resolver := func(ctx context.Context) (http.Header, error) { return nil, want } + go func() { + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[{"type":"text","text":"err"}]}}`) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + lines, auth := runProxyForTest(t, tr, resolver, stdin) + if !auth.Attempted() { + t.Errorf("Attempted() = false") + } + if len(lines) != 1 || !bytes.Contains(lines[0], []byte(`"isError":true`)) { + t.Fatalf("stdout did not contain original error: %q", lines) + } +} + +func TestProxy_BodyError_RetrySendFails_ForwardsOriginal(t *testing.T) { + tr := newScriptedTransport() + resolver := func(ctx context.Context) (http.Header, error) { + return http.Header{"X-Token": []string{"ok"}}, nil + } + go func() { + // Wait for runLoop to send the original request before flipping + // sendErr; otherwise the initial Send fails and the retry path + // is not exercised. + for tr.sentCount() < 1 { + time.Sleep(5 * time.Millisecond) + } + tr.mu.Lock() + tr.sendErr = errors.New("connect fail") + tr.mu.Unlock() + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[{"type":"text","text":"err"}]}}`) + time.Sleep(50 * time.Millisecond) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + lines, auth := runProxyForTest(t, tr, resolver, stdin) + if !auth.Attempted() { + t.Errorf("Attempted() = false") + } + if len(lines) != 1 || !bytes.Contains(lines[0], []byte(`"isError":true`)) { + t.Fatalf("stdout did not contain original error: %q", lines) + } +} + +func TestProxy_BodyError_RetryAlsoErrors_ForwardsRetryResponse(t *testing.T) { + tr := newScriptedTransport() + resolver := func(ctx context.Context) (http.Header, error) { + return http.Header{"X-Token": []string{"ok"}}, nil + } + go func() { + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[{"type":"text","text":"err1"}]}}`) + for tr.sentCount() < 2 { + time.Sleep(5 * time.Millisecond) + } + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[{"type":"text","text":"err2"}]}}`) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + lines, _ := runProxyForTest(t, tr, resolver, stdin) + if len(lines) != 1 || !bytes.Contains(lines[0], []byte(`"err2"`)) { + t.Fatalf("client should see retry response only; got: %q", lines) + } +} + +func TestProxy_ConcurrentInFlight_RetriesCorrectId(t *testing.T) { + tr := newScriptedTransport() + resolver := func(ctx context.Context) (http.Header, error) { + return http.Header{"X-Token": []string{"ok"}}, nil + } + go func() { + tr.push(`{"jsonrpc":"2.0","id":2,"result":{"isError":true,"content":[{"type":"text","text":"e2"}]}}`) + for tr.sentCount() < 3 { + time.Sleep(5 * time.Millisecond) + } + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":false,"content":[{"type":"text","text":"ok1"}]}}`) + tr.push(`{"jsonrpc":"2.0","id":2,"result":{"isError":false,"content":[{"type":"text","text":"ok2"}]}}`) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\"}\n" + lines, _ := runProxyForTest(t, tr, resolver, stdin) + if len(lines) != 2 { + t.Fatalf("expected 2 forwarded responses, got %d: %q", len(lines), lines) + } + if !bytes.Contains(lines[0], []byte(`"id":1`)) || !bytes.Contains(lines[1], []byte(`"id":2`)) { + t.Errorf("unexpected line order: %q", lines) + } + if !bytes.Contains(lines[1], []byte(`"ok2"`)) { + t.Errorf("id=2 line should be retry success: %q", lines[1]) + } + for _, l := range lines { + if bytes.Contains(l, []byte(`"e2"`)) { + t.Errorf("original error leaked: %q", l) + } + } +} + +func TestProxy_OnlyFirstErrorTriggers(t *testing.T) { + tr := newScriptedTransport() + var calls atomic.Int32 + resolver := func(ctx context.Context) (http.Header, error) { + calls.Add(1) + return http.Header{"X-Token": []string{"ok"}}, nil + } + go func() { + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":false,"content":[]}}`) + tr.push(`{"jsonrpc":"2.0","id":2,"result":{"isError":true,"content":[]}}`) + for tr.sentCount() < 3 { + time.Sleep(5 * time.Millisecond) + } + tr.push(`{"jsonrpc":"2.0","id":2,"result":{"isError":false,"content":[{"type":"text","text":"ok"}]}}`) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\"}\n" + lines, auth := runProxyForTest(t, tr, resolver, stdin) + if calls.Load() != 1 { + t.Errorf("resolver calls = %d, want 1", calls.Load()) + } + if !auth.Attempted() { + t.Errorf("Attempted() = false") + } + if len(lines) != 2 { + t.Fatalf("expected 2 forwarded responses, got %d: %q", len(lines), lines) + } + if !bytes.Contains(lines[1], []byte(`"ok"`)) { + t.Errorf("id=2 line should be retry success: %q", lines[1]) + } +} + +func TestProxy_AfterAttempted_NoParse(t *testing.T) { + tr := newScriptedTransport() + var calls atomic.Int32 + resolver := func(ctx context.Context) (http.Header, error) { + calls.Add(1) + return http.Header{"X-Token": []string{"ok"}}, nil + } + go func() { + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[]}}`) + for tr.sentCount() < 2 { + time.Sleep(5 * time.Millisecond) + } + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":false,"content":[]}}`) + tr.push(`{"jsonrpc":"2.0","id":2,"result":{"isError":true,"content":[]}}`) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\"}\n" + _, _ = runProxyForTest(t, tr, resolver, stdin) + if got := calls.Load(); got != 1 { + t.Errorf("resolver calls = %d, want 1", got) + } +} + +func TestProxy_Notification_NotTracked(t *testing.T) { + tr := newScriptedTransport() + resolver := func(ctx context.Context) (http.Header, error) { + return http.Header{"X-Token": []string{"ok"}}, nil + } + go func() { + tr.push(`{"jsonrpc":"2.0","method":"notifications/cancelled","params":{}}`) + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":true,"content":[]}}`) + for tr.sentCount() < 2 { + time.Sleep(5 * time.Millisecond) + } + tr.push(`{"jsonrpc":"2.0","id":1,"result":{"isError":false,"content":[]}}`) + tr.Close() + }() + stdin := "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/call\"}\n" + lines, _ := runProxyForTest(t, tr, resolver, stdin) + if len(lines) != 2 { + t.Fatalf("expected 2 lines (notification + retry success), got %d: %q", len(lines), lines) + } +} From ce3705f7252c6999bc2c74576afb07785f1c9d02 Mon Sep 17 00:00:00 2001 From: Lorem Dev Date: Fri, 15 May 2026 22:38:15 +0200 Subject: [PATCH 6/7] docs: describe body-level JSON-RPC error trigger for lazy auth --- CHANGES.md | 7 +++++++ docs/architecture.md | 14 ++++++++++++++ docs/configuration.md | 8 ++++++++ 3 files changed, 29 insertions(+) diff --git a/CHANGES.md b/CHANGES.md index 2b7745e..629128e 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -2,6 +2,13 @@ ## Development +- `locksmith mcp run --url` now triggers templated-header resolution + not only on HTTP `401/403` but also when the remote MCP server + returns `200 OK` with a JSON-RPC `error` field or a tool-level + `result.isError: true`. Resolution and retry happen exactly once + per session; if the retry also fails, the error is forwarded to the + AI client unchanged. Detection looks only at structural fields - no + keyword matching against `result.content[].text`. - docs(cli): drop Touch ID wording from CLI help, README, configuration and architecture docs, and plugin READMEs; keychain authorization is described as OS-delegated without naming a specific auth method diff --git a/docs/architecture.md b/docs/architecture.md index 470bed3..054a6d5 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -115,6 +115,20 @@ sends its first MCP request: the MCP `initialize` handshake without auth therefore never trigger a vault prompt for that connection. +The same `resolveOnce` lifecycle also triggers on **body-level errors**: +while the resolver has not yet been invoked, the proxy run loop +inspects each server response. If the response is a JSON-RPC envelope +with a non-empty top-level `error` field, or a `result.isError: true` +tool error, the proxy treats it as an auth-failure signal: it calls +`resolveOnce`, re-sends the original request bytes (keyed by JSON-RPC +`id`), and waits for the response with the same id before forwarding +to the client. Any other responses that arrive during the retry wait +are forwarded as-is. After the first attempt (success or failure), no +further body inspection happens; subsequent errors propagate verbatim. +Free-form text inside `result.content[].text` is never matched against +keywords - only structural signals (`error`, `isError`) trigger the +retry. + Each mode resolves its secrets exactly once - lazily, but not repeatedly. Subsequent client requests reuse the env vars or HTTP headers established on the first message. diff --git a/docs/configuration.md b/docs/configuration.md index 461c763..abbe5ec 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -256,6 +256,14 @@ first request. The auth-deferral is automatic and has no config knob; servers that do not require auth on the MCP handshake therefore never trigger a vault prompt for that connection. +The same one-shot resolution also fires on **JSON-RPC body errors**: +if the remote server returns HTTP `200 OK` but the response carries a +JSON-RPC `error` field or a tool-level `result.isError: true`, +locksmith treats it as an auth-failure signal, resolves the templated +headers, and retries the failing request once. If the retry also +fails, the response is forwarded to the AI client unchanged. No +configuration knob - the behaviour is automatic. + ### mcp.servers.\.command **Required (local mode).** List of strings: executable followed by From c80677eb8d81dc4d262de3e61eb14e1d5a29df72 Mon Sep 17 00:00:00 2001 From: Lorem Dev Date: Fri, 15 May 2026 22:41:37 +0200 Subject: [PATCH 7/7] feat(mcp): log retry-wait shutdown and clarify clear() doc --- internal/mcp/proxy.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/internal/mcp/proxy.go b/internal/mcp/proxy.go index 4163f59..3f57e8c 100644 --- a/internal/mcp/proxy.go +++ b/internal/mcp/proxy.go @@ -351,6 +351,7 @@ func handleBodyError( for { next, ok := recv() if !ok { + log.Debug().Str("id", id).Msg("mcp proxy: retry wait interrupted by shutdown") return true } forward(next) @@ -436,7 +437,10 @@ func (s *proxyState) take(id string) []byte { } // clear releases the inFlight map. Subsequent record/take calls are -// no-ops. Called once authState.Attempted() flips to true. +// no-ops. Called from the body-error retry handshake on each exit +// branch (resolve failure, send failure, retry completed) - once any +// of those fires, authState.Attempted() is true and tracking is no +// longer needed. func (s *proxyState) clear() { s.mu.Lock() defer s.mu.Unlock()