1from __future__ import annotations

2

3import json

4import os

5import queue

6import subprocess

7import threading

8import time

9import uuid

10from collections import deque

11from dataclasses import dataclass

12from pathlib import Path

13from typing import Callable, TypeAlias, TypeVar

14

15from pydantic import BaseModel

16

17from .errors import JsonRpcError, TransportClosedError

18from .models import IncomingRequest, InitializeResponse, JsonObject, JsonValue, Notification

19

20ModelT = TypeVar("ModelT", bound=BaseModel)

21NotificationFilter: TypeAlias = Callable[[Notification], bool]

22

23

24@dataclass(slots=True)

25class HarnessConfig:

26 """Configuration for launching the local DeepSeek Harness SDK runtime."""

27

28 dsh_bin: str | None = None

29 profile: str = "sdk"

30 patches: tuple[str, ...] = ()

31 dsh_home: str | None = None

32 cwd: str | None = None

33 env: dict[str, str] | None = None

34 initialize_timeout_seconds: float = 30.0

35 request_timeout_seconds: float | None = None

36 shutdown_timeout_seconds: float | None = 1.0

37

38

39class HarnessClient:

40 """Synchronous JSON-RPC client for the DeepSeek Harness SDK runtime over stdio."""

41

42 def __init__(

43 self,

44 config: HarnessConfig | None = None,

45 *,

46 _launch_args: tuple[str, ...] | None = None,

47 ) -> None:

48 self.config = config or HarnessConfig()

49 self._launch_args = _launch_args

50 self._proc: subprocess.Popen[str] | None = None

51 self._lock = threading.Lock()

52 self._write_lock = threading.Lock()

53 self._responses: dict[str, queue.Queue[JsonValue | BaseException]] = {}

54 self._notifications: queue.Queue[Notification | BaseException] = queue.Queue()

55 self._notification_subscribers: dict[

56 str, tuple[queue.Queue[Notification | BaseException], NotificationFilter | None]

57 ] = {}

58 self._session_parents: dict[str, str] = {}

59 self._requests: queue.Queue[IncomingRequest | BaseException] = queue.Queue()

60 self._stderr_lines: deque[str] = deque(maxlen=400)

61 self._reader_thread: threading.Thread | None = None

62 self._stderr_thread: threading.Thread | None = None

63

64 def __enter__(self) -> "HarnessClient":

65 self.start()

66 return self

67

68 def __exit__(self, _exc_type, _exc, _tb) -> None:

69 self.close()

70

71 def start(self) -> None:

72 if self._proc is not None:

73 return

74 with self._lock:

75 self._session_parents.clear()

76 env = os.environ.copy()

77 if self.config.env:

78 env.update(self.config.env)

79 args = list(self._launch_args or self._default_launch_args(env))

80 self._proc = subprocess.Popen(

81 args,

82 stdin=subprocess.PIPE,

83 stdout=subprocess.PIPE,

84 stderr=subprocess.PIPE,

85 text=True,

86 encoding="utf-8",

87 cwd=None if self.config.cwd is None else str(Path(self.config.cwd).resolve()),

88 env=env,

89 bufsize=1,

90 )

91 self._start_reader_thread()

92 self._start_stderr_thread()

93

94 def close(self) -> None:

95 """Close the runtime after a bounded opportunity to flush durable state."""

96 proc = self._proc

97 if proc is None:

98 return

99 shutdown_completed = False

100 try:

101 self.request("shutdown", None, response_model=_ShutdownResponse, timeout_seconds=self.config.shutdown_timeout_seconds)

102 shutdown_completed = True

103 except Exception as exc:

104 self._stderr_lines.append(f"shutdown request failed: {exc}")

105 if proc.stdin:

106 try:

107 proc.stdin.close()

108 except Exception as exc:

109 self._stderr_lines.append(f"stdin close failed: {exc}")

110 if shutdown_completed:

111 try:

112 proc.wait(timeout=self.config.shutdown_timeout_seconds)

113 except subprocess.TimeoutExpired:

114 pass

115 if proc.poll() is None:

116 try:

117 proc.terminate()

118 except ProcessLookupError:

119 pass

120 if proc.poll() is None:

121 try:

122 proc.wait(timeout=self.config.shutdown_timeout_seconds)

123 except subprocess.TimeoutExpired:

124 proc.kill()

125 proc.wait()

126 self._proc = None

127 self._fail_waiters(self._runtime_closed_error("DeepSeek Harness runtime closed"))

128 if self._reader_thread and self._reader_thread.is_alive():

129 self._reader_thread.join(timeout=0.5)

130 if self._stderr_thread and self._stderr_thread.is_alive():

131 self._stderr_thread.join(timeout=0.5)

132

133 def initialize(

134 self,

135 *,

136 cwd: str,

137 provider: str,

138 model: str,

139 reasoning_effort: str | None = None,

140 max_tokens: int | None = None,

141 ) -> InitializeResponse:

142 payload: JsonObject = {

143 "cwd": str(Path(cwd).resolve()),

144 "provider": provider,

145 "model": model,

146 }

147 if reasoning_effort is not None:

148 payload["reasoningEffort"] = reasoning_effort

149 if max_tokens is not None:

150 payload["maxTokens"] = max_tokens

151 try:

152 return self.request(

153 "initialize",

154 payload,

155 response_model=InitializeResponse,

156 timeout_seconds=self.config.initialize_timeout_seconds,

157 )

158 except TimeoutError as error:

159 self.close()

160 raise TimeoutError(f"{error}\nselected dsh profile {self.config.profile!r}") from error

161 except BaseException as error:

162 self.close()

163 diagnostics = self._runtime_diagnostics()

164 if isinstance(error, JsonRpcError) and diagnostics:

165 raise JsonRpcError(

166 error.code,

167 f"{error.message}\n{diagnostics}",

168 error.data,

169 ) from error

170 raise

171

172 def session_prompt(

173 self,

174 session_id: str,

175 content_blocks: list[JsonObject],

176 *,

177 on_notification: Callable[[Notification], None] | None = None,

178 notification_subscription: "NotificationSubscription | None" = None,

179 ) -> str:

180 payload: JsonObject = {"sessionId": session_id, "contentBlocks": content_blocks}

181 response = self.request(

182 "session/prompt",

183 payload,

184 response_model=_SessionPromptResponse,

185 on_notification=on_notification,

186 notification_filter=self._notification_belongs_to_session_tree(session_id),

187 notification_subscription=notification_subscription,

188 )

189 return response.messageId

190

191 def request(

192 self,

193 method: str,

194 params: JsonObject | None,

195 *,

196 response_model: type[ModelT],

197 timeout_seconds: float | None = None,

198 on_notification: Callable[[Notification], None] | None = None,

199 notification_filter: NotificationFilter | None = None,

200 notification_subscription: "NotificationSubscription | None" = None,

201 ) -> ModelT:

202 result = self._request_raw(

203 method,

204 params,

205 timeout_seconds=timeout_seconds,

206 on_notification=on_notification,

207 notification_filter=notification_filter,

208 notification_subscription=notification_subscription,

209 )

210 if not isinstance(result, dict):

211 raise TypeError(f"{method} response must be a JSON object")

212 return response_model.model_validate(result)

213

214 def notify(self, method: str, params: JsonObject | None = None) -> None:

215 message: JsonObject = {"jsonrpc": "2.0", "method": method}

216 if params is not None:

217 message["params"] = params

218 self._write_message(message)

219

220 def next_notification(self) -> Notification:

221 item = self._notifications.get()

222 if isinstance(item, BaseException):

223 raise item

224 return item

225

226 def subscribe_notifications(

227 self,

228 notification_filter: NotificationFilter | None = None,

229 ) -> "NotificationSubscription":

230 subscription_id = str(uuid.uuid4())

231 notifications: queue.Queue[Notification | BaseException] = queue.Queue()

232 with self._lock:

233 self._notification_subscribers[subscription_id] = (notifications, notification_filter)

234 return NotificationSubscription(self, subscription_id, notifications)

235

236 def subscribe_session_notifications(self, session_id: str) -> "NotificationSubscription":

237 """Subscribe to a session and descendants discovered from subagent lifecycle edges."""

238 return self.subscribe_notifications(self._notification_belongs_to_session_tree(session_id))

239

240 def next_request(self) -> IncomingRequest:

241 item = self._requests.get()

242 if isinstance(item, BaseException):

243 raise item

244 return item

245

246 def respond(self, request_id: str | int, result: JsonValue) -> None:

247 self._write_message({"jsonrpc": "2.0", "id": request_id, "result": result})

248

249 def respond_error(

250 self,

251 request_id: str | int,

252 *,

253 code: int,

254 message: str,

255 data: JsonValue | None = None,

256 ) -> None:

257 error: JsonObject = {"code": code, "message": message}

258 if data is not None:

259 error["data"] = data

260 self._write_message({"jsonrpc": "2.0", "id": request_id, "error": error})

261

262 def _request_raw(

263 self,

264 method: str,

265 params: JsonObject | None = None,

266 *,

267 timeout_seconds: float | None = None,

268 on_notification: Callable[[Notification], None] | None = None,

269 notification_filter: NotificationFilter | None = None,

270 notification_subscription: "NotificationSubscription | None" = None,

271 ) -> JsonValue:

272 request_id = str(uuid.uuid4())

273 waiter: queue.Queue[JsonValue | BaseException] = queue.Queue(maxsize=1)

274 temp_subscription: NotificationSubscription | None = None

275 subscription = notification_subscription

276 with self._lock:

277 self._responses[request_id] = waiter

278 if on_notification is not None and subscription is None:

279 temp_subscription = self.subscribe_notifications(notification_filter)

280 subscription = temp_subscription

281 try:

282 message: JsonObject = {"jsonrpc": "2.0", "id": request_id, "method": method}

283 if params is not None:

284 message["params"] = params

285 self._write_message(message)

286 except BaseException:

287 with self._lock:

288 self._responses.pop(request_id, None)

289 if temp_subscription is not None:

290 temp_subscription.close()

291 raise

292 timeout = self.config.request_timeout_seconds if timeout_seconds is None else timeout_seconds

293 deadline = None if timeout is None else time.monotonic() + timeout

294 try:

295 while True:

296 if on_notification is not None and subscription is not None:

297 subscription.drain(on_notification)

298 wait_timeout = None

299 if on_notification is not None:

300 wait_timeout = 0.05

301 if deadline is not None:

302 remaining = deadline - time.monotonic()

303 if remaining <= 0:

304 with self._lock:

305 self._responses.pop(request_id, None)

306 diagnostics = self._runtime_diagnostics()

307 suffix = f"\n{diagnostics}" if diagnostics else ""

308 raise TimeoutError(

309 f"{method} timed out waiting for DeepSeek Harness runtime{suffix}"

310 )

311 wait_timeout = remaining if wait_timeout is None else min(wait_timeout, remaining)

312 try:

313 item = waiter.get(timeout=wait_timeout)

314 if on_notification is not None and subscription is not None:

315 subscription.drain(on_notification)

316 break

317 except queue.Empty:

318 continue

319 except BaseException:

320 with self._lock:

321 self._responses.pop(request_id, None)

322 if temp_subscription is not None:

323 temp_subscription.close()

324 raise

325 finally:

326 if temp_subscription is not None:

327 temp_subscription.close()

328 if isinstance(item, BaseException):

329 raise item

330 return item

331

332 def _write_message(self, message: JsonObject) -> None:

333 proc = self._proc

334 if proc is None or proc.stdin is None:

335 raise TransportClosedError("DeepSeek Harness runtime is not running")

336 try:

337 payload = json.dumps(message, separators=(",", ":")) + "\n"

338 with self._write_lock:

339 proc.stdin.write(payload)

340 proc.stdin.flush()

341 except Exception as exc:

342 raise self._runtime_closed_error("Failed to write to DeepSeek Harness runtime") from exc

343

344 def _start_reader_thread(self) -> None:

345 self._reader_thread = threading.Thread(target=self._reader_loop, name="dsh-runtime-reader", daemon=True)

346 self._reader_thread.start()

347

348 def _start_stderr_thread(self) -> None:

349 self._stderr_thread = threading.Thread(target=self._stderr_loop, name="dsh-runtime-stderr", daemon=True)

350 self._stderr_thread.start()

351

352 def _reader_loop(self) -> None:

353 proc = self._proc

354 if proc is None or proc.stdout is None:

355 return

356 try:

357 for line in proc.stdout:

358 if not line.strip():

359 continue

360 try:

361 message = json.loads(line)

362 except json.JSONDecodeError:

363 continue

364 self._handle_message(message)

365 except BaseException as exc:

366 self._fail_waiters(exc)

367 finally:

368 self._fail_waiters(self._runtime_closed_error("DeepSeek Harness runtime stdout closed"))

369

370 def _stderr_loop(self) -> None:

371 proc = self._proc

372 if proc is None or proc.stderr is None:

373 return

374 for line in proc.stderr:

375 self._stderr_lines.append(line.rstrip())

376

377 def _handle_message(self, message: object) -> None:

378 if not isinstance(message, dict):

379 return

380 msg_id = message.get("id")

381 method = message.get("method")

382 if isinstance(msg_id, (str, int)) and isinstance(method, str):

383 params = message.get("params")

384 self._requests.put(IncomingRequest(id=msg_id, method=method, payload=params if isinstance(params, dict) else {}))

385 return

386 if isinstance(msg_id, (str, int)):

387 with self._lock:

388 waiter = self._responses.pop(str(msg_id), None)

389 if waiter is None:

390 return

391 if isinstance(message.get("error"), dict):

392 err = message["error"]

393 waiter.put(JsonRpcError(_int_or_none(err.get("code")), str(err.get("message", "JSON-RPC error")), err.get("data")))

394 else:

395 waiter.put(message.get("result"))

396 return

397 if isinstance(method, str):

398 params = message.get("params")

399 notification = Notification(method=method, payload=params if isinstance(params, dict) else {})

400 with self._lock:

401 self._record_session_relationship_locked(notification)

402 subscribers = list(self._notification_subscribers.items())

403 delivered = False

404 for subscription_id, (subscriber, predicate) in subscribers:

405 try:

406 matches = predicate is None or predicate(notification)

407 except BaseException as exc:

408 with self._lock:

409 current = self._notification_subscribers.get(subscription_id)

410 if current is not None and current[0] is subscriber:

411 self._notification_subscribers.pop(subscription_id, None)

412 subscriber.put(exc)

413 continue

414 if matches:

415 subscriber.put(notification)

416 delivered = True

417 if not delivered:

418 self._notifications.put(notification)

419

420 def _fail_waiters(self, exc: BaseException) -> None:

421 with self._lock:

422 waiters = list(self._responses.values())

423 self._responses.clear()

424 subscribers = list(self._notification_subscribers.values())

425 self._notification_subscribers.clear()

426 for waiter in waiters:

427 waiter.put(exc)

428 for subscriber, _predicate in subscribers:

429 subscriber.put(exc)

430 self._notifications.put(exc)

431 self._requests.put(exc)

432

433 def _runtime_closed_error(self, reason: str) -> TransportClosedError:

434 diagnostics = self._runtime_diagnostics()

435 return TransportClosedError(f"{reason}\n{diagnostics}" if diagnostics else reason)

436

437 def _runtime_diagnostics(self) -> str:

438 """Return available subprocess state for transport failures and timeouts."""

439 proc = self._proc

440 if (

441 proc is not None

442 and proc.poll() is not None

443 and self._stderr_thread is not None

444 and self._stderr_thread.is_alive()

445 and threading.current_thread() is not self._stderr_thread

446 ):

447 self._stderr_thread.join(timeout=0.1)

448

449 parts: list[str] = []

450 if proc is not None:

451 exit_code = proc.poll()

452 if exit_code is not None:

453 parts.append(f"exit code: {exit_code}")

454 if self._stderr_lines:

455 parts.append("stderr tail:\n" + "\n".join(self._stderr_lines))

456 return "\n".join(parts)

457

458 def _default_launch_args(self, env: dict[str, str]) -> tuple[str, ...]:

459 if self.config.dsh_bin is None:

460 try:

461 from deepseek_harness_runtime import resolve_bundled_launch_args

462 except ImportError as exc:

463 raise FileNotFoundError(

464 "Unable to locate the bundled DeepSeek Harness dsh runtime. "

465 "Install deepseek-harness-runtime-bin."

466 ) from exc

467 base = resolve_bundled_launch_args()

468 else:

469 base = (str(Path(self.config.dsh_bin).expanduser().resolve()),)

470

471 if self.config.dsh_home is not None:

472 if not self.config.dsh_home.strip():

473 raise ValueError("HarnessConfig requires a non-empty dsh_home")

474 env["DSH_HOME"] = str(Path(self.config.dsh_home).expanduser().resolve())

475 elif not env.get("DSH_HOME", "").strip():

476 raise ValueError(

477 "HarnessConfig requires an explicit dsh_home or non-empty DSH_HOME; "

478 "the Python SDK never uses ~/.dsh implicitly"

479 )

480

481 patches = tuple(

482 argument

483 for patch in self.config.patches

484 for argument in ("--patch", str(Path(patch).expanduser().resolve()))

485 )

486 return (*base, "--profile", self.config.profile, *patches)

487

488 def _unsubscribe_notifications(self, subscription_id: str) -> None:

489 with self._lock:

490 self._notification_subscribers.pop(subscription_id, None)

491

492 def _record_session_relationship_locked(self, notification: Notification) -> None:

493 if notification.method != "subagent.started":

494 return

495 parent_id = notification.payload.get("parentSessionId")

496 child_id = notification.payload.get("childSessionId")

497 if (

498 isinstance(parent_id, str)

499 and parent_id

500 and isinstance(child_id, str)

501 and child_id

502 and parent_id != child_id

503 ):

504 self._session_parents[child_id] = parent_id

505

506 def _notification_belongs_to_session_tree(self, session_id: str) -> NotificationFilter:

507 def belongs(notification: Notification) -> bool:

508 payload = notification.payload

509 if notification.method in {"subagent.started", "subagent.finished"}:

510 parent_id = payload.get("parentSessionId")

511 if (

512 isinstance(parent_id, str)

513 and self._session_is_descendant_of(parent_id, session_id)

514 ):

515 return True

516 return payload.get("childSessionId") == session_id

517 related_id = payload.get("sessionId")

518 return (

519 isinstance(related_id, str)

520 and self._session_is_descendant_of(related_id, session_id)

521 )

522

523 return belongs

524

525 def _session_is_descendant_of(self, session_id: str, root_session_id: str) -> bool:

526 current = session_id

527 visited: set[str] = set()

528 while current not in visited:

529 if current == root_session_id:

530 return True

531 visited.add(current)

532 parent = self._session_parents.get(current)

533 if parent is None:

534 return False

535 current = parent

536 return False

537

538

539class NotificationSubscription:

540 def __init__(

541 self,

542 client: HarnessClient,

543 subscription_id: str,

544 notifications: queue.Queue[Notification | BaseException],

545 ) -> None:

546 self._client = client

547 self._subscription_id = subscription_id

548 self._notifications = notifications

549 self._closed = False

550

551 def __enter__(self) -> "NotificationSubscription":

552 return self

553

554 def __exit__(self, _exc_type, _exc, _tb) -> None:

555 self.close()

556

557 def close(self) -> None:

558 if self._closed:

559 return

560 self._closed = True

561 self._client._unsubscribe_notifications(self._subscription_id)

562

563 def next(self) -> Notification:

564 item = self._notifications.get()

565 if isinstance(item, BaseException):

566 raise item

567 return item

568

569 def drain(self, on_notification: Callable[[Notification], None]) -> None:

570 while True:

571 try:

572 item = self._notifications.get_nowait()

573 except queue.Empty:

574 return

575 if isinstance(item, BaseException):

576 raise item

577 on_notification(item)

578

579

580class _SessionPromptResponse(BaseModel):

581 messageId: str

582

583

584class _ShutdownResponse(BaseModel):

585 pass

586

587

588def _int_or_none(value: object) -> int | None:

589 return value if isinstance(value, int) else None