[Server] Send progress with the token of the gateway's own request

CHANGELOG.md

@@ -6,6 +6,7 @@ All notable changes to `mcp/sdk` will be documented in this file.
 -----
 
 * Fix concurrent requests of one session overwriting each other's session changes: `Session::save()` now writes only the keys it changed, onto what the store holds at that point.
+* Fix `ClientGateway::progress()` sending another call's progress token when calls of one session run in parallel: the token now comes from the gateway's own request.
 * [BC Break] `SchemaValidator` takes an optional `Opis\JsonSchema\Validator` as its first constructor argument, moving `$logger` to second. Pass `logger:` by name.
 * [BC Break] `SchemaValidator::validateAgainstJsonSchema()` no longer validates an empty array as an object. Pass `new \stdClass()` for an empty object.
 * Add `Builder::setSchemaValidator()` to configure the validator used for `tools/call` input and output, e.g. with a resolver for external `$ref` schemas.

src/Capability/Registry/InjectableParameters.php

@@ -55,7 +55,7 @@ public static function resolve(string $typeName, array $arguments): ?object
         }
 
         if (ClientGateway::class === $typeName && isset($arguments['_session'])) {
-            return new ClientGateway($arguments['_session']);
+            return new ClientGateway($arguments['_session'], $arguments['_request'] ?? null);
         }
 
         return null;

src/Capability/Registry/Loader/ExplicitElementLoader.php

@@ -54,7 +54,7 @@ public function load(RegistryInterface $registry): void
             $handler = $entry['handler'];
             $registry->registerTool($entry['definition'], $this->boundClosure(
                 static function (array $arguments) use ($handler): mixed {
-                    $gateway = new ClientGateway($arguments['_session']);
+                    $gateway = new ClientGateway($arguments['_session'], $arguments['_request'] ?? null);
                     unset($arguments['_session'], $arguments['_request']);
 
                     return $handler->execute($arguments, $gateway);
@@ -67,7 +67,7 @@ static function (array $arguments) use ($handler): mixed {
             $registry->registerResource($entry['definition'], $this->boundClosure(
                 static fn (array $arguments): mixed => $handler->read(
                     $arguments['uri'],
-                    new ClientGateway($arguments['_session']),
+                    new ClientGateway($arguments['_session'], $arguments['_request'] ?? null),
                 ),
             ));
         }
@@ -76,7 +76,7 @@ static function (array $arguments) use ($handler): mixed {
             $handler = $entry['handler'];
             $registry->registerResourceTemplate($entry['definition'], $this->boundClosure(
                 static function (array $arguments) use ($handler): mixed {
-                    $gateway = new ClientGateway($arguments['_session']);
+                    $gateway = new ClientGateway($arguments['_session'], $arguments['_request'] ?? null);
                     $uri = $arguments['uri'];
                     unset($arguments['_session'], $arguments['_request'], $arguments['uri']);
 
@@ -89,7 +89,7 @@ static function (array $arguments) use ($handler): mixed {
             $handler = $entry['handler'];
             $registry->registerPrompt($entry['definition'], $this->boundClosure(
                 static function (array $arguments) use ($handler): mixed {
-                    $gateway = new ClientGateway($arguments['_session']);
+                    $gateway = new ClientGateway($arguments['_session'], $arguments['_request'] ?? null);
                     unset($arguments['_session'], $arguments['_request']);
 
                     return $handler->get($arguments, $gateway);

src/Server/ClientGateway.php

@@ -80,8 +80,12 @@
  */
 class ClientGateway
 {
+    /**
+     * @param Request|null $request the request the handler serves; its meta carries the progress token
+     */
     public function __construct(
         private readonly SessionInterface $session,
+        private readonly ?Request $request = null,
     ) {
     }
 
@@ -122,7 +126,11 @@ public function log(LoggingLevel $level, mixed $data, ?string $logger = null): v
      */
     public function progress(float $progress, ?float $total = null, ?string $message = null): void
     {
-        $meta = $this->session->get(Protocol::SESSION_ACTIVE_REQUEST_META, []);
+        // The session holds the meta of whichever request of the session came last,
+        // so it only stands in when the gateway does not know its own request.
+        $meta = null !== $this->request
+            ? $this->request->getMeta()
+            : $this->session->get(Protocol::SESSION_ACTIVE_REQUEST_META, []);
         $progressToken = $meta['progressToken'] ?? null;
 
         if (null === $progressToken) {

src/Server/RequestContext.php

@@ -94,7 +94,7 @@ public function getAccessToken(): ?AccessToken
     public function getClientGateway(): ClientGateway
     {
         if (null == $this->clientGateway) {
-            $this->clientGateway = new ClientGateway($this->session);
+            $this->clientGateway = new ClientGateway($this->session, $this->request);
         }
 
         return $this->clientGateway;

tests/Unit/Capability/Registry/Loader/ExplicitElementLoaderTest.php

@@ -18,6 +18,7 @@
 use Mcp\Capability\RegistryInterface;
 use Mcp\Exception\InvalidArgumentException;
 use Mcp\Schema\Prompt;
+use Mcp\Schema\Request\PingRequest;
 use Mcp\Schema\ResourceDefinition;
 use Mcp\Schema\ResourceTemplate;
 use Mcp\Schema\Tool;
@@ -68,7 +69,7 @@ public function execute(array $arguments, ClientGateway $gateway): mixed
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'foo' => 'bar',
         ]);
 
@@ -111,7 +112,7 @@ public function read(string $uri, ClientGateway $gateway): mixed
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'uri' => 'config://demo',
         ]);
 
@@ -154,7 +155,7 @@ public function read(string $uri, array $variables, ClientGateway $gateway): mix
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'uri' => 'config://abc',
             'key' => 'abc',
         ]);
@@ -197,7 +198,7 @@ public function get(array $arguments, ClientGateway $gateway): mixed
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'topic' => 'php',
         ]);
 

tests/Unit/Capability/Registry/ReferenceHandlerTest.php

@@ -15,6 +15,7 @@
 use Mcp\Capability\Registry\ReferenceHandler;
 use Mcp\Exception\InvalidArgumentException;
 use Mcp\Exception\RegistryException;
+use Mcp\Schema\Request\PingRequest;
 use Mcp\Server\ClientGateway;
 use Mcp\Server\Handler\ResourceHandlerInterface;
 use Mcp\Server\Handler\ToolHandlerInterface;
@@ -57,7 +58,7 @@ static function (array $arguments) use ($toolHandler): mixed {
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'kept' => 'value',
             'other' => 'value2',
         ]);
@@ -100,7 +101,7 @@ public function read(string $uri, ClientGateway $gateway): mixed
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'uri' => 'config://x',
         ]);
 
@@ -127,7 +128,7 @@ public function search(string $query, ClientGateway $gateway): string
 
         $result = (new ReferenceHandler())->handle(new ElementReference([$handler, 'search']), [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'query' => 'foo',
         ]);
 
@@ -142,7 +143,7 @@ public function testHandleResolvesAClassStringAndMethodPairToAnInstance(): void
 
         $result = (new ReferenceHandler())->handle(new ElementReference([ArrayHandlerFixture::class, 'greet']), [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'who' => 'bob',
         ]);
 
@@ -158,7 +159,7 @@ public function testHandleRefusesAnArrayHandlerWhoseMethodIsNotCallable(): void
 
         (new ReferenceHandler())->handle(new ElementReference([new ArrayHandlerFixture(), 'hidden']), [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
         ]);
     }
 
@@ -177,7 +178,7 @@ public function testHandleStillReflectsOrdinaryClosuresAndDoesNotInjectArgumentB
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $session,
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'kept' => 'value',
         ]);
 
@@ -269,7 +270,7 @@ public function testHandleKeepsArgumentOrderWhenAnInjectableParameterPrecedesAVa
 
         $result = (new ReferenceHandler())->handle($reference, [
             '_session' => $this->createMock(SessionInterface::class),
-            '_request' => new \stdClass(),
+            '_request' => new PingRequest(),
             'parts' => ['a', 'b'],
         ]);
 

tests/Unit/Server/ClientGatewayTest.php

@@ -24,8 +24,10 @@
 use Mcp\Schema\JsonRpc\Request;
 use Mcp\Schema\JsonRpc\Response;
 use Mcp\Schema\Notification\LoggingMessageNotification;
+use Mcp\Schema\Notification\ProgressNotification;
 use Mcp\Schema\Request\ElicitRequest;
 use Mcp\Schema\Request\ListRootsRequest;
+use Mcp\Schema\Request\PingRequest;
 use Mcp\Schema\Result\ElicitResult;
 use Mcp\Schema\Result\ListRootsResult;
 use Mcp\Server\ClientGateway;
@@ -213,6 +215,62 @@ public function testLogDefaultsToWarningWhenTheClientSetNoLevel(): void
         $this->assertInstanceOf(LoggingMessageNotification::class, $this->logInFiber(null, LoggingLevel::Warning));
     }
 
+    public function testProgressUsesTheTokenOfItsOwnRequest(): void
+    {
+        $session = $this->createMock(SessionInterface::class);
+        $session->method('getId')->willReturn(Uuid::v4());
+        // The session holds the meta of whichever call of the session came last.
+        $session->method('get')->willReturnCallback(
+            static fn (string $key, mixed $default = null): mixed => Protocol::SESSION_ACTIVE_REQUEST_META === $key ? ['progressToken' => 'token-b'] : $default,
+        );
+
+        $gateway = new ClientGateway($session, (new PingRequest())->withId(1)->withMeta(['progressToken' => 'token-a']));
+
+        $this->assertSame('token-a', $this->progressInFiber($gateway)?->progressToken);
+    }
+
+    public function testProgressIsDroppedWhenItsRequestAskedForNone(): void
+    {
+        $session = $this->createMock(SessionInterface::class);
+        $session->method('getId')->willReturn(Uuid::v4());
+        $session->method('get')->willReturnCallback(
+            static fn (string $key, mixed $default = null): mixed => Protocol::SESSION_ACTIVE_REQUEST_META === $key ? ['progressToken' => 'token-b'] : $default,
+        );
+
+        $gateway = new ClientGateway($session, (new PingRequest())->withId(1));
+
+        $this->assertNull($this->progressInFiber($gateway));
+    }
+
+    public function testProgressFallsBackToTheSessionWithoutARequest(): void
+    {
+        $session = $this->createMock(SessionInterface::class);
+        $session->method('getId')->willReturn(Uuid::v4());
+        $session->method('get')->willReturnCallback(
+            static fn (string $key, mixed $default = null): mixed => Protocol::SESSION_ACTIVE_REQUEST_META === $key ? ['progressToken' => 'token-b'] : $default,
+        );
+
+        $this->assertSame('token-b', $this->progressInFiber(new ClientGateway($session))?->progressToken);
+    }
+
+    /**
+     * Reports progress inside a Fiber and returns the notification the gateway suspended with, if any.
+     */
+    private function progressInFiber(ClientGateway $gateway): ?ProgressNotification
+    {
+        $fiber = new \Fiber(static fn () => $gateway->progress(1, 2));
+        $suspend = $fiber->start();
+
+        if (null === $suspend) {
+            return null;
+        }
+
+        $this->assertInstanceOf(NotificationSuspension::class, $suspend);
+        $this->assertInstanceOf(ProgressNotification::class, $suspend->notification);
+
+        return $suspend->notification;
+    }
+
     /**
      * Logs inside a Fiber and returns the notification the gateway suspended with, if any.
      */

tests/Unit/Server/ProtocolTest.php

@@ -24,12 +24,14 @@
 use Mcp\Schema\JsonRpc\Request;
 use Mcp\Schema\JsonRpc\Response;
 use Mcp\Schema\Notification\LoggingMessageNotification;
+use Mcp\Schema\Notification\ProgressNotification;
 use Mcp\Schema\Request\CallToolRequest;
 use Mcp\Schema\Request\PingRequest;
 use Mcp\Server\ClientGateway;
 use Mcp\Server\Handler\Notification\NotificationHandlerInterface;
 use Mcp\Server\Handler\Request\RequestHandlerInterface;
 use Mcp\Server\Protocol;
+use Mcp\Server\RequestContext;
 use Mcp\Server\Session\InMemorySessionStore;
 use Mcp\Server\Session\Psr16SessionStore;
 use Mcp\Server\Session\Session;
@@ -2218,6 +2220,67 @@ public function testSuspendedFibersDispatchResultUnderTheirOwnRequest(): void
         $this->assertSame(1, $results[1]->getRequest()->getId());
     }
 
+    #[TestDox('Progress of a suspended call keeps its own token while a parallel call of the session saves in between')]
+    public function testProgressKeepsItsTokenAcrossAParallelCall(): void
+    {
+        $handler = $this->createMock(RequestHandlerInterface::class);
+        $handler->method('supports')->willReturn(true);
+        $handler->method('handle')->willReturnCallback(static function (Request $request, SessionInterface $session): Response {
+            $gateway = (new RequestContext($session, $request))->getClientGateway();
+            $gateway->progress(1, 2);
+            $gateway->progress(2, 2);
+
+            return new Response($request->getId(), []);
+        });
+
+        // Runs call B right after call A's first write, between the save after
+        // A's fiber suspended and the one closing A's processInput().
+        $store = new class extends InMemorySessionStore {
+            public ?\Closure $afterWrite = null;
+
+            public function write(Uuid $id, string $data): bool
+            {
+                $written = parent::write($id, $data);
+                $hook = $this->afterWrite;
+                $this->afterWrite = null;
+                $hook?->__invoke();
+
+                return $written;
+            }
+        };
+
+        $session = (new SessionManager($store))->create();
+        $session->save();
+
+        $fibers = [];
+        $protocols = [];
+        $transports = [];
+        foreach (['a', 'b'] as $call) {
+            $transports[$call] = $this->createMock(TransportInterface::class);
+            $transports[$call]->method('attachFiberToSession')->willReturnCallback(static function (\Fiber $fiber) use (&$fibers, $call): void {
+                $fibers[$call] = $fiber;
+            });
+            $protocols[$call] = new Protocol(
+                requestHandlers: [$handler],
+                notificationHandlers: [],
+                messageFactory: MessageFactory::make(),
+                sessionManager: new SessionManager($store),
+            );
+        }
+
+        $store->afterWrite = static fn () => $protocols['b']->processInput($transports['b'], '{"jsonrpc": "2.0", "id": 2, "method": "ping", "params": {"_meta": {"progressToken": "token-b"}}}', $session->getId());
+        $protocols['a']->processInput($transports['a'], '{"jsonrpc": "2.0", "id": 1, "method": "ping", "params": {"_meta": {"progressToken": "token-a"}}}', $session->getId());
+
+        $this->assertCount(2, $fibers);
+
+        foreach (['a' => 'token-a', 'b' => 'token-b'] as $call => $token) {
+            $suspension = $fibers[$call]->resume();
+            $this->assertInstanceOf(NotificationSuspension::class, $suspension);
+            $this->assertInstanceOf(ProgressNotification::class, $suspension->notification);
+            $this->assertSame($token, $suspension->notification->progressToken, \sprintf('Call %s sent its second progress step with the wrong token.', strtoupper($call)));
+        }
+    }
+
     #[TestDox('A throwing result listener of a resumed fiber is answered with an internal error')]
     public function testResumedFiberTurnsThrowingResponseListenerIntoInternalError(): void
     {