refactor: make progress token dependent on request ID

Author: CodeWithKyrian

src/Client/Client.php

@@ -54,7 +54,6 @@ class Client
     private Protocol $protocol;
     private ClientSessionInterface $session;
     private ?ClientTransportInterface $transport = null;
-    private int $progressTokenCounter = 0;
 
     /**
      * @param NotificationHandlerInterface[] $notificationHandlers
@@ -143,11 +142,8 @@ public function listTools(?string $cursor = null): ListToolsResult
      *        Optional callback for progress updates. If provided, a progress token
      *        is automatically generated and attached to the request.
      */
-    public function callTool(
-        string $name,
-        array $arguments = [],
-        ?callable $onProgress = null,
-    ): CallToolResult {
+    public function callTool(string $name, array $arguments = [], ?callable $onProgress = null): CallToolResult
+    {
         $this->ensureConnected();
 
         $request = new CallToolRequest($name, $arguments);
@@ -209,11 +205,8 @@ public function listPrompts(?string $cursor = null): ListPromptsResult
      * @param (callable(float $progress, ?float $total, ?string $message): void)|null $onProgress
      *        Optional callback for progress updates.
      */
-    public function getPrompt(
-        string $name,
-        array $arguments = [],
-        ?callable $onProgress = null,
-    ): GetPromptResult {
+    public function getPrompt(string $name, array $arguments = [], ?callable $onProgress = null): GetPromptResult
+    {
         $this->ensureConnected();
 
         $request = new GetPromptRequest($name, $arguments);
@@ -261,10 +254,13 @@ public function disconnect(): void
      *
      * @throws RequestException
      */
-    private function doRequest(object $request, ?string $resultClass = null, ?callable $onProgress = null): mixed
+    private function doRequest(Request $request, ?string $resultClass = null, ?callable $onProgress = null): mixed
     {
-        if (null !== $onProgress && $request instanceof Request) {
-            $progressToken = $this->generateProgressToken();
+        $requestId = $this->session->nextRequestId();
+        $request = $request->withId($requestId);
+
+        if (null !== $onProgress) {
+            $progressToken = 'prog-' . $requestId;
             $request = $request->withMeta(['progressToken' => $progressToken]);
         }
 
@@ -283,14 +279,6 @@ private function doRequest(object $request, ?string $resultClass = null, ?callab
         return $resultClass::fromArray($response->result);
     }
 
-    /**
-     * Generate a unique progress token for a request.
-     */
-    private function generateProgressToken(): string
-    {
-        return 'prog-' . (++$this->progressTokenCounter);
-    }
-
     private function ensureConnected(): void
     {
         if (!$this->isConnected()) {

src/Client/Protocol.php

@@ -66,7 +66,7 @@ public function connect(ClientTransportInterface $transport): void
     {
         $this->transport = $transport;
         $transport->setSession($this->session);
-        $transport->onInitialize(fn() => $this->performInitialize());
+        $transport->onInitialize($this->initialize(...));
         $transport->onMessage($this->processMessage(...));
         $transport->onError(fn(\Throwable $e) => $this->logger->error('Transport error', ['exception' => $e]));
 
@@ -80,14 +80,17 @@ public function connect(ClientTransportInterface $transport): void
      *
      * @return Response<array<string, mixed>>|Error
      */
-    public function performInitialize(): Response|Error
+    public function initialize(): Response|Error
     {
         $request = new InitializeRequest(
             $this->config->protocolVersion,
             $this->config->capabilities,
             $this->config->clientInfo,
         );
 
+        $requestId = $this->session->nextRequestId();
+        $request = $request->withId($requestId);
+
         $response = $this->request($request, $this->config->initTimeout);
 
         if ($response instanceof Response) {
@@ -114,15 +117,14 @@ public function performInitialize(): Response|Error
      */
     public function request(Request $request, int $timeout): Response|Error
     {
-        $requestId = $this->session->nextRequestId();
-        $requestWithId = $request->withId($requestId);
+        $requestId = $request->getId();
 
         $this->logger->debug('Sending request', [
             'id' => $requestId,
             'method' => $request::getMethod(),
         ]);
 
-        $encoded = json_encode($requestWithId, \JSON_THROW_ON_ERROR);
+        $encoded = json_encode($request, \JSON_THROW_ON_ERROR);
         $this->session->queueOutgoing($encoded, ['type' => 'request']);
         $this->session->addPendingRequest($requestId, $timeout);