Merge branch 'main' into mcp-server-req-auth-2

Author: jba

.github/workflows/test.yml

@@ -3,7 +3,7 @@ on:
   # Manual trigger
   workflow_dispatch:
   push:
-    branches: main
+    branches: [main]
   pull_request:
 
 permissions:
@@ -13,43 +13,51 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-    - name: Check out code
-      uses: actions/checkout@v4
-    - name: Set up Go
-      uses: actions/setup-go@v5
-    - name: Check formatting
-      run: |
-        unformatted=$(gofmt -l .)
-        if [ -n "$unformatted" ]; then
-          echo "The following files are not properly formatted:"
-          echo "$unformatted"
-          exit 1
-        fi
-        echo "All Go files are properly formatted"
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Set up Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: "^1.23"
+      - name: Check formatting
+        run: |
+          unformatted=$(gofmt -l .)
+          if [ -n "$unformatted" ]; then
+            echo "The following files are not properly formatted:"
+            echo "$unformatted"
+            exit 1
+          fi
+          echo "All Go files are properly formatted"
+      - name: Run Go vet
+        run: go vet ./...
+      - name: Run staticcheck
+        uses: dominikh/staticcheck-action@v1
+        with:
+          version: "latest"
 
   test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        go: [ '1.23', '1.24', '1.25.0-rc.3' ]
+        go: ["1.23", "1.24", "1.25.0-rc.3"]
     steps:
-    - name: Check out code
-      uses: actions/checkout@v4
-    - name: Set up Go
-      uses: actions/setup-go@v5
-      with:
-        go-version: ${{ matrix.go }}
-    - name: Test
-      run: go test -v ./...
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Set up Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ matrix.go }}
+      - name: Test
+        run: go test -v ./...
 
   race-test:
     runs-on: ubuntu-latest
     steps:
-    - name: Check out code
-      uses: actions/checkout@v4
-    - name: Set up Go
-      uses: actions/setup-go@v5
-      with:
-        go-version: '1.24'
-    - name: Test with -race
-      run: go test -v -race ./...
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Set up Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: "1.24"
+      - name: Test with -race
+        run: go test -v -race ./...

internal/jsonrpc2/conn.go

@@ -374,6 +374,46 @@ func (c *Connection) Call(ctx context.Context, method string, params any) *Async
 	return ac
 }
 
+// Async, signals that the current jsonrpc2 request may be handled
+// asynchronously to subsequent requests, when ctx is the request context.
+//
+// Async must be called at most once on each request's context (and its
+// descendants).
+func Async(ctx context.Context) {
+	if r, ok := ctx.Value(asyncKey).(*releaser); ok {
+		r.release(false)
+	}
+}
+
+type asyncKeyType struct{}
+
+var asyncKey = asyncKeyType{}
+
+// A releaser implements concurrency safe 'releasing' of async requests. (A
+// request is released when it is allowed to run concurrent with other
+// requests, via a call to [Async].)
+type releaser struct {
+	mu       sync.Mutex
+	ch       chan struct{}
+	released bool
+}
+
+// release closes the associated channel. If soft is set, multiple calls to
+// release are allowed.
+func (r *releaser) release(soft bool) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if r.released {
+		if !soft {
+			panic("jsonrpc2.Async called multiple times")
+		}
+	} else {
+		close(r.ch)
+		r.released = true
+	}
+}
+
 type AsyncCall struct {
 	id       ID
 	ready    chan struct{} // closed after response has been set
@@ -425,28 +465,6 @@ func (ac *AsyncCall) Await(ctx context.Context, result any) error {
 	return json.Unmarshal(ac.response.Result, result)
 }
 
-// Respond delivers a response to an incoming Call.
-//
-// Respond must be called exactly once for any message for which a handler
-// returns ErrAsyncResponse. It must not be called for any other message.
-func (c *Connection) Respond(id ID, result any, err error) error {
-	var req *incomingRequest
-	c.updateInFlight(func(s *inFlightState) {
-		req = s.incomingByID[id]
-	})
-	if req == nil {
-		return c.internalErrorf("Request not found for ID %v", id)
-	}
-
-	if err == ErrAsyncResponse {
-		// Respond is supposed to supply the asynchronous response, so it would be
-		// confusing to call Respond with an error that promises to call Respond
-		// again.
-		err = c.internalErrorf("Respond called with ErrAsyncResponse for %q", req.Method)
-	}
-	return c.processResult("Respond", req, result, err)
-}
-
 // Cancel cancels the Context passed to the Handle call for the inbound message
 // with the given ID.
 //
@@ -576,11 +594,6 @@ func (c *Connection) acceptRequest(ctx context.Context, msg *Request, preempter
 	if preempter != nil {
 		result, err := preempter.Preempt(req.ctx, req.Request)
 
-		if req.IsCall() && errors.Is(err, ErrAsyncResponse) {
-			// This request will remain in flight until Respond is called for it.
-			return
-		}
-
 		if !errors.Is(err, ErrNotHandled) {
 			c.processResult("Preempt", req, result, err)
 			return
@@ -655,19 +668,20 @@ func (c *Connection) handleAsync() {
 			continue
 		}
 
-		result, err := c.handler.Handle(req.ctx, req.Request)
-		c.processResult(c.handler, req, result, err)
+		releaser := &releaser{ch: make(chan struct{})}
+		ctx := context.WithValue(req.ctx, asyncKey, releaser)
+		go func() {
+			defer releaser.release(true)
+			result, err := c.handler.Handle(ctx, req.Request)
+			c.processResult(c.handler, req, result, err)
+		}()
+		<-releaser.ch
 	}
 }
 
 // processResult processes the result of a request and, if appropriate, sends a response.
 func (c *Connection) processResult(from any, req *incomingRequest, result any, err error) error {
 	switch err {
-	case ErrAsyncResponse:
-		if !req.IsCall() {
-			return c.internalErrorf("%#v returned ErrAsyncResponse for a %q Request without an ID", from, req.Method)
-		}
-		return nil // This request is still in flight, so don't record the result yet.
 	case ErrNotHandled, ErrMethodNotFound:
 		// Add detail describing the unhandled method.
 		err = fmt.Errorf("%w: %q", ErrMethodNotFound, req.Method)
@@ -705,10 +719,10 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
 		} else if err != nil {
 			err = fmt.Errorf("%w: %q notification failed: %v", ErrInternal, req.Method, err)
 		}
-		if err != nil {
-			// TODO: can/should we do anything with this error beyond writing it to the event log?
-			// (Is this the right label to attach to the log?)
-		}
+	}
+	if err != nil {
+		// TODO: can/should we do anything with this error beyond writing it to the event log?
+		// (Is this the right label to attach to the log?)
 	}
 
 	// Cancel the request to free any associated resources.
@@ -725,7 +739,23 @@ func (c *Connection) processResult(from any, req *incomingRequest, result any, e
 // write is used by all things that write outgoing messages, including replies.
 // it makes sure that writes are atomic
 func (c *Connection) write(ctx context.Context, msg Message) error {
-	err := c.writer.Write(ctx, msg)
+	var err error
+	// Fail writes immediately if the connection is shutting down.
+	//
+	// TODO(rfindley): should we allow cancellation notifications through? It
+	// could be the case that writes can still succeed.
+	c.updateInFlight(func(s *inFlightState) {
+		err = s.shuttingDown(ErrServerClosing)
+	})
+	if err == nil {
+		err = c.writer.Write(ctx, msg)
+	}
+
+	// For rejected requests, we don't set the writeErr (which would break the
+	// connection). They can just be returned to the caller.
+	if errors.Is(err, ErrRejected) {
+		return err
+	}
 
 	if err != nil && ctx.Err() == nil {
 		// The call to Write failed, and since ctx.Err() is nil we can't attribute

internal/jsonrpc2/jsonrpc2.go

@@ -22,13 +22,6 @@ var (
 	// If a Handler returns ErrNotHandled, the server replies with
 	// ErrMethodNotFound.
 	ErrNotHandled = errors.New("JSON RPC not handled")
-
-	// ErrAsyncResponse is returned from a handler to indicate it will generate a
-	// response asynchronously.
-	//
-	// ErrAsyncResponse must not be returned for notifications,
-	// which do not receive responses.
-	ErrAsyncResponse = errors.New("JSON RPC asynchronous response")
 )
 
 // Preempter handles messages on a connection before they are queued to the main

internal/jsonrpc2/jsonrpc2_test.go

@@ -371,16 +371,14 @@ func (h *handler) Handle(ctx context.Context, req *jsonrpc2.Request) (any, error
 		if err := json.Unmarshal(req.Params, &name); err != nil {
 			return nil, fmt.Errorf("%w: %s", jsonrpc2.ErrParse, err)
 		}
+		jsonrpc2.Async(ctx)
 		waitFor := h.waiter(name)
-		go func() {
-			select {
-			case <-waitFor:
-				h.conn.Respond(req.ID, true, nil)
-			case <-ctx.Done():
-				h.conn.Respond(req.ID, nil, ctx.Err())
-			}
-		}()
-		return nil, jsonrpc2.ErrAsyncResponse
+		select {
+		case <-waitFor:
+			return true, nil
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
 	default:
 		return nil, jsonrpc2.ErrNotHandled
 	}

internal/jsonrpc2/wire.go

@@ -37,6 +37,17 @@ var (
 	ErrServerClosing = NewError(-32004, "server is closing")
 	// ErrClientClosing is a dummy error returned for calls initiated while the client is closing.
 	ErrClientClosing = NewError(-32003, "client is closing")
+
+	// The following errors have special semantics for MCP transports
+
+	// ErrRejected may be wrapped to return errors from calls to Writer.Write
+	// that signal that the request was rejected by the transport layer as
+	// invalid.
+	//
+	// Such failures do not indicate that the connection is broken, but rather
+	// should be returned to the caller to indicate that the specific request is
+	// invalid in the current context.
+	ErrRejected = NewError(-32004, "rejected by transport")
 )
 
 const wireVersion = "2.0"

mcp/client.go

@@ -103,8 +103,7 @@ func (e unsupportedProtocolVersionError) Error() string {
 }
 
 // ClientSessionOptions is reserved for future use.
-type ClientSessionOptions struct {
-}
+type ClientSessionOptions struct{}
 
 // Connect begins an MCP session by connecting to a server over the given
 // transport, and initializing the session.
@@ -177,6 +176,8 @@ type clientSessionState struct {
 	InitializeResult *InitializeResult
 }
 
+func (cs *ClientSession) InitializeResult() *InitializeResult { return cs.state.InitializeResult }
+
 func (cs *ClientSession) ID() string {
 	if c, ok := cs.mcpConn.(hasSessionID); ok {
 		return c.SessionID()
@@ -323,16 +324,19 @@ func (cs *ClientSession) receivingMethodInfos() map[string]methodInfo {
 }
 
 func (cs *ClientSession) handle(ctx context.Context, req *jsonrpc.Request) (any, error) {
+	if req.IsCall() {
+		jsonrpc2.Async(ctx)
+	}
 	return handleReceive(ctx, cs, req)
 }
 
-func (cs *ClientSession) sendingMethodHandler() methodHandler {
+func (cs *ClientSession) sendingMethodHandler() MethodHandler {
 	cs.client.mu.Lock()
 	defer cs.client.mu.Unlock()
 	return cs.client.sendingMethodHandler_
 }
 
-func (cs *ClientSession) receivingMethodHandler() methodHandler {
+func (cs *ClientSession) receivingMethodHandler() MethodHandler {
 	cs.client.mu.Lock()
 	defer cs.client.mu.Unlock()
 	return cs.client.receivingMethodHandler_
@@ -392,7 +396,7 @@ func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) (
 	return handleSend[*CallToolResult](ctx, methodCallTool, newClientRequest(cs, orZero[Params](params)))
 }
 
-func (cs *ClientSession) SetLevel(ctx context.Context, params *SetLevelParams) error {
+func (cs *ClientSession) SetLoggingLevel(ctx context.Context, params *SetLoggingLevelParams) error {
 	_, err := handleSend[*emptyResult](ctx, methodSetLevel, newClientRequest(cs, orZero[Params](params)))
 	return err
 }

mcp/conformance_test.go

@@ -183,7 +183,7 @@ func runServerTest(t *testing.T, test *conformanceTest) {
 				return nil, err, false
 			}
 			serverMessages = append(serverMessages, msg)
-			if req, ok := msg.(*jsonrpc.Request); ok && req.ID.IsValid() {
+			if req, ok := msg.(*jsonrpc.Request); ok && req.IsCall() {
 				// Pair up the next outgoing response with this request.
 				// We assume requests arrive in the same order every time.
 				if len(outResponses) == 0 {
@@ -201,8 +201,8 @@ func runServerTest(t *testing.T, test *conformanceTest) {
 	// Synthetic peer interacts with real peer.
 	for _, req := range outRequests {
 		writeMsg(req)
-		if req.ID.IsValid() {
-			// A request (as opposed to a notification). Wait for the response.
+		if req.IsCall() {
+			// A call (as opposed to a notification). Wait for the response.
 			res, err, ok := nextResponse()
 			if err != nil {
 				t.Fatalf("reading server messages failed: %v", err)

mcp/content.go

@@ -253,7 +253,7 @@ func contentsFromWire(wires []*wireContent, allow map[string]bool) ([]Content, e
 
 func contentFromWire(wire *wireContent, allow map[string]bool) (Content, error) {
 	if wire == nil {
-		return nil, fmt.Errorf("content wire is nil")
+		return nil, fmt.Errorf("nil content")
 	}
 	if allow != nil && !allow[wire.Type] {
 		return nil, fmt.Errorf("invalid content type %q", wire.Type)