@neoline/hostbridge 1.4.3 → 2.0.0

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.
@@ -1,922 +0,0 @@
1
- from __future__ import annotations
2
-
3
- import base64
4
- import hashlib
5
- import re
6
- import shlex
7
- import threading
8
- import time
9
- import uuid
10
- from collections.abc import Callable
11
- from contextlib import suppress
12
- from dataclasses import dataclass, field
13
- from pathlib import Path
14
- from typing import Protocol, cast
15
-
16
- import pexpect
17
-
18
- from .hosts import HostRegistry
19
- from .policy import AuditSink, PolicyConfig, PolicyDecision, audit_event
20
-
21
- DEFAULT_CONNECT_TIMEOUT = 30
22
- DEFAULT_COMMAND_TIMEOUT = 60
23
- DEFAULT_REAPER_INTERVAL_SECONDS = 60
24
- MAX_OUTPUT_CHARS = 120_000
25
- REMOTE_JOB_ROOT = "$HOME/.hostbridge/jobs"
26
- REMOTE_SHELL = "bash"
27
- MAX_SHELL_TRANSFER_BYTES = 100 * 1024 * 1024
28
- DEFAULT_MAX_TRANSFER_BYTES = MAX_SHELL_TRANSFER_BYTES
29
- UPLOAD_BASE64_CHUNK_CHARS = 256 * 1024
30
- DOWNLOAD_CHUNK_BYTES = 5 * 1024 * 1024
31
-
32
- class PolicyDeniedError(RuntimeError):
33
- """Raised when a command is blocked by the configured policy."""
34
-
35
-
36
- class ChildProcess(Protocol):
37
- before: str
38
- after: object
39
-
40
- def sendline(self, line: str) -> None: ...
41
-
42
- def send(self, text: str) -> None: ...
43
-
44
- def expect(self, patterns: object, timeout: int | float = -1) -> int: ...
45
-
46
- def isalive(self) -> bool: ...
47
-
48
- def close(self, force: bool = False) -> None: ...
49
-
50
- def read_nonblocking(self, size: int = 1, timeout: int | float = -1) -> str: ...
51
-
52
-
53
- ChildFactory = Callable[[str, list[str], int | float], ChildProcess]
54
- TokenFactory = Callable[[], str]
55
-
56
-
57
- class RemoteCommandError(RuntimeError):
58
- """Raised when a remote session cannot execute a requested operation."""
59
-
60
-
61
- @dataclass(slots=True)
62
- class RemoteSession:
63
- id: str
64
- host_id: str
65
- label: str
66
- command: str
67
- args: list[str]
68
- child: ChildProcess
69
- created_at: float = field(default_factory=time.time)
70
- last_used_at: float = field(default_factory=time.time)
71
- busy: bool = False
72
- owner_agent_id: str | None = None
73
- last_output: str = ""
74
-
75
-
76
- @dataclass(slots=True)
77
- class CommandResult:
78
- session_id: str
79
- host_id: str
80
- command: str
81
- exit_code: int | None
82
- output: str
83
- timed_out: bool = False
84
-
85
-
86
- @dataclass(slots=True)
87
- class RemoteTask:
88
- id: str
89
- session_id: str
90
- host_id: str
91
- command: str
92
- job_dir: str
93
- remote_pid: int | None
94
- owner_agent_id: str | None = None
95
- created_at: float = field(default_factory=time.time)
96
-
97
-
98
- def _new_token() -> str:
99
- return uuid.uuid4().hex[:12]
100
-
101
-
102
- def _default_child_factory(command: str, args: list[str], timeout: int | float) -> ChildProcess:
103
- return pexpect.spawn(
104
- command,
105
- args,
106
- encoding="utf-8",
107
- codec_errors="replace",
108
- timeout=timeout,
109
- echo=False,
110
- )
111
-
112
-
113
- def _tail(value: str, limit: int = MAX_OUTPUT_CHARS) -> str:
114
- if len(value) <= limit:
115
- return value
116
- return value[-limit:]
117
-
118
-
119
- def _extract_between_markers(raw: str, start_marker: str) -> str:
120
- raw = raw.replace("\x1b[?2004h", "").replace("\x1b[?2004l", "")
121
- if start_marker in raw:
122
- raw = raw.split(start_marker, 1)[1]
123
- lines = raw.replace("\r\n", "\n").replace("\r", "\n").split("\n")
124
- cleaned = [
125
- line
126
- for line in lines
127
- if not line.startswith("printf '") and "__SCM_START_" not in line and "__SCM_EXIT_" not in line
128
- ]
129
- return "\n".join(cleaned).strip("\n")
130
-
131
-
132
- def _parse_exit_marker(marker: object, token: str) -> int | None:
133
- text = marker.group(0) if hasattr(marker, "group") else str(marker)
134
- match = re.search(rf"__SCM_EXIT_{re.escape(token)}:(\d+)", text)
135
- return int(match.group(1)) if match else None
136
-
137
-
138
- def _contains_terminal_control(text: str) -> bool:
139
- return "\x1b" in text or "\x07" in text or "\\x1b" in text or "\\u001b" in text
140
-
141
-
142
- def _looks_shell_prompt(text: str) -> bool:
143
- return re.search(r"[#$](?:\\s\*)?$", text.strip()) is not None
144
-
145
-
146
- def _runtime_expect_pattern(pattern: str) -> str:
147
- if _contains_terminal_control(pattern) and _looks_shell_prompt(pattern):
148
- return r"(?m)[^\r\n]*[#$]\s*$"
149
- return pattern
150
-
151
-
152
- class SessionManager:
153
- """Manages long-lived interactive PTYs and remote background jobs."""
154
-
155
- def __init__(
156
- self,
157
- hosts: HostRegistry,
158
- *,
159
- child_factory: ChildFactory = _default_child_factory,
160
- remote_shell: str = REMOTE_SHELL,
161
- policy: PolicyConfig | None = None,
162
- audit_sink: AuditSink | None = None,
163
- max_sessions_per_host: int = 0,
164
- ):
165
- self.hosts = hosts
166
- self.child_factory = child_factory
167
- self.remote_shell = remote_shell
168
- self._sessions: dict[str, RemoteSession] = {}
169
- self._tasks: dict[str, RemoteTask] = {}
170
- self._policy = policy
171
- self._audit = audit_sink
172
- self._max_sessions_per_host = max(0, int(max_sessions_per_host))
173
- self._reaper_stop = threading.Event()
174
- self._reaper_thread: threading.Thread | None = None
175
-
176
- @property
177
- def policy(self) -> PolicyConfig | None:
178
- return self._policy
179
-
180
- def start_idle_reaper(self, *, interval: int = DEFAULT_REAPER_INTERVAL_SECONDS) -> None:
181
- """Spawn a background thread that closes idle sessions past policy timeout."""
182
-
183
- if self._policy is None or not self._policy.enabled or self._policy.idle_timeout_seconds <= 0:
184
- return
185
- if self._reaper_thread is not None and self._reaper_thread.is_alive():
186
- return
187
- self._reaper_stop.clear()
188
- timeout = self._policy.idle_timeout_seconds
189
- thread = threading.Thread(
190
- target=self._reap_loop,
191
- kwargs={"interval": interval, "timeout": timeout},
192
- daemon=True,
193
- name="scm-idle-reaper",
194
- )
195
- self._reaper_thread = thread
196
- thread.start()
197
-
198
- def stop_idle_reaper(self) -> None:
199
- self._reaper_stop.set()
200
- thread = self._reaper_thread
201
- if thread is not None and thread.is_alive():
202
- thread.join(timeout=2.0)
203
- self._reaper_thread = None
204
-
205
- def _reap_loop(self, *, interval: int, timeout: int) -> None:
206
- while not self._reaper_stop.wait(interval):
207
- self._reap_idle(timeout)
208
-
209
- def _reap_idle(self, timeout: int) -> int:
210
- now = time.time()
211
- expired = [
212
- sid for sid, session in list(self._sessions.items())
213
- if not session.busy and (now - session.last_used_at) > timeout
214
- ]
215
- for sid in expired:
216
- self._safe_close(sid)
217
- self._emit("session_idle_closed", session_id=sid, detail=f"idle > {timeout}s")
218
- return len(expired)
219
-
220
- def _safe_close(self, session_id: str) -> None:
221
- session = self._sessions.pop(session_id, None)
222
- if session is not None:
223
- with suppress(Exception):
224
- session.child.close(force=True)
225
-
226
- def _emit(self, event: str, **fields: object) -> None:
227
- if self._audit is None:
228
- return
229
- audit_event(self._audit, event=event, **{k: v for k, v in fields.items() if v is not None}) # type: ignore[arg-type]
230
-
231
- def _check_policy(self, command: str) -> PolicyDecision | None:
232
- if self._policy is None:
233
- return None
234
- return self._policy.evaluate(command)
235
-
236
- def open_session(
237
- self,
238
- host_id: str,
239
- *,
240
- connect_timeout: int = DEFAULT_CONNECT_TIMEOUT,
241
- token_factory: TokenFactory = _new_token,
242
- owner_agent_id: str | None = None,
243
- ) -> RemoteSession:
244
- host = self.hosts.get(host_id)
245
- self._check_host_session_limit(host.id)
246
- child = self.child_factory(host.command, list(host.args), connect_timeout)
247
- token = token_factory()
248
- marker = f"__SCM_READY_{token}__"
249
- try:
250
- self._run_login_steps(child, host, connect_timeout)
251
- if getattr(host, "login_steps", []):
252
- self._wait_for_ready_prompt(child, host, connect_timeout)
253
- child.send(f"printf '\\n{marker}\\n'\r")
254
- child.expect(re.escape(marker), timeout=connect_timeout)
255
- except Exception as exc: # pexpect.TIMEOUT/EOF and fake child errors are surfaced uniformly.
256
- try:
257
- child.close(force=True)
258
- finally:
259
- raise RemoteCommandError(f"failed to establish ready shell for {host_id}: {exc}") from exc
260
-
261
- session = RemoteSession(
262
- id=uuid.uuid4().hex,
263
- host_id=host.id,
264
- label=host.label,
265
- command=host.command,
266
- args=list(host.args),
267
- child=child,
268
- owner_agent_id=owner_agent_id,
269
- last_output=_tail(str(child.before)),
270
- )
271
- self._sessions[session.id] = session
272
- self._emit("session_opened", session_id=session.id, host_id=host.id, outcome="ok")
273
- return session
274
-
275
- def _run_login_steps(self, child: ChildProcess, host: object, timeout: int) -> None:
276
- for step in getattr(host, "login_steps", []):
277
- child.expect(_runtime_expect_pattern(step.expect), timeout=timeout)
278
- if step.send_secret is not None:
279
- secrets = getattr(host, "secrets", {})
280
- if step.send_secret not in secrets:
281
- raise RemoteCommandError(f"login step references missing secret '{step.send_secret}'")
282
- self._send_login_text(child, secrets[step.send_secret])
283
- else:
284
- self._send_login_text(child, step.send or "")
285
-
286
- def _send_login_text(self, child: ChildProcess, text: str) -> None:
287
- # JumpServer-style menus often require carriage return rather than LF.
288
- child.send(f"{text}\r")
289
-
290
- def _wait_for_ready_prompt(self, child: ChildProcess, host: object, timeout: int) -> None:
291
- ready_expect = getattr(host, "ready_expect", None)
292
- pattern = ready_expect if isinstance(ready_expect, str) and ready_expect.strip() else r"(?m)[^\r\n]*[#$]\s*$"
293
- child.expect(_runtime_expect_pattern(pattern), timeout=timeout)
294
-
295
- def list_sessions(self) -> list[dict[str, object]]:
296
- return [self._describe_session(session) for session in self._sessions.values()]
297
-
298
- def close_all_sessions(self) -> dict[str, object]:
299
- """Close every live session. Used for graceful shutdown."""
300
-
301
- closed: list[str] = []
302
- for session_id in list(self._sessions.keys()):
303
- self._safe_close(session_id)
304
- self._emit("session_closed", session_id=session_id, outcome="shutdown")
305
- closed.append(session_id)
306
- return {"closed": closed, "count": len(closed)}
307
-
308
- def close_session(
309
- self,
310
- session_id: str,
311
- *,
312
- owner_agent_id: str | None = None,
313
- force: bool = False,
314
- ) -> dict[str, object]:
315
- session = self._get_session(session_id)
316
- self._check_owner(session, owner_agent_id, force=force)
317
- if session.busy and not force:
318
- raise RemoteCommandError(f"session {session_id} is busy; retry later or close with force")
319
- session.child.close(force=True)
320
- self._sessions.pop(session_id, None)
321
- self._emit("session_closed", session_id=session_id, host_id=session.host_id, outcome="forced" if force else "user")
322
- return {"session_id": session_id, "closed": True}
323
-
324
- def stop_task(self, session_id: str, task_id: str, *, owner_agent_id: str | None = None, force: bool = False) -> dict[str, object]:
325
- """Best-effort terminate a background task by killing its remote PID."""
326
-
327
- task = self._tasks.get(task_id)
328
- if task is None:
329
- raise RemoteCommandError(f"unknown task {task_id}")
330
- if task.session_id != session_id:
331
- raise RemoteCommandError(f"task {task_id} does not belong to session {session_id}")
332
- self._check_task_owner(task, owner_agent_id, force=force)
333
- if task.remote_pid is None:
334
- raise RemoteCommandError(f"task {task_id} has no recorded remote pid")
335
- kill_script = f'if [ -f "{task.job_dir}/exit_code" ]; then echo already; else kill -TERM {task.remote_pid} 2>/dev/null && echo killed || echo noop; fi'
336
- result = self.run_command(session_id, kill_script, timeout=10, owner_agent_id=owner_agent_id)
337
- outcome = result.output.strip().splitlines()[-1] if result.output.strip() else "unknown"
338
- self._emit(
339
- "task_stopped",
340
- session_id=session_id,
341
- task_id=task_id,
342
- host_id=task.host_id,
343
- outcome=outcome,
344
- )
345
- return {"task_id": task_id, "stopped": outcome in ("killed", "already"), "outcome": outcome}
346
-
347
- def run_command(
348
- self,
349
- session_id: str,
350
- command: str,
351
- *,
352
- timeout: int = DEFAULT_COMMAND_TIMEOUT,
353
- token_factory: TokenFactory = _new_token,
354
- owner_agent_id: str | None = None,
355
- output_limit: int = MAX_OUTPUT_CHARS,
356
- ) -> CommandResult:
357
- if not command.strip():
358
- raise RemoteCommandError("command must not be empty")
359
- decision = self._check_policy(command)
360
- if decision is not None and not decision.allowed:
361
- self._emit(
362
- "command_denied",
363
- session_id=session_id,
364
- command=command,
365
- outcome=decision.reason,
366
- detail=decision.matched_pattern,
367
- )
368
- raise PolicyDeniedError(decision.reason)
369
- session = self._get_session(session_id)
370
- self._check_owner(session, owner_agent_id)
371
- if session.busy:
372
- raise RemoteCommandError(f"session {session_id} is busy")
373
- if not session.child.isalive():
374
- raise RemoteCommandError(f"session {session_id} is not alive")
375
-
376
- token = token_factory()
377
- start_marker = f"__SCM_START_{token}__"
378
- exit_pattern = re.compile(rf"__SCM_EXIT_{re.escape(token)}:(\d+)")
379
- remote_line = self._build_remote_line(command, token, start_marker)
380
- session.busy = True
381
- try:
382
- session.child.sendline(remote_line)
383
- session.child.expect(exit_pattern, timeout=timeout)
384
- output = _extract_between_markers(str(session.child.before), start_marker)
385
- exit_code = _parse_exit_marker(session.child.after, token)
386
- result = CommandResult(session.id, session.host_id, command, exit_code, _tail(output, max(1, output_limit)))
387
- self._emit(
388
- "command_run",
389
- session_id=session.id,
390
- host_id=session.host_id,
391
- command=command,
392
- outcome=f"exit={exit_code}",
393
- )
394
- except pexpect.TIMEOUT:
395
- output = _tail(str(session.child.before))
396
- result = CommandResult(session.id, session.host_id, command, None, output, timed_out=True)
397
- self._emit(
398
- "command_run",
399
- session_id=session.id,
400
- host_id=session.host_id,
401
- command=command,
402
- outcome="timeout",
403
- )
404
- finally:
405
- session.busy = False
406
- session.last_used_at = time.time()
407
- session.last_output = _tail(str(session.child.before))
408
- return result
409
-
410
- def start_task(
411
- self,
412
- session_id: str,
413
- command: str,
414
- *,
415
- token_factory: TokenFactory = _new_token,
416
- owner_agent_id: str | None = None,
417
- ) -> RemoteTask:
418
- if not command.strip():
419
- raise RemoteCommandError("command must not be empty")
420
- decision = self._check_policy(command)
421
- if decision is not None and not decision.allowed:
422
- self._emit(
423
- "task_denied",
424
- session_id=session_id,
425
- command=command,
426
- outcome=decision.reason,
427
- detail=decision.matched_pattern,
428
- )
429
- raise PolicyDeniedError(decision.reason)
430
- task_id = token_factory()
431
- job_dir = f"{REMOTE_JOB_ROOT}/{task_id}"
432
- quoted_command = shlex.quote(command)
433
- script = (
434
- f'job_dir="{job_dir}"; '
435
- 'mkdir -p "$job_dir"; '
436
- 'date -u +%Y-%m-%dT%H:%M:%SZ > "$job_dir/started_at"; '
437
- f'({self.remote_shell} -lc {quoted_command} > "$job_dir/stdout.log" 2>&1; '
438
- 'status=$?; printf "%s\\n" "$status" > "$job_dir/exit_code"; '
439
- 'date -u +%Y-%m-%dT%H:%M:%SZ > "$job_dir/finished_at") & '
440
- 'pid=$!; printf "%s\\n" "$pid" > "$job_dir/pid"; printf "%s" "$pid"'
441
- )
442
- result = self.run_command(session_id, script, timeout=10, token_factory=lambda: f"task_{task_id}", owner_agent_id=owner_agent_id)
443
- pid_match = re.search(r"(\d+)", result.output)
444
- task = RemoteTask(
445
- id=task_id,
446
- session_id=session_id,
447
- host_id=self._get_session(session_id).host_id,
448
- command=command,
449
- job_dir=job_dir,
450
- remote_pid=int(pid_match.group(1)) if pid_match else None,
451
- owner_agent_id=owner_agent_id,
452
- )
453
- self._tasks[task.id] = task
454
- self._emit("task_started", session_id=session_id, task_id=task.id, host_id=task.host_id, command=command, outcome="ok")
455
- return task
456
-
457
- def task_status(
458
- self,
459
- session_id: str,
460
- task_id: str,
461
- *,
462
- tail_lines: int = 80,
463
- owner_agent_id: str | None = None,
464
- ) -> dict[str, object]:
465
- task = self._tasks.get(task_id)
466
- if task is None:
467
- raise RemoteCommandError(f"unknown task {task_id}")
468
- if task.session_id != session_id:
469
- raise RemoteCommandError(f"task {task_id} does not belong to session {session_id}")
470
- self._check_task_owner(task, owner_agent_id)
471
- safe_tail_lines = max(1, min(int(tail_lines), 500))
472
- script = (
473
- f'job_dir="{task.job_dir}"; '
474
- 'pid=$(cat "$job_dir/pid" 2>/dev/null || true); '
475
- 'if [ -f "$job_dir/exit_code" ]; then state=finished; exit_code=$(cat "$job_dir/exit_code"); '
476
- 'elif [ -n "$pid" ] && kill -0 "$pid" 2>/dev/null; then state=running; exit_code=""; '
477
- 'else state=unknown; exit_code=""; fi; '
478
- 'printf "STATE=%s\\nPID=%s\\nEXIT=%s\\n---LOG---\\n" "$state" "$pid" "$exit_code"; '
479
- f'tail -n {safe_tail_lines} "$job_dir/stdout.log" 2>/dev/null || true'
480
- )
481
- result = self.run_command(session_id, script, timeout=15, owner_agent_id=owner_agent_id)
482
- metadata, _, log = result.output.partition("---LOG---\n")
483
- fields: dict[str, str] = {}
484
- for line in metadata.splitlines():
485
- key, sep, value = line.partition("=")
486
- if sep:
487
- fields[key] = value
488
- exit_text = fields.get("EXIT") or ""
489
- return {
490
- "task_id": task.id,
491
- "session_id": session_id,
492
- "host_id": task.host_id,
493
- "command": task.command,
494
- "job_dir": task.job_dir,
495
- "state": fields.get("STATE", "unknown"),
496
- "pid": fields.get("PID") or task.remote_pid,
497
- "exit_code": int(exit_text) if exit_text.isdigit() else None,
498
- "log_tail": log,
499
- }
500
-
501
-
502
- def file_write_text(
503
- self,
504
- session_id: str,
505
- remote_path: str,
506
- content: str,
507
- *,
508
- mode: str = "0644",
509
- max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
510
- token_factory: TokenFactory = _new_token,
511
- owner_agent_id: str | None = None,
512
- ) -> dict[str, object]:
513
- data = content.encode("utf-8")
514
- return self._upload_bytes(
515
- session_id,
516
- data,
517
- remote_path,
518
- mode=mode,
519
- max_bytes=max_bytes,
520
- token_factory=token_factory,
521
- owner_agent_id=owner_agent_id,
522
- )
523
-
524
- def file_read_text(
525
- self,
526
- session_id: str,
527
- remote_path: str,
528
- *,
529
- max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
530
- token_factory: TokenFactory = _new_token,
531
- owner_agent_id: str | None = None,
532
- ) -> dict[str, object]:
533
- result = self._download_bytes(
534
- session_id,
535
- remote_path,
536
- max_bytes=max_bytes,
537
- token_factory=token_factory,
538
- owner_agent_id=owner_agent_id,
539
- )
540
- data = cast(bytes, result.pop("data_bytes"))
541
- result["content"] = data.decode("utf-8")
542
- return result
543
-
544
- def file_write_bytes(
545
- self,
546
- session_id: str,
547
- remote_path: str,
548
- data: bytes,
549
- *,
550
- mode: str = "0644",
551
- max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
552
- token_factory: TokenFactory = _new_token,
553
- owner_agent_id: str | None = None,
554
- ) -> dict[str, object]:
555
- return self._upload_bytes(
556
- session_id,
557
- data,
558
- remote_path,
559
- mode=mode,
560
- max_bytes=max_bytes,
561
- token_factory=token_factory,
562
- owner_agent_id=owner_agent_id,
563
- )
564
-
565
- def file_read_bytes(
566
- self,
567
- session_id: str,
568
- remote_path: str,
569
- *,
570
- max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
571
- token_factory: TokenFactory = _new_token,
572
- owner_agent_id: str | None = None,
573
- ) -> dict[str, object]:
574
- return self._download_bytes(
575
- session_id,
576
- remote_path,
577
- max_bytes=max_bytes,
578
- token_factory=token_factory,
579
- owner_agent_id=owner_agent_id,
580
- )
581
-
582
- def file_upload(
583
- self,
584
- session_id: str,
585
- local_path: Path | str,
586
- remote_path: str,
587
- *,
588
- mode: str = "0644",
589
- max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
590
- token_factory: TokenFactory = _new_token,
591
- owner_agent_id: str | None = None,
592
- ) -> dict[str, object]:
593
- path = Path(local_path).expanduser()
594
- if not path.is_absolute():
595
- raise RemoteCommandError("local_path must be absolute")
596
- if not path.is_file():
597
- raise RemoteCommandError(f"local file not found: {path}")
598
- size = path.stat().st_size
599
- self._check_transfer_size(size, max_bytes)
600
- return self._upload_bytes(
601
- session_id,
602
- path.read_bytes(),
603
- remote_path,
604
- mode=mode,
605
- max_bytes=max_bytes,
606
- token_factory=token_factory,
607
- owner_agent_id=owner_agent_id,
608
- local_path=str(path),
609
- )
610
-
611
- def file_download(
612
- self,
613
- session_id: str,
614
- remote_path: str,
615
- local_path: Path | str,
616
- *,
617
- max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
618
- token_factory: TokenFactory = _new_token,
619
- owner_agent_id: str | None = None,
620
- ) -> dict[str, object]:
621
- path = Path(local_path).expanduser()
622
- if not path.is_absolute():
623
- raise RemoteCommandError("local_path must be absolute")
624
- result = self._download_bytes(
625
- session_id,
626
- remote_path,
627
- max_bytes=max_bytes,
628
- token_factory=token_factory,
629
- owner_agent_id=owner_agent_id,
630
- )
631
- data = cast(bytes, result.pop("data_bytes"))
632
- path.parent.mkdir(parents=True, exist_ok=True)
633
- path.write_bytes(data)
634
- result["local_path"] = str(path)
635
- return result
636
-
637
- def _upload_bytes(
638
- self,
639
- session_id: str,
640
- data: bytes,
641
- remote_path: str,
642
- *,
643
- mode: str,
644
- max_bytes: int,
645
- token_factory: TokenFactory,
646
- owner_agent_id: str | None,
647
- local_path: str | None = None,
648
- ) -> dict[str, object]:
649
- self._check_transfer_size(len(data), max_bytes)
650
- self._validate_remote_path(remote_path)
651
- b64 = base64.b64encode(data).decode("ascii")
652
- sha256 = hashlib.sha256(data).hexdigest()
653
- token = token_factory()
654
- remote_q = shlex.quote(remote_path)
655
- mode_q = shlex.quote(mode)
656
- tmp_prefix = f"/tmp/hostbridge.upload.{token}"
657
- tmp_b64 = f"{tmp_prefix}.b64"
658
- tmp_out = f"{tmp_prefix}.out"
659
- cleanup_script = f"rm -f {shlex.quote(tmp_b64)} {shlex.quote(tmp_out)}"
660
- init_script = (
661
- "set -e; "
662
- f"target={remote_q}; "
663
- 'dir=$(dirname "$target"); mkdir -p "$dir"; '
664
- f"rm -f {shlex.quote(tmp_b64)} {shlex.quote(tmp_out)}; "
665
- f": > {shlex.quote(tmp_b64)}"
666
- )
667
- result = self.run_command(session_id, init_script, timeout=60, token_factory=lambda: token, owner_agent_id=owner_agent_id)
668
- if result.exit_code != 0:
669
- raise RemoteCommandError(result.output or f"failed to prepare upload for {remote_path}")
670
- try:
671
- self._stream_base64_to_remote_file(
672
- session_id,
673
- b64,
674
- tmp_b64,
675
- token=token,
676
- owner_agent_id=owner_agent_id,
677
- )
678
- finalize_script = (
679
- "set -e; "
680
- f"target={remote_q}; mode={mode_q}; tmp_b64={shlex.quote(tmp_b64)}; tmp_out={shlex.quote(tmp_out)}; "
681
- '{ base64 -d < "$tmp_b64" 2>/dev/null || base64 -D < "$tmp_b64"; } > "$tmp_out"; '
682
- 'chmod "$mode" "$tmp_out"; mv "$tmp_out" "$target"; rm -f "$tmp_b64"; '
683
- f"printf 'BYTES=%s\\nSHA256=%s\\n' {len(data)} {shlex.quote(sha256)}"
684
- )
685
- result = self.run_command(
686
- session_id,
687
- finalize_script,
688
- timeout=60,
689
- token_factory=lambda: token,
690
- owner_agent_id=owner_agent_id,
691
- )
692
- if result.exit_code != 0:
693
- raise RemoteCommandError(result.output or f"failed to finalize upload for {remote_path}")
694
- except Exception:
695
- with suppress(Exception):
696
- self.run_command(session_id, cleanup_script, timeout=10, token_factory=lambda: token, owner_agent_id=owner_agent_id)
697
- raise
698
- payload: dict[str, object] = {
699
- "session_id": session_id,
700
- "remote_path": remote_path,
701
- "bytes": len(data),
702
- "sha256": sha256,
703
- "strategy": "shell_base64_chunked",
704
- }
705
- if local_path is not None:
706
- payload["local_path"] = local_path
707
- self._emit("file_uploaded", session_id=session_id, command=remote_path, outcome=f"bytes={len(data)}")
708
- return payload
709
-
710
- def _stream_base64_to_remote_file(
711
- self,
712
- session_id: str,
713
- encoded: str,
714
- remote_tmp_b64: str,
715
- *,
716
- token: str,
717
- owner_agent_id: str | None,
718
- ) -> None:
719
- session = self._get_session(session_id)
720
- self._check_owner(session, owner_agent_id)
721
- if session.busy:
722
- raise RemoteCommandError(f"session {session_id} is busy")
723
- if not session.child.isalive():
724
- raise RemoteCommandError(f"session {session_id} is not alive")
725
-
726
- heredoc = f"__HOSTBRIDGE_UPLOAD_{token}__"
727
- done_pattern = re.compile(rf"__HOSTBRIDGE_UPLOAD_DONE_{re.escape(token)}__:(\d+)")
728
- chunk_size = max(4, int(UPLOAD_BASE64_CHUNK_CHARS))
729
- session.busy = True
730
- try:
731
- session.child.send(f"cat >> {shlex.quote(remote_tmp_b64)} <<'{heredoc}'\n")
732
- for offset in range(0, len(encoded), chunk_size):
733
- session.child.send(encoded[offset : offset + chunk_size])
734
- session.child.send(f"\n{heredoc}\nprintf '\\n__HOSTBRIDGE_UPLOAD_DONE_{token}__:%s\\n' \"$?\"\n")
735
- session.child.expect(done_pattern, timeout=60)
736
- marker = session.child.after
737
- text = marker.group(0) if hasattr(marker, "group") else str(marker)
738
- match = re.search(r":(\d+)", text)
739
- if match is None or int(match.group(1)) != 0:
740
- raise RemoteCommandError(f"failed to stream upload chunk for {session_id}")
741
- finally:
742
- session.busy = False
743
- session.last_used_at = time.time()
744
- session.last_output = _tail(str(session.child.before))
745
-
746
- def _download_bytes(
747
- self,
748
- session_id: str,
749
- remote_path: str,
750
- *,
751
- max_bytes: int,
752
- token_factory: TokenFactory,
753
- owner_agent_id: str | None,
754
- ) -> dict[str, object]:
755
- self._validate_remote_path(remote_path)
756
- token = token_factory()
757
- remote_q = shlex.quote(remote_path)
758
- metadata_script = (
759
- "set -e; "
760
- f"target={remote_q}; "
761
- 'bytes=$(wc -c < "$target" | tr -d " "); '
762
- f"if [ \"$bytes\" -gt {int(max_bytes)} ]; then echo 'FILE_TOO_LARGE remote pull required' >&2; exit 73; fi; "
763
- "if command -v sha256sum >/dev/null 2>&1; then sha=$(sha256sum \"$target\" | awk '{print $1}'); "
764
- "else sha=$(shasum -a 256 \"$target\" | awk '{print $1}'); fi; "
765
- "printf 'BYTES=%s\\nSHA256=%s\\n' \"$bytes\" \"$sha\""
766
- )
767
- result = self.run_command(
768
- session_id,
769
- metadata_script,
770
- timeout=60,
771
- token_factory=lambda: token,
772
- owner_agent_id=owner_agent_id,
773
- )
774
- if result.exit_code == 73:
775
- self._raise_remote_pull_required(max_bytes)
776
- if result.exit_code != 0:
777
- raise RemoteCommandError(result.output or f"failed to download {remote_path}")
778
- fields: dict[str, str] = {}
779
- for line in result.output.splitlines():
780
- key, _, value = line.partition("=")
781
- if key:
782
- fields[key] = value
783
- byte_count = int(fields.get("BYTES") or 0)
784
- self._check_transfer_size(byte_count, max_bytes)
785
- chunks: list[bytes] = []
786
- chunk_size = max(1, int(DOWNLOAD_CHUNK_BYTES))
787
- for index in range((byte_count + chunk_size - 1) // chunk_size):
788
- chunk_script = (
789
- "set -e; "
790
- f"dd if={remote_q} bs={chunk_size} skip={index} count=1 2>/dev/null | base64"
791
- )
792
- output_limit = max(MAX_OUTPUT_CHARS, chunk_size * 2 + 4096)
793
- result = self.run_command(
794
- session_id,
795
- chunk_script,
796
- timeout=60,
797
- token_factory=lambda: token,
798
- owner_agent_id=owner_agent_id,
799
- output_limit=output_limit,
800
- )
801
- if result.exit_code != 0:
802
- raise RemoteCommandError(result.output or f"failed to download chunk {index} from {remote_path}")
803
- chunks.append(base64.b64decode("".join(result.output.split())))
804
- data = b"".join(chunks)
805
- sha256 = hashlib.sha256(data).hexdigest()
806
- remote_sha = fields.get("SHA256") or sha256
807
- if remote_sha != "ignored" and sha256 != remote_sha:
808
- raise RemoteCommandError("download checksum mismatch")
809
- self._emit("file_downloaded", session_id=session_id, command=remote_path, outcome=f"bytes={len(data)}")
810
- return {
811
- "session_id": session_id,
812
- "remote_path": remote_path,
813
- "bytes": len(data),
814
- "sha256": remote_sha,
815
- "strategy": "shell_base64_chunked",
816
- "data_bytes": data,
817
- }
818
-
819
- def _validate_remote_path(self, remote_path: str) -> None:
820
- if not remote_path or remote_path.strip() in {"", "/"}:
821
- raise RemoteCommandError("remote_path must not be empty or root")
822
-
823
- def _check_transfer_size(self, size: int, max_bytes: int) -> None:
824
- if size > int(max_bytes):
825
- self._raise_remote_pull_required(max_bytes)
826
-
827
- def _raise_remote_pull_required(self, max_bytes: int) -> None:
828
- raise RemoteCommandError(
829
- f"file exceeds HostBridge shell transfer limit ({int(max_bytes)} bytes); use remote pull (curl/wget/modelscope/git) for large files"
830
- )
831
-
832
- def write(self, session_id: str, text: str, *, owner_agent_id: str | None = None) -> dict[str, object]:
833
- return self.send(session_id, f"{text}\n", owner_agent_id=owner_agent_id)
834
-
835
- def send(self, session_id: str, text: str, *, owner_agent_id: str | None = None) -> dict[str, object]:
836
- session = self._get_session(session_id)
837
- self._check_owner(session, owner_agent_id)
838
- if session.busy:
839
- raise RemoteCommandError(f"session {session_id} is busy")
840
- if not session.child.isalive():
841
- raise RemoteCommandError(f"session {session_id} is not alive")
842
- session.child.send(text)
843
- session.last_used_at = time.time()
844
- return {"session_id": session_id, "written": True}
845
-
846
- def read(
847
- self,
848
- session_id: str,
849
- *,
850
- timeout: float = 1.0,
851
- max_chars: int = 20_000,
852
- owner_agent_id: str | None = None,
853
- ) -> dict[str, object]:
854
- session = self._get_session(session_id)
855
- self._check_owner(session, owner_agent_id)
856
- chunks: list[str] = []
857
- deadline = time.time() + max(0.0, timeout)
858
- while len("".join(chunks)) < max_chars:
859
- remaining = max(0.0, deadline - time.time())
860
- if remaining <= 0 and chunks:
861
- break
862
- try:
863
- chunk = session.child.read_nonblocking(size=4096, timeout=remaining)
864
- except (pexpect.TIMEOUT, pexpect.EOF):
865
- break
866
- if not chunk:
867
- break
868
- chunks.append(str(chunk))
869
- output = _tail("".join(chunks), max_chars)
870
- session.last_output = output or session.last_output
871
- session.last_used_at = time.time()
872
- return {"session_id": session_id, "output": output, "alive": session.child.isalive()}
873
-
874
- def _get_session(self, session_id: str) -> RemoteSession:
875
- try:
876
- return self._sessions[session_id]
877
- except KeyError as exc:
878
- raise RemoteCommandError(f"unknown session {session_id}") from exc
879
-
880
- def _check_owner(self, session: RemoteSession, owner_agent_id: str | None, *, force: bool = False) -> None:
881
- if force or owner_agent_id is None or session.owner_agent_id is None:
882
- return
883
- if owner_agent_id != session.owner_agent_id:
884
- raise RemoteCommandError(f"session {session.id} is owned by another agent")
885
-
886
- def _check_task_owner(self, task: RemoteTask, owner_agent_id: str | None, *, force: bool = False) -> None:
887
- if force or owner_agent_id is None or task.owner_agent_id is None:
888
- return
889
- if owner_agent_id != task.owner_agent_id:
890
- raise RemoteCommandError(f"task {task.id} is owned by another agent")
891
-
892
- def _check_host_session_limit(self, host_id: str) -> None:
893
- if self._max_sessions_per_host <= 0:
894
- return
895
- open_count = sum(1 for session in self._sessions.values() if session.host_id == host_id and session.child.isalive())
896
- if open_count >= self._max_sessions_per_host:
897
- raise RemoteCommandError(f"host {host_id} reached session limit {self._max_sessions_per_host}")
898
-
899
- def _build_remote_line(self, command: str, token: str, start_marker: str) -> str:
900
- encoded_command = base64.b64encode(command.encode("utf-8")).decode("ascii")
901
- return (
902
- f"__scm_cmd=$(printf %s {shlex.quote(encoded_command)} | {{ base64 --decode 2>/dev/null || base64 -D; }}); "
903
- f"printf '\\n{start_marker}\\n'; "
904
- f"{self.remote_shell} -lc \"$__scm_cmd\"; "
905
- "__scm_status=$?; "
906
- f"printf '\\n__SCM_EXIT_{token}:%s\\n' \"$__scm_status\""
907
- )
908
-
909
- @staticmethod
910
- def _describe_session(session: RemoteSession) -> dict[str, object]:
911
- return {
912
- "session_id": session.id,
913
- "host_id": session.host_id,
914
- "label": session.label,
915
- "command": session.command,
916
- "args": session.args,
917
- "alive": session.child.isalive(),
918
- "busy": session.busy,
919
- "owner_agent_id": session.owner_agent_id,
920
- "created_at": session.created_at,
921
- "last_used_at": session.last_used_at,
922
- }