@neoline/hostbridge 2.0.4 → 2.0.5

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (40) hide show
  1. package/README.md +11 -9
  2. package/docs/ARCHITECTURE.md +62 -0
  3. package/package.json +2 -1
  4. package/pyproject.toml +1 -1
  5. package/src/server_control_mcp/__init__.py +4 -3
  6. package/src/server_control_mcp/async_lifecycle.py +41 -0
  7. package/src/server_control_mcp/cli.py +1 -1
  8. package/src/server_control_mcp/client.py +302 -287
  9. package/src/server_control_mcp/config.py +2 -1
  10. package/src/server_control_mcp/daemon.py +233 -526
  11. package/src/server_control_mcp/doctor.py +34 -2
  12. package/src/server_control_mcp/hosts.py +136 -2
  13. package/src/server_control_mcp/mock_ssh.py +52 -14
  14. package/src/server_control_mcp/mock_ssh_exec.py +147 -58
  15. package/src/server_control_mcp/mock_ssh_session.py +3 -2
  16. package/src/server_control_mcp/mock_ssh_sftp.py +126 -28
  17. package/src/server_control_mcp/mock_ssh_tunnel.py +141 -444
  18. package/src/server_control_mcp/mux_connection.py +326 -0
  19. package/src/server_control_mcp/mux_daemon.py +186 -0
  20. package/src/server_control_mcp/mux_protocol.py +279 -0
  21. package/src/server_control_mcp/mux_records.py +71 -0
  22. package/src/server_control_mcp/mux_rpc.py +1779 -0
  23. package/src/server_control_mcp/mux_service.py +506 -0
  24. package/src/server_control_mcp/mux_stream.py +283 -0
  25. package/src/server_control_mcp/mux_sync_client.py +224 -0
  26. package/src/server_control_mcp/remote_agent_bundle.py +56 -0
  27. package/src/server_control_mcp/remote_mux_agent.py +769 -0
  28. package/src/server_control_mcp/runtime_services.py +862 -0
  29. package/src/server_control_mcp/server.py +33 -18
  30. package/src/server_control_mcp/services.py +2 -666
  31. package/src/server_control_mcp/transports/__init__.py +4 -8
  32. package/src/server_control_mcp/transports/base.py +60 -38
  33. package/src/server_control_mcp/transports/ssh.py +206 -259
  34. package/src/server_control_mcp/tunnel_manager.py +1245 -0
  35. package/src/server_control_mcp/tunnel_native_ssh.py +454 -0
  36. package/src/server_control_mcp/tunnel_providers.py +20 -0
  37. package/src/server_control_mcp/tunnel_pty_agent.py +909 -0
  38. package/src/server_control_mcp/protocol.py +0 -363
  39. package/src/server_control_mcp/remote_tunnel_agent.py +0 -143
  40. package/src/server_control_mcp/transports/pty.py +0 -515
@@ -1,37 +1,24 @@
1
- """Typed synchronous client for the HostBridge daemon protocol."""
1
+ """Typed synchronous client for the HostBridge multiplexed daemon protocol."""
2
2
 
3
3
  from __future__ import annotations
4
4
 
5
- import base64
6
- import binascii
7
5
  import hashlib
8
- import json
9
- import socket
10
6
  import threading
7
+ import time
11
8
  import uuid
12
9
  from collections.abc import Iterable, Mapping
13
10
  from contextlib import suppress
14
11
  from dataclasses import dataclass
15
12
  from pathlib import Path
13
+ from types import TracebackType
16
14
 
17
15
  from . import __version__
18
- from .protocol import (
19
- MAX_FRAME_BYTES,
20
- PROTOCOL_VERSION,
21
- BinaryFrame,
22
- ErrorInfo,
23
- FrameFlags,
24
- FrameSequenceValidator,
25
- ProtocolError,
26
- Request,
27
- Response,
28
- encode_frame,
29
- read_frame,
30
- )
16
+ from .mux_protocol import MUX_PROTOCOL_VERSION
17
+ from .mux_records import MuxRecordChannel, MuxRecordDecoder, MuxRecordError
18
+ from .mux_rpc import MuxCapacityError, MuxConnectionClosed, MuxRemoteError
19
+ from .mux_sync_client import MuxSyncClient, MuxSyncStream
31
20
  from .runtime import runtime_paths
32
- from .transports.base import ShellReadResult, TransferResult
33
-
34
- MAX_CONTROL_LINE_BYTES = 4 * 1024 * 1024
21
+ from .transports.base import ShellReadResult, ShellTerminalResult, TransferResult
35
22
 
36
23
 
37
24
  def _required_text(value: object, field: str) -> str:
@@ -40,23 +27,6 @@ def _required_text(value: object, field: str) -> str:
40
27
  return value
41
28
 
42
29
 
43
- def _encode_inline_bytes(value: bytes | None) -> dict[str, str] | None:
44
- if value is None:
45
- return None
46
- if not isinstance(value, bytes):
47
- raise TypeError("stdin must be bytes or None")
48
- return {"encoding": "base64", "data": base64.b64encode(value).decode("ascii")}
49
-
50
-
51
- def _decode_inline_bytes(value: object, field: str) -> bytes:
52
- if not isinstance(value, dict) or value.get("encoding") != "base64" or not isinstance(value.get("data"), str):
53
- raise HostBridgeError("protocol_mismatch", f"response field {field} must contain base64 data")
54
- try:
55
- return base64.b64decode(value["data"], validate=True)
56
- except (binascii.Error, ValueError) as exc:
57
- raise HostBridgeError("protocol_mismatch", f"response field {field} contains invalid base64") from exc
58
-
59
-
60
30
  @dataclass(frozen=True, slots=True)
61
31
  class SessionInfo:
62
32
  session_id: str
@@ -83,21 +53,70 @@ class ExecResult:
83
53
  stderr: bytes
84
54
  timed_out: bool = False
85
55
 
86
- @classmethod
87
- def from_dict(cls, value: Mapping[str, object]) -> ExecResult:
88
- exit_code = value.get("exit_code")
89
- if exit_code is not None and (not isinstance(exit_code, int) or isinstance(exit_code, bool)):
90
- raise HostBridgeError("protocol_mismatch", "response field exit_code must be an integer or null")
91
- timed_out = value.get("timed_out", False)
92
- if not isinstance(timed_out, bool):
93
- raise HostBridgeError("protocol_mismatch", "response field timed_out must be a boolean")
94
- return cls(
95
- session_id=_required_text(value.get("session_id"), "session_id"),
96
- exit_code=exit_code,
97
- stdout=_decode_inline_bytes(value.get("stdout"), "stdout"),
98
- stderr=_decode_inline_bytes(value.get("stderr"), "stderr"),
99
- timed_out=timed_out,
100
- )
56
+
57
+ class HostBridgeShell:
58
+ def __init__(self, stream: MuxSyncStream) -> None:
59
+ self._stream = stream
60
+ self._state_lock = threading.Lock()
61
+ self._closed = False
62
+ self._input_closed = False
63
+ self._terminal: ShellTerminalResult | None = None
64
+
65
+ @property
66
+ def closed(self) -> bool:
67
+ with self._state_lock:
68
+ return self._closed
69
+
70
+ @property
71
+ def exit_code(self) -> int | None:
72
+ with self._state_lock:
73
+ return None if self._terminal is None else self._terminal.exit_code
74
+
75
+ @property
76
+ def exit_signal(self) -> str | None:
77
+ with self._state_lock:
78
+ return None if self._terminal is None else self._terminal.signal
79
+
80
+ def receive(self, max_bytes: int = 64 * 1024) -> bytes | None:
81
+ if max_bytes <= 0:
82
+ raise ValueError("max_bytes must be positive")
83
+ with self._state_lock:
84
+ if self._closed:
85
+ return None
86
+ chunk = self._stream.receive_unbounded(max_bytes)
87
+ if chunk is not None:
88
+ return chunk
89
+ result = self._stream.finish(timeout=1)
90
+ try:
91
+ terminal = ShellTerminalResult.from_mapping(result)
92
+ except ValueError as exc:
93
+ raise HostBridgeError("protocol_mismatch", f"invalid shell terminal result: {exc}") from exc
94
+ with self._state_lock:
95
+ self._terminal = terminal
96
+ self._closed = True
97
+ return None
98
+
99
+ def send(self, data: bytes) -> None:
100
+ if not isinstance(data, bytes):
101
+ raise TypeError("shell data must be bytes")
102
+ with self._state_lock:
103
+ if self._closed or self._input_closed:
104
+ raise HostBridgeError("connection_lost", "shell input is closed")
105
+ self._stream.send(data)
106
+
107
+ def send_eof(self) -> None:
108
+ with self._state_lock:
109
+ if self._closed or self._input_closed:
110
+ return
111
+ self._input_closed = True
112
+ self._stream.send_eof()
113
+
114
+ def close(self) -> None:
115
+ with self._state_lock:
116
+ if self._closed:
117
+ return
118
+ self._closed = True
119
+ self._stream.reset("shell output consumer closed", timeout=1)
101
120
 
102
121
 
103
122
  class HostBridgeError(RuntimeError):
@@ -114,10 +133,6 @@ class HostBridgeError(RuntimeError):
114
133
  self.retryable = retryable
115
134
  self.details = details or {}
116
135
 
117
- @classmethod
118
- def from_info(cls, error: ErrorInfo) -> HostBridgeError:
119
- return cls(error.code, error.message, retryable=error.retryable, details=error.details)
120
-
121
136
 
122
137
  class HostBridgeClient:
123
138
  def __init__(
@@ -130,6 +145,25 @@ class HostBridgeClient:
130
145
  raise ValueError("default_timeout must be positive")
131
146
  self.socket_path = runtime_paths(socket_path).socket
132
147
  self.default_timeout = float(default_timeout)
148
+ self._mux = MuxSyncClient(self.socket_path, default_timeout=self.default_timeout)
149
+
150
+ @property
151
+ def closed(self) -> bool:
152
+ return self._mux.closed
153
+
154
+ def close(self) -> None:
155
+ self._mux.close()
156
+
157
+ def __enter__(self) -> HostBridgeClient:
158
+ return self
159
+
160
+ def __exit__(
161
+ self,
162
+ exc_type: type[BaseException] | None,
163
+ exc: BaseException | None,
164
+ traceback: TracebackType | None,
165
+ ) -> None:
166
+ self.close()
133
167
 
134
168
  def request(
135
169
  self,
@@ -138,19 +172,16 @@ class HostBridgeClient:
138
172
  *,
139
173
  timeout: float | None = None,
140
174
  ) -> dict[str, object]:
141
- request = Request.new(method, dict(params))
142
- response = self._exchange(request, timeout=self.default_timeout if timeout is None else timeout)
143
- if not response.ok:
144
- if response.error is None:
145
- raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
146
- raise HostBridgeError.from_info(response.error)
147
- return response.result
175
+ try:
176
+ return self._mux.request(method, params, timeout=timeout)
177
+ except BaseException as exc:
178
+ raise _translate_error(exc) from exc
148
179
 
149
180
  def hello(self) -> dict[str, object]:
150
181
  result = self.request("system.hello", {})
151
182
  if (
152
183
  result.get("service") != "hostbridge"
153
- or result.get("protocol") != PROTOCOL_VERSION
184
+ or result.get("protocol") != MUX_PROTOCOL_VERSION
154
185
  or result.get("version") != __version__
155
186
  ):
156
187
  raise HostBridgeError(
@@ -159,6 +190,25 @@ class HostBridgeClient:
159
190
  )
160
191
  return result
161
192
 
193
+ def list_hosts(self) -> list[dict[str, object]]:
194
+ hosts = self.request("host.list", {}).get("hosts")
195
+ if not isinstance(hosts, list) or any(not isinstance(host, dict) for host in hosts):
196
+ raise HostBridgeError("protocol_mismatch", "host list response must contain objects")
197
+ return hosts
198
+
199
+ def tunnel_health(self) -> dict[str, object]:
200
+ result = self.request("tunnel.health", {})
201
+ hosts = result.get("hosts")
202
+ if not isinstance(hosts, list) or any(not isinstance(host, dict) for host in hosts):
203
+ raise HostBridgeError("protocol_mismatch", "tunnel health response must contain host objects")
204
+ return result
205
+
206
+ def list_sessions(self) -> list[dict[str, object]]:
207
+ sessions = self.request("session.list", {}).get("sessions")
208
+ if not isinstance(sessions, list) or any(not isinstance(session, dict) for session in sessions):
209
+ raise HostBridgeError("protocol_mismatch", "session list response must contain objects")
210
+ return sessions
211
+
162
212
  def open_session(
163
213
  self,
164
214
  host_id: str,
@@ -185,19 +235,17 @@ class HostBridgeClient:
185
235
  output_limit: int | None = None,
186
236
  owner: str | None = None,
187
237
  ) -> ExecResult:
188
- params: dict[str, object] = {
189
- "session_id": session_id,
190
- "command": command,
191
- "stdin": _encode_inline_bytes(stdin),
192
- }
193
- if timeout is not None:
194
- params["timeout"] = timeout
195
- if output_limit is not None:
196
- params["output_limit"] = output_limit
197
- if owner is not None:
198
- params["owner"] = owner
199
- result = self.request("exec.start", params, timeout=timeout)
200
- return ExecResult.from_dict(result)
238
+ if stdin is not None and not isinstance(stdin, bytes):
239
+ raise TypeError("stdin must be bytes or None")
240
+ chunks: Iterable[bytes] = () if stdin is None else (stdin,)
241
+ return self.exec_stream(
242
+ session_id,
243
+ command,
244
+ chunks,
245
+ timeout=timeout,
246
+ output_limit=output_limit,
247
+ owner=owner,
248
+ )
201
249
 
202
250
  def exec_stream(
203
251
  self,
@@ -210,58 +258,43 @@ class HostBridgeClient:
210
258
  owner: str | None = None,
211
259
  cancel_event: threading.Event | None = None,
212
260
  ) -> ExecResult:
213
- stream_id = uuid.uuid4().hex
214
- params: dict[str, object] = {
215
- "session_id": session_id,
216
- "command": command,
217
- "stream_id": stream_id,
218
- }
261
+ params: dict[str, object] = {"session_id": session_id, "command": command}
219
262
  if timeout is not None:
220
263
  params["timeout"] = timeout
221
264
  if output_limit is not None:
222
265
  params["output_limit"] = output_limit
223
266
  if owner is not None:
224
267
  params["owner"] = owner
225
- request = Request.new("exec.stream", params)
226
- request_timeout = self.default_timeout if timeout is None else timeout
227
- sequence = 0
268
+ deadline = _deadline(self.default_timeout if timeout is None else timeout)
269
+ stream = self._open_stream("exec.start", params, deadline)
270
+ decoder = MuxRecordDecoder()
271
+ stdout = bytearray()
272
+ stderr = bytearray()
228
273
  try:
229
- with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
230
- connection.settimeout(request_timeout)
231
- connection.connect(str(self.socket_path))
232
- connection.sendall(request.to_json().encode("utf-8") + b"\n")
233
- try:
234
- for chunk in stdin:
235
- if cancel_event is not None and cancel_event.is_set():
236
- raise HostBridgeError("cancelled", "exec stdin cancelled")
237
- if not isinstance(chunk, bytes):
238
- raise TypeError("exec stdin chunks must be bytes")
239
- for offset in range(0, len(chunk), MAX_FRAME_BYTES):
240
- payload = chunk[offset : offset + MAX_FRAME_BYTES]
241
- connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.DATA, payload)))
242
- sequence += 1
243
- if cancel_event is not None and cancel_event.is_set():
244
- raise HostBridgeError("cancelled", "exec stdin cancelled")
245
- connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.EOF)))
246
- except BaseException as exc:
247
- metadata = json.dumps(
248
- {"code": "cancelled", "message": str(exc) or "exec stdin failed"},
249
- separators=(",", ":"),
250
- ).encode("utf-8")
251
- with suppress(OSError):
252
- connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.CANCEL, metadata)))
253
- raise
254
- with connection.makefile("rb") as reader:
255
- response = self._read_response(reader, request)
256
- except HostBridgeError:
257
- raise
258
- except TimeoutError as exc:
259
- raise HostBridgeError("timeout", "HostBridge exec stream timed out", retryable=True) from exc
260
- except OSError as exc:
261
- raise HostBridgeError(
262
- "connection_lost", "HostBridge exec stream connection failed", retryable=True
263
- ) from exc
264
- return ExecResult.from_dict(response.result)
274
+ self._send_input(stream, stdin, deadline, cancel_event=cancel_event, label="exec stdin")
275
+ while chunk := stream.receive(timeout=_remaining(deadline)):
276
+ for record in decoder.feed(chunk):
277
+ target = stdout if record.channel is MuxRecordChannel.STDOUT else stderr
278
+ target.extend(record.payload)
279
+ if decoder.buffered_bytes:
280
+ raise HostBridgeError("protocol_mismatch", "exec output ended with an incomplete binary record")
281
+ result = stream.finish(timeout=_remaining(deadline))
282
+ except BaseException as exc:
283
+ self._reset_stream(stream, str(exc))
284
+ raise _translate_error(exc) from exc
285
+ exit_code = result.get("exit_code")
286
+ if exit_code is not None and (not isinstance(exit_code, int) or isinstance(exit_code, bool)):
287
+ raise HostBridgeError("protocol_mismatch", "response field exit_code must be an integer or null")
288
+ timed_out = result.get("timed_out", False)
289
+ if not isinstance(timed_out, bool):
290
+ raise HostBridgeError("protocol_mismatch", "response field timed_out must be a boolean")
291
+ return ExecResult(
292
+ _required_text(result.get("session_id"), "session_id"),
293
+ exit_code,
294
+ bytes(stdout),
295
+ bytes(stderr),
296
+ timed_out,
297
+ )
265
298
 
266
299
  def close_session(self, session_id: str, *, owner: str | None = None) -> None:
267
300
  params: dict[str, object] = {"session_id": session_id}
@@ -269,6 +302,9 @@ class HostBridgeClient:
269
302
  params["owner"] = owner
270
303
  self.request("session.close", params)
271
304
 
305
+ def release_mock_ssh_session(self, session_id: str) -> None:
306
+ self.request("session.release_mock_ssh", {"session_id": session_id})
307
+
272
308
  def write_text(
273
309
  self,
274
310
  session_id: str,
@@ -288,8 +324,7 @@ class HostBridgeClient:
288
324
  }
289
325
  if owner is not None:
290
326
  params["owner"] = owner
291
- result = self.request("file.write_text", params)
292
- return self._transfer_result(result)
327
+ return self._transfer_result(self.request("file.write_text", params))
293
328
 
294
329
  def read_text(
295
330
  self,
@@ -299,28 +334,55 @@ class HostBridgeClient:
299
334
  max_bytes: int = 4 * 1024 * 1024,
300
335
  owner: str | None = None,
301
336
  ) -> str:
302
- params: dict[str, object] = {"session_id": session_id, "remote_path": remote_path, "max_bytes": max_bytes}
337
+ params: dict[str, object] = {
338
+ "session_id": session_id,
339
+ "remote_path": remote_path,
340
+ "max_bytes": max_bytes,
341
+ }
303
342
  if owner is not None:
304
343
  params["owner"] = owner
305
- result = self.request("file.read_text", params)
306
- content = result.get("content")
344
+ content = self.request("file.read_text", params).get("content")
307
345
  if not isinstance(content, str):
308
346
  raise HostBridgeError("protocol_mismatch", "text response content must be a string")
309
347
  return content
310
348
 
311
349
  def shell_write(self, session_id: str, data: bytes, *, owner: str | None = None) -> int:
312
- params: dict[str, object] = {
313
- "session_id": session_id,
314
- "data": _encode_inline_bytes(data),
315
- }
350
+ if not isinstance(data, bytes):
351
+ raise TypeError("shell data must be bytes")
352
+ params: dict[str, object] = {"session_id": session_id, "max_bytes": max(1, len(data))}
316
353
  if owner is not None:
317
354
  params["owner"] = owner
318
- result = self.request("shell.write", params)
355
+ deadline = _deadline(self.default_timeout)
356
+ stream = self._open_stream("shell.write", params, deadline)
357
+ try:
358
+ self._send_input(stream, (data,), deadline, cancel_event=None, label="shell input")
359
+ while stream.receive(timeout=_remaining(deadline)) is not None:
360
+ pass
361
+ result = stream.finish(timeout=_remaining(deadline))
362
+ except BaseException as exc:
363
+ self._reset_stream(stream, str(exc))
364
+ raise _translate_error(exc) from exc
319
365
  written = result.get("written")
320
366
  if not isinstance(written, int) or isinstance(written, bool) or written < 0:
321
367
  raise HostBridgeError("protocol_mismatch", "shell write response contains invalid byte count")
322
368
  return written
323
369
 
370
+ def open_shell(
371
+ self,
372
+ session_id: str,
373
+ *,
374
+ owner: str | None = None,
375
+ max_bytes: int = 64 * 1024,
376
+ ) -> HostBridgeShell:
377
+ if max_bytes <= 0:
378
+ raise ValueError("max_bytes must be positive")
379
+ params: dict[str, object] = {"session_id": session_id, "max_bytes": max_bytes}
380
+ if owner is not None:
381
+ params["owner"] = owner
382
+ deadline = _deadline(self.default_timeout)
383
+ stream = self._open_stream("shell.attach", params, deadline)
384
+ return HostBridgeShell(stream)
385
+
324
386
  def shell_read(
325
387
  self,
326
388
  session_id: str,
@@ -329,18 +391,24 @@ class HostBridgeClient:
329
391
  max_bytes: int = 4096,
330
392
  owner: str | None = None,
331
393
  ) -> ShellReadResult:
332
- params: dict[str, object] = {
333
- "session_id": session_id,
334
- "timeout": timeout,
335
- "max_bytes": max_bytes,
336
- }
394
+ params: dict[str, object] = {"session_id": session_id, "timeout": timeout, "max_bytes": max_bytes}
337
395
  if owner is not None:
338
396
  params["owner"] = owner
339
- result = self.request("shell.read", params, timeout=max(self.default_timeout, timeout + 1))
397
+ deadline = _deadline(max(self.default_timeout, timeout + 1))
398
+ stream = self._open_stream("shell.read", params, deadline)
399
+ output = bytearray()
400
+ try:
401
+ stream.send_eof(timeout=_remaining(deadline))
402
+ while chunk := stream.receive(timeout=_remaining(deadline)):
403
+ output.extend(chunk)
404
+ result = stream.finish(timeout=_remaining(deadline))
405
+ except BaseException as exc:
406
+ self._reset_stream(stream, str(exc))
407
+ raise _translate_error(exc) from exc
340
408
  alive = result.get("alive")
341
409
  if not isinstance(alive, bool):
342
410
  raise HostBridgeError("protocol_mismatch", "shell read response contains invalid alive flag")
343
- return ShellReadResult(_decode_inline_bytes(result.get("data"), "data"), alive)
411
+ return ShellReadResult(bytes(output), alive)
344
412
 
345
413
  def upload_file(
346
414
  self,
@@ -356,55 +424,37 @@ class HostBridgeClient:
356
424
  max_bytes: int | None = None,
357
425
  ) -> TransferResult:
358
426
  self._validate_chunk_size(chunk_size)
427
+ self._validate_optional_byte_limit(max_bytes)
359
428
  source = Path(local_path)
360
429
  expected_size = source.stat().st_size
361
- self._validate_optional_byte_limit(max_bytes)
362
430
  if max_bytes is not None and expected_size > max_bytes:
363
431
  raise HostBridgeError("resource_limit", f"upload exceeds {max_bytes} bytes")
364
- stream_id = uuid.uuid4().hex
365
432
  params: dict[str, object] = {
366
433
  "session_id": session_id,
367
434
  "remote_path": remote_path,
368
- "stream_id": stream_id,
369
435
  "mode": mode,
370
436
  "expected_size": expected_size,
371
437
  }
372
438
  if owner is not None:
373
439
  params["owner"] = owner
374
- request = Request.new("transfer.upload", params)
440
+ deadline = _deadline(self.default_timeout if timeout is None else timeout)
441
+ stream = self._open_stream("transfer.upload", params, deadline)
375
442
  digest = hashlib.sha256()
376
- sequence = 0
377
- request_timeout = self.default_timeout if timeout is None else timeout
443
+
444
+ def chunks() -> Iterable[bytes]:
445
+ with source.open("rb") as handle:
446
+ while chunk := handle.read(chunk_size):
447
+ digest.update(chunk)
448
+ yield chunk
449
+
378
450
  try:
379
- with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
380
- connection.settimeout(request_timeout)
381
- connection.connect(str(self.socket_path))
382
- connection.sendall(request.to_json().encode("utf-8") + b"\n")
383
- with source.open("rb") as local_file:
384
- while True:
385
- if cancel_event is not None and cancel_event.is_set():
386
- payload = json.dumps(
387
- {"code": "cancelled", "message": "upload cancelled by client"},
388
- separators=(",", ":"),
389
- ).encode("utf-8")
390
- connection.sendall(
391
- encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.CANCEL, payload))
392
- )
393
- break
394
- chunk = local_file.read(chunk_size)
395
- if not chunk:
396
- connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.EOF)))
397
- break
398
- digest.update(chunk)
399
- connection.sendall(encode_frame(BinaryFrame(stream_id, sequence, FrameFlags.DATA, chunk)))
400
- sequence += 1
401
- with connection.makefile("rb") as reader:
402
- response = self._read_response(reader, request)
403
- except TimeoutError as exc:
404
- raise HostBridgeError("timeout", "HostBridge upload timed out", retryable=True) from exc
405
- except OSError as exc:
406
- raise HostBridgeError("connection_lost", "HostBridge upload connection failed", retryable=True) from exc
407
- result = self._transfer_result(response.result)
451
+ self._send_input(stream, chunks(), deadline, cancel_event=cancel_event, label="upload")
452
+ while stream.receive(timeout=_remaining(deadline)) is not None:
453
+ pass
454
+ result = self._transfer_result(stream.finish(timeout=_remaining(deadline)))
455
+ except BaseException as exc:
456
+ self._reset_stream(stream, str(exc))
457
+ raise _translate_error(exc) from exc
408
458
  if result.bytes_transferred != expected_size or result.sha256 != digest.hexdigest():
409
459
  raise HostBridgeError("integrity_error", "HostBridge upload integrity check failed")
410
460
  return result
@@ -425,66 +475,70 @@ class HostBridgeClient:
425
475
  destination = Path(local_path)
426
476
  destination.parent.mkdir(parents=True, exist_ok=True)
427
477
  temporary = destination.with_name(f".{destination.name}.hostbridge-{uuid.uuid4().hex}.tmp")
428
- stream_id = uuid.uuid4().hex
429
478
  params: dict[str, object] = {
430
479
  "session_id": session_id,
431
480
  "remote_path": remote_path,
432
- "stream_id": stream_id,
433
481
  "chunk_size": chunk_size,
434
482
  }
435
483
  if owner is not None:
436
484
  params["owner"] = owner
437
- request = Request.new("transfer.download", params)
485
+ deadline = _deadline(self.default_timeout if timeout is None else timeout)
486
+ stream = self._open_stream("transfer.download", params, deadline)
438
487
  digest = hashlib.sha256()
439
488
  transferred = 0
440
- validator = FrameSequenceValidator()
441
- request_timeout = self.default_timeout if timeout is None else timeout
442
489
  try:
443
- with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
444
- connection.settimeout(request_timeout)
445
- connection.connect(str(self.socket_path))
446
- connection.sendall(request.to_json().encode("utf-8") + b"\n")
447
- with connection.makefile("rb") as reader:
448
- response = self._read_response(reader, request)
449
- if response.result.get("stream_id") != stream_id:
450
- raise HostBridgeError("protocol_mismatch", "download stream_id does not match")
451
- with temporary.open("wb") as local_file:
452
- while True:
453
- frame = read_frame(reader)
454
- validator.accept(frame)
455
- if frame.stream_id != stream_id:
456
- raise HostBridgeError("protocol_mismatch", "download frame stream_id does not match")
457
- if frame.flags & FrameFlags.CANCEL:
458
- self._raise_cancel_frame(frame)
459
- if frame.flags & FrameFlags.DATA:
460
- local_file.write(frame.payload)
461
- digest.update(frame.payload)
462
- transferred += len(frame.payload)
463
- if max_bytes is not None and transferred > max_bytes:
464
- raise HostBridgeError("resource_limit", f"download exceeds {max_bytes} bytes")
465
- if frame.flags & FrameFlags.EOF:
466
- metadata = self._frame_metadata(frame.payload)
467
- break
468
- except HostBridgeError:
469
- temporary.unlink(missing_ok=True)
470
- raise
471
- except TimeoutError as exc:
490
+ stream.send_eof(timeout=_remaining(deadline))
491
+ with temporary.open("wb") as handle:
492
+ while chunk := stream.receive(timeout=_remaining(deadline)):
493
+ transferred += len(chunk)
494
+ if max_bytes is not None and transferred > max_bytes:
495
+ raise HostBridgeError("resource_limit", f"download exceeds {max_bytes} bytes")
496
+ handle.write(chunk)
497
+ digest.update(chunk)
498
+ result = self._transfer_result(stream.finish(timeout=_remaining(deadline)))
499
+ if result.bytes_transferred != transferred or result.sha256 != digest.hexdigest():
500
+ raise HostBridgeError("integrity_error", "HostBridge download integrity check failed")
501
+ temporary.replace(destination)
502
+ return result
503
+ except BaseException as exc:
472
504
  temporary.unlink(missing_ok=True)
473
- raise HostBridgeError("timeout", "HostBridge download timed out", retryable=True) from exc
474
- except (OSError, ProtocolError) as exc:
475
- temporary.unlink(missing_ok=True)
476
- raise HostBridgeError("connection_lost", "HostBridge download connection failed", retryable=True) from exc
477
- expected = self._transfer_result(metadata)
478
- if expected.bytes_transferred != transferred or expected.sha256 != digest.hexdigest():
479
- temporary.unlink(missing_ok=True)
480
- raise HostBridgeError("integrity_error", "HostBridge download integrity check failed")
481
- temporary.replace(destination)
482
- return expected
505
+ self._reset_stream(stream, str(exc))
506
+ raise _translate_error(exc) from exc
507
+
508
+ def _open_stream(self, method: str, params: Mapping[str, object], deadline: float) -> MuxSyncStream:
509
+ try:
510
+ return self._mux.open_stream(method, params, timeout=_remaining(deadline))
511
+ except BaseException as exc:
512
+ raise _translate_error(exc) from exc
513
+
514
+ @staticmethod
515
+ def _send_input(
516
+ stream: MuxSyncStream,
517
+ chunks: Iterable[bytes],
518
+ deadline: float,
519
+ *,
520
+ cancel_event: threading.Event | None,
521
+ label: str,
522
+ ) -> None:
523
+ for chunk in chunks:
524
+ if cancel_event is not None and cancel_event.is_set():
525
+ raise HostBridgeError("cancelled", f"{label} cancelled")
526
+ if not isinstance(chunk, bytes):
527
+ raise TypeError(f"{label} chunks must be bytes")
528
+ stream.send(chunk, timeout=_remaining(deadline))
529
+ if cancel_event is not None and cancel_event.is_set():
530
+ raise HostBridgeError("cancelled", f"{label} cancelled")
531
+ stream.send_eof(timeout=_remaining(deadline))
532
+
533
+ @staticmethod
534
+ def _reset_stream(stream: MuxSyncStream, message: str) -> None:
535
+ with suppress(BaseException):
536
+ stream.reset(message or "client operation failed", timeout=1)
483
537
 
484
538
  @staticmethod
485
539
  def _validate_chunk_size(chunk_size: int) -> None:
486
- if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= MAX_FRAME_BYTES:
487
- raise ValueError(f"chunk_size must be between 4096 and {MAX_FRAME_BYTES} bytes")
540
+ if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or not 4096 <= chunk_size <= 1024 * 1024:
541
+ raise ValueError("chunk_size must be between 4096 and 1048576 bytes")
488
542
 
489
543
  @staticmethod
490
544
  def _validate_optional_byte_limit(max_bytes: int | None) -> None:
@@ -501,70 +555,31 @@ class HostBridgeClient:
501
555
  raise HostBridgeError("protocol_mismatch", "transfer response contains invalid SHA-256")
502
556
  return TransferResult(byte_count, sha256.lower())
503
557
 
504
- def _read_response(self, reader, request: Request) -> Response:
505
- raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
506
- if not raw_response:
507
- raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
508
- if len(raw_response) > MAX_CONTROL_LINE_BYTES:
509
- raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
510
- try:
511
- response = Response.from_json(raw_response.decode("utf-8"))
512
- except (ProtocolError, UnicodeDecodeError) as exc:
513
- raise HostBridgeError("protocol_mismatch", "HostBridge daemon returned an invalid response") from exc
514
- if response.request_id != request.request_id:
515
- raise HostBridgeError("protocol_mismatch", "HostBridge daemon response request_id does not match")
516
- if not response.ok:
517
- if response.error is None:
518
- raise HostBridgeError("protocol_mismatch", "failed response omitted error details")
519
- raise HostBridgeError.from_info(response.error)
520
- return response
521
-
522
- @staticmethod
523
- def _frame_metadata(payload: bytes) -> dict[str, object]:
524
- try:
525
- value = json.loads(payload.decode("utf-8"))
526
- except (UnicodeDecodeError, json.JSONDecodeError) as exc:
527
- raise HostBridgeError("protocol_mismatch", "download EOF contains invalid metadata") from exc
528
- if not isinstance(value, dict):
529
- raise HostBridgeError("protocol_mismatch", "download EOF metadata must be an object")
530
- return value
531
-
532
- @staticmethod
533
- def _raise_cancel_frame(frame: BinaryFrame) -> None:
534
- metadata = HostBridgeClient._frame_metadata(frame.payload)
535
- code = metadata.get("code")
536
- message = metadata.get("message")
537
- raise HostBridgeError(
538
- code if isinstance(code, str) else "cancelled",
539
- message if isinstance(message, str) else "HostBridge transfer cancelled",
540
- )
541
-
542
- def _exchange(self, request: Request, *, timeout: float) -> Response:
543
- if timeout <= 0:
544
- raise ValueError("timeout must be positive")
545
- try:
546
- with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
547
- connection.settimeout(timeout)
548
- connection.connect(str(self.socket_path))
549
- connection.sendall(request.to_json().encode("utf-8") + b"\n")
550
- with connection.makefile("rb") as reader:
551
- raw_response = reader.readline(MAX_CONTROL_LINE_BYTES + 1)
552
- except TimeoutError as exc:
553
- raise HostBridgeError("timeout", "HostBridge daemon request timed out", retryable=True) from exc
554
- except OSError as exc:
555
- raise HostBridgeError("connection_lost", "HostBridge daemon connection failed", retryable=True) from exc
556
- if not raw_response:
557
- raise HostBridgeError("connection_lost", "HostBridge daemon closed before responding", retryable=True)
558
- if len(raw_response) > MAX_CONTROL_LINE_BYTES:
559
- raise HostBridgeError("protocol_mismatch", "HostBridge daemon response exceeds the control limit")
560
- return self._read_response_from_bytes(raw_response, request)
561
558
 
562
- @staticmethod
563
- def _read_response_from_bytes(raw_response: bytes, request: Request) -> Response:
564
- try:
565
- response = Response.from_json(raw_response.decode("utf-8"))
566
- except (ProtocolError, UnicodeDecodeError) as exc:
567
- raise HostBridgeError("protocol_mismatch", "HostBridge daemon returned an invalid response") from exc
568
- if response.request_id != request.request_id:
569
- raise HostBridgeError("protocol_mismatch", "HostBridge daemon response request_id does not match")
570
- return response
559
+ def _deadline(timeout: float) -> float:
560
+ if timeout <= 0:
561
+ raise ValueError("timeout must be positive")
562
+ return time.monotonic() + timeout
563
+
564
+
565
+ def _remaining(deadline: float) -> float:
566
+ remaining = deadline - time.monotonic()
567
+ if remaining <= 0:
568
+ raise HostBridgeError("timeout", "HostBridge operation timed out", retryable=True)
569
+ return remaining
570
+
571
+
572
+ def _translate_error(error: BaseException) -> BaseException:
573
+ if isinstance(error, MuxRecordError):
574
+ return HostBridgeError("protocol_mismatch", str(error))
575
+ if isinstance(error, HostBridgeError | TypeError | ValueError):
576
+ return error
577
+ if isinstance(error, MuxRemoteError):
578
+ return HostBridgeError(error.code, str(error), retryable=error.retryable, details=error.details)
579
+ if isinstance(error, MuxCapacityError):
580
+ return HostBridgeError("busy", str(error), retryable=True)
581
+ if isinstance(error, MuxConnectionClosed | RuntimeError):
582
+ return HostBridgeError("connection_lost", str(error), retryable=True)
583
+ if isinstance(error, TimeoutError):
584
+ return HostBridgeError("timeout", str(error), retryable=True)
585
+ return error