[Server] Stop concurrent requests from overwriting each other's session changes

CHANGELOG.md

@@ -5,6 +5,7 @@ All notable changes to `mcp/sdk` will be documented in this file.
 0.9.0
 -----
 
+* 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.
 * [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/Server/Session/Session.php

@@ -15,6 +15,9 @@
 use Symfony\Component\Uid\UuidV4;
 
 /**
+ * Saving writes only the keys changed since the last save, onto what the store holds by then:
+ * concurrent requests of one session keep each other's changes as long as they change different keys.
+ *
  * @author Kyrian Obikwelu <koshnawaza@gmail.com>
  */
 class Session implements SessionInterface
@@ -31,6 +34,13 @@ class Session implements SessionInterface
      */
     private array $data;
 
+    /**
+     * Keys changed since the last save, or null once clear() or hydrate() replaced all of them.
+     *
+     * @var array<string, true>|null
+     */
+    private ?array $changes = [];
+
     public function __construct(
         private SessionStoreInterface $store,
         private Uuid $id = new UuidV4(),
@@ -44,7 +54,23 @@ public function getId(): Uuid
 
     public function save(): bool
     {
-        return $this->store->write($this->id, json_encode($this->readData(), \JSON_THROW_ON_ERROR));
+        $data = $this->readData();
+
+        if (null !== $this->changes) {
+            $data = $this->load();
+            foreach (array_keys($this->changes) as $key) {
+                $this->apply($data, (string) $key);
+            }
+        }
+
+        if (!$this->store->write($this->id, json_encode($data, \JSON_THROW_ON_ERROR))) {
+            return false;
+        }
+
+        $this->data = $data;
+        $this->changes = [];
+
+        return true;
     }
 
     public function get(string $key, mixed $default = null): mixed
@@ -79,6 +105,7 @@ public function set(string $key, mixed $value, bool $overwrite = true): void
 
         if ($overwrite || !isset($data[$lastKey])) {
             $data[$lastKey] = $value;
+            $this->trackChange($key);
         }
     }
 
@@ -114,12 +141,16 @@ public function forget(string $key): void
             $data = &$data[$segment];
         }
 
-        unset($data[$lastKey]);
+        if (\array_key_exists($lastKey, $data)) {
+            unset($data[$lastKey]);
+            $this->trackChange($key);
+        }
     }
 
     public function clear(): void
     {
         $this->data = [];
+        $this->changes = null;
     }
 
     public function pull(string $key, mixed $default = null): mixed
@@ -138,6 +169,7 @@ public function all(): array
     public function hydrate(array $attributes): void
     {
         $this->data = $attributes;
+        $this->changes = null;
     }
 
     /** @return array<string, mixed> */
@@ -151,23 +183,61 @@ public function jsonSerialize(): array
      */
     private function readData(): array
     {
-        if (isset($this->data)) {
-            return $this->data;
-        }
+        return $this->data ??= $this->load();
+    }
 
+    /**
+     * @return array<string, mixed>
+     */
+    private function load(): array
+    {
         $rawData = $this->store->read($this->id);
 
         if (false === $rawData) {
-            return $this->data = [];
+            return [];
         }
 
         // Empty session should not throw
         $decoded = json_decode($rawData, true);
 
-        if (!\is_array($decoded)) {
-            return $this->data = [];
+        return \is_array($decoded) ? $decoded : [];
+    }
+
+    private function trackChange(string $key): void
+    {
+        if (null !== $this->changes) {
+            $this->changes[$key] = true;
         }
+    }
 
-        return $this->data = $decoded;
+    /**
+     * Carries the current value of a changed key over to $data, or removes it there if it was forgotten.
+     *
+     * @param array<string, mixed> $data
+     */
+    private function apply(array &$data, string $key): void
+    {
+        $segments = explode('.', $key);
+        $lastKey = array_pop($segments);
+        $source = $this->data;
+        $target = &$data;
+
+        foreach ($segments as $segment) {
+            $source = \is_array($source) ? $source[$segment] ?? null : null;
+
+            if (!isset($target[$segment]) || !\is_array($target[$segment])) {
+                if (!\is_array($source)) {
+                    return;
+                }
+                $target[$segment] = [];
+            }
+            $target = &$target[$segment];
+        }
+
+        if (\is_array($source) && \array_key_exists($lastKey, $source)) {
+            $target[$lastKey] = $source[$lastKey];
+        } else {
+            unset($target[$lastKey]);
+        }
     }
 }

tests/Unit/Server/Session/SessionTest.php

@@ -365,4 +365,120 @@ public function testAllReturnsEmptyArrayForUndecodablePayload(): void
 
         $this->assertSame([], $session->all());
     }
+
+    public function testSaveKeepsWhatAConcurrentSessionSaved(): void
+    {
+        $id = new UuidV4();
+        $this->store->write($id, json_encode(['initialized' => true], \JSON_THROW_ON_ERROR));
+
+        $first = new Session($this->store, $id);
+        $second = new Session($this->store, $id);
+        $first->get('initialized');
+        $second->get('initialized');
+
+        $first->set('_mcp.logging_level', 'debug');
+        $first->save();
+        $second->set('counter', 1);
+        $second->save();
+
+        $this->assertSame(
+            ['initialized' => true, '_mcp' => ['logging_level' => 'debug'], 'counter' => 1],
+            (new Session($this->store, $id))->all(),
+        );
+    }
+
+    public function testSaveMergesNestedKeysOfConcurrentSessions(): void
+    {
+        $id = new UuidV4();
+        $first = new Session($this->store, $id);
+        $second = new Session($this->store, $id);
+        $first->all();
+        $second->all();
+
+        $first->set('_mcp.responses.1', ['result' => 'a']);
+        $first->save();
+        $second->set('_mcp.responses.2', ['result' => 'b']);
+        $second->save();
+
+        $this->assertSame(
+            ['1' => ['result' => 'a'], '2' => ['result' => 'b']],
+            (new Session($this->store, $id))->get('_mcp.responses'),
+        );
+    }
+
+    public function testSaveAppliesForgetOntoConcurrentChanges(): void
+    {
+        $id = new UuidV4();
+        $this->store->write($id, json_encode(['stale' => true], \JSON_THROW_ON_ERROR));
+
+        $first = new Session($this->store, $id);
+        $second = new Session($this->store, $id);
+        $first->all();
+        $second->all();
+
+        $first->forget('stale');
+        $first->save();
+        $second->set('fresh', true);
+        $second->save();
+
+        $this->assertSame(['fresh' => true], (new Session($this->store, $id))->all());
+    }
+
+    public function testSaveDoesNotWriteBackUnchangedKeys(): void
+    {
+        $id = new UuidV4();
+        $this->store->write($id, json_encode(['level' => 'info'], \JSON_THROW_ON_ERROR));
+
+        $first = new Session($this->store, $id);
+        $second = new Session($this->store, $id);
+        $first->all();
+        $second->all();
+
+        $first->set('level', 'debug');
+        $first->save();
+        $second->set('level', 'error', overwrite: false);
+        $second->save();
+
+        $this->assertSame('debug', (new Session($this->store, $id))->get('level'));
+    }
+
+    public function testSaveSeesConcurrentChangesAfterwards(): void
+    {
+        $id = new UuidV4();
+        $first = new Session($this->store, $id);
+        $second = new Session($this->store, $id);
+        $first->all();
+
+        $second->set('other', true);
+        $second->save();
+        $first->set('mine', true);
+        $first->save();
+
+        $this->assertSame(['other' => true, 'mine' => true], $first->all());
+    }
+
+    public function testSaveAfterClearReplacesStoredData(): void
+    {
+        $id = new UuidV4();
+        $this->store->write($id, json_encode(['stale' => true], \JSON_THROW_ON_ERROR));
+
+        $session = new Session($this->store, $id);
+        $session->clear();
+        $session->set('fresh', true);
+        $session->save();
+
+        $this->assertSame(['fresh' => true], (new Session($this->store, $id))->all());
+    }
+
+    public function testSaveAfterHydrateReplacesStoredData(): void
+    {
+        $id = new UuidV4();
+        $this->store->write($id, json_encode(['stale' => true], \JSON_THROW_ON_ERROR));
+
+        $session = new Session($this->store, $id);
+        $session->hydrate(['fresh' => true]);
+        $session->save();
+
+        $this->assertSame(['fresh' => true], (new Session($this->store, $id))->all());
+    }
 }