feat(sdk): route native commands through run

python/sdk/src/deepseek_harness/api.py

@@ -1,14 +1,11 @@
 from __future__ import annotations
 
 import uuid
-import threading
-from collections import deque
-from collections.abc import Iterator
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Callable, Literal
+from typing import Callable, Literal, overload
 
-from .client import HarnessClient, HarnessConfig, NotificationSubscription, UserQuestionHandler, _PendingRequest
+from .client import HarnessClient, HarnessConfig, UserQuestionHandler
 from .errors import SdkProtocolError, TransportClosedError
 from .models import CommandExecution, JsonObject, Notification
 
@@ -84,101 +81,6 @@ class _PromptActivity:
         )
 
 
-class _QueryResponseIterator(Iterator[Notification | RunResult | CommandExecution]):
-    def __init__(
-        self, client: HarnessClient, session_id: str, request: _PendingRequest,
-        subscription: NotificationSubscription,
-    ) -> None:
-        self._client = client
-        self._session_id = session_id
-        self._request = request
-        self._subscription = subscription
-        self._lock = threading.Lock()
-        self._closed = False
-        self._released = False
-        self._routed = False
-        self._activity: _PromptActivity | None = None
-        self._buffer: deque[Notification] = deque()
-        self._terminal: RunResult | CommandExecution | None = None
-
-    def _release(self, error: BaseException | None = None) -> None:
-        with self._lock:
-            if self._released:
-                return
-            self._released = True
-        self._client._discard_request(self._request, error)
-        if error is None:
-            self._subscription.close()
-        else:
-            self._subscription._cancel(error)
-
-    def close(self) -> None:
-        """Stop observing this query; the runtime operation is not cancelled."""
-        with self._lock:
-            if self._closed:
-                return
-            self._closed = True
-        self._release(TransportClosedError("query response iterator closed"))
-
-    def _ensure_open(self) -> None:
-        if self._closed:
-            raise StopIteration
-
-    def __next__(self) -> Notification | RunResult | CommandExecution:
-        try:
-            self._ensure_open()
-            if not self._routed:
-                response = self._client._wait_request(self._request)
-                self._ensure_open()
-                self._routed = True
-                if not isinstance(response, dict):
-                    raise SdkProtocolError("session/query response must be a JSON object")
-                if response.get("kind") == "command":
-                    self._terminal = CommandExecution.model_validate(response.get("execution"))
-                    self._release()
-                    while True:
-                        try:
-                            notification = self._subscription._next_queued()
-                        except TransportClosedError:
-                            break
-                        if notification is None:
-                            break
-                        self._buffer.append(notification)
-                        if _is_command_done(notification, self._session_id, self._terminal.command_id):
-                            break
-                elif response.get("kind") == "prompt" and isinstance(response.get("messageId"), str) and response["messageId"]:
-                    self._activity = _PromptActivity(self._session_id, response["messageId"])
-                else:
-                    raise SdkProtocolError("session/query requires a prompt messageId or command execution result")
-            self._ensure_open()
-            if self._buffer:
-                return self._buffer.popleft()
-            if self._terminal is not None:
-                terminal = self._terminal
-                self._terminal = None
-                self._release()
-                self._closed = True
-                return terminal
-            assert self._activity is not None
-            while True:
-                notification = self._subscription.next()
-                self._ensure_open()
-                if not self._activity.record(notification):
-                    continue
-                if (notification.method == "session.status"
-                    and notification.payload.get("sessionId") == self._session_id
-                    and notification.payload.get("status") == "idle"):
-                    self._terminal = self._activity.result()
-                    self._release()
-                return notification
-        except BaseException as error:
-            locally_closed = self._closed
-            self.close()
-            if locally_closed and isinstance(error, (TransportClosedError, StopIteration)):
-                raise StopIteration
-            raise
-
-
 class DeepSeekHarness:
     """Reusable synchronous SDK for running DeepSeek Harness agent turns.
 
@@ -265,55 +167,35 @@ class DeepSeekHarness:
         self._client.resume_session(session_id)
         return Session(self, session_id)
 
+    @overload
+    def run(
+        self, input: list[JsonObject], *, session_id: str | None = None,
+        on_notification: Callable[[Notification], None] | None = None,
+    ) -> RunResult: ...
+
+    @overload
+    def run(
+        self, input: str | list[JsonObject], *, session_id: str | None = None,
+        on_notification: Callable[[Notification], None] | None = None,
+        delivery: Literal["queue", "steer"] = "queue",
+    ) -> RunResult | CommandExecution: ...
+
     def run(
         self,
         input: str | list[JsonObject],
         *,
         session_id: str | None = None,
         on_notification: Callable[[Notification], None] | None = None,
-    ) -> RunResult:
-        return self.start_session(session_id).run(input, on_notification=on_notification)
+        delivery: Literal["queue", "steer"] = "queue",
+    ) -> RunResult | CommandExecution:
+        """Run text or a native command on a fresh or named Session."""
+        return self.start_session(session_id).run(input, on_notification=on_notification, delivery=delivery)
 
 
 class Session:
     def __init__(self, harness: DeepSeekHarness, session_id: str) -> None:
         self.harness = harness
         self.id = session_id
-        self._queries: deque[_QueryResponseIterator] = deque()
-        self._query_lock = threading.Lock()
-
-    def query(self, text: str, *, delivery: Literal["queue", "steer"] = "queue") -> None:
-        """Send text for runtime prompt or command routing without waiting for the response.
-
-        Pair each call with receive_response(). Delivery selects where ordinary
-        prompts enter the inbox; commands retain their native execution timing.
-        """
-        if not isinstance(text, str):
-            raise TypeError("query text must be a string")
-        if delivery not in ("queue", "steer"):
-            raise ValueError("query delivery must be queue or steer")
-        self.harness.start()
-        client = self.harness.client
-        with self._query_lock:
-            subscription = client.subscribe_session_notifications(self.id)
-            try:
-                request = client._begin_request("session/query", {"sessionId": self.id, "text": text, "delivery": delivery})
-            except BaseException:
-                subscription.close()
-                raise
-            self._queries.append(_QueryResponseIterator(client, self.id, request, subscription))
-
-    def receive_response(self) -> _QueryResponseIterator:
-        """Claim the next query in send order and iterate its notifications and terminal result.
-
-        Prompts collect the inbox-receipt-to-idle interval used by run(); commands
-        finish on their RPC result. Close an unfinished iterator to release its
-        local observation resources without cancelling the runtime operation.
-        """
-        with self._query_lock:
-            if not self._queries:
-                raise RuntimeError("receive_response requires a preceding query")
-            return self._queries.popleft()
 
     def get_working_directory(self) -> str:
         """Read the effective directory, recovering a missing directory to the Session origin."""
@@ -325,20 +207,68 @@ class Session:
         self.harness.start()
         return self.harness.client.set_working_directory(self.id, path)
 
+    @overload
+    def run(
+        self, input: list[JsonObject], *,
+        on_notification: Callable[[Notification], None] | None = None,
+    ) -> RunResult: ...
+
+    @overload
+    def run(
+        self, input: str | list[JsonObject], *,
+        on_notification: Callable[[Notification], None] | None = None,
+        delivery: Literal["queue", "steer"] = "queue",
+    ) -> RunResult | CommandExecution: ...
+
     def run(
         self,
         input: str | list[JsonObject],
         *,
         on_notification: Callable[[Notification], None] | None = None,
-    ) -> RunResult:
-        content_blocks = normalize_input(input)
-        with self.harness.client.subscribe_session_notifications(self.id) as subscription:
-            message_id = self.harness.client.session_prompt(
-                self.id,
-                content_blocks,
-                notification_subscription=subscription,
-            )
+        delivery: Literal["queue", "steer"] = "queue",
+    ) -> RunResult | CommandExecution:
+        """Run text through native input routing, or queue literal content blocks.
 
+        Text returns the prompt's receipt-to-idle RunResult or a settled
+        CommandExecution; command completion does not wait for later model work.
+        Content blocks stay literal and accept only queue delivery. Notifications
+        retain wire order; command notifications are delivered after its response,
+        up to its own command/done event. Handler errors remain command results;
+        routing, transport and observer failures raise without rolling back work.
+        """
+        if delivery not in ("queue", "steer"):
+            raise ValueError("run delivery must be queue or steer")
+        if not isinstance(input, str) and delivery != "queue":
+            raise ValueError("content-block input requires queue delivery")
+        self.harness.start()
+        with self.harness.client.subscribe_session_notifications(self.id) as subscription:
+            if isinstance(input, str):
+                response = self.harness.client._request_raw(
+                    "session/query", {"sessionId": self.id, "text": input, "delivery": delivery},
+                )
+                if not isinstance(response, dict):
+                    raise SdkProtocolError("session/query response must be a JSON object")
+                if response.get("kind") == "command":
+                    execution = CommandExecution.model_validate(response.get("execution"))
+                    while True:
+                        try:
+                            notification = subscription._next_queued()
+                        except TransportClosedError:
+                            break
+                        if notification is None:
+                            break
+                        if on_notification is not None:
+                            on_notification(notification)
+                        if _is_command_done(notification, self.id, execution.command_id):
+                            break
+                    return execution
+                if response.get("kind") != "prompt" or not isinstance(response.get("messageId"), str) or not response["messageId"]:
+                    raise SdkProtocolError("session/query requires a prompt messageId or command execution result")
+                message_id = response["messageId"]
+            else:
+                message_id = self.harness.client.session_prompt(
+                    self.id, input, notification_subscription=subscription,
+                )
             activity = _PromptActivity(self.id, message_id)
             while True:
                 notification = subscription.next()

python/sdk/src/deepseek_harness/client.py

@@ -39,14 +39,6 @@ class HarnessConfig:
     on_user_question: UserQuestionHandler | None = None
 
 
-@dataclass(slots=True)
-class _PendingRequest:
-    request_id: str
-    method: str
-    waiter: queue.Queue[JsonValue | BaseException]
-    deadline: float | None
-
-
 class HarnessClient:
     """Synchronous JSON-RPC client for the DeepSeek Harness SDK runtime over stdio."""
 
@@ -340,29 +332,15 @@ class HarnessClient:
         notification_filter: NotificationFilter | None = None,
         notification_subscription: "NotificationSubscription | None" = None,
     ) -> JsonValue:
+        request_id = str(uuid.uuid4())
+        waiter: queue.Queue[JsonValue | BaseException] = queue.Queue(maxsize=1)
         temp_subscription: NotificationSubscription | None = None
         subscription = notification_subscription
+        with self._lock:
+            self._responses[request_id] = waiter
         if on_notification is not None and subscription is None:
             temp_subscription = self.subscribe_notifications(notification_filter)
             subscription = temp_subscription
-        try:
-            pending = self._begin_request(method, params, timeout_seconds=timeout_seconds)
-            return self._wait_request(pending, on_notification=on_notification, notification_subscription=subscription)
-        finally:
-            if temp_subscription is not None:
-                temp_subscription.close()
-
-    def _begin_request(
-        self,
-        method: str,
-        params: JsonObject | None = None,
-        *,
-        timeout_seconds: float | None = None,
-    ) -> _PendingRequest:
-        request_id = str(uuid.uuid4())
-        waiter: queue.Queue[JsonValue | BaseException] = queue.Queue(maxsize=1)
-        with self._lock:
-            self._responses[request_id] = waiter
         try:
             message: JsonObject = {"jsonrpc": "2.0", "id": request_id, "method": method}
             if params is not None:
@@ -371,59 +349,48 @@ class HarnessClient:
         except BaseException:
             with self._lock:
                 self._responses.pop(request_id, None)
+            if temp_subscription is not None:
+                temp_subscription.close()
             raise
         timeout = self.config.request_timeout_seconds if timeout_seconds is None else timeout_seconds
         deadline = None if timeout is None else time.monotonic() + timeout
-        return _PendingRequest(request_id, method, waiter, deadline)
-
-    def _discard_request(self, pending: _PendingRequest, error: BaseException | None = None) -> None:
-        with self._lock:
-            if self._responses.get(pending.request_id) is pending.waiter:
-                self._responses.pop(pending.request_id)
-                if error is not None:
-                    pending.waiter.put_nowait(error)
-
-    def _wait_request(
-        self,
-        pending: _PendingRequest,
-        *,
-        on_notification: Callable[[Notification], None] | None = None,
-        notification_subscription: "NotificationSubscription | None" = None,
-    ) -> JsonValue:
-        subscription = notification_subscription
         try:
             while True:
                 if on_notification is not None and subscription is not None:
                     subscription.drain(on_notification)
-                try:
-                    item = pending.waiter.get_nowait()
-                    break
-                except queue.Empty:
-                    pass
                 wait_timeout = None
                 if on_notification is not None:
                     wait_timeout = 0.05
-                if pending.deadline is not None:
-                    remaining = pending.deadline - time.monotonic()
+                if deadline is not None:
+                    remaining = deadline - time.monotonic()
                     if remaining <= 0:
+                        with self._lock:
+                            self._responses.pop(request_id, None)
                         diagnostics = self._runtime_diagnostics()
                         suffix = f"\n{diagnostics}" if diagnostics else ""
                         raise TimeoutError(
-                            f"{pending.method} timed out waiting for DeepSeek Harness runtime{suffix}"
+                            f"{method} timed out waiting for DeepSeek Harness runtime{suffix}"
                         )
                     wait_timeout = remaining if wait_timeout is None else min(wait_timeout, remaining)
                 try:
-                    item = pending.waiter.get(timeout=wait_timeout)
+                    item = waiter.get(timeout=wait_timeout)
+                    if on_notification is not None and subscription is not None:
+                        subscription.drain(on_notification)
                     break
                 except queue.Empty:
                     continue
-            if on_notification is not None and subscription is not None:
-                subscription.drain(on_notification)
-            if isinstance(item, BaseException):
-                raise item
-            return item
+        except BaseException:
+            with self._lock:
+                self._responses.pop(request_id, None)
+            if temp_subscription is not None:
+                temp_subscription.close()
+            raise
         finally:
-            self._discard_request(pending)
+            if temp_subscription is not None:
+                temp_subscription.close()
+        if isinstance(item, BaseException):
+            raise item
+        return item
 
     def _write_message(self, message: JsonObject) -> None:
         proc = self._proc
@@ -716,10 +683,6 @@ class NotificationSubscription:
         self._closed = True
         self._client._unsubscribe_notifications(self._subscription_id)
 
-    def _cancel(self, error: BaseException) -> None:
-        self.close()
-        self._notifications.put(error)
-
     def next(self, timeout_seconds: float | None = None) -> Notification:
         """Read one notification; an optional timeout stops only this local wait."""
         try: