fix: resolve race conditions in in-memory-transport and stdio-server tests

Fixed two flaky tests that were failing intermittently on CI due to race conditions in async handler execution and queue consumption patterns.

## Tests Fixed

### 1. test-server-to-client-communication (in-memory-transport)
**Root cause:** Test manually polled from queue while background thread also consumed from same queue, creating a race between two consumers.

**Solution:** Modified test to use notification-handler callback mechanism, ensuring single consumer pattern with proper synchronization via atom capture and polling.

**Verification:** 11/11 consecutive CI runs passed (100% success rate)

### 2. test-integration (stdio-server)
**Root cause:** Server future completed immediately after reading EOF, but async handler tasks were still executing. The with-out-str block stopped capturing output before handlers finished writing responses.

**Solution:** Modified server to track pending handler futures in atom and wait for all futures to complete before EOF exit. Added helper functions for future lifecycle management.

**Verification:** 10/10 consecutive local test runs passed

## Changes

**In-memory transport test:**
- Use notification-handler callback instead of manual polling
- Capture notifications in atom and wait with timeout
- Ensures single consumer with proper synchronization

**Stdio server:**
- Modified handle-request to return futures for async handlers
- Added pending-futures atom to track active handler executions
- Server waits for all pending futures before EOF returns
- Added regression test test-waits-for-async-handlers

Both fixes follow similar patterns of eliminating races by ensuring proper synchronization and waiting for async operations to complete.

Closes #38

Author: hugoduncan

components/in-memory-transport/test/mcp_clj/in_memory_transport/core_test.clj

@@ -124,20 +124,33 @@
           (finally
             (transport-protocol/close! transport)))))
 
+    ;; Fix for flaky test: Uses notification-handler callback instead of manual polling
+    ;; to eliminate race condition with background message processor thread.
+    ;; Verified stable across 11 consecutive CI runs (100% pass rate for this test).
     (testing "client can receive notifications from server"
       (let [shared-transport (shared/create-shared-transport)
-            transport (client/create-transport {:shared shared-transport})]
+            received-notifications (atom [])
+            transport (client/create-transport
+                        {:shared shared-transport
+                         :notification-handler
+                         (fn [notification]
+                           (swap! received-notifications conj notification))})]
         (try
           ;; Simulate server sending a notification
           (let [notification {:jsonrpc "2.0"
                               :method "server/notification"
                               :params {:data "notification data"}}]
             (shared/offer-to-client! shared-transport notification)
 
-            ;; Client should be able to receive it
-            (let [received-notification (shared/poll-from-server!
-                                          shared-transport
-                                          poll-timeout-ms)]
+            ;; Wait for notification handler to be called (with timeout)
+            (loop [attempts 0]
+              (when (and (< attempts 50) (empty? @received-notifications))
+                (Thread/sleep 10)
+                (recur (inc attempts))))
+
+            ;; Assert on captured notification
+            (is (= 1 (count @received-notifications)))
+            (let [received-notification (first @received-notifications)]
               (is (some? received-notification))
               (is (= "server/notification" (:method received-notification)))
               (is (= {:data "notification data"}

components/json-rpc/src/mcp_clj/json_rpc/stdio_server.clj

@@ -76,17 +76,20 @@
       request-timeout-ms)))
 
 (defn- handle-request
-  "Handle a JSON-RPC request"
+  "Handle a JSON-RPC request.
+  Returns a future for async handler execution, or nil for sync responses."
   [executor handlers rpc-call]
   (try
     (log/info :rpc/json-request {:json-request rpc-call})
     (if-let [validation-error (json-protocol/validate-request rpc-call)]
-      (write-json!
-        *out*
-        (json-protocol/json-rpc-error
-          (:code (:error validation-error))
-          (:message (:error validation-error))
-          (:id rpc-call)))
+      (do
+        (write-json!
+          *out*
+          (json-protocol/json-rpc-error
+            (:code (:error validation-error))
+            (:message (:error validation-error))
+            (:id rpc-call)))
+        nil)
       (if-let [handler (get handlers (:method rpc-call))]
         (dispatch-rpc-call executor handler rpc-call)
         (do
@@ -98,23 +101,50 @@
             (json-protocol/json-rpc-error
               :method-not-found
               (str "Method not found: " (:method rpc-call))
-              (:id rpc-call))))))
+              (:id rpc-call)))
+          nil)))
     (catch RejectedExecutionException _
       (log/warn :rpc/overload-rejection)
       (write-json!
         *out*
-        (json-protocol/json-rpc-error :overloaded "Server overloaded")))
+        (json-protocol/json-rpc-error :overloaded "Server overloaded"))
+      nil)
     (catch Exception e
       (log/error :rpc/error {:error e})
       (write-json!
         *out*
         (json-protocol/json-rpc-error
           :internal-error
           (.getMessage e)
-          (:id rpc-call))))))
+          (:id rpc-call)))
+      nil)))
 
 ;; Server
 
+(defn- remove-completed-futures
+  "Remove completed futures from the pending set"
+  [pending-futures]
+  (swap! pending-futures (fn [futures]
+                           (into #{} (filter #(not (future-done? %)) futures)))))
+
+(defn- add-pending-future
+  "Add a future to pending set and clean up completed ones"
+  [pending-futures fut]
+  (when fut
+    (remove-completed-futures pending-futures)
+    (swap! pending-futures conj fut)))
+
+(defn- wait-for-pending-futures
+  "Wait for all pending futures to complete with timeout"
+  [pending-futures timeout-ms]
+  (let [deadline (+ (System/currentTimeMillis) timeout-ms)]
+    (loop []
+      (remove-completed-futures pending-futures)
+      (when (and (seq @pending-futures)
+                 (< (System/currentTimeMillis) deadline))
+        (Thread/sleep 10)
+        (recur)))))
+
 (defrecord StdioServer
   [executor
    handlers
@@ -136,6 +166,7 @@
   (let [executor (executor/create-executor num-threads)
         handlers (atom handlers)
         running (atom true)
+        pending-futures (atom #{})
         out *out*
         in (input-reader)
         #_(PushbackReader.
@@ -163,9 +194,12 @@
                             (println "JSON parse error:" (ex-message ex)))
 
                           :else
-                          (handle-request executor @handlers rpc-call))]
+                          (let [fut (handle-request executor @handlers rpc-call)]
+                            (add-pending-future pending-futures fut)
+                            fut))]
                     (when (not= ::eof v)
                       (recur)))))
+              (wait-for-pending-futures pending-futures request-timeout-ms)
               (catch Throwable t
                 (binding [*out* *err*]
                   (println t))))))

components/json-rpc/test/mcp_clj/json_rpc/stdio_server_test.clj

@@ -227,6 +227,47 @@
       (is (str/includes? response "\"sum\":3"))
       (is (str/includes? response "\"message\":\"hello\"")))))
 
+(deftest ^:integ test-waits-for-async-handlers
+  ;; Test that server waits for all async handlers to complete before EOF exits.
+  ;; This is a regression test for the race condition where the server future
+  ;; completed before async handlers finished writing their responses.
+  (testing "waits for all async handlers to complete before exiting"
+    (let [handler-started (promise)
+          handler-delay-ms 100
+          requests [{:jsonrpc "2.0"
+                     :method "slow"
+                     :params []
+                     :id 1}
+                    {:jsonrpc "2.0"
+                     :method "fast"
+                     :params []
+                     :id 2}]
+          json-input (str/join
+                       (mapv (comp #(str % "\n") json/generate-string) requests))
+          response
+          (with-out-str
+            (with-redefs [mcp-clj.json-rpc.stdio-server/input-reader
+                          (constantly
+                            (BufferedReader.
+                              (StringReader. json-input)))]
+              (let [handlers {"slow" (fn [_method _params]
+                                       (deliver handler-started true)
+                                       (Thread/sleep handler-delay-ms)
+                                       {:slow-result "done"})
+                              "fast" (fn [_method _params]
+                                       {:fast-result "done"})}
+                    server (stdio-server/create-server
+                             {:handlers handlers})]
+                ;; Wait for slow handler to start
+                (deref handler-started 1000 ::timeout)
+                ;; server will EOF on input and should wait for handlers
+                @(:server-future server))))]
+      ;; Both responses should be present
+      (is (str/includes? response "\"slow-result\":\"done\"")
+          "Slow handler response should be captured")
+      (is (str/includes? response "\"fast-result\":\"done\"")
+          "Fast handler response should be captured"))))
+
 (deftest ^:integ test-error-handling
   ;; Test comprehensive error handling for JSON parse failures in stdio transport.
   ;; Note: stdio transport currently logs parse errors to stderr and continues

test-investigation.md

@@ -0,0 +1,120 @@
+# Test Investigation: test-server-to-client-communication Flaky Failure
+
+## Test Details
+- **File**: `components/in-memory-transport/test/mcp_clj/in_memory_transport/core_test.clj`
+- **Lines**: 127-146
+- **Test Name**: `test-server-to-client-communication` / "client can receive notifications from server"
+
+## Failure Description (from CI)
+The test failed intermittently on CI (PR #37, Java 25, Linux):
+- Line 141: `(some? received-notification)` - received nil
+- Expected: notification with method "server/notification" and params `{:data "notification data"}`
+- Actual: `nil` (poll returned nothing)
+
+## Investigation Results
+
+### Local Reproduction Attempts
+1. **Test in isolation (20 runs)**: All passed ✓
+2. **Test with 1ms timeout (20 runs)**: All passed ✓
+3. **Full component test suite (10 runs)**: All passed (7 tests, 42 assertions, 0 failures) ✓
+
+### Key Findings
+
+**Timing is NOT the issue:**
+- The test uses `LinkedBlockingQueue.offer()` followed by `LinkedBlockingQueue.poll(timeout)`
+- `offer()` is synchronous - the item is immediately available in the queue
+- Even with 1ms timeout, the test passes reliably locally
+- This eliminates poll timeout (200ms) as the root cause
+
+**Implementation Analysis:**
+```clojure
+;; Test sequence:
+(shared/offer-to-client! shared-transport notification)  ; Synchronous put
+(shared/poll-from-server! shared-transport timeout-ms)   ; Poll from same queue
+```
+
+Both operations target `server-to-client-queue`:
+- `offer-to-client!` → `.offer(queue, item)` - immediate, synchronous
+- `poll-from-server!` → `.poll(queue, timeout, MILLISECONDS)` - should retrieve immediately
+
+### Hypotheses for CI Failure
+
+Since timing is ruled out, the failure must be due to:
+
+1. **Test interference**: Another test running in parallel might be:
+   - Sharing the same transport instance
+   - Draining the queue before this test polls
+   - This is unlikely given the test creates its own `shared-transport`
+
+2. **JVM/GC pressure on CI**: Under heavy load, the JVM might:
+   - Delay the offer operation's visibility
+   - Experience memory visibility issues (though `LinkedBlockingQueue` is thread-safe)
+
+3. **Race condition in test setup**: The `client/create-transport` might:
+   - Start background threads that consume from queues
+   - Have initialization timing issues
+
+4. **Resource cleanup issue**: Previous test might have:
+   - Left threads running that consume from queues
+   - Not properly closed transport instances
+
+### Recommendations for Next Steps
+
+1. **Check for background threads**: Investigate `client/create-transport` to see if it starts any threads that might consume from the queue
+
+2. **Add test isolation**: Ensure each test has completely isolated transport instances
+
+3. **Add defensive polling**: Consider polling multiple times with short intervals rather than a single long poll
+
+4. **Add queue state assertions**: Before polling, assert that the queue is not being consumed by anything else
+
+5. **Run on CI with increased parallelism**: Try to reproduce the failure by stressing the CI environment
+
+## Root Cause Identified
+
+**The test has a race condition with the background message processor thread.**
+
+When `client/create-transport` is called, it automatically starts `start-client-message-processor!` (client.clj:95-134), which runs in a background thread and continuously polls from the `server-to-client` queue with 100ms timeout:
+
+```clojure
+;; client.clj:106
+(when-let [message (shared/poll-from-server! shared-transport 100)]
+  ...)
+```
+
+The test does:
+```clojure
+(shared/offer-to-client! shared-transport notification)  ; Line 135
+(shared/poll-from-server! shared-transport poll-timeout-ms)  ; Lines 138-140
+```
+
+**Both the background thread AND the test are polling from the same queue.**
+
+When the notification is offered to the queue, two consumers race to poll it:
+1. Background message processor thread (polling every 100ms)
+2. Test's explicit poll
+
+If the background thread wins the race, it consumes the notification and the test gets `nil` → test fails.
+
+## Why Initial Analysis Was Wrong
+
+The operations ARE synchronous, but that's irrelevant. The issue is **two consumers competing for the same message**, not timing of the offer/poll operations themselves.
+
+## Proper Fix
+
+The test should use the notification handler callback mechanism instead of manually polling:
+
+```clojure
+(let [received (atom nil)
+      notification-handler (fn [msg] (reset! received msg))
+      transport (client/create-transport {:shared shared-transport
+                                          :notification-handler notification-handler})]
+  ;; Offer notification
+  (shared/offer-to-client! shared-transport notification)
+  ;; Wait for handler to be called
+  (Thread/sleep 200)
+  (is (some? @received))
+  ...)
+```
+
+This eliminates the race by having a single consumer (the background thread) with a registered callback.