mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-07-09 17:08:29 +00:00
7298a88fda
Add host_browser_selection config normalization, Browser Settings UI choices from CLI-advertised inventory, and browser_selection forwarding through connector browser operations. Expose host browser ids, labels, and available_browsers through A0 connector metadata and the browser_runtime API, with regression coverage for selection normalization.
925 lines
28 KiB
Python
925 lines
28 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import binascii
|
|
import copy
|
|
import json
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
|
|
@dataclass
|
|
class PendingFileOperation:
|
|
sid: str
|
|
loop: asyncio.AbstractEventLoop
|
|
future: asyncio.Future[dict[str, Any]]
|
|
context_id: str | None = None
|
|
chunk_count: int | None = None
|
|
chunks: dict[int, bytes] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class PendingExecOperation:
|
|
sid: str
|
|
loop: asyncio.AbstractEventLoop
|
|
future: asyncio.Future[dict[str, Any]]
|
|
context_id: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class PendingComputerUseOperation:
|
|
sid: str
|
|
loop: asyncio.AbstractEventLoop
|
|
future: asyncio.Future[dict[str, Any]]
|
|
context_id: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class PendingBrowserOperation:
|
|
sid: str
|
|
loop: asyncio.AbstractEventLoop
|
|
future: asyncio.Future[dict[str, Any]]
|
|
context_id: str | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteTreeSnapshot:
|
|
sid: str
|
|
payload: dict[str, Any]
|
|
updated_at: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ComputerUseMetadata:
|
|
supported: bool
|
|
enabled: bool
|
|
trust_mode: str
|
|
status: str
|
|
last_error: str
|
|
restore_token_present: bool
|
|
artifact_root: str
|
|
backend_id: str
|
|
backend_family: str
|
|
features: tuple[str, ...]
|
|
contract_version: int
|
|
capabilities: dict[str, Any]
|
|
support_reason: str
|
|
updated_at: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HostBrowserMetadata:
|
|
supported: bool
|
|
can_prepare: bool
|
|
enabled: bool
|
|
status: str
|
|
browser_family: str
|
|
profile_label: str
|
|
profile_path: str
|
|
cdp_endpoint: str
|
|
browser_id: str
|
|
browser_label: str
|
|
available_browsers: tuple[dict[str, Any], ...]
|
|
content_helper_sha256: str
|
|
features: tuple[str, ...]
|
|
support_reason: str
|
|
updated_at: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteFileMetadata:
|
|
enabled: bool
|
|
write_enabled: bool
|
|
mode: str
|
|
updated_at: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RemoteExecMetadata:
|
|
enabled: bool
|
|
updated_at: float
|
|
|
|
|
|
_context_subscriptions: dict[str, set[str]] = {}
|
|
_sid_contexts: dict[str, set[str]] = {}
|
|
_pending_file_ops: dict[str, PendingFileOperation] = {}
|
|
_pending_exec_ops: dict[str, PendingExecOperation] = {}
|
|
_pending_computer_use_ops: dict[str, PendingComputerUseOperation] = {}
|
|
_pending_browser_ops: dict[str, PendingBrowserOperation] = {}
|
|
_remote_tree_snapshots: dict[str, RemoteTreeSnapshot] = {}
|
|
_sid_computer_use_metadata: dict[str, ComputerUseMetadata] = {}
|
|
_sid_host_browser_metadata: dict[str, HostBrowserMetadata] = {}
|
|
_sid_remote_file_metadata: dict[str, RemoteFileMetadata] = {}
|
|
_sid_remote_exec_metadata: dict[str, RemoteExecMetadata] = {}
|
|
_state_lock = threading.RLock()
|
|
|
|
|
|
def register_sid(sid: str) -> None:
|
|
with _state_lock:
|
|
_sid_contexts.setdefault(sid, set())
|
|
|
|
|
|
def unregister_sid(sid: str) -> set[str]:
|
|
with _state_lock:
|
|
contexts = _sid_contexts.pop(sid, set())
|
|
_remote_tree_snapshots.pop(sid, None)
|
|
_sid_computer_use_metadata.pop(sid, None)
|
|
_sid_host_browser_metadata.pop(sid, None)
|
|
_sid_remote_file_metadata.pop(sid, None)
|
|
_sid_remote_exec_metadata.pop(sid, None)
|
|
for context_id in contexts:
|
|
subscribers = _context_subscriptions.get(context_id)
|
|
if not subscribers:
|
|
continue
|
|
subscribers.discard(sid)
|
|
if not subscribers:
|
|
_context_subscriptions.pop(context_id, None)
|
|
return contexts
|
|
|
|
|
|
def subscribe_sid_to_context(sid: str, context_id: str) -> None:
|
|
with _state_lock:
|
|
_sid_contexts.setdefault(sid, set()).add(context_id)
|
|
_context_subscriptions.setdefault(context_id, set()).add(sid)
|
|
|
|
|
|
def unsubscribe_sid_from_context(sid: str, context_id: str) -> None:
|
|
with _state_lock:
|
|
contexts = _sid_contexts.get(sid)
|
|
if contexts is not None:
|
|
contexts.discard(context_id)
|
|
if not contexts:
|
|
_sid_contexts.pop(sid, None)
|
|
|
|
subscribers = _context_subscriptions.get(context_id)
|
|
if subscribers is not None:
|
|
subscribers.discard(sid)
|
|
if not subscribers:
|
|
_context_subscriptions.pop(context_id, None)
|
|
|
|
|
|
def subscribed_contexts_for_sid(sid: str) -> set[str]:
|
|
with _state_lock:
|
|
return set(_sid_contexts.get(sid, set()))
|
|
|
|
|
|
def subscribed_sids_for_context(context_id: str) -> set[str]:
|
|
with _state_lock:
|
|
return set(_context_subscriptions.get(context_id, set()))
|
|
|
|
|
|
def connected_sids() -> set[str]:
|
|
with _state_lock:
|
|
return set(_sid_contexts.keys())
|
|
|
|
|
|
def _candidate_sids_for_context_locked(context_id: str) -> list[str]:
|
|
context_sids = sorted(_context_subscriptions.get(context_id, set()))
|
|
context_set = set(context_sids)
|
|
global_sids = sorted(sid for sid in _sid_contexts if sid not in context_set)
|
|
return context_sids + global_sids
|
|
|
|
|
|
def remote_tool_sids_for_context(context_id: str) -> list[str]:
|
|
"""Return connected CLI candidates, preferring clients subscribed to context_id."""
|
|
with _state_lock:
|
|
return _candidate_sids_for_context_locked(context_id)
|
|
|
|
|
|
def store_remote_tree_snapshot(
|
|
sid: str,
|
|
payload: dict[str, Any],
|
|
) -> RemoteTreeSnapshot:
|
|
snapshot = RemoteTreeSnapshot(
|
|
sid=sid,
|
|
payload=dict(payload),
|
|
updated_at=time.time(),
|
|
)
|
|
with _state_lock:
|
|
_remote_tree_snapshots[sid] = snapshot
|
|
return snapshot
|
|
|
|
|
|
def clear_remote_tree_snapshot(sid: str) -> None:
|
|
with _state_lock:
|
|
_remote_tree_snapshots.pop(sid, None)
|
|
|
|
|
|
def latest_remote_tree_for_context(
|
|
context_id: str,
|
|
*,
|
|
max_age_seconds: float = 90.0,
|
|
) -> dict[str, Any] | None:
|
|
now = time.time()
|
|
with _state_lock:
|
|
context_sids = sorted(_context_subscriptions.get(context_id, set()))
|
|
context_set = set(context_sids)
|
|
global_sids = sorted(sid for sid in _sid_contexts if sid not in context_set)
|
|
snapshot_groups = [
|
|
[
|
|
_remote_tree_snapshots[sid]
|
|
for sid in context_sids
|
|
if sid in _remote_tree_snapshots
|
|
],
|
|
[
|
|
_remote_tree_snapshots[sid]
|
|
for sid in global_sids
|
|
if sid in _remote_tree_snapshots
|
|
],
|
|
]
|
|
|
|
for snapshots in snapshot_groups:
|
|
snapshots.sort(key=lambda item: item.updated_at, reverse=True)
|
|
for snapshot in snapshots:
|
|
if max_age_seconds > 0 and now - snapshot.updated_at > max_age_seconds:
|
|
continue
|
|
payload = dict(snapshot.payload)
|
|
payload["sid"] = snapshot.sid
|
|
payload["updated_at"] = snapshot.updated_at
|
|
return payload
|
|
return None
|
|
|
|
|
|
def select_target_sid(context_id: str) -> str | None:
|
|
with _state_lock:
|
|
subscribers = _context_subscriptions.get(context_id, set())
|
|
if not subscribers:
|
|
return None
|
|
return sorted(subscribers)[0]
|
|
|
|
|
|
def store_sid_remote_file_metadata(sid: str, payload: dict[str, Any]) -> RemoteFileMetadata:
|
|
write_enabled = bool(payload.get("write_enabled"))
|
|
mode = str(payload.get("mode", "") or "").strip().lower()
|
|
if mode not in {"read_only", "read_write"}:
|
|
mode = "read_write" if write_enabled else "read_only"
|
|
metadata = RemoteFileMetadata(
|
|
enabled=bool(payload.get("enabled", True)),
|
|
write_enabled=write_enabled,
|
|
mode=mode,
|
|
updated_at=time.time(),
|
|
)
|
|
with _state_lock:
|
|
_sid_remote_file_metadata[sid] = metadata
|
|
return metadata
|
|
|
|
|
|
def clear_sid_remote_file_metadata(sid: str) -> None:
|
|
with _state_lock:
|
|
_sid_remote_file_metadata.pop(sid, None)
|
|
|
|
|
|
def remote_file_metadata_for_sid(sid: str) -> dict[str, Any] | None:
|
|
with _state_lock:
|
|
metadata = _sid_remote_file_metadata.get(sid)
|
|
if metadata is None:
|
|
return None
|
|
return {
|
|
"enabled": metadata.enabled,
|
|
"write_enabled": metadata.write_enabled,
|
|
"mode": metadata.mode,
|
|
"updated_at": metadata.updated_at,
|
|
}
|
|
|
|
|
|
def select_remote_file_target_sid(context_id: str, *, require_writes: bool = False) -> str | None:
|
|
with _state_lock:
|
|
for sid in _candidate_sids_for_context_locked(context_id):
|
|
metadata = _sid_remote_file_metadata.get(sid)
|
|
if metadata is None:
|
|
continue
|
|
if not metadata.enabled:
|
|
continue
|
|
if require_writes and not metadata.write_enabled:
|
|
continue
|
|
return sid
|
|
return None
|
|
|
|
|
|
def store_sid_remote_exec_metadata(sid: str, payload: dict[str, Any]) -> RemoteExecMetadata:
|
|
metadata = RemoteExecMetadata(
|
|
enabled=bool(payload.get("enabled")),
|
|
updated_at=time.time(),
|
|
)
|
|
with _state_lock:
|
|
_sid_remote_exec_metadata[sid] = metadata
|
|
return metadata
|
|
|
|
|
|
def clear_sid_remote_exec_metadata(sid: str) -> None:
|
|
with _state_lock:
|
|
_sid_remote_exec_metadata.pop(sid, None)
|
|
|
|
|
|
def remote_exec_metadata_for_sid(sid: str) -> dict[str, Any] | None:
|
|
with _state_lock:
|
|
metadata = _sid_remote_exec_metadata.get(sid)
|
|
if metadata is None:
|
|
return None
|
|
return {
|
|
"enabled": metadata.enabled,
|
|
"updated_at": metadata.updated_at,
|
|
}
|
|
|
|
|
|
def select_remote_exec_target_sid(context_id: str, *, require_writes: bool = False) -> str | None:
|
|
with _state_lock:
|
|
for sid in _candidate_sids_for_context_locked(context_id):
|
|
metadata = _sid_remote_exec_metadata.get(sid)
|
|
if metadata is None:
|
|
continue
|
|
if metadata.enabled:
|
|
if require_writes:
|
|
file_metadata = _sid_remote_file_metadata.get(sid)
|
|
if file_metadata is None or (
|
|
not file_metadata.enabled or not file_metadata.write_enabled
|
|
):
|
|
continue
|
|
return sid
|
|
return None
|
|
|
|
|
|
def store_sid_computer_use_metadata(sid: str, payload: dict[str, Any]) -> ComputerUseMetadata:
|
|
features_value = payload.get("features")
|
|
if isinstance(features_value, (list, tuple)):
|
|
features = tuple(str(item).strip() for item in features_value if str(item).strip())
|
|
else:
|
|
features = ()
|
|
capabilities_value = payload.get("capabilities")
|
|
capabilities = copy.deepcopy(capabilities_value) if isinstance(capabilities_value, dict) else {}
|
|
try:
|
|
contract_version = int(payload.get("contract_version") or 0)
|
|
except (TypeError, ValueError):
|
|
contract_version = 0
|
|
metadata = ComputerUseMetadata(
|
|
supported=bool(payload.get("supported")),
|
|
enabled=bool(payload.get("supported")) and bool(payload.get("enabled")),
|
|
trust_mode=str(payload.get("trust_mode", "") or "").strip(),
|
|
status=str(payload.get("status", "") or "").strip(),
|
|
last_error=str(payload.get("last_error", "") or "").strip(),
|
|
restore_token_present=bool(payload.get("restore_token_present")),
|
|
artifact_root=str(payload.get("artifact_root", "") or "").strip(),
|
|
backend_id=str(payload.get("backend_id", "") or "").strip(),
|
|
backend_family=str(payload.get("backend_family", "") or "").strip(),
|
|
features=features,
|
|
contract_version=contract_version,
|
|
capabilities=capabilities,
|
|
support_reason=str(payload.get("support_reason", "") or "").strip(),
|
|
updated_at=time.time(),
|
|
)
|
|
with _state_lock:
|
|
_sid_computer_use_metadata[sid] = metadata
|
|
return metadata
|
|
|
|
|
|
def clear_sid_computer_use_metadata(sid: str) -> None:
|
|
with _state_lock:
|
|
_sid_computer_use_metadata.pop(sid, None)
|
|
|
|
|
|
def computer_use_metadata_for_sid(sid: str) -> dict[str, Any] | None:
|
|
with _state_lock:
|
|
metadata = _sid_computer_use_metadata.get(sid)
|
|
if metadata is None:
|
|
return None
|
|
return {
|
|
"supported": metadata.supported,
|
|
"enabled": metadata.enabled,
|
|
"trust_mode": metadata.trust_mode,
|
|
"status": metadata.status,
|
|
"last_error": metadata.last_error,
|
|
"restore_token_present": metadata.restore_token_present,
|
|
"artifact_root": metadata.artifact_root,
|
|
"backend_id": metadata.backend_id,
|
|
"backend_family": metadata.backend_family,
|
|
"features": list(metadata.features),
|
|
"contract_version": metadata.contract_version,
|
|
"capabilities": copy.deepcopy(metadata.capabilities),
|
|
"support_reason": metadata.support_reason,
|
|
"updated_at": metadata.updated_at,
|
|
}
|
|
|
|
|
|
def store_sid_host_browser_metadata(sid: str, payload: dict[str, Any]) -> HostBrowserMetadata:
|
|
features_value = payload.get("features")
|
|
if isinstance(features_value, (list, tuple)):
|
|
features = tuple(str(item).strip() for item in features_value if str(item).strip())
|
|
else:
|
|
features = ()
|
|
support_reason = str(payload.get("support_reason", "") or "").strip()
|
|
metadata = HostBrowserMetadata(
|
|
supported=bool(payload.get("supported")),
|
|
can_prepare=_host_browser_can_prepare(payload, features=features, support_reason=support_reason),
|
|
enabled=bool(payload.get("supported")) and bool(payload.get("enabled")),
|
|
status=str(payload.get("status", "") or "").strip(),
|
|
browser_family=str(payload.get("browser_family", "") or "").strip(),
|
|
profile_label=str(payload.get("profile_label", "") or "").strip(),
|
|
profile_path=str(payload.get("profile_path", "") or "").strip(),
|
|
cdp_endpoint=str(payload.get("cdp_endpoint", "") or "").strip(),
|
|
browser_id=str(payload.get("browser_id", payload.get("browser_selection", "")) or "").strip(),
|
|
browser_label=str(payload.get("browser_label", "") or "").strip(),
|
|
available_browsers=_normalize_available_host_browsers(payload.get("available_browsers")),
|
|
content_helper_sha256=str(payload.get("content_helper_sha256", "") or "").strip().lower(),
|
|
features=features,
|
|
support_reason=support_reason,
|
|
updated_at=time.time(),
|
|
)
|
|
with _state_lock:
|
|
_sid_host_browser_metadata[sid] = metadata
|
|
return metadata
|
|
|
|
|
|
def _host_browser_can_prepare(
|
|
payload: dict[str, Any],
|
|
*,
|
|
features: tuple[str, ...],
|
|
support_reason: str,
|
|
) -> bool:
|
|
if "can_prepare" in payload:
|
|
return bool(payload.get("can_prepare"))
|
|
if "ensure" not in features:
|
|
return False
|
|
reason = support_reason.lower()
|
|
return (
|
|
"python playwright" in reason
|
|
or "a0-controlled local profile" in reason
|
|
or "chrome-a0" in reason
|
|
or "remote debugging" in reason
|
|
)
|
|
|
|
|
|
def _normalize_available_host_browsers(value: Any) -> tuple[dict[str, Any], ...]:
|
|
if not isinstance(value, (list, tuple)):
|
|
return ()
|
|
browsers: list[dict[str, Any]] = []
|
|
for item in value:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
browser_id = str(item.get("id", item.get("browser_id", item.get("selection", ""))) or "").strip()
|
|
family = str(item.get("family", item.get("browser_family", "")) or "").strip()
|
|
label = str(item.get("label", item.get("name", "")) or "").strip()
|
|
cdp_endpoint = str(item.get("cdp_endpoint", "") or "").strip()
|
|
status = str(item.get("status", "") or "").strip()
|
|
enabled = bool(item.get("enabled", True))
|
|
if not any((browser_id, family, label, cdp_endpoint)):
|
|
continue
|
|
browsers.append({
|
|
"id": browser_id or family or cdp_endpoint,
|
|
"family": family,
|
|
"label": label or family or browser_id or cdp_endpoint,
|
|
"cdp_endpoint": cdp_endpoint,
|
|
"status": status,
|
|
"enabled": enabled,
|
|
})
|
|
return tuple(browsers)
|
|
|
|
|
|
def clear_sid_host_browser_metadata(sid: str) -> None:
|
|
with _state_lock:
|
|
_sid_host_browser_metadata.pop(sid, None)
|
|
|
|
|
|
def host_browser_metadata_for_sid(sid: str) -> dict[str, Any] | None:
|
|
with _state_lock:
|
|
metadata = _sid_host_browser_metadata.get(sid)
|
|
if metadata is None:
|
|
return None
|
|
return {
|
|
"supported": metadata.supported,
|
|
"can_prepare": metadata.can_prepare,
|
|
"enabled": metadata.enabled,
|
|
"status": metadata.status,
|
|
"browser_family": metadata.browser_family,
|
|
"profile_label": metadata.profile_label,
|
|
"profile_path": metadata.profile_path,
|
|
"cdp_endpoint": metadata.cdp_endpoint,
|
|
"browser_id": metadata.browser_id,
|
|
"browser_label": metadata.browser_label,
|
|
"available_browsers": copy.deepcopy(list(metadata.available_browsers)),
|
|
"content_helper_sha256": metadata.content_helper_sha256,
|
|
"features": list(metadata.features),
|
|
"support_reason": metadata.support_reason,
|
|
"updated_at": metadata.updated_at,
|
|
}
|
|
|
|
|
|
def select_host_browser_target_sid(context_id: str) -> str | None:
|
|
with _state_lock:
|
|
fallback: str | None = None
|
|
for sid in _candidate_sids_for_context_locked(context_id):
|
|
metadata = _sid_host_browser_metadata.get(sid)
|
|
if not metadata:
|
|
continue
|
|
if not metadata.supported:
|
|
continue
|
|
if metadata.enabled and metadata.status in {"ready", "active"}:
|
|
return sid
|
|
if metadata.enabled and fallback is None:
|
|
fallback = sid
|
|
return fallback
|
|
|
|
|
|
def select_host_browser_candidate_sid(context_id: str) -> str | None:
|
|
with _state_lock:
|
|
fallback: str | None = None
|
|
for sid in _candidate_sids_for_context_locked(context_id):
|
|
metadata = _sid_host_browser_metadata.get(sid)
|
|
if not metadata or not (metadata.supported or metadata.can_prepare):
|
|
continue
|
|
if metadata.enabled and metadata.status in {"ready", "active"}:
|
|
return sid
|
|
if metadata.enabled and fallback is None:
|
|
fallback = sid
|
|
elif fallback is None:
|
|
fallback = sid
|
|
return fallback
|
|
|
|
|
|
def host_browser_metadata_for_context(context_id: str) -> list[dict[str, Any]]:
|
|
with _state_lock:
|
|
candidates = _candidate_sids_for_context_locked(context_id)
|
|
rows: list[dict[str, Any]] = []
|
|
for sid in candidates:
|
|
metadata = host_browser_metadata_for_sid(sid)
|
|
if metadata is not None:
|
|
metadata["sid"] = sid
|
|
rows.append(metadata)
|
|
return rows
|
|
|
|
|
|
def all_host_browser_metadata() -> list[dict[str, Any]]:
|
|
with _state_lock:
|
|
items = sorted(_sid_host_browser_metadata.items())
|
|
rows: list[dict[str, Any]] = []
|
|
for sid, metadata in items:
|
|
rows.append(
|
|
{
|
|
"sid": sid,
|
|
"supported": metadata.supported,
|
|
"enabled": metadata.enabled,
|
|
"status": metadata.status,
|
|
"browser_family": metadata.browser_family,
|
|
"profile_label": metadata.profile_label,
|
|
"profile_path": metadata.profile_path,
|
|
"cdp_endpoint": metadata.cdp_endpoint,
|
|
"browser_id": metadata.browser_id,
|
|
"browser_label": metadata.browser_label,
|
|
"available_browsers": copy.deepcopy(list(metadata.available_browsers)),
|
|
"content_helper_sha256": metadata.content_helper_sha256,
|
|
"features": list(metadata.features),
|
|
"support_reason": metadata.support_reason,
|
|
"updated_at": metadata.updated_at,
|
|
}
|
|
)
|
|
return rows
|
|
|
|
|
|
def select_computer_use_target_sid(context_id: str) -> str | None:
|
|
with _state_lock:
|
|
for sid in _candidate_sids_for_context_locked(context_id):
|
|
metadata = _sid_computer_use_metadata.get(sid)
|
|
if metadata and metadata.supported and metadata.enabled:
|
|
return sid
|
|
return None
|
|
|
|
|
|
def store_pending_file_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
future: asyncio.Future[dict[str, Any]],
|
|
loop: asyncio.AbstractEventLoop,
|
|
context_id: str | None = None,
|
|
) -> None:
|
|
with _state_lock:
|
|
_pending_file_ops[op_id] = PendingFileOperation(
|
|
sid=sid,
|
|
loop=loop,
|
|
future=future,
|
|
context_id=context_id,
|
|
)
|
|
|
|
|
|
def clear_pending_file_op(op_id: str) -> None:
|
|
with _state_lock:
|
|
_pending_file_ops.pop(op_id, None)
|
|
|
|
|
|
def resolve_pending_file_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
payload: dict[str, Any],
|
|
) -> bool:
|
|
if payload.get("chunked") is True:
|
|
return _resolve_pending_file_chunk(op_id, sid=sid, payload=payload)
|
|
return _resolve_pending(_pending_file_ops, op_id, sid=sid, payload=payload)
|
|
|
|
|
|
def _resolve_pending_file_chunk(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
payload: dict[str, Any],
|
|
) -> bool:
|
|
error = _validate_file_chunk_payload(payload)
|
|
if error:
|
|
return _fail_pending(
|
|
_pending_file_ops,
|
|
op_id,
|
|
sid=sid,
|
|
error=f"Invalid chunked file operation result: {error}",
|
|
)
|
|
|
|
chunk_index = int(payload["chunk_index"])
|
|
chunk_count = int(payload["chunk_count"])
|
|
encoded = str(payload.get("data") or "")
|
|
try:
|
|
chunk = base64.b64decode(encoded.encode("ascii"), validate=True)
|
|
except (UnicodeEncodeError, binascii.Error) as exc:
|
|
return _fail_pending(
|
|
_pending_file_ops,
|
|
op_id,
|
|
sid=sid,
|
|
error=f"Invalid chunked file operation result: {exc}",
|
|
)
|
|
|
|
with _state_lock:
|
|
pending = _pending_file_ops.get(op_id)
|
|
if pending is None or pending.sid != sid:
|
|
return False
|
|
|
|
if pending.chunk_count is None:
|
|
pending.chunk_count = chunk_count
|
|
elif pending.chunk_count != chunk_count:
|
|
_pending_file_ops.pop(op_id, None)
|
|
pending.loop.call_soon_threadsafe(
|
|
_set_future_result,
|
|
pending.future,
|
|
{
|
|
"op_id": op_id,
|
|
"ok": False,
|
|
"error": "Invalid chunked file operation result: chunk_count changed",
|
|
},
|
|
)
|
|
return True
|
|
|
|
pending.chunks[chunk_index] = chunk
|
|
if len(pending.chunks) < chunk_count:
|
|
return True
|
|
|
|
ordered = [pending.chunks[index] for index in range(chunk_count)]
|
|
_pending_file_ops.pop(op_id, None)
|
|
|
|
try:
|
|
assembled = b"".join(ordered).decode("utf-8")
|
|
result = json.loads(assembled)
|
|
if not isinstance(result, dict):
|
|
raise ValueError("decoded result is not an object")
|
|
except Exception as exc:
|
|
result = {
|
|
"op_id": op_id,
|
|
"ok": False,
|
|
"error": f"Invalid chunked file operation result: {exc}",
|
|
}
|
|
|
|
pending.loop.call_soon_threadsafe(_set_future_result, pending.future, result)
|
|
return True
|
|
|
|
|
|
def _validate_file_chunk_payload(payload: dict[str, Any]) -> str:
|
|
if payload.get("encoding") != "json+base64":
|
|
return "encoding must be json+base64"
|
|
|
|
try:
|
|
chunk_index = int(payload.get("chunk_index"))
|
|
chunk_count = int(payload.get("chunk_count"))
|
|
except (TypeError, ValueError):
|
|
return "chunk_index and chunk_count must be integers"
|
|
|
|
if chunk_count <= 0:
|
|
return "chunk_count must be positive"
|
|
if chunk_index < 0 or chunk_index >= chunk_count:
|
|
return "chunk_index out of range"
|
|
if not isinstance(payload.get("data"), str):
|
|
return "data must be a string"
|
|
return ""
|
|
|
|
|
|
def fail_pending_file_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str | None = None,
|
|
error: str,
|
|
) -> bool:
|
|
return _fail_pending(_pending_file_ops, op_id, sid=sid, error=error)
|
|
|
|
|
|
def fail_pending_file_ops_for_sid(sid: str, *, error: str) -> None:
|
|
_fail_pending_for_sid(_pending_file_ops, sid=sid, error=error)
|
|
|
|
|
|
def store_pending_exec_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
future: asyncio.Future[dict[str, Any]],
|
|
loop: asyncio.AbstractEventLoop,
|
|
context_id: str | None = None,
|
|
) -> None:
|
|
with _state_lock:
|
|
_pending_exec_ops[op_id] = PendingExecOperation(
|
|
sid=sid,
|
|
loop=loop,
|
|
future=future,
|
|
context_id=context_id,
|
|
)
|
|
|
|
|
|
def clear_pending_exec_op(op_id: str) -> None:
|
|
with _state_lock:
|
|
_pending_exec_ops.pop(op_id, None)
|
|
|
|
|
|
def resolve_pending_exec_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
payload: dict[str, Any],
|
|
) -> bool:
|
|
return _resolve_pending(_pending_exec_ops, op_id, sid=sid, payload=payload)
|
|
|
|
|
|
def fail_pending_exec_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str | None = None,
|
|
error: str,
|
|
) -> bool:
|
|
return _fail_pending(_pending_exec_ops, op_id, sid=sid, error=error)
|
|
|
|
|
|
def fail_pending_exec_ops_for_sid(sid: str, *, error: str) -> None:
|
|
_fail_pending_for_sid(_pending_exec_ops, sid=sid, error=error)
|
|
|
|
|
|
def store_pending_computer_use_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
future: asyncio.Future[dict[str, Any]],
|
|
loop: asyncio.AbstractEventLoop,
|
|
context_id: str | None = None,
|
|
) -> None:
|
|
with _state_lock:
|
|
_pending_computer_use_ops[op_id] = PendingComputerUseOperation(
|
|
sid=sid,
|
|
loop=loop,
|
|
future=future,
|
|
context_id=context_id,
|
|
)
|
|
|
|
|
|
def clear_pending_computer_use_op(op_id: str) -> None:
|
|
with _state_lock:
|
|
_pending_computer_use_ops.pop(op_id, None)
|
|
|
|
|
|
def resolve_pending_computer_use_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
payload: dict[str, Any],
|
|
) -> bool:
|
|
return _resolve_pending(_pending_computer_use_ops, op_id, sid=sid, payload=payload)
|
|
|
|
|
|
def fail_pending_computer_use_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str | None = None,
|
|
error: str,
|
|
) -> bool:
|
|
return _fail_pending(_pending_computer_use_ops, op_id, sid=sid, error=error)
|
|
|
|
|
|
def fail_pending_computer_use_ops_for_sid(sid: str, *, error: str) -> None:
|
|
_fail_pending_for_sid(_pending_computer_use_ops, sid=sid, error=error)
|
|
|
|
|
|
def store_pending_browser_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
future: asyncio.Future[dict[str, Any]],
|
|
loop: asyncio.AbstractEventLoop,
|
|
context_id: str | None = None,
|
|
) -> None:
|
|
with _state_lock:
|
|
_pending_browser_ops[op_id] = PendingBrowserOperation(
|
|
sid=sid,
|
|
loop=loop,
|
|
future=future,
|
|
context_id=context_id,
|
|
)
|
|
|
|
|
|
def clear_pending_browser_op(op_id: str) -> None:
|
|
with _state_lock:
|
|
_pending_browser_ops.pop(op_id, None)
|
|
|
|
|
|
def resolve_pending_browser_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
payload: dict[str, Any],
|
|
) -> bool:
|
|
return _resolve_pending(_pending_browser_ops, op_id, sid=sid, payload=payload)
|
|
|
|
|
|
def fail_pending_browser_op(
|
|
op_id: str,
|
|
*,
|
|
sid: str | None = None,
|
|
error: str,
|
|
) -> bool:
|
|
return _fail_pending(_pending_browser_ops, op_id, sid=sid, error=error)
|
|
|
|
|
|
def fail_pending_browser_ops_for_sid(sid: str, *, error: str) -> None:
|
|
_fail_pending_for_sid(_pending_browser_ops, sid=sid, error=error)
|
|
|
|
|
|
def _resolve_pending(
|
|
registry: dict[str, PendingFileOperation | PendingExecOperation | PendingComputerUseOperation | PendingBrowserOperation],
|
|
op_id: str,
|
|
*,
|
|
sid: str,
|
|
payload: dict[str, Any],
|
|
) -> bool:
|
|
with _state_lock:
|
|
pending = registry.get(op_id)
|
|
if pending is None or pending.sid != sid:
|
|
return False
|
|
registry.pop(op_id, None)
|
|
|
|
pending.loop.call_soon_threadsafe(_set_future_result, pending.future, dict(payload))
|
|
return True
|
|
|
|
|
|
def _fail_pending(
|
|
registry: dict[str, PendingFileOperation | PendingExecOperation | PendingComputerUseOperation | PendingBrowserOperation],
|
|
op_id: str,
|
|
*,
|
|
sid: str | None,
|
|
error: str,
|
|
) -> bool:
|
|
with _state_lock:
|
|
pending = registry.get(op_id)
|
|
if pending is None:
|
|
return False
|
|
if sid is not None and pending.sid != sid:
|
|
return False
|
|
registry.pop(op_id, None)
|
|
|
|
pending.loop.call_soon_threadsafe(
|
|
_set_future_result,
|
|
pending.future,
|
|
{"op_id": op_id, "ok": False, "error": error},
|
|
)
|
|
return True
|
|
|
|
|
|
def _fail_pending_for_sid(
|
|
registry: dict[str, PendingFileOperation | PendingExecOperation | PendingComputerUseOperation | PendingBrowserOperation],
|
|
*,
|
|
sid: str,
|
|
error: str,
|
|
) -> None:
|
|
with _state_lock:
|
|
matches = [
|
|
(op_id, pending)
|
|
for op_id, pending in registry.items()
|
|
if pending.sid == sid
|
|
]
|
|
for op_id, _pending in matches:
|
|
registry.pop(op_id, None)
|
|
|
|
for op_id, pending in matches:
|
|
pending.loop.call_soon_threadsafe(
|
|
_set_future_result,
|
|
pending.future,
|
|
{"op_id": op_id, "ok": False, "error": error},
|
|
)
|
|
|
|
|
|
def _set_future_result(
|
|
future: asyncio.Future[dict[str, Any]],
|
|
payload: dict[str, Any],
|
|
) -> None:
|
|
if not future.done():
|
|
future.set_result(payload)
|