Fix draft release detection (attempt 3) (#1176)

## Motivation

Previous fix still failed in CI. Suspecting permissions issue with
GITHUB_TOKEN not being able to see draft releases via API.

## Changes

1. Add explicit `permissions: contents: write` to the job
2. Use `gh release list` first to check if draft exists (this uses a
different code path that might work better)
3. Add debug echo statements

## Test Plan

Delete v1.0.63 tag and re-push after merging.

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

.github/workflows/build-app.yml

@@ -1,5 +1,16 @@
 name: Build EXO macOS DMG
 
+# Release workflow:
+# 1. Create a draft GitHub Release with the tag name (e.g. v1.0.0) and write release notes in markdown
+# 2. Push the tag: git tag v1.0.0 && git push origin v1.0.0
+# 3. This workflow builds, signs, and notarizes the DMG
+# 4. Release notes are embedded in appcast.xml for Sparkle (rendered as markdown)
+# 5. DMG and appcast.xml are uploaded to S3
+# 6. The draft GitHub Release is published with the DMG attached
+#
+# For alpha releases (e.g. v1.0.0-alpha.1): draft release and notes are optional.
+# If no draft exists, a release is auto-created with generated notes.
+
 on:
   workflow_dispatch:
   push:
@@ -11,8 +22,10 @@ on:
 jobs:
   build-macos-app:
     runs-on: "macos-26"
+    permissions:
+      contents: write
     env:
-      SPARKLE_VERSION: 2.8.1
+      SPARKLE_VERSION: 2.9.0-beta.1
       SPARKLE_DOWNLOAD_PREFIX: ${{ secrets.SPARKLE_DOWNLOAD_PREFIX }}
       SPARKLE_FEED_URL: ${{ secrets.SPARKLE_FEED_URL }}
       SPARKLE_ED25519_PUBLIC: ${{ secrets.SPARKLE_ED25519_PUBLIC }}
@@ -87,6 +100,52 @@ jobs:
             exit 1
           fi
 
+      - name: Fetch and validate release notes
+        if: github.ref_type == 'tag'
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          # Find draft release by name using gh release list (more reliable with default token)
+          echo "Looking for draft release named '$GITHUB_REF_NAME'..."
+          DRAFT_EXISTS=$(gh release list --json name,isDraft --jq ".[] | select(.isDraft == true) | select(.name == \"$GITHUB_REF_NAME\") | .name" 2>/dev/null || echo "")
+
+          if [[ -z "$DRAFT_EXISTS" ]]; then
+            if [[ "$IS_ALPHA" == "true" ]]; then
+              echo "No draft release found for alpha tag $GITHUB_REF_NAME (optional for alphas)"
+              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
+              exit 0
+            fi
+            echo "ERROR: No draft release found for tag $GITHUB_REF_NAME"
+            echo "Please create a draft release with release notes before pushing the tag."
+            exit 1
+          fi
+
+          # Fetch full release details via API to get body and ID
+          echo "Found draft release, fetching details..."
+          RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq ".[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")" 2>/dev/null || echo "")
+
+          # Extract release notes
+          NOTES=$(echo "$RELEASE_JSON" | jq -r '.body // ""')
+          if [[ -z "$NOTES" || "$NOTES" == "null" ]]; then
+            if [[ "$IS_ALPHA" == "true" ]]; then
+              echo "Draft release has no notes (optional for alphas)"
+              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
+              exit 0
+            fi
+            echo "ERROR: Draft release exists but has no release notes"
+            echo "Please add release notes to the draft release before pushing the tag."
+            exit 1
+          fi
+
+          # Save release ID for later publishing
+          RELEASE_ID=$(echo "$RELEASE_JSON" | jq -r '.id')
+          echo "DRAFT_RELEASE_ID=$RELEASE_ID" >> $GITHUB_ENV
+          echo "HAS_RELEASE_NOTES=true" >> $GITHUB_ENV
+
+          echo "Found draft release (ID: $RELEASE_ID), saving release notes..."
+          echo "$NOTES" > /tmp/release_notes.md
+          echo "RELEASE_NOTES_FILE=/tmp/release_notes.md" >> $GITHUB_ENV
+
       # ============================================================
       # Install dependencies
       # ============================================================
@@ -304,6 +363,28 @@ jobs:
             $CHANNEL_FLAG \
             .
 
+      - name: Inject release notes into appcast
+        if: github.ref_type == 'tag' && env.HAS_RELEASE_NOTES == 'true'
+        env:
+          RELEASE_VERSION: ${{ env.RELEASE_VERSION }}
+        run: |
+          # Inject markdown release notes with sparkle:format="markdown" (Sparkle 2.9+)
+          export NOTES=$(cat "$RELEASE_NOTES_FILE")
+
+          # Insert description after the enclosure tag for this version
+          awk '
+            /<enclosure[^>]*>/ && index($0, ENVIRON["RELEASE_VERSION"]) {
+              print
+              print "            <description sparkle:format=\"markdown\"><![CDATA["
+              print ENVIRON["NOTES"]
+              print "            ]]></description>"
+              next
+            }
+            { print }
+          ' output/appcast.xml > output/appcast.xml.tmp && mv output/appcast.xml.tmp output/appcast.xml
+
+          echo "Injected markdown release notes for version $RELEASE_VERSION"
+
       # ============================================================
       # Upload artifacts
       # ============================================================
@@ -336,3 +417,26 @@ jobs:
             aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg"
             aws s3 cp appcast.xml "s3://${SPARKLE_S3_BUCKET}/${PREFIX}appcast.xml" --content-type application/xml --cache-control no-cache
           fi
+
+      - name: Publish GitHub Release
+        if: github.ref_type == 'tag'
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+        run: |
+          DMG_PATH="output/EXO-${RELEASE_VERSION}.dmg"
+
+          if [[ "$HAS_RELEASE_NOTES" == "true" ]]; then
+            # Update the draft release with the tag and upload DMG
+            gh api --method PATCH "repos/${{ github.repository }}/releases/$DRAFT_RELEASE_ID" \
+              -f tag_name="$GITHUB_REF_NAME" \
+              -F draft=false
+            gh release upload "$GITHUB_REF_NAME" "$DMG_PATH" --clobber
+            echo "Published release $GITHUB_REF_NAME with DMG attached"
+          else
+            # Alpha without draft release - create one with auto-generated notes
+            gh release create "$GITHUB_REF_NAME" "$DMG_PATH" \
+              --title "$GITHUB_REF_NAME" \
+              --generate-notes \
+              --prerelease
+            echo "Created alpha release $GITHUB_REF_NAME with auto-generated notes"
+          fi

.mlx_typings/mlx_lm/generate.pyi

@@ -276,23 +276,24 @@ class BatchGenerator:
         logprobs: mx.array
         finish_reason: Optional[str]
 
-    unprocessed_prompts: List[Any]
-
     def __init__(
         self,
-        model: nn.Module,
+        model,
         max_tokens: int = ...,
-        stop_tokens: Optional[set[int]] = ...,
+        stop_tokens: Optional[set] = ...,
         sampler: Optional[Callable[[mx.array], mx.array]] = ...,
         completion_batch_size: int = ...,
         prefill_batch_size: int = ...,
         prefill_step_size: int = ...,
     ) -> None: ...
     def insert(
-        self, prompts: List[List[int]], max_tokens: Union[List[int], int, None] = ...
-    ) -> List[int]: ...
-    def stats(self) -> BatchStats: ...
-    def next(self) -> List[Response]: ...
+        self, prompts, max_tokens: Union[List[int], int, None] = ...
+    ):  # -> list[Any]:
+        ...
+    def stats(self):  # -> BatchStats:
+        ...
+    def next(self):  # -> list[Any]:
+        ...
 
 def batch_generate(
     model,

.mlx_typings/mlx_lm/tokenizer_utils.pyi

@@ -39,18 +39,12 @@ class StreamingDetokenizer:
     """
 
     __slots__ = ...
-    tokens: list[int]
-    def reset(self) -> None: ...
-    def add_token(self, token: int) -> None: ...
-    def finalize(self) -> None: ...
+    def reset(self): ...
+    def add_token(self, token): ...
+    def finalize(self): ...
     @property
-    def text(self) -> str:
-        """The full text decoded so far."""
-        ...
-    @property
-    def last_segment(self) -> str:
+    def last_segment(self):
         """Return the last segment of readable text since last time this property was accessed."""
-        ...
 
 class NaiveStreamingDetokenizer(StreamingDetokenizer):
     """NaiveStreamingDetokenizer relies on the underlying tokenizer
@@ -114,7 +108,6 @@ class TokenizerWrapper:
     _tokenizer: PreTrainedTokenizerFast
     eos_token_id: int | None
     eos_token: str | None
-    eos_token_ids: list[int] | None
     bos_token_id: int | None
     bos_token: str | None
     vocab_size: int

AGENTS.md

@@ -91,45 +91,6 @@ From .cursorrules:
 - Catch exceptions only where you can handle them meaningfully
 - Use `@final` and immutability wherever applicable
 
-## Model Storage
-
-Downloaded models are stored in `~/.exo/models/` (not the standard HuggingFace cache location).
-
-## Creating Model Instances via API
-
-When testing with the API, you must first create a model instance before sending chat completions:
-
-```bash
-# 1. Get instance previews for a model
-curl "http://localhost:52415/instance/previews?model_id=llama-3.2-1b"
-
-# 2. Create an instance from the first valid preview
-INSTANCE=$(curl -s "http://localhost:52415/instance/previews?model_id=llama-3.2-1b" | jq -c '.previews[] | select(.error == null) | .instance' | head -n1)
-curl -X POST http://localhost:52415/instance -H 'Content-Type: application/json' -d "{\"instance\": $INSTANCE}"
-
-# 3. Wait for the runner to become ready (check logs for "runner ready")
-
-# 4. Send chat completions using the full model ID
-curl -X POST http://localhost:52415/v1/chat/completions \
-  -H "Content-Type: application/json" \
-  -d '{"model": "mlx-community/Llama-3.2-1B-Instruct-4bit", "messages": [{"role": "user", "content": "Hello"}], "max_tokens": 50}'
-```
-
-## Logs
-
-Exo logs are stored in `~/.exo/exo.log`. This is useful for debugging runner crashes and distributed issues.
-
 ## Testing
 
 Tests use pytest-asyncio with `asyncio_mode = "auto"`. Tests are in `tests/` subdirectories alongside the code they test. The `EXO_TESTS=1` env var is set during tests.
-
-### Distributed Testing
-
-When running distributed tests across multiple machines, use `EXO_LIBP2P_NAMESPACE` to isolate your test cluster from other exo instances on the same network:
-
-```bash
-# On each machine in the test cluster, use the same unique namespace
-EXO_LIBP2P_NAMESPACE=my-test-cluster uv run exo
-```
-
-This prevents your test cluster from discovering and interfering with production or other developers' exo clusters.

app/EXO/EXO.xcodeproj/project.pbxproj

@@ -585,7 +585,7 @@
 			repositoryURL = "https://github.com/sparkle-project/Sparkle.git";
 			requirement = {
 				kind = upToNextMajorVersion;
-				minimumVersion = 2.8.1;
+				minimumVersion = 2.9.0-beta.1;
 			};
 		};
 /* End XCRemoteSwiftPackageReference section */

app/EXO/EXO.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -6,8 +6,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/sparkle-project/Sparkle.git",
       "state" : {
-        "revision" : "5581748cef2bae787496fe6d61139aebe0a451f6",
-        "version" : "2.8.1"
+        "revision" : "e641adb41915a8409895e2e30666aa64e487b637",
+        "version" : "2.9.0-beta.1"
       }
     }
   ],

app/EXO/EXO/ContentView.swift

@@ -56,6 +56,11 @@ struct ContentView: View {
     }
 
     private var shouldShowLocalNetworkWarning: Bool {
+        // Show warning if local network is not working and EXO is running.
+        // The checker uses a longer timeout on first launch to allow time for
+        // the permission prompt, so this correctly handles both:
+        // 1. User denied permission on first launch
+        // 2. Permission broke after restart (macOS TCC bug)
         if case .notWorking = localNetworkChecker.status {
             return controller.status != .stopped
         }

app/EXO/EXO/Services/LocalNetworkChecker.swift

@@ -5,8 +5,8 @@ import os.log
 /// Checks if the app's local network permission is actually functional.
 ///
 /// macOS local network permission can appear enabled in System Preferences but not
-/// actually work after a restart. This service detects this by creating a UDP
-/// connection to the mDNS multicast address (224.0.0.251:5353).
+/// actually work after a restart. This service uses NWConnection to mDNS multicast
+/// to verify actual connectivity.
 @MainActor
 final class LocalNetworkChecker: ObservableObject {
     enum Status: Equatable {
@@ -35,30 +35,43 @@ final class LocalNetworkChecker: ObservableObject {
     }
 
     private static let logger = Logger(subsystem: "io.exo.EXO", category: "LocalNetworkChecker")
+    private static let hasCompletedInitialCheckKey = "LocalNetworkChecker.hasCompletedInitialCheck"
 
     @Published private(set) var status: Status = .unknown
-    @Published private(set) var lastConnectionState: String = "none"
 
     private var connection: NWConnection?
     private var checkTask: Task<Void, Never>?
 
+    /// Whether we've completed at least one check (stored in UserDefaults)
+    private var hasCompletedInitialCheck: Bool {
+        get { UserDefaults.standard.bool(forKey: Self.hasCompletedInitialCheckKey) }
+        set { UserDefaults.standard.set(newValue, forKey: Self.hasCompletedInitialCheckKey) }
+    }
+
     /// Checks if local network access is working.
     func check() {
         checkTask?.cancel()
         status = .checking
-        lastConnectionState = "connecting"
+
+        // Use longer timeout on first launch to allow time for permission prompt
+        let isFirstCheck = !hasCompletedInitialCheck
+        let timeout: UInt64 = isFirstCheck ? 30_000_000_000 : 3_000_000_000
 
         checkTask = Task { [weak self] in
             guard let self else { return }
-            let result = await self.performCheck()
+
+            Self.logger.info("Checking local network connectivity (first check: \(isFirstCheck))")
+            let result = await self.checkConnectivity(timeout: timeout)
             self.status = result
+            self.hasCompletedInitialCheck = true
+
             Self.logger.info("Local network check complete: \(result.displayText)")
         }
     }
 
-    private func performCheck() async -> Status {
-        Self.logger.info("Checking local network access via UDP multicast")
-
+    /// Checks connectivity using NWConnection to mDNS multicast.
+    /// The connection attempt triggers the permission prompt if not yet shown.
+    private func checkConnectivity(timeout: UInt64) async -> Status {
         connection?.cancel()
         connection = nil
 
@@ -84,22 +97,7 @@ final class LocalNetworkChecker: ObservableObject {
                 continuation.resume(returning: status)
             }
 
-            conn.stateUpdateHandler = { [weak self] state in
-                let stateStr: String
-                switch state {
-                case .setup: stateStr = "setup"
-                case .preparing: stateStr = "preparing"
-                case .ready: stateStr = "ready"
-                case .waiting(let e): stateStr = "waiting(\(e))"
-                case .failed(let e): stateStr = "failed(\(e))"
-                case .cancelled: stateStr = "cancelled"
-                @unknown default: stateStr = "unknown"
-                }
-
-                Task { @MainActor in
-                    self?.lastConnectionState = stateStr
-                }
-
+            conn.stateUpdateHandler = { state in
                 switch state {
                 case .ready:
                     resumeOnce(.working)
@@ -108,6 +106,7 @@ final class LocalNetworkChecker: ObservableObject {
                     if errorStr.contains("54") || errorStr.contains("ECONNRESET") {
                         resumeOnce(.notWorking(reason: "Connection blocked"))
                     }
+                // Otherwise keep waiting - might be showing permission prompt
                 case .failed(let error):
                     let errorStr = "\(error)"
                     if errorStr.contains("65") || errorStr.contains("EHOSTUNREACH")
@@ -127,7 +126,7 @@ final class LocalNetworkChecker: ObservableObject {
             conn.start(queue: .main)
 
             Task {
-                try? await Task.sleep(nanoseconds: 3_000_000_000)
+                try? await Task.sleep(nanoseconds: timeout)
                 let state = conn.state
                 switch state {
                 case .ready:

bench/exo_bench.py

@@ -241,6 +241,9 @@ class PromptSizer:
             ids = tokenizer.apply_chat_template(
                 messages, tokenize=True, add_generation_prompt=True
             )
+            # Fix for transformers 5.x
+            if hasattr(ids, "input_ids"):
+                ids = ids.input_ids
             return int(len(ids))
 
         return count_fn

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
+    "httpx>=0.28.1",
 ]
 
 [project.scripts]

src/exo/shared/logging.py

@@ -29,6 +29,11 @@ class _InterceptHandler(logging.Handler):
 
 def logger_setup(log_file: Path | None, verbosity: int = 0):
     """Set up logging for this process - formatting, file handles, verbosity and output"""
+
+    logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
+    logging.getLogger("httpx").setLevel(logging.WARNING)
+    logging.getLogger("httpcore").setLevel(logging.WARNING)
+
     logger.remove()
 
     # replace all stdlib loggers with _InterceptHandlers that log to loguru

src/exo/shared/types/worker/runners.py

@@ -50,9 +50,7 @@ class RunnerReady(BaseRunnerStatus):
 
 
 class RunnerRunning(BaseRunnerStatus):
-    """Runner is processing requests and can accept more (continuous batching)."""
-
-    active_requests: int = 0
+    pass
 
 
 class RunnerShuttingDown(BaseRunnerStatus):

src/exo/worker/engines/mlx/generator/batch_engine.py

@@ -1,302 +0,0 @@
-"""Batch generation engine using mlx_lm's BatchGenerator for continuous batching."""
-
-import time
-from dataclasses import dataclass, field
-
-import mlx.core as mx
-from mlx_lm.generate import BatchGenerator
-from mlx_lm.sample_utils import make_sampler
-from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
-
-from exo.shared.types.api import FinishReason, GenerationStats
-from exo.shared.types.common import CommandId
-from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import ChatCompletionTaskParams, TaskId
-from exo.shared.types.worker.runner_response import GenerationResponse
-from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import MAX_TOKENS
-from exo.worker.engines.mlx.generator.distributed_sync import share_object
-from exo.worker.engines.mlx.utils_mlx import apply_chat_template
-from exo.worker.runner.bootstrap import logger
-
-
-@dataclass
-class ActiveRequest:
-    """Tracks an active request in the batch."""
-
-    command_id: CommandId
-    task_id: TaskId
-    uid: int  # BatchGenerator's internal ID
-    detokenizer: StreamingDetokenizer
-    tokens_generated: int = 0
-    prompt_tokens: int = 0
-    start_time: float = field(default_factory=time.perf_counter)
-
-
-@dataclass
-class BatchedGenerationResponse:
-    """Response from batch engine, tagged with command_id and task_id."""
-
-    command_id: CommandId
-    task_id: TaskId
-    response: GenerationResponse
-
-
-class BatchGenerationEngine:
-    """Manages continuous batching using mlx_lm's BatchGenerator."""
-
-    def __init__(
-        self,
-        model: Model,
-        tokenizer: TokenizerWrapper,
-        group: mx.distributed.Group | None = None,
-        max_tokens: int = MAX_TOKENS,
-        completion_batch_size: int = 32,
-        prefill_batch_size: int = 8,
-        prefill_step_size: int = 2048,
-    ):
-        self.model = model
-        self.tokenizer = tokenizer
-        self.max_tokens = max_tokens
-        self.active_requests: dict[int, ActiveRequest] = {}
-        self._pending_inserts: list[
-            tuple[CommandId, TaskId, ChatCompletionTaskParams]
-        ] = []
-        self._pending_completions: list[
-            int
-        ] = []  # UIDs completed but not yet synced/removed
-
-        self.group = group
-        self.rank = group.rank() if group else 0
-        self.is_distributed = group is not None and group.size() > 1
-
-        sampler = make_sampler(temp=0.7, top_p=1.0)
-
-        eos_tokens: set[int] = set(tokenizer.eos_token_ids or [])
-
-        self.batch_gen: BatchGenerator = BatchGenerator(
-            model=model,
-            max_tokens=max_tokens,
-            stop_tokens=eos_tokens,
-            sampler=sampler,
-            completion_batch_size=completion_batch_size,
-            prefill_batch_size=prefill_batch_size,
-            prefill_step_size=prefill_step_size,
-        )
-
-        logger.info(
-            f"BatchGenerationEngine initialized with completion_batch_size={completion_batch_size}, "
-            f"prefill_batch_size={prefill_batch_size}, distributed={self.is_distributed}"
-        )
-
-    def queue_request(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams,
-    ) -> None:
-        """Queue a request for insertion. Only rank 0 should call this.
-
-        In distributed mode, rank 0 receives tasks from the control plane and
-        queues them here. The actual insertion happens in sync_and_insert_pending()
-        which ensures all ranks insert the same requests together.
-        """
-        assert self.rank == 0, "Only rank 0 should queue requests"
-        self._pending_inserts.append((command_id, task_id, task_params))
-        logger.info(
-            f"Queued request {command_id} for insertion (pending={len(self._pending_inserts)})"
-        )
-
-    def sync_and_insert_pending(self) -> list[int]:
-        """Sync pending inserts across ranks and insert them. Returns UIDs.
-
-        This method ensures all ranks insert the same requests in the same order.
-        In non-distributed mode, it simply inserts all pending requests.
-        In distributed mode, it broadcasts pending requests from rank 0 to all ranks.
-
-        Batches all pending inserts into a single batch_gen.insert() call for
-        efficient prefill batching.
-        """
-        inserts_to_process: list[tuple[CommandId, TaskId, ChatCompletionTaskParams]]
-
-        if not self.is_distributed:
-            # Non-distributed: just insert directly from pending
-            inserts_to_process = list(self._pending_inserts)
-        else:
-            # Distributed: broadcast pending inserts from rank 0 to all ranks
-            assert self.group is not None
-            pending_data = self._pending_inserts if self.rank == 0 else None
-            synced_data = share_object(pending_data, self.rank, self.group)
-
-            if synced_data is None:
-                self._pending_inserts.clear()
-                return []
-
-            inserts_to_process = synced_data
-
-        if not inserts_to_process:
-            self._pending_inserts.clear()
-            return []
-
-        # Prepare all requests for batched insertion
-        all_tokens: list[list[int]] = []
-        all_max_tokens: list[int] = []
-        all_prompt_tokens: list[int] = []
-        request_info: list[tuple[CommandId, TaskId]] = []
-
-        for cmd_id, task_id, params in inserts_to_process:
-            prompt_str = apply_chat_template(self.tokenizer, params)
-            tokens: list[int] = self.tokenizer.encode(
-                prompt_str, add_special_tokens=False
-            )
-            max_tokens = params.max_tokens or self.max_tokens
-
-            all_tokens.append(tokens)
-            all_max_tokens.append(max_tokens)
-            all_prompt_tokens.append(len(tokens))
-            request_info.append((cmd_id, task_id))
-
-        # Single batched insert for efficient prefill
-        uids = self.batch_gen.insert(all_tokens, max_tokens=all_max_tokens)
-
-        # Track all inserted requests
-        for i, uid in enumerate(uids):
-            cmd_id, task_id = request_info[i]
-            self.active_requests[uid] = ActiveRequest(
-                command_id=cmd_id,
-                task_id=task_id,
-                uid=uid,
-                detokenizer=self.tokenizer.detokenizer,
-                prompt_tokens=all_prompt_tokens[i],
-            )
-            logger.info(
-                f"Inserted request {cmd_id} with uid={uid}, prompt_tokens={all_prompt_tokens[i]}, max_tokens={all_max_tokens[i]}"
-            )
-
-        self._pending_inserts.clear()
-        return uids
-
-    def step(self) -> list[BatchedGenerationResponse]:
-        """Run one decode step. Tracks completions but does not sync - call sync_completions() at budget boundaries."""
-        responses = self.batch_gen.next()
-        if not responses:
-            return []
-
-        results: list[BatchedGenerationResponse] = []
-
-        for r in responses:
-            uid: int = r.uid
-            req = self.active_requests.get(uid)
-            if req is None:
-                logger.warning(f"Received response for unknown uid={uid}")
-                continue
-
-            req.tokens_generated += 1
-
-            # Decode the token
-            token: int = r.token
-            req.detokenizer.add_token(token)
-            text: str = req.detokenizer.last_segment
-
-            stats: GenerationStats | None = None
-            finish_reason: FinishReason | None = None
-
-            raw_finish_reason: str | None = r.finish_reason
-            if raw_finish_reason is not None:
-                # Finalize to get remaining text
-                req.detokenizer.finalize()
-                text = req.detokenizer.last_segment
-
-                elapsed = time.perf_counter() - req.start_time
-                generation_tps = req.tokens_generated / elapsed if elapsed > 0 else 0.0
-
-                stats = GenerationStats(
-                    prompt_tps=0.0,  # Not tracked per-request in batch mode
-                    generation_tps=generation_tps,
-                    prompt_tokens=req.prompt_tokens,
-                    generation_tokens=req.tokens_generated,
-                    peak_memory_usage=Memory.from_gb(mx.get_peak_memory() / 1e9),
-                )
-
-                if raw_finish_reason == "stop":
-                    finish_reason = "stop"
-                elif raw_finish_reason == "length":
-                    finish_reason = "length"
-                else:
-                    logger.warning(f"Unknown finish_reason: {raw_finish_reason}")
-                    finish_reason = "stop"
-
-                # Track completion but don't remove yet - wait for sync_completions()
-                self._pending_completions.append(uid)
-                logger.info(
-                    f"Request {req.command_id} completed: {req.tokens_generated} tokens, {generation_tps:.2f} tps, reason={finish_reason}"
-                )
-
-            results.append(
-                BatchedGenerationResponse(
-                    command_id=req.command_id,
-                    task_id=req.task_id,
-                    response=GenerationResponse(
-                        text=text, token=token, finish_reason=finish_reason, stats=stats
-                    ),
-                )
-            )
-
-        # In non-distributed mode, clean up completions immediately
-        if not self.is_distributed:
-            self._remove_completed()
-
-        return results
-
-    def sync_completions(self) -> None:
-        """Sync and remove completed requests. Call at time budget boundaries in distributed mode."""
-        if not self.is_distributed:
-            # Non-distributed: early return if nothing to do
-            if not self._pending_completions:
-                return
-            self._remove_completed()
-            return
-
-        # Distributed mode: ALWAYS sync to ensure all ranks participate in collective op
-        # This prevents deadlock if one rank has completions and another doesn't
-        assert self.group is not None
-        synced_uids = share_object(
-            self._pending_completions if self.rank == 0 else None,
-            self.rank,
-            self.group,
-        )
-        if synced_uids:
-            self._pending_completions = synced_uids
-
-        self._remove_completed()
-
-    def _remove_completed(self) -> None:
-        """Remove completed requests from tracking."""
-        for uid in self._pending_completions:
-            if uid in self.active_requests:
-                del self.active_requests[uid]
-        self._pending_completions.clear()
-
-    @property
-    def has_active_requests(self) -> bool:
-        return bool(self.active_requests or self.batch_gen.unprocessed_prompts)
-
-    @property
-    def has_pending_inserts(self) -> bool:
-        return bool(self._pending_inserts)
-
-    @property
-    def active_count(self) -> int:
-        return len(self.active_requests)
-
-    @property
-    def pending_count(self) -> int:
-        return len(self.batch_gen.unprocessed_prompts)
-
-    @property
-    def pending_insert_count(self) -> int:
-        return len(self._pending_inserts)
-
-    @property
-    def has_pending_completions(self) -> bool:
-        return bool(self._pending_completions)

src/exo/worker/engines/mlx/generator/distributed_sync.py

@@ -1,30 +0,0 @@
-"""Distributed sync utilities using mx.distributed.all_sum() to broadcast from rank 0."""
-
-# pyright: reportAny=false
-
-import pickle
-from typing import TypeVar, cast
-
-import mlx.core as mx
-
-T = TypeVar("T")
-
-
-def share_object(obj: T | None, rank: int, group: mx.distributed.Group) -> T | None:
-    """Broadcast object from rank 0 to all ranks. Two-phase: size then data."""
-    if rank == 0:
-        if obj is None:
-            mx.eval(mx.distributed.all_sum(mx.array([0]), group=group))
-            return None
-        data = mx.array(list(pickle.dumps(obj)), dtype=mx.uint8)
-        mx.eval(mx.distributed.all_sum(mx.array([data.size]), group=group))
-        mx.eval(mx.distributed.all_sum(data, group=group))
-        return obj
-    else:
-        size = int(mx.distributed.all_sum(mx.array([0]), group=group).item())
-        if size == 0:
-            return None
-        data = mx.zeros(size, dtype=mx.uint8)
-        data = mx.distributed.all_sum(data, group=group)
-        mx.eval(data)
-        return cast(T, pickle.loads(bytes(cast(list[int], data.tolist()))))

src/exo/worker/engines/mlx/generator/time_budget.py

@@ -1,104 +0,0 @@
-"""Time budget iterator for controlling generation loop timing in distributed mode.
-
-Based on mlx-lm's TimeBudget pattern - runs for a time budget then syncs,
-rather than syncing every token. This reduces distributed sync overhead.
-"""
-
-import time
-from typing import Iterator
-
-import mlx.core as mx
-
-from exo.worker.runner.bootstrap import logger
-
-generation_stream = mx.new_stream(mx.default_device())
-
-
-class TimeBudget(Iterator[None]):
-    """Controls generation loop timing, syncing across ranks periodically.
-
-    In distributed mode, periodically syncs timing across all ranks to
-    dynamically adjust iteration count based on actual performance.
-
-    In non-distributed mode, simply runs for the time budget.
-
-    Usage:
-        for _ in TimeBudget(budget=0.5):
-            batch_engine.step()
-            # ... process responses ...
-    """
-
-    def __init__(
-        self,
-        budget: float = 0.5,
-        iterations: int = 25,
-        sync_frequency: int = 10,
-        group: mx.distributed.Group | None = None,
-    ):
-        """Initialize TimeBudget.
-
-        Args:
-            budget: Time budget in seconds before yielding control
-            iterations: Initial number of iterations per budget period (distributed only)
-            sync_frequency: How often to sync timing across ranks (distributed only)
-            group: Distributed group, or None for non-distributed mode
-        """
-        self._budget = budget
-        self._iterations = iterations
-        self._sync_frequency = sync_frequency
-        self._group = group
-        self._is_distributed = group is not None and group.size() > 1
-
-        # Runtime state
-        self._start: float = 0.0
-        self._current_iterations: int = 0
-        self._loops: int = 0
-        self._time_spent: float = 0.0
-
-    def __iter__(self) -> "TimeBudget":
-        self._start = time.perf_counter()
-        self._current_iterations = 0
-        return self
-
-    def __next__(self) -> None:
-        if not self._is_distributed:
-            # Non-distributed: just check time budget
-            if time.perf_counter() - self._start > self._budget:
-                raise StopIteration()
-            return None
-
-        # Distributed mode: iteration-based with periodic timing sync
-        self._current_iterations += 1
-        if self._current_iterations > self._iterations:
-            self._loops += 1
-            self._time_spent += time.perf_counter() - self._start
-
-            if self._loops % self._sync_frequency == 0:
-                # Sync timing across all ranks
-                assert self._group is not None
-                with mx.stream(generation_stream):
-                    time_array = mx.array([self._time_spent], dtype=mx.float32)
-                    total_time = mx.distributed.all_sum(time_array, group=self._group)
-                    mx.eval(total_time)
-                    loop_time = float(total_time.item())
-
-                avg_loop_time = loop_time / (self._group.size() * self._sync_frequency)
-
-                if avg_loop_time > 0:
-                    factor = self._budget / avg_loop_time
-                    self._iterations = max(round(self._iterations * factor), 1)
-                    logger.debug(
-                        f"TimeBudget adjusted iterations to {self._iterations}"
-                    )
-
-                self._loops = 0
-                self._time_spent = 0.0
-
-            raise StopIteration()
-
-        return None
-
-    @property
-    def iterations(self) -> int:
-        """Current iterations per budget period."""
-        return self._iterations

src/exo/worker/plan.py

@@ -277,14 +277,12 @@ def _pending_tasks(
             # I have a design point here; this is a state race in disguise as the task status doesn't get updated to completed fast enough
             # however, realistically the task status should be set to completed by the LAST runner, so this is a true race
             # the actual solution is somewhat deeper than this bypass - TODO!
-            # Also skip tasks in pending to prevent duplicate forwarding with continuous batching
-            if task.task_id in runner.completed or task.task_id in runner.pending:
+            if task.task_id in runner.completed:
                 continue
 
             # TODO: Check ordering aligns with MLX distributeds expectations.
 
-            # Allow forwarding tasks when runner is Ready or Running (for continuous batching)
-            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
+            if isinstance(runner.status, RunnerReady) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/runner.py

@@ -1,8 +1,15 @@
-import gc
 import time
+from collections.abc import Generator
+from functools import cache
 
 import mlx.core as mx
-from anyio import WouldBlock
+from mlx_lm.models.gpt_oss import Model as GptOssModel
+from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
+    HarmonyEncodingName,
+    Role,
+    StreamableParser,
+    load_harmony_encoding,
+)
 
 from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.chunks import TokenChunk
@@ -23,6 +30,9 @@ from exo.shared.types.tasks import (
     TaskStatus,
 )
 from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.runner_response import (
+    GenerationResponse,
+)
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
@@ -38,9 +48,7 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx.generator.batch_engine import BatchGenerationEngine
-from exo.worker.engines.mlx.generator.generate import warmup_inference
-from exo.worker.engines.mlx.generator.time_budget import TimeBudget
+from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
     load_mlx_items,
@@ -70,318 +78,187 @@ def main(
     model = None
     tokenizer = None
     group = None
-    batch_engine: BatchGenerationEngine | None = None
-    pending_shutdown: Shutdown | None = None
 
     current_status: RunnerStatus = RunnerIdle()
-
-    def send_status(status: RunnerStatus) -> None:
-        event_sender.send(
-            RunnerStatusUpdated(runner_id=runner_id, runner_status=status)
-        )
-
     logger.info("runner created")
-    send_status(current_status)
-
-    def handle_task(task: Task, is_deferred: bool = False) -> bool:
-        nonlocal current_status, model, tokenizer, group, batch_engine, pending_shutdown
-
-        # For Shutdown, check if we need to defer BEFORE sending Running/Acknowledged
-        if (
-            isinstance(task, Shutdown)
-            and not is_deferred
-            and batch_engine is not None
-            and (batch_engine.has_active_requests or batch_engine.has_pending_inserts)
-        ):
-            logger.info("deferring shutdown until active requests complete")
-            pending_shutdown = task
-            return True
-
-        event_sender.send(
-            TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
-        )
-        event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-        match task:
-            case ConnectToGroup() if isinstance(
-                current_status, (RunnerIdle, RunnerFailed)
-            ):
-                logger.info("runner connecting")
-                current_status = RunnerConnecting()
-                send_status(current_status)
-                group = initialize_mlx(bound_instance)
-
-                logger.info("runner connected")
-                current_status = RunnerConnected()
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
+    event_sender.send(
+        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+    )
+    with task_receiver as tasks:
+        for task in tasks:
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
+            )
+            event_sender.send(TaskAcknowledged(task_id=task.task_id))
+            match task:
+                case ConnectToGroup() if isinstance(
+                    current_status, (RunnerIdle, RunnerFailed)
+                ):
+                    logger.info("runner connecting")
+                    current_status = RunnerConnecting()
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
                     )
-                )
-                send_status(current_status)
+                    group = initialize_mlx(bound_instance)
 
-            case LoadModel() if (
-                isinstance(current_status, RunnerConnected) and group is not None
-            ) or (isinstance(current_status, RunnerIdle) and group is None):
-                current_status = RunnerLoading()
-                logger.info("runner loading")
-                send_status(current_status)
+                    logger.info("runner connected")
+                    current_status = RunnerConnected()
 
-                model, tokenizer = load_mlx_items(bound_instance, group)
-
-                current_status = RunnerLoaded()
-                logger.info("runner loaded")
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
+                # we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
+                case LoadModel() if (
+                    isinstance(current_status, RunnerConnected) and group is not None
+                ) or (isinstance(current_status, RunnerIdle) and group is None):
+                    current_status = RunnerLoading()
+                    logger.info("runner loading")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
                     )
-                )
-                send_status(current_status)
 
-            case StartWarmup() if isinstance(current_status, RunnerLoaded):
-                assert model is not None
-                assert tokenizer is not None
-                current_status = RunnerWarmingUp()
-                logger.info("runner warming up")
-                send_status(current_status)
+                    model, tokenizer = load_mlx_items(bound_instance, group)
 
-                logger.info(f"warming up inference for instance: {instance}")
-                toks = warmup_inference(model=model, tokenizer=tokenizer)
-                logger.info(f"warmed up by generating {toks} tokens")
-                logger.info(
-                    f"runner initialized in {time.time() - setup_start_time} seconds"
-                )
-
-                batch_engine = BatchGenerationEngine(
-                    model=model, tokenizer=tokenizer, group=group
-                )
-
-                current_status = RunnerReady()
-                logger.info("runner ready")
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
+                    current_status = RunnerLoaded()
+                    logger.info("runner loaded")
+                case StartWarmup() if isinstance(current_status, RunnerLoaded):
+                    assert model
+                    assert tokenizer
+                    current_status = RunnerWarmingUp()
+                    logger.info("runner warming up")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
                     )
-                )
-                send_status(current_status)
 
-            case ChatCompletion(task_params=task_params, command_id=command_id) if (
-                isinstance(current_status, (RunnerReady, RunnerRunning))
-            ):
-                assert batch_engine is not None
-
-                # In distributed mode, only rank 0 should queue requests
-                # Other ranks should skip - they'll participate in sync_and_insert_pending()
-                is_distributed_mode = group is not None and group.size() > 1
-                if is_distributed_mode and shard_metadata.device_rank != 0:
-                    logger.debug(
-                        f"Rank {shard_metadata.device_rank} skipping ChatCompletionTask (only rank 0 queues)"
+                    logger.info(f"warming up inference for instance: {instance}")
+                    toks = warmup_inference(
+                        model=model,
+                        tokenizer=tokenizer,
+                        # kv_prefix_cache=kv_prefix_cache,  # supply for warmup-time prefix caching
                     )
-                    return True
-
-                if task_params.messages and task_params.messages[0].content is not None:
+                    logger.info(f"warmed up by generating {toks} tokens")
+                    logger.info(
+                        f"runner initialized in {time.time() - setup_start_time} seconds"
+                    )
+                    current_status = RunnerReady()
+                    logger.info("runner ready")
+                case ChatCompletion(task_params=task_params, command_id=command_id) if (
+                    isinstance(current_status, RunnerReady)
+                ):
+                    assert model
+                    assert tokenizer
+                    logger.info(f"received chat request: {str(task)[:500]}")
+                    current_status = RunnerRunning()
+                    logger.info("runner running")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
+                        )
+                    )
+                    assert task_params.messages[0].content is not None
                     _check_for_debug_prompts(task_params.messages[0].content)
 
-                # Queue the request - actual insertion happens in sync_and_insert_pending()
-                batch_engine.queue_request(
-                    command_id=command_id, task_id=task.task_id, task_params=task_params
-                )
-
-                # Status will be updated after actual insertion in the main loop
-                # For now, set to RunnerRunning to indicate we're processing
-                current_status = RunnerRunning(
-                    active_requests=batch_engine.active_count
-                    + batch_engine.pending_insert_count
-                )
-                send_status(current_status)
-
-            case Shutdown():
-                current_status = RunnerShuttingDown()
-                logger.info("runner shutting down")
-                send_status(current_status)
-                event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
+                    # Generate responses using the actual MLX generation
+                    mlx_generator = mlx_generate(
+                        model=model,
+                        tokenizer=tokenizer,
+                        task=task_params,
                     )
-                )
-                current_status = RunnerShutdown()
-                send_status(current_status)
-                return False
 
-            case _:
-                raise ValueError(
-                    f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
-                )
+                    # GPT-OSS specific parsing to match other model formats.
+                    if isinstance(model, GptOssModel):
+                        mlx_generator = parse_gpt_oss(mlx_generator)
 
-        return True
+                    # TODO: Add tool call parser here
 
-    with task_receiver as tasks:
-        running = True
-        is_rank_0 = shard_metadata.device_rank == 0
-
-        while running:
-            # Use batch_engine.is_distributed since it's set correctly after group initialization
-            # (the group variable is None at loop start, but set by ConnectToGroup task)
-            if batch_engine is not None and batch_engine.is_distributed:
-                assert group is not None
-                assert batch_engine is not None
-
-                # Distributed mode: tasks wake up all ranks, then we sync and generate
-
-                # Check deferred shutdown FIRST - all ranks must check and process together
-                # This must run before any collective operations to prevent deadlock
-                if (
-                    pending_shutdown is not None
-                    and not batch_engine.has_active_requests
-                    and not batch_engine.has_pending_inserts
-                ):
-                    handle_task(pending_shutdown, is_deferred=True)
-                    running = False
-                    continue
-
-                # When idle, block waiting for task (exo sends tasks to all ranks)
-                # When active, poll non-blocking to batch incoming requests
-                if (
-                    not batch_engine.has_active_requests
-                    and not batch_engine.has_pending_inserts
-                ):
-                    # IDLE: Block until task arrives (all ranks receive the same task)
-                    task = tasks.receive()
-                    task_result = handle_task(task)
-                    if not task_result:
-                        running = False
-                        continue
-                else:
-                    # ACTIVE: Poll for new tasks without blocking
-                    while True:
-                        try:
-                            task = tasks.receive_nowait()
-                            task_result = handle_task(task)
-                            if not task_result:
-                                running = False
-                                break
-                        except WouldBlock:
-                            break
-                    if not running:
-                        continue
-
-                # Sync and insert pending requests (collective operation)
-                # Rank 0 broadcasts its pending to all ranks
-                inserted = batch_engine.sync_and_insert_pending()
-                if is_rank_0 and inserted:
-                    current_status = RunnerRunning(
-                        active_requests=batch_engine.active_count
-                    )
-                    send_status(current_status)
-
-                # Run generation for time budget
-                if batch_engine.has_active_requests:
-                    time_budget = TimeBudget(budget=0.5, group=group)
-                    for _ in time_budget:
-                        if not batch_engine.has_active_requests:
-                            break
-                        for resp in batch_engine.step():
-                            # Send token IMMEDIATELY for smooth streaming (only rank 0)
-                            if is_rank_0:
-                                event_sender.send(
-                                    ChunkGenerated(
-                                        command_id=resp.command_id,
-                                        chunk=TokenChunk(
-                                            idx=resp.response.token,
-                                            model=shard_metadata.model_meta.model_id,
-                                            text=resp.response.text,
-                                            token_id=resp.response.token,
-                                            finish_reason=resp.response.finish_reason,
-                                            stats=resp.response.stats,
-                                        ),
-                                    )
-                                )
-                                if resp.response.finish_reason is not None:
+                    for response in mlx_generator:
+                        match response:
+                            case GenerationResponse():
+                                if shard_metadata.device_rank == 0:
                                     event_sender.send(
-                                        TaskStatusUpdated(
-                                            task_id=resp.task_id,
-                                            task_status=TaskStatus.Complete,
+                                        ChunkGenerated(
+                                            command_id=command_id,
+                                            chunk=TokenChunk(
+                                                idx=response.token,
+                                                model=shard_metadata.model_meta.model_id,
+                                                text=response.text,
+                                                token_id=response.token,
+                                                finish_reason=response.finish_reason,
+                                                stats=response.stats,
+                                            ),
                                         )
                                     )
+                                # case TokenizedResponse():
+                                # TODO: something here ig
 
-                # Sync completions at budget boundary (always call - it's a collective operation)
-                batch_engine.sync_completions()
-
-                # Update status after budget
-                if is_rank_0:
-                    current_status = (
-                        RunnerRunning(active_requests=batch_engine.active_count)
-                        if batch_engine.has_active_requests
-                        else RunnerReady()
-                    )
-                    send_status(current_status)
-
-            else:
-                # Non-distributed mode: original logic with queue + insert
-                while True:
-                    try:
-                        task = tasks.receive_nowait()
-                        running = handle_task(task)
-                        if not running:
-                            break
-                    except WouldBlock:
-                        break
-
-                if not running:
-                    break
-
-                # Insert any queued requests (non-distributed just inserts directly)
-                # Status was already sent in handle_task when queueing
-                if batch_engine is not None and batch_engine.has_pending_inserts:
-                    batch_engine.sync_and_insert_pending()
-
-                if batch_engine is not None and batch_engine.has_active_requests:
-                    for resp in batch_engine.step():
-                        if shard_metadata.device_rank == 0:
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=resp.command_id,
-                                    chunk=TokenChunk(
-                                        idx=resp.response.token,
-                                        model=shard_metadata.model_meta.model_id,
-                                        text=resp.response.text,
-                                        token_id=resp.response.token,
-                                        finish_reason=resp.response.finish_reason,
-                                        stats=resp.response.stats,
-                                    ),
-                                )
-                            )
-                        if resp.response.finish_reason is not None:
-                            event_sender.send(
-                                TaskStatusUpdated(
-                                    task_id=resp.task_id,
-                                    task_status=TaskStatus.Complete,
-                                )
-                            )
-
-                    if batch_engine.has_active_requests:
-                        current_status = RunnerRunning(
-                            active_requests=batch_engine.active_count
+                    current_status = RunnerReady()
+                    logger.info("runner ready")
+                case Shutdown():
+                    current_status = RunnerShuttingDown()
+                    logger.info("runner shutting down")
+                    event_sender.send(
+                        RunnerStatusUpdated(
+                            runner_id=runner_id, runner_status=current_status
                         )
-                    else:
-                        current_status = RunnerReady()
-                    send_status(current_status)
+                    )
+                    current_status = RunnerShutdown()
+                case _:
+                    raise ValueError(
+                        f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
+                    )
+            event_sender.send(
+                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
+            )
+            event_sender.send(
+                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
+            )
+            if isinstance(current_status, RunnerShutdown):
+                del model, tokenizer, group
+                mx.clear_cache()
+                import gc
 
-                    # Process deferred shutdown after all requests complete
-                    if (
-                        pending_shutdown is not None
-                        and not batch_engine.has_active_requests
-                        and not batch_engine.has_pending_inserts
-                    ):
-                        running = handle_task(pending_shutdown, is_deferred=True)
-                else:
-                    task = tasks.receive()
-                    running = handle_task(task)
+                gc.collect()
+                break
 
-    # Cleanup
-    del model, tokenizer, group, batch_engine
-    mx.clear_cache()
-    gc.collect()
+
+@cache
+def get_gpt_oss_encoding():
+    encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+    return encoding
+
+
+def parse_gpt_oss(
+    responses: Generator[GenerationResponse],
+) -> Generator[GenerationResponse]:
+    encoding = get_gpt_oss_encoding()
+    stream = StreamableParser(encoding, role=Role.ASSISTANT)
+    thinking = False
+
+    for response in responses:
+        stream.process(response.token)
+
+        delta = stream.last_content_delta
+        ch = stream.current_channel
+
+        if ch == "analysis" and not thinking:
+            thinking = True
+            yield response.model_copy(update={"text": "<think>"})
+
+        if ch != "analysis" and thinking:
+            thinking = False
+            yield response.model_copy(update={"text": "</think>"})
+
+        if delta:
+            yield response.model_copy(update={"text": delta})
+
+        if response.finish_reason is not None:
+            if thinking:
+                yield response.model_copy(update={"text": "</think>"})
+            yield response
+            break
 
 
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"

src/exo/worker/runner/runner_supervisor.py

@@ -105,7 +105,7 @@ class RunnerSupervisor:
             return
 
         # This is overkill but it's not technically bad, just unnecessary.
-        logger.warning("Runner process didn't shutdown successfully, terminating")
+        logger.warning("Runner process didn't shutdown succesfully, terminating")
         self.runner_process.terminate()
         await to_thread.run_sync(self.runner_process.join, 5)
         if not self.runner_process.is_alive():
@@ -128,11 +128,9 @@ class RunnerSupervisor:
 
     async def start_task(self, task: Task):
         if task.task_id in self.completed:
-            logger.info(f"Skipping task {task.task_id} - already completed")
-            return
-        if task.task_id in self.pending:
-            logger.info(f"Skipping task {task.task_id} - already pending")
-            return
+            logger.info(
+                f"Skipping invalid task {task} as it has already been completed"
+            )
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
@@ -151,17 +149,13 @@ class RunnerSupervisor:
                     if isinstance(event, RunnerStatusUpdated):
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
-                        # Just set the event to unblock start_task, but keep in pending
-                        # to prevent duplicate forwarding until completion
-                        if event.task_id in self.pending:
-                            self.pending[event.task_id].set()
+                        self.pending.pop(event.task_id).set()
                         continue
-                    if isinstance(event, TaskStatusUpdated) and event.task_status in (
-                        TaskStatus.Complete,
-                        TaskStatus.TimedOut,
-                        TaskStatus.Failed,
+                    if (
+                        isinstance(event, TaskStatusUpdated)
+                        and event.task_status == TaskStatus.Complete
                     ):
-                        # If a task has just finished, we should be working on it.
+                        # If a task has just been completed, we should be working on it.
                         assert isinstance(
                             self.status,
                             (
@@ -172,8 +166,6 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
-                        # Now safe to remove from pending and add to completed
-                        self.pending.pop(event.task_id, None)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
             except (ClosedResourceError, BrokenResourceError) as e:

src/exo/worker/tests/unittests/conftest.py

@@ -20,7 +20,6 @@ class FakeRunnerSupervisor:
     bound_instance: BoundInstance
     status: RunnerStatus
     completed: set[TaskId] = field(default_factory=set)
-    pending: dict[TaskId, object] = field(default_factory=dict)
 
 
 class OtherTask(BaseTask):

src/exo/worker/tests/unittests/test_runner/test_continuous_batching.py

@@ -1,319 +0,0 @@
-"""
-Tests for continuous batching behavior in the runner.
-
-These tests verify that:
-1. Single requests work through the batch path
-2. Multiple concurrent requests batch together
-3. Tokens are routed to the correct requests
-4. Requests complete at different times appropriately
-"""
-
-# pyright: reportAny=false
-# pyright: reportUnknownArgumentType=false
-# pyright: reportUnknownMemberType=false
-# pyright: reportAttributeAccessIssue=false
-# pyright: reportInvalidTypeVarUse=false
-
-from typing import Any
-from unittest.mock import MagicMock
-
-import pytest
-
-import exo.worker.runner.runner as mlx_runner
-from exo.shared.types.api import ChatCompletionMessage
-from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.events import (
-    Event,
-    RunnerStatusUpdated,
-    TaskStatusUpdated,
-)
-from exo.shared.types.tasks import (
-    ChatCompletion,
-    ChatCompletionTaskParams,
-    ConnectToGroup,
-    LoadModel,
-    Shutdown,
-    StartWarmup,
-    Task,
-    TaskId,
-    TaskStatus,
-)
-from exo.shared.types.worker.runner_response import GenerationResponse
-from exo.shared.types.worker.runners import RunnerRunning
-from exo.utils.channels import mp_channel
-from exo.worker.engines.mlx.generator.batch_engine import (
-    BatchedGenerationResponse,
-)
-from exo.worker.tests.constants import (
-    INSTANCE_1_ID,
-    MODEL_A_ID,
-    NODE_A,
-    RUNNER_1_ID,
-)
-from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
-
-
-class FakeBatchEngineWithTokens:
-    """
-    Fake batch engine that generates a specified number of tokens per request.
-
-    This simulates realistic batch generation behavior where:
-    - Requests are queued on insert
-    - Each step() call generates one token for all active requests
-    - Requests complete when they've generated all their tokens
-    """
-
-    def __init__(self, *_args: Any, **_kwargs: Any):
-        self._active_requests: dict[int, tuple[CommandId, TaskId, int, int]] = {}
-        self._pending_inserts: list[
-            tuple[CommandId, TaskId, ChatCompletionTaskParams]
-        ] = []
-        self._uid_counter = 0
-        self._tokens_per_request = 3  # Default: generate 3 tokens before completing
-        self.rank = 0  # Fake rank for testing
-
-    def queue_request(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams,
-    ) -> None:
-        """Queue a request for insertion."""
-        self._pending_inserts.append((command_id, task_id, task_params))
-
-    def sync_and_insert_pending(self) -> list[int]:
-        """Insert all pending requests."""
-        uids: list[int] = []
-        for command_id, task_id, task_params in self._pending_inserts:
-            uid = self._do_insert(command_id, task_id, task_params)
-            uids.append(uid)
-        self._pending_inserts.clear()
-        return uids
-
-    @property
-    def has_pending_inserts(self) -> bool:
-        return len(self._pending_inserts) > 0
-
-    def _do_insert(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams | None,
-    ) -> int:
-        uid = self._uid_counter
-        self._uid_counter += 1
-        # Track: (command_id, task_id, tokens_generated, max_tokens)
-        max_tokens = task_params.max_tokens if task_params else self._tokens_per_request
-        self._active_requests[uid] = (command_id, task_id, 0, max_tokens or 3)
-        return uid
-
-    def step(self) -> list[BatchedGenerationResponse]:
-        results: list[BatchedGenerationResponse] = []
-        uids_to_remove: list[int] = []
-
-        for uid, (command_id, task_id, tokens_gen, max_tokens) in list(
-            self._active_requests.items()
-        ):
-            tokens_gen += 1
-            finish_reason = "stop" if tokens_gen >= max_tokens else None
-            text = f"token{tokens_gen}"
-
-            if finish_reason:
-                uids_to_remove.append(uid)
-            else:
-                self._active_requests[uid] = (
-                    command_id,
-                    task_id,
-                    tokens_gen,
-                    max_tokens,
-                )
-
-            results.append(
-                BatchedGenerationResponse(
-                    command_id=command_id,
-                    task_id=task_id,
-                    response=GenerationResponse(
-                        token=tokens_gen,
-                        text=text,
-                        finish_reason=finish_reason,
-                    ),
-                )
-            )
-
-        for uid in uids_to_remove:
-            del self._active_requests[uid]
-
-        return results
-
-    @property
-    def has_active_requests(self) -> bool:
-        return len(self._active_requests) > 0
-
-    @property
-    def active_count(self) -> int:
-        return len(self._active_requests)
-
-    @property
-    def pending_insert_count(self) -> int:
-        return len(self._pending_inserts)
-
-
-def make_nothin[T, U, V](res: T):
-    def nothin(*_1: U, **_2: V) -> T:
-        return res
-
-    return nothin
-
-
-@pytest.fixture
-def patch_batch_engine(monkeypatch: pytest.MonkeyPatch):
-    """Patch MLX dependencies and use FakeBatchEngineWithTokens."""
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MagicMock()))
-    monkeypatch.setattr(
-        mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MagicMock()))
-    )
-    monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
-    monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", make_nothin(None))
-    monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngineWithTokens)
-
-
-def _run_with_tasks(tasks: list[Task]) -> list[Event]:
-    """
-    Run tasks through the runner, adding shutdown at the end.
-
-    Tasks are sent in order, with shutdown sent last.
-    The batch engine processes between task handling.
-    """
-    bound_instance = get_bound_mlx_ring_instance(
-        instance_id=INSTANCE_1_ID,
-        model_id=MODEL_A_ID,
-        runner_id=RUNNER_1_ID,
-        node_id=NodeId(NODE_A),
-    )
-
-    task_sender, task_receiver = mp_channel[Task]()
-    event_sender, event_receiver = mp_channel[Event]()
-
-    shutdown_task = Shutdown(
-        task_id=TaskId("shutdown"),
-        instance_id=INSTANCE_1_ID,
-        runner_id=RUNNER_1_ID,
-    )
-
-    with task_sender, event_receiver:
-        # Send all tasks including shutdown
-        for t in tasks:
-            task_sender.send(t)
-        task_sender.send(shutdown_task)
-
-        # Disable cleanup methods to prevent issues
-        event_sender.close = lambda: None
-        event_sender.join = lambda: None
-        task_receiver.close = lambda: None
-        task_receiver.join = lambda: None
-
-        mlx_runner.main(bound_instance, event_sender, task_receiver)
-
-        return event_receiver.collect()
-
-
-INIT_TASK = ConnectToGroup(task_id=TaskId("init"), instance_id=INSTANCE_1_ID)
-LOAD_TASK = LoadModel(task_id=TaskId("load"), instance_id=INSTANCE_1_ID)
-WARMUP_TASK = StartWarmup(task_id=TaskId("warmup"), instance_id=INSTANCE_1_ID)
-
-
-def make_chat_task(
-    task_id: str, command_id: str, max_tokens: int = 3
-) -> ChatCompletion:
-    return ChatCompletion(
-        task_id=TaskId(task_id),
-        command_id=CommandId(command_id),
-        task_params=ChatCompletionTaskParams(
-            model=str(MODEL_A_ID),
-            messages=[ChatCompletionMessage(role="user", content="hello")],
-            stream=True,
-            max_tokens=max_tokens,
-        ),
-        instance_id=INSTANCE_1_ID,
-    )
-
-
-def test_single_request_generates_tokens(patch_batch_engine: None):
-    """
-    Verify a single request generates the expected tokens through the batch path.
-
-    Note: With the current non-blocking design, shutdown is processed before
-    batch steps run when all tasks are queued together. This test verifies
-    the runner status reflects active requests.
-    """
-    chat_task = make_chat_task("chat1", "cmd1", max_tokens=3)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
-
-    # Find RunnerRunning status events - this shows the request was inserted
-    running_events = [
-        e
-        for e in events
-        if isinstance(e, RunnerStatusUpdated)
-        and isinstance(e.runner_status, RunnerRunning)
-    ]
-
-    assert len(running_events) >= 1, "Expected at least one RunnerRunning event"
-    assert running_events[0].runner_status.active_requests == 1
-
-
-def test_runner_status_reflects_active_requests(patch_batch_engine: None):
-    """Verify RunnerRunning status includes active_requests count."""
-    chat_task = make_chat_task("chat1", "cmd1", max_tokens=2)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
-
-    # Find RunnerRunning status events
-    running_events = [
-        e
-        for e in events
-        if isinstance(e, RunnerStatusUpdated)
-        and isinstance(e.runner_status, RunnerRunning)
-    ]
-
-    assert len(running_events) > 0, "Expected at least one RunnerRunning event"
-    assert running_events[0].runner_status.active_requests == 1
-
-
-def test_chat_task_acknowledged(patch_batch_engine: None):
-    """Verify chat completion task is acknowledged with proper status updates."""
-    chat_task = make_chat_task("chat1", "cmd1", max_tokens=2)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat_task])
-
-    # Find the chat task status events
-    chat_running = [
-        e
-        for e in events
-        if isinstance(e, TaskStatusUpdated)
-        and e.task_id == TaskId("chat1")
-        and e.task_status == TaskStatus.Running
-    ]
-
-    assert len(chat_running) == 1, "Expected exactly one chat task Running status"
-
-
-def test_multiple_requests_tracked(patch_batch_engine: None):
-    """Verify multiple concurrent requests are tracked in active_requests."""
-    chat1 = make_chat_task("chat1", "cmd1", max_tokens=2)
-    chat2 = make_chat_task("chat2", "cmd2", max_tokens=2)
-    events = _run_with_tasks([INIT_TASK, LOAD_TASK, WARMUP_TASK, chat1, chat2])
-
-    # Find RunnerRunning status events
-    running_events = [
-        e
-        for e in events
-        if isinstance(e, RunnerStatusUpdated)
-        and isinstance(e.runner_status, RunnerRunning)
-    ]
-
-    # Should have at least 2 RunnerRunning events (one per request inserted)
-    assert len(running_events) >= 2, (
-        f"Expected at least 2 RunnerRunning events, got {len(running_events)}"
-    )
-
-    # First should have 1 active request, second should have 2
-    assert running_events[0].runner_status.active_requests == 1
-    assert running_events[1].runner_status.active_requests == 2

src/exo/worker/tests/unittests/test_runner/test_event_ordering.py

@@ -1,17 +1,12 @@
 # Check tasks are complete before runner is ever ready.
-
-# pyright: reportAny=false
-
 from collections.abc import Iterable
-from typing import Any, Callable
-from unittest.mock import MagicMock
+from typing import Callable
 
 import pytest
 
 import exo.worker.runner.runner as mlx_runner
 from exo.shared.types.api import ChatCompletionMessage
 from exo.shared.types.chunks import TokenChunk
-from exo.shared.types.common import CommandId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
@@ -27,7 +22,6 @@ from exo.shared.types.tasks import (
     Shutdown,
     StartWarmup,
     Task,
-    TaskId,
     TaskStatus,
 )
 from exo.shared.types.worker.runner_response import GenerationResponse
@@ -44,9 +38,6 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import mp_channel
-from exo.worker.engines.mlx.generator.batch_engine import (
-    BatchedGenerationResponse,
-)
 
 from ...constants import (
     CHAT_COMPLETION_TASK_ID,
@@ -116,89 +107,18 @@ def assert_events_equal(test_events: Iterable[Event], true_events: Iterable[Even
         assert test_event == true_event, f"{test_event} != {true_event}"
 
 
-class FakeBatchEngine:
-    """
-    Fake batch engine for testing.
-
-    Queues requests on insert, returns one token per step.
-    The runner's non-blocking loop drains all tasks before running batch steps,
-    so this engine queues requests and has_active_requests returns True only
-    after at least one request has been inserted.
-    """
-
-    def __init__(self, *_args: Any, **_kwargs: Any):
-        self._active_requests: dict[int, tuple[CommandId, TaskId]] = {}
-        self._pending_inserts: list[
-            tuple[CommandId, TaskId, ChatCompletionTaskParams]
-        ] = []
-        self._uid_counter = 0
-        self.rank = 0  # Fake rank for testing
-
-    def queue_request(
-        self,
-        command_id: CommandId,
-        task_id: TaskId,
-        task_params: ChatCompletionTaskParams,
-    ) -> None:
-        """Queue a request for insertion."""
-        self._pending_inserts.append((command_id, task_id, task_params))
-
-    def sync_and_insert_pending(self) -> list[int]:
-        """Insert all pending requests."""
-        uids: list[int] = []
-        for command_id, task_id, _task_params in self._pending_inserts:
-            uid = self._uid_counter
-            self._uid_counter += 1
-            self._active_requests[uid] = (command_id, task_id)
-            uids.append(uid)
-        self._pending_inserts.clear()
-        return uids
-
-    @property
-    def has_pending_inserts(self) -> bool:
-        return len(self._pending_inserts) > 0
-
-    def step(self) -> list[BatchedGenerationResponse]:
-        results: list[BatchedGenerationResponse] = []
-        # Process all active requests - return one token and complete
-        for uid, (command_id, task_id) in list(self._active_requests.items()):
-            results.append(
-                BatchedGenerationResponse(
-                    command_id=command_id,
-                    task_id=task_id,
-                    response=GenerationResponse(
-                        token=0,
-                        text="hi",
-                        finish_reason="stop",
-                    ),
-                )
-            )
-            del self._active_requests[uid]
-        return results
-
-    @property
-    def has_active_requests(self) -> bool:
-        return len(self._active_requests) > 0
-
-    @property
-    def active_count(self) -> int:
-        return len(self._active_requests)
-
-    @property
-    def pending_insert_count(self) -> int:
-        return len(self._pending_inserts)
-
-
 @pytest.fixture
 def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
-    # initialize_mlx returns a fake "group" (non-None for state machine)
-    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MagicMock()))
-    monkeypatch.setattr(
-        mlx_runner, "load_mlx_items", make_nothin((MagicMock(), MagicMock()))
-    )
+    # initialize_mlx returns a "group" equal to 1
+    monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(1))
+    monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, 1)))
     monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
-    monkeypatch.setattr(mlx_runner, "BatchGenerationEngine", FakeBatchEngine)
+
+    def fake_generate(*_1: object, **_2: object):
+        yield GenerationResponse(token=0, text="hi", finish_reason="stop")
+
+    monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
 
 
 def _run(tasks: Iterable[Task]):
@@ -228,8 +148,7 @@ def _run(tasks: Iterable[Task]):
         return event_receiver.collect()
 
 
-def test_chat_completion_generates_and_completes(patch_out_mlx: pytest.MonkeyPatch):
-    """Verify chat completion generates tokens, completes, and runner returns to Ready."""
+def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
     events = _run([INIT_TASK, LOAD_TASK, WARMUP_TASK, CHAT_TASK, SHUTDOWN_TASK])
 
     expected_chunk = ChunkGenerated(
@@ -272,9 +191,7 @@ def test_chat_completion_generates_and_completes(patch_out_mlx: pytest.MonkeyPat
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Running
             ),
             TaskAcknowledged(task_id=CHAT_COMPLETION_TASK_ID),
-            RunnerStatusUpdated(
-                runner_id=RUNNER_1_ID, runner_status=RunnerRunning(active_requests=1)
-            ),
+            RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerRunning()),
             expected_chunk,
             TaskStatusUpdated(
                 task_id=CHAT_COMPLETION_TASK_ID, task_status=TaskStatus.Complete
@@ -289,6 +206,7 @@ def test_chat_completion_generates_and_completes(patch_out_mlx: pytest.MonkeyPat
             TaskStatusUpdated(
                 task_id=SHUTDOWN_TASK_ID, task_status=TaskStatus.Complete
             ),
+            # SPECIAL EXCEPTION FOR RUNNER SHUTDOWN
             RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerShutdown()),
         ],
     )

src/exo/worker/utils/net_profile.py

@@ -1,62 +1,64 @@
-import http.client
-import time
-
-from anyio import create_task_group, to_thread
+import anyio
+import httpx
+from anyio import create_task_group
 from loguru import logger
 
 from exo.shared.topology import Topology
 from exo.shared.types.common import NodeId
 
-BAD_STATUSLINE_ATTEMPTS = 3
+REACHABILITY_ATTEMPTS = 3
 
 
 async def check_reachability(
     target_ip: str,
     expected_node_id: NodeId,
-    self_node_id: NodeId,
     out: dict[NodeId, set[str]],
+    client: httpx.AsyncClient,
 ) -> None:
     """Check if a node is reachable at the given IP and verify its identity."""
+    if ":" in target_ip:
+        # TODO: use real IpAddress types
+        target_ip = f"[{target_ip}]"
+    url = f"http://{target_ip}:52415/node_id"
 
-    # TODO: use an async http client
-    def _fetch_remote_node_id(*, attempt: int = 1) -> NodeId | None:
-        connection = http.client.HTTPConnection(target_ip, 52415, timeout=3)
+    remote_node_id = None
+    last_error = None
+
+    for _ in range(REACHABILITY_ATTEMPTS):
         try:
-            connection.request("GET", "/node_id")
-            response = connection.getresponse()
-            if response.status != 200:
-                return None
+            r = await client.get(url)
+            if r.status_code != 200:
+                await anyio.sleep(1)
+                continue
 
-            body = response.read().decode("utf-8").strip()
+            body = r.text.strip().strip('"')
+            if not body:
+                await anyio.sleep(1)
+                continue
 
-            # Strip quotes if present (JSON string response)
-            if body.startswith('"') and body.endswith('"') and len(body) >= 2:
-                body = body[1:-1]
+            remote_node_id = NodeId(body)
+            break
 
-            return NodeId(body) or None
-        except OSError:
-            return None
-        except http.client.BadStatusLine:
-            if attempt >= BAD_STATUSLINE_ATTEMPTS:
-                logger.warning(
-                    f"BadStatusLine from {target_ip}, after {attempt} attempts, assuming connection to {expected_node_id} has dropped"
-                )
-                return None
-            time.sleep(1)
-            return _fetch_remote_node_id(attempt=attempt + 1)
-        except http.client.HTTPException as e:
-            logger.warning(f"HTTPException from {target_ip}: {type(e).__name__}: {e}")
-            return None
-        finally:
-            connection.close()
+        # expected failure cases
+        except (
+            httpx.TimeoutException,
+            httpx.NetworkError,
+        ):
+            await anyio.sleep(1)
+
+        # other failures should be logged on last attempt
+        except httpx.HTTPError as e:
+            last_error = e
+            await anyio.sleep(1)
+
+    if last_error is not None:
+        logger.warning(
+            f"connect error {type(last_error).__name__} from {target_ip} after {REACHABILITY_ATTEMPTS} attempts; treating as down"
+        )
 
-    remote_node_id = await to_thread.run_sync(_fetch_remote_node_id)
     if remote_node_id is None:
         return
 
-    if remote_node_id == self_node_id:
-        return
-
     if remote_node_id != expected_node_id:
         logger.warning(
             f"Discovered node with unexpected node_id; "
@@ -74,18 +76,33 @@ async def check_reachable(
     topology: Topology, self_node_id: NodeId
 ) -> dict[NodeId, set[str]]:
     """Check which nodes are reachable and return their IPs."""
+
     reachable: dict[NodeId, set[str]] = {}
-    async with create_task_group() as tg:
+
+    # these are intentionally httpx's defaults so we can tune them later
+    timeout = httpx.Timeout(timeout=5.0)
+    limits = httpx.Limits(
+        max_connections=100,
+        max_keepalive_connections=20,
+        keepalive_expiry=5,
+    )
+
+    async with (
+        httpx.AsyncClient(timeout=timeout, limits=limits) as client,
+        create_task_group() as tg,
+    ):
         for node in topology.list_nodes():
             if not node.node_profile:
                 continue
+            if node.node_id == self_node_id:
+                continue
             for iface in node.node_profile.network_interfaces:
                 tg.start_soon(
                     check_reachability,
                     iface.ip_address,
                     node.node_id,
-                    self_node_id,
                     reachable,
+                    client,
                 )
 
     return reachable