From d4dc83ba787d2bbb8c4413c0aac41e8bdd6cde84 Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Tue, 26 May 2026 20:05:47 +0200 Subject: [PATCH 001/231] Expose installed plugin toggles Advertise the installed_plugins connector capability and add a protected API endpoint that lists already-installed Agent Zero plugins and toggles supported plugins only. The endpoint normalizes plugin metadata, preserves the installed-only safety boundary, and refuses changes to protected plugins such as _a0_connector so the CLI cannot disconnect itself. --- plugins/_a0_connector/api/v1/capabilities.py | 1 + .../_a0_connector/api/v1/installed_plugins.py | 159 ++++++++++++++++++ 2 files changed, 160 insertions(+) create mode 100644 plugins/_a0_connector/api/v1/installed_plugins.py diff --git a/plugins/_a0_connector/api/v1/capabilities.py b/plugins/_a0_connector/api/v1/capabilities.py index 327188119..7a337831e 100644 --- a/plugins/_a0_connector/api/v1/capabilities.py +++ b/plugins/_a0_connector/api/v1/capabilities.py @@ -37,6 +37,7 @@ _OPTIONAL_FEATURES: dict[str, tuple[str, ...]] = { "skills_list": ("helpers.skills", "helpers.files", "helpers.projects", "helpers.runtime"), "skills_activate": ("helpers.skills", "helpers.persist_chat"), "skills_delete": ("helpers.skills", "helpers.files", "helpers.projects", "helpers.runtime"), + "installed_plugins": ("helpers.plugins",), "model_presets": ("plugins._model_config.helpers.model_config",), "model_switcher": ("plugins._model_config.helpers.model_config",), "browser_runtime_config": ("plugins._browser.helpers.config", "helpers.plugins"), diff --git a/plugins/_a0_connector/api/v1/installed_plugins.py b/plugins/_a0_connector/api/v1/installed_plugins.py new file mode 100644 index 000000000..f2bc44be0 --- /dev/null +++ b/plugins/_a0_connector/api/v1/installed_plugins.py @@ -0,0 +1,159 @@ +"""POST /api/plugins/_a0_connector/v1/installed_plugins.""" +from __future__ import annotations + +from typing import Any, Mapping + +from helpers.api import Request, Response +import plugins._a0_connector.api.v1.base as connector_base + + +_PROTECTED_PLUGIN_REASONS = { + "_a0_connector": "The A0 Connector plugin keeps this CLI session connected.", +} + + +def _clean_text(value: object) -> str: + return str(value or "").strip() + + +def _plugin_mapping(plugin: object) -> dict[str, Any]: + if isinstance(plugin, Mapping): + return dict(plugin) + + model_dump = getattr(plugin, "model_dump", None) + if callable(model_dump): + data = model_dump(mode="json") + return data if isinstance(data, dict) else {} + + result: dict[str, Any] = {} + for key in ( + "name", + "display_name", + "description", + "version", + "author", + "repo", + "always_enabled", + "is_custom", + "has_main_screen", + "has_config_screen", + "toggle_state", + ): + if hasattr(plugin, key): + result[key] = getattr(plugin, key) + return result + + +def _plugin_payload(plugin: object) -> dict[str, Any]: + data = _plugin_mapping(plugin) + name = _clean_text(data.get("name")) + display_name = _clean_text(data.get("display_name")) or name + toggle_state = _clean_text(data.get("toggle_state")).lower() + always_enabled = bool(data.get("always_enabled")) + enabled = always_enabled or toggle_state == "enabled" + protected_reason = _PROTECTED_PLUGIN_REASONS.get(name, "") + if always_enabled and not protected_reason: + protected_reason = "Agent Zero marks this plugin as always enabled." + + return { + "name": name, + "display_name": display_name, + "description": _clean_text(data.get("description")), + "version": _clean_text(data.get("version")), + "author": _clean_text(data.get("author")), + "repo": _clean_text(data.get("repo")), + "source": "custom" if bool(data.get("is_custom")) else "builtin", + "is_custom": bool(data.get("is_custom")), + "always_enabled": always_enabled, + "enabled": enabled, + "toggle_state": "enabled" if enabled else "disabled", + "toggleable": not bool(protected_reason), + "protected_reason": protected_reason, + "has_main_screen": bool(data.get("has_main_screen")), + "has_config_screen": bool(data.get("has_config_screen")), + } + + +def _installed_plugin_payloads() -> list[dict[str, Any]]: + from helpers import plugins + + items = plugins.get_enhanced_plugins_list(custom=True, builtin=True) + payloads = [_plugin_payload(item) for item in items] + return [payload for payload in payloads if payload["name"]] + + +def _find_installed_plugin(plugin_name: str) -> dict[str, Any] | None: + normalized = plugin_name.strip() + if not normalized: + return None + for payload in _installed_plugin_payloads(): + if payload["name"] == normalized: + return payload + return None + + +def _parse_enabled(value: object) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"1", "true", "yes", "on", "enable", "enabled"}: + return True + if normalized in {"0", "false", "no", "off", "disable", "disabled"}: + return False + return None + + +class InstalledPlugins(connector_base.ProtectedConnectorApiHandler): + """List and toggle already-installed Agent Zero Core plugins only.""" + + async def process(self, input: dict, request: Request) -> dict | Response: + del request + + action = _clean_text(input.get("action") or "list").lower() + if action == "list": + plugins = _installed_plugin_payloads() + enabled_count = sum(1 for plugin in plugins if plugin["enabled"]) + return { + "ok": True, + "plugins": plugins, + "installed_count": len(plugins), + "enabled_count": enabled_count, + } + + if action != "set_enabled": + return Response(status=400, response=f"Unknown action: {action}") + + plugin_name = _clean_text(input.get("plugin_name")) + if not plugin_name: + return Response(status=400, response="Missing plugin_name") + + enabled = _parse_enabled(input.get("enabled")) + if enabled is None: + return Response(status=400, response="Missing or invalid enabled state") + + plugin = _find_installed_plugin(plugin_name) + if plugin is None: + return Response(status=404, response="Plugin not found") + + if not plugin["toggleable"]: + return Response( + status=400, + response=plugin["protected_reason"] or "Plugin cannot be toggled.", + ) + + from helpers import plugins as plugin_helpers + + plugin_helpers.toggle_plugin( + plugin_name, + enabled, + project_name="", + agent_profile="", + clear_overrides=False, + ) + updated = _find_installed_plugin(plugin_name) or plugin + return { + "ok": True, + "plugin": updated, + "plugins": _installed_plugin_payloads(), + } From 8af14fcd93b62776a1b18a477ad56fa4a2637e78 Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Wed, 27 May 2026 16:25:18 +0200 Subject: [PATCH 002/231] Clean desktop SSH agent state during self-update Remove stale runtime entries from the desktop SSH agent directory when a self-update request is consumed. Keep the cleanup best-effort so missing paths, non-directory paths, and unexpected cleanup failures do not block update startup. Cover successful cleanup, missing-directory skips, and failure fallback with focused self-update manager tests. --- docker/run/fs/exe/self_update_manager.py | 62 ++++++++++++++++ tests/test_self_update_tag_filter.py | 94 ++++++++++++++++++++++++ 2 files changed, 156 insertions(+) diff --git a/docker/run/fs/exe/self_update_manager.py b/docker/run/fs/exe/self_update_manager.py index a1917bc31..3e5301909 100644 --- a/docker/run/fs/exe/self_update_manager.py +++ b/docker/run/fs/exe/self_update_manager.py @@ -46,6 +46,9 @@ DEFAULT_BACKUP_CONFLICT_POLICY = "rename" BACKUP_CONFLICT_POLICIES = {"rename", "overwrite", "fail"} MIN_SELECTOR_VERSION = (1, 0) LATEST_SELECTOR_TAG = "latest" +TRANSIENT_DESKTOP_SSH_AGENT_RELATIVE_DIR = Path( + "usr/plugins/_desktop/profiles/agent-zero-desktop/.ssh/agent" +) def now_iso() -> str: @@ -448,6 +451,61 @@ def should_include_usr_backup_entry(source_file: Path, logger: AttemptLogger) -> return True +def clean_transient_desktop_ssh_agent_dir( + repo_dir: Path, + logger: AttemptLogger, +) -> None: + agent_dir = repo_dir / TRANSIENT_DESKTOP_SSH_AGENT_RELATIVE_DIR + try: + agent_stat = agent_dir.lstat() + except FileNotFoundError: + logger.log( + f"Transient desktop SSH agent directory not found, skipping: {agent_dir}" + ) + return + except OSError as exc: + logger.log( + f"Transient desktop SSH agent directory could not be inspected: {agent_dir}: {exc}" + ) + return + + if stat.S_ISLNK(agent_stat.st_mode) or not stat.S_ISDIR(agent_stat.st_mode): + logger.log( + f"Transient desktop SSH agent path is not a directory, skipping: {agent_dir}" + ) + return + + removed = 0 + try: + entries = list(agent_dir.iterdir()) + except OSError as exc: + logger.log( + f"Transient desktop SSH agent directory could not be listed: {agent_dir}: {exc}" + ) + return + + for entry in entries: + try: + if entry.is_symlink(): + entry.unlink(missing_ok=True) + elif entry.is_dir(): + shutil.rmtree(entry) + else: + entry.unlink(missing_ok=True) + removed += 1 + except FileNotFoundError: + continue + except OSError as exc: + logger.log( + f"Skipping transient desktop SSH agent entry after error: {entry}: {exc}" + ) + + if removed: + logger.log(f"Removed {removed} transient desktop SSH agent entries from {agent_dir}") + else: + logger.log(f"Transient desktop SSH agent directory already empty: {agent_dir}") + + def run_command( command: list[str], *, @@ -1266,6 +1324,10 @@ def docker_run_ui() -> int: logger.log(f"Consumed update file at {TRIGGER_FILE}") logger.log_block("Trigger file content", raw_text) clean_uv_cache(logger) + try: + clean_transient_desktop_ssh_agent_dir(REPO_DIR, logger) + except Exception as exc: + logger.log(f"Transient desktop SSH agent cleanup skipped after error: {exc}") try: current = get_repo_version_info(REPO_DIR) diff --git a/tests/test_self_update_tag_filter.py b/tests/test_self_update_tag_filter.py index 056225659..88844b3c4 100644 --- a/tests/test_self_update_tag_filter.py +++ b/tests/test_self_update_tag_filter.py @@ -876,6 +876,100 @@ def test_self_update_manager_usr_backup_skips_transient_desktop_ssh_agent_dir(tm ) +def test_self_update_manager_cleans_transient_desktop_ssh_agent_dir(): + manager = load_self_update_manager() + with tempfile.TemporaryDirectory(prefix="a0su-", dir="/tmp") as temp_root: + repo_dir = Path(temp_root) / "repo" + agent_dir = repo_dir / manager.TRANSIENT_DESKTOP_SSH_AGENT_RELATIVE_DIR + nested_dir = agent_dir / "nested" + agent_dir.mkdir(parents=True) + nested_dir.mkdir() + (agent_dir / "socket").write_text("ephemeral\n", encoding="utf-8") + (nested_dir / "token").write_text("ephemeral\n", encoding="utf-8") + (agent_dir / "broken-link").symlink_to("/missing/ssh-agent-socket") + socket_path = agent_dir / "runtime.sock" + messages = [] + + class ListLogger: + def log(self, message=""): + messages.append(message) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as runtime_socket: + runtime_socket.bind(str(socket_path)) + + manager.clean_transient_desktop_ssh_agent_dir(repo_dir, ListLogger()) + + assert agent_dir.exists() + assert list(agent_dir.iterdir()) == [] + assert any( + f"Removed 4 transient desktop SSH agent entries from {agent_dir}" in message + for message in messages + ) + + +def test_self_update_manager_cleans_transient_desktop_ssh_agent_dir_skips_missing( + tmp_path, +): + manager = load_self_update_manager() + repo_dir = tmp_path / "repo" + messages = [] + + class ListLogger: + def log(self, message=""): + messages.append(message) + + manager.clean_transient_desktop_ssh_agent_dir(repo_dir, ListLogger()) + + assert any( + "Transient desktop SSH agent directory not found, skipping:" in message + for message in messages + ) + + +def test_self_update_manager_desktop_ssh_cleanup_failure_does_not_block_startup( + monkeypatch, + tmp_path, +): + manager = load_self_update_manager() + request_data = {"branch": "main", "tag": "v1.2", "requested_at": "now"} + current_info = { + "branch": "main", + "describe": "v1.2", + "short_tag": "v1.2", + "commit": "abc1234", + "short_commit": "abc1234", + } + launched = [] + fake_process = object() + + monkeypatch.setattr(manager, "LOG_FILE", tmp_path / "a0-self-update.log") + monkeypatch.setattr( + manager, + "load_request_file", + lambda: (request_data, "branch: main\ntag: v1.2\n"), + ) + monkeypatch.setattr(manager, "clean_uv_cache", lambda logger: None) + monkeypatch.setattr( + manager, + "clean_transient_desktop_ssh_agent_dir", + lambda repo_dir, logger: (_ for _ in ()).throw(RuntimeError("cleanup boom")), + ) + monkeypatch.setattr(manager, "get_repo_version_info", lambda repo_dir: current_info) + monkeypatch.setattr(manager, "record_result", lambda **kwargs: None) + monkeypatch.setattr( + manager, + "launch_ui_process", + lambda repo_dir, logger: launched.append(repo_dir) or fake_process, + ) + monkeypatch.setattr(manager, "wait_for_process", lambda process: 0) + + assert manager.docker_run_ui() == 0 + assert launched == [manager.REPO_DIR] + assert "Transient desktop SSH agent cleanup skipped after error: cleanup boom" in ( + tmp_path / "a0-self-update.log" + ).read_text(encoding="utf-8") + + def test_self_update_manager_clean_uv_cache_uses_uv_when_available(monkeypatch): manager = load_self_update_manager() commands = [] From 1bbbfa44acd87b9702176dfb88be410bcd861554 Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Tue, 26 May 2026 20:18:38 +0530 Subject: [PATCH 003/231] Improve Telegram integration command and streaming UX --- helpers/integration_commands.py | 314 +++++++++++++- plugins/_telegram_integration/README.md | 24 +- .../python/job_loop/_10_telegram_bot.py | 5 +- .../_45_telegram_draft_start.py | 17 + .../process_chain_end/_55_telegram_reply.py | 3 + .../_45_telegram_draft_response.py | 30 ++ .../system_prompt/_20_telegram_context.py | 17 + .../_50_telegram_response.py | 8 +- .../_45_telegram_draft_tool.py | 18 + .../helpers/bot_manager.py | 22 +- .../helpers/command_ui.py | 401 ++++++++++++++++++ .../helpers/constants.py | 8 + .../helpers/draft_stream.py | 361 ++++++++++++++++ .../_telegram_integration/helpers/handler.py | 165 ++++--- .../helpers/telegram_client.py | 84 ++++ .../_whatsapp_integration/helpers/handler.py | 18 + 16 files changed, 1405 insertions(+), 90 deletions(-) create mode 100644 plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py create mode 100644 plugins/_telegram_integration/extensions/python/response_stream/_45_telegram_draft_response.py create mode 100644 plugins/_telegram_integration/extensions/python/tool_execute_before/_45_telegram_draft_tool.py create mode 100644 plugins/_telegram_integration/helpers/command_ui.py create mode 100644 plugins/_telegram_integration/helpers/draft_stream.py diff --git a/helpers/integration_commands.py b/helpers/integration_commands.py index d78f3ff6e..655a3cc87 100644 --- a/helpers/integration_commands.py +++ b/helpers/integration_commands.py @@ -1,6 +1,7 @@ from __future__ import annotations import re +from dataclasses import dataclass from typing import TYPE_CHECKING from helpers import message_queue as mq @@ -14,7 +15,83 @@ if TYPE_CHECKING: _CLEAR_VALUES = {"", "default", "none", "clear", "off"} -_SUPPORTED_COMMANDS = {"/send", "/queue", "/project", "/config", "/preset"} + + +@dataclass(frozen=True) +class IntegrationCommandDef: + name: str + description: str + category: str + aliases: tuple[str, ...] = () + args_hint: str = "" + menu: bool = True + + +COMMAND_REGISTRY: tuple[IntegrationCommandDef, ...] = ( + IntegrationCommandDef("commands", "Show all integration commands.", "Info", aliases=("help",)), + IntegrationCommandDef( + "status", + "Show this chat's project, model, agent, and queue state.", + "Info", + ), + IntegrationCommandDef("new", "Start a fresh chat context.", "Session"), + IntegrationCommandDef("clear", "Reset the current chat context.", "Session", aliases=("reset",)), + IntegrationCommandDef( + "queue", + "Show or manage queued messages.", + "Session", + args_hint="[send|clear]", + ), + IntegrationCommandDef("send", "Send queued messages now.", "Session", aliases=("push",)), + IntegrationCommandDef( + "steer", + "Intervene in the currently running task.", + "Session", + args_hint="", + ), + IntegrationCommandDef("pause", "Pause the active run.", "Session"), + IntegrationCommandDef("resume", "Resume a paused run.", "Session"), + IntegrationCommandDef("nudge", "Nudge the active run.", "Session"), + IntegrationCommandDef( + "stream", + "Enable or disable Telegram response streaming.", + "Configuration", + args_hint="[on|off]", + ), + IntegrationCommandDef( + "tools", + "Show or hide Telegram tool progress.", + "Configuration", + args_hint="[on|off]", + ), + IntegrationCommandDef( + "project", + "Show or switch the active project.", + "Configuration", + args_hint="[name|none]", + ), + IntegrationCommandDef( + "model", + "Show or switch the chat model preset.", + "Configuration", + aliases=("config", "preset"), + args_hint="[preset|default]", + ), + IntegrationCommandDef( + "agent", + "Show or switch the agent profile.", + "Configuration", + aliases=("profile",), + args_hint="[profile]", + ), +) + + +_COMMAND_LOOKUP = { + f"/{name}": command + for command in COMMAND_REGISTRY + for name in (command.name, *command.aliases) +} def extract_command_line(text: str) -> str: @@ -32,11 +109,53 @@ def parse_command(text: str) -> tuple[str, str] | None: return None command, _, args = line.partition(" ") - command = command.strip().lower() - if command not in _SUPPORTED_COMMANDS: + command = _normalize_command_token(command) + resolved = resolve_command(command) + if not resolved: return None - return command, args.strip() + return f"/{resolved.name}", args.strip() + + +def resolve_command(command: str) -> IntegrationCommandDef | None: + normalized = _normalize_command_token(command) + if not normalized.startswith("/"): + normalized = f"/{normalized}" + return _COMMAND_LOOKUP.get(normalized) + + +def telegram_menu_commands() -> list[tuple[str, str]]: + return [ + (command.name, _telegram_description(command)) + for command in COMMAND_REGISTRY + if command.menu + ] + + +def command_names(include_aliases: bool = True) -> list[str]: + names: list[str] = [] + for command in COMMAND_REGISTRY: + names.append(command.name) + if include_aliases: + names.extend(command.aliases) + return names + + +def help_text(*, full: bool = False) -> str: + commands = COMMAND_REGISTRY if full else tuple(c for c in COMMAND_REGISTRY if c.menu) + lines = ["Available commands:"] + for command in commands: + args = f" {command.args_hint}" if command.args_hint else "" + alias_text = "" + if command.aliases: + alias_text = f" (alias: {', '.join('/' + alias for alias in command.aliases)})" + lines.append(f"/{command.name}{args} - {command.description}{alias_text}") + return "\n".join(lines) + + +def unknown_command_text(command: str) -> str: + token = _normalize_command_token(command).split(" ", 1)[0] + return f"Unknown command: {token}\n\n{help_text(full=True)}" def try_handle_command(context: "AgentContext", text: str) -> str | None: @@ -45,17 +164,47 @@ def try_handle_command(context: "AgentContext", text: str) -> str | None: return None command, args = parsed + if command == "/commands": + return help_text(full=True) + if command == "/status": + return _handle_status(context) + if command in {"/new", "/clear"}: + return _handle_clear(context, new_chat=(command == "/new")) if command == "/send": return _handle_queue(context, "send") if command == "/queue": return _handle_queue(context, args) + if command == "/steer": + return _handle_steer(context, args) + if command == "/pause": + return _handle_pause(context) + if command == "/resume": + return _handle_resume(context) + if command == "/nudge": + return _handle_nudge(context) + if command == "/stream": + return _handle_toggle(context, args, "telegram_stream_enabled", "Response streaming") + if command == "/tools": + return _handle_toggle(context, args, "telegram_tools_enabled", "Tool progress") if command == "/project": return _handle_project(context, args) - if command in {"/config", "/preset"}: - return _handle_config(context, args) + if command == "/model": + return _handle_model(context, args) + if command == "/agent": + return _handle_agent(context, args) return None +def _normalize_command_token(command: str) -> str: + normalized = command.strip().lower() + if not normalized: + return "" + token, *rest = normalized.split(" ", 1) + if "@" in token: + token = token.split("@", 1)[0] + return f"{token} {rest[0]}".strip() if rest else token + + def _handle_queue(context: "AgentContext", args: str) -> str: queue = mq.get_queue(context) count = len(queue) @@ -68,8 +217,13 @@ def _handle_queue(context: "AgentContext", args: str) -> str: "Use /send or /queue send to send everything as one batch." ) + if action in {"clear", "reset"}: + mq.remove(context) + mark_dirty_for_context(context.id, reason="integration_commands.queue_clear") + return "Queue cleared." + if action not in {"send", "all"}: - return "Unknown queue action. Use /queue send to flush the queue." + return "Unknown queue action. Use /queue send to flush or /queue clear to clear." if count == 0: return "Queue is empty." @@ -80,6 +234,73 @@ def _handle_queue(context: "AgentContext", args: str) -> str: return f"Sent {sent_count} queued {noun} as one batch." +def _handle_status(context: "AgentContext") -> str: + project_name = context.get_data("project") or "none" + override = context.get_data("chat_model_override") + agent_profile = getattr(context.agent0.config, "profile", "default") + running = "running" if context.is_running() else "idle" + if getattr(context, "paused", False): + running = "paused" + queue_count = len(mq.get_queue(context)) + return ( + f"Status: {running}\n" + f"Project: {project_name}\n" + f"Model: {_describe_override(override)}\n" + f"Agent: {agent_profile}\n" + f"Queued messages: {queue_count}" + ) + + +def _handle_clear(context: "AgentContext", *, new_chat: bool) -> str: + context.reset() + mq.remove(context) + save_tmp_chat(context) + reason = "integration_commands.new" if new_chat else "integration_commands.clear" + mark_dirty_for_context(context.id, reason=reason) + return "Started a fresh chat." if new_chat else "Chat cleared." + + +def _handle_steer(context: "AgentContext", args: str) -> str: + message = args.strip() + if not message: + return "Usage: /steer " + from agent import UserMessage + + context.communicate(UserMessage(message=message)) + if context.is_running(): + return "Steering message sent to the active run." + return "Message sent." + + +def _handle_pause(context: "AgentContext") -> str: + if not context.is_running(): + return "No active run is currently running." + context.paused = True + return "Agent paused." + + +def _handle_resume(context: "AgentContext") -> str: + context.paused = False + return "Agent resumed." + + +def _handle_nudge(context: "AgentContext") -> str: + context.nudge() + return "Agent nudged." + + +def _handle_toggle(context: "AgentContext", args: str, key: str, label: str) -> str: + value = _parse_toggle(args) + current = _get_toggle(context, key) + if value is None: + state = "on" if current else "off" + return f"{label}: {state}. Use /{key.split('_')[1]} on or /{key.split('_')[1]} off." + context.set_data(key, value) + save_tmp_chat(context) + mark_dirty_for_context(context.id, reason=f"integration_commands.{key}") + return f"{label} {'enabled' if value else 'disabled'}." + + def _handle_project(context: "AgentContext", args: str) -> str: items = projects.get_active_projects_list() or [] current_name = context.get_data("project") or "" @@ -115,7 +336,7 @@ def _handle_project(context: "AgentContext", args: str) -> str: return f"Switched project to {match.get('title') or match['name']}." -def _handle_config(context: "AgentContext", args: str) -> str: +def _handle_model(context: "AgentContext", args: str) -> str: allowed = model_config.is_chat_override_allowed(context.agent0) presets = [preset for preset in model_config.get_presets() if preset.get("name")] current_override = context.get_data("chat_model_override") @@ -123,12 +344,12 @@ def _handle_config(context: "AgentContext", args: str) -> str: if not args: current_label = _describe_override(current_override) available = ", ".join(preset["name"] for preset in presets) or "none" - suffix = "Use /config to switch, or /config default to clear it." + suffix = "Use /model to switch, or /model default to clear it." if not allowed: suffix = "Per-chat config switching is disabled in Model Configuration." return ( - f"Current config: {current_label}\n" - f"Available configs: {available}\n" + f"Current model: {current_label}\n" + f"Available presets: {available}\n" f"{suffix}" ) @@ -159,7 +380,49 @@ def _handle_config(context: "AgentContext", args: str) -> str: context.set_data("chat_model_override", {"preset_name": preset_name}) save_tmp_chat(context) mark_dirty_for_context(context.id, reason="integration_commands.config_set") - return f"Switched config to {preset_name}." + return f"Switched model preset to {preset_name}." + + +def _handle_agent(context: "AgentContext", args: str) -> str: + from agent import Agent + from helpers import subagents + from initialize import initialize_agent + + items = subagents.get_all_agents_list() + current = getattr(context.agent0.config, "profile", "default") + if not args: + available = ", ".join(_format_agent_entry(item) for item in items) or "none" + return ( + f"Current agent: {current}\n" + f"Available agents: {available}\n" + "Use /agent to switch after the current run finishes." + ) + + if context.is_running(): + return "Agent profile can be changed after the current run finishes." + + desired = _strip_quotes(args) + match, ambiguous = _match_named_item(items, desired, keys=("key", "label")) + if ambiguous: + names = ", ".join(_format_agent_entry(item) for item in ambiguous) + return f"Agent profile is ambiguous. Matches: {names}" + if not match: + available = ", ".join(_format_agent_entry(item) for item in items) or "none" + return f"Agent profile '{desired}' was not found. Available agents: {available}" + + profile = str(match["key"]) + if profile == current: + return f"Already using agent {match.get('label') or profile}." + + config = initialize_agent(override_settings={"agent_profile": profile}) + context.config = config + agent = context.agent0 + while agent: + agent.config = config + agent = agent.get_data(Agent.DATA_NAME_SUBORDINATE) + save_tmp_chat(context) + mark_dirty_for_context(context.id, reason="integration_commands.agent_set") + return f"Switched agent to {match.get('label') or profile}." def _format_project_entry(item: dict) -> str: @@ -170,6 +433,19 @@ def _format_project_entry(item: dict) -> str: return name or title +def _format_agent_entry(item: dict) -> str: + key = str(item.get("key", "") or "").strip() + label = str(item.get("label", "") or "").strip() + if label and label.lower() != key.lower(): + return f"{label} ({key})" + return key or label + + +def _telegram_description(command: IntegrationCommandDef) -> str: + description = command.description.strip() + return description[:255] if len(description) > 255 else description + + def _describe_project(items: list[dict], current_name: str) -> str: if not current_name: return "none" @@ -201,6 +477,20 @@ def _normalize_lookup(value: str) -> str: return lowered.strip() +def _get_toggle(context: "AgentContext", key: str) -> bool: + value = context.get_data(key) + return True if value is None else bool(value) + + +def _parse_toggle(args: str) -> bool | None: + value = _normalize_lookup(args) + if value in {"on", "enable", "enabled", "yes", "true", "1"}: + return True + if value in {"off", "disable", "disabled", "no", "false", "0"}: + return False + return None + + def _match_named_item( items: list[dict], desired: str, diff --git a/plugins/_telegram_integration/README.md b/plugins/_telegram_integration/README.md index 8201ab882..bf486184e 100644 --- a/plugins/_telegram_integration/README.md +++ b/plugins/_telegram_integration/README.md @@ -16,10 +16,16 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - Supports both long-polling and webhook delivery modes. - **Per-user chat sessions** - Each Telegram user gets a dedicated `AgentContext`, persisted across restarts via a JSON state file. - - `/start` creates a context; `/clear` resets it. + - `/start` creates a context; `/new` starts fresh; `/clear` resets the current context. + - `/commands` shows the shared integration command menu (`/help` is an alias). + - `/status` shows project, model, agent profile, and queue state. - `/project ` switches the active project for the current chat. - - `/config ` switches the active model preset for the current chat. - - `/send` or `/queue send` flushes the queued messages for the current chat. + - `/model ` switches the active model preset for the current chat (`/config` remains an alias). + - `/agent ` switches the active agent profile when the current run is idle. + - `/model`, `/project`, and `/agent` show Telegram inline keyboard pickers when used without arguments. + - `/stream` toggles live response streaming; `/tools` toggles tool progress messages. + - `/send` or `/queue send` flushes queued messages for the current chat. + - `/steer ` sends an intervention to the active run. - **Group support** - Three modes: `mention` (respond only when @mentioned or replied to), `all` (respond to every message), `off` (private only). - Optional welcome message for new members. @@ -27,6 +33,8 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - Extracts text, captions, locations, contacts, stickers, and attachment indicators. - Downloads photos, documents, audio, voice, and video into `usr/uploads/` with configurable auto-cleanup. - **Reply delivery** + - Streams tool progress and response text through real Telegram messages updated with `editMessageText`. + - Tool progress and the AI response are separate messages; only the AI response replies to the user message. - `tool_execute_after` intercepts the `response` tool — sends inline progress for `break_loop=false`. - `process_chain_end` auto-sends the final response, with retry logic on failure. - **Formatting** @@ -36,8 +44,11 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - Agent can attach a `keyboard` array to the `response` tool; button presses feed back as user messages. - **Typing indicator** - Persistent "typing…" action while the agent is processing, cancelled on reply. +- **Busy-run queue** + - Normal user messages received while a run is active are queued automatically. + - `/steer ` is still delivered immediately as an intervention. - **Notifications** - - Optional WebUI notifications for incoming Telegram messages and `/clear` events. + - Optional WebUI notifications for incoming Telegram messages. - **Access control** - Per-bot allow-list by Telegram user ID or @username. Empty list = open to all. - **Project binding** @@ -50,8 +61,13 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - `helpers/handler.py` — Central message routing, context lifecycle, user auth, attachment download, reply sending, typing indicator. - `helpers/bot_manager.py` — Bot creation, polling/webhook lifecycle, bot registry. - `helpers/telegram_client.py` — Low-level Telegram API wrapper: send text/file/photo, Markdown→HTML converter, keyboard builder, message splitting. + - `helpers/command_ui.py` — Telegram inline keyboard command pickers and callback handling. + - `helpers/draft_stream.py` — Editable Telegram message streaming for tool progress and response previews. - **Extensions** - `extensions/python/job_loop/_10_telegram_bot.py` — Bot lifecycle manager, starts/stops bots on each tick. + - `extensions/python/message_loop_start/_45_telegram_draft_start.py` — Initializes Telegram response streaming. + - `extensions/python/tool_execute_before/_45_telegram_draft_tool.py` — Streams tool-start progress. + - `extensions/python/response_stream/_45_telegram_draft_response.py` — Streams final response previews. - `extensions/python/system_prompt/_20_telegram_context.py` — Injects Telegram-specific system prompt. - `extensions/python/tool_execute_after/_50_telegram_response.py` — Intercepts `response` tool for inline delivery. - `extensions/python/process_chain_end/_55_telegram_reply.py` — Auto-sends final reply with retry. diff --git a/plugins/_telegram_integration/extensions/python/job_loop/_10_telegram_bot.py b/plugins/_telegram_integration/extensions/python/job_loop/_10_telegram_bot.py index 6355a9a96..dff9222ef 100644 --- a/plugins/_telegram_integration/extensions/python/job_loop/_10_telegram_bot.py +++ b/plugins/_telegram_integration/extensions/python/job_loop/_10_telegram_bot.py @@ -31,13 +31,13 @@ class TelegramBotManager(Extension): get_all_bots, create_bot, cache_bot_info, + register_bot_commands, start_polling, setup_webhook, stop_bot, ) from plugins._telegram_integration.helpers.handler import ( handle_start, - handle_clear, handle_message, handle_callback_query, handle_new_members, @@ -73,7 +73,6 @@ class TelegramBotManager(Extension): try: # Create handler closures that capture bot_name and config _on_start = partial(_make_handler(handle_start), bot_name=name, bot_cfg=bot_cfg) - _on_clear = partial(_make_handler(handle_clear), bot_name=name, bot_cfg=bot_cfg) _on_message = partial(_make_handler(handle_message), bot_name=name, bot_cfg=bot_cfg) _on_callback = partial(_make_handler(handle_callback_query), bot_name=name, bot_cfg=bot_cfg) _on_new_members = partial(_make_handler(handle_new_members), bot_name=name, bot_cfg=bot_cfg) @@ -83,7 +82,6 @@ class TelegramBotManager(Extension): token=bot_cfg["token"], on_message=_on_message, on_command_start=_on_start, - on_command_clear=_on_clear, on_command_control=_on_message, on_callback_query=_on_callback, on_new_members=_on_new_members, @@ -91,6 +89,7 @@ class TelegramBotManager(Extension): ) await cache_bot_info(instance) + await register_bot_commands(instance) mode = bot_cfg.get("mode", "polling") if mode == "webhook": diff --git a/plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py b/plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py new file mode 100644 index 000000000..af773207d --- /dev/null +++ b/plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py @@ -0,0 +1,17 @@ +from agent import LoopData +from helpers.extension import Extension +from plugins._telegram_integration.helpers.constants import CTX_TG_BOT + + +class TelegramDraftStart(Extension): + + async def execute(self, loop_data: LoopData = LoopData(), **kwargs): + if not self.agent or self.agent.number != 0: + return + context = self.agent.context + if not context.data.get(CTX_TG_BOT): + return + + from plugins._telegram_integration.helpers import draft_stream + + await draft_stream.start(context) diff --git a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py index d0730a6dc..dcd28665e 100644 --- a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py +++ b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py @@ -41,6 +41,9 @@ class TelegramAutoReply(Extension): typing_stop = context.data.pop(CTX_TG_TYPING_STOP, None) if typing_stop: typing_stop.set() + from plugins._telegram_integration.helpers import draft_stream + + draft_stream.clear(context) context.data.pop(CTX_TG_REPLY_TO, None) async def _send_reply( diff --git a/plugins/_telegram_integration/extensions/python/response_stream/_45_telegram_draft_response.py b/plugins/_telegram_integration/extensions/python/response_stream/_45_telegram_draft_response.py new file mode 100644 index 000000000..47e647d36 --- /dev/null +++ b/plugins/_telegram_integration/extensions/python/response_stream/_45_telegram_draft_response.py @@ -0,0 +1,30 @@ +from agent import LoopData +from helpers.extension import Extension +from plugins._telegram_integration.helpers.constants import CTX_TG_BOT + + +class TelegramDraftResponse(Extension): + + async def execute( + self, + loop_data: LoopData = LoopData(), + parsed: dict | None = None, + **kwargs, + ): + if not self.agent or self.agent.number != 0: + return + context = self.agent.context + if not context.data.get(CTX_TG_BOT): + return + + parsed = parsed or {} + if parsed.get("tool_name") != "response": + return + tool_args = parsed.get("tool_args") or {} + text = tool_args.get("text") or tool_args.get("message") or "" + if not text: + return + + from plugins._telegram_integration.helpers import draft_stream + + await draft_stream.update_response(context, str(text)) diff --git a/plugins/_telegram_integration/extensions/python/system_prompt/_20_telegram_context.py b/plugins/_telegram_integration/extensions/python/system_prompt/_20_telegram_context.py index bf580a43c..9d337cf8d 100644 --- a/plugins/_telegram_integration/extensions/python/system_prompt/_20_telegram_context.py +++ b/plugins/_telegram_integration/extensions/python/system_prompt/_20_telegram_context.py @@ -1,5 +1,6 @@ from helpers.extension import Extension from agent import LoopData +from helpers import integration_commands from plugins._telegram_integration.helpers.constants import CTX_TG_BOT, CTX_TG_BOT_CFG @@ -18,6 +19,7 @@ class TelegramContextPrompt(Extension): system_prompt.append( self.agent.read_prompt("fw.telegram.system_context_reply.md") ) + system_prompt.append(_telegram_commands_prompt()) # Inject per-bot agent instructions (once in system prompt) bot_cfg = self.agent.context.data.get(CTX_TG_BOT_CFG, {}) @@ -29,3 +31,18 @@ class TelegramContextPrompt(Extension): instructions=instructions, ) ) + + +def _telegram_commands_prompt() -> str: + lines = [ + "Telegram slash commands are handled by the integration before you see the message.", + "Do not invent command help, unknown-command replies, or pretend to execute slash commands.", + "If the user asks what commands exist, refer them to /commands.", + "Current integration commands:", + ] + for name in integration_commands.command_names(include_aliases=False): + definition = integration_commands.resolve_command(name) + if definition: + args = f" {definition.args_hint}" if definition.args_hint else "" + lines.append(f"- /{definition.name}{args}: {definition.description}") + return "\n".join(lines) diff --git a/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py b/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py index 14968c623..43fcbae24 100644 --- a/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py +++ b/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py @@ -13,14 +13,18 @@ class TelegramResponseIntercept(Extension): async def execute( self, tool_name: str = "", response: Response | None = None, **kwargs, ): - if tool_name != "response": - return if not self.agent: return context = self.agent.context if not context.data.get(CTX_TG_BOT): return + from plugins._telegram_integration.helpers import draft_stream + + if tool_name != "response": + await draft_stream.add_tool_done(context, tool_name, ok=response is not None) + return + tool = self.agent.loop_data.current_tool if not tool: return diff --git a/plugins/_telegram_integration/extensions/python/tool_execute_before/_45_telegram_draft_tool.py b/plugins/_telegram_integration/extensions/python/tool_execute_before/_45_telegram_draft_tool.py new file mode 100644 index 000000000..4496cd422 --- /dev/null +++ b/plugins/_telegram_integration/extensions/python/tool_execute_before/_45_telegram_draft_tool.py @@ -0,0 +1,18 @@ +from helpers.extension import Extension +from plugins._telegram_integration.helpers.constants import CTX_TG_BOT + + +class TelegramDraftToolStart(Extension): + + async def execute(self, tool_name: str = "", tool_args: dict | None = None, **kwargs): + if not self.agent or self.agent.number != 0: + return + if (tool_name or "").strip().lower() == "response": + return + context = self.agent.context + if not context.data.get(CTX_TG_BOT): + return + + from plugins._telegram_integration.helpers import draft_stream + + await draft_stream.add_tool_start(context, tool_name, tool_args or {}) diff --git a/plugins/_telegram_integration/helpers/bot_manager.py b/plugins/_telegram_integration/helpers/bot_manager.py index b59c68664..757804c21 100644 --- a/plugins/_telegram_integration/helpers/bot_manager.py +++ b/plugins/_telegram_integration/helpers/bot_manager.py @@ -6,7 +6,8 @@ from aiogram import Bot, Dispatcher, Router, F from aiogram.client.default import DefaultBotProperties from aiogram.enums import ParseMode, ChatType, ContentType from aiogram.filters import Command, CommandStart -from aiogram.types import Message +from aiogram.types import BotCommand, Message +from helpers import integration_commands from helpers.errors import format_error from helpers.print_style import PrintStyle @@ -44,7 +45,6 @@ def create_bot( token: str, on_message: Callable[..., Awaitable], on_command_start: Callable[..., Awaitable], - on_command_clear: Callable[..., Awaitable], on_command_control: Callable[..., Awaitable] | None = None, on_callback_query: Callable[..., Awaitable] | None = None, on_new_members: Callable[..., Awaitable] | None = None, @@ -56,11 +56,10 @@ def create_bot( # Register command handlers router.message.register(on_command_start, CommandStart()) - router.message.register(on_command_clear, Command("clear")) if on_command_control: router.message.register( on_command_control, - Command(commands=["project", "config", "preset", "queue", "send"]), + Command(commands=integration_commands.command_names()), ) if on_callback_query: @@ -93,6 +92,21 @@ def create_bot( return instance +async def register_bot_commands(instance: BotInstance) -> None: + """Register Telegram's native / command menu from the shared integration registry.""" + commands = [ + BotCommand(command=name, description=description) + for name, description in integration_commands.telegram_menu_commands() + ] + if not commands: + return + try: + await instance.bot.set_my_commands(commands) + PrintStyle.info(f"Telegram ({instance.name}): registered {len(commands)} bot commands") + except Exception as e: + PrintStyle.error(f"Telegram ({instance.name}): failed to register bot commands: {format_error(e)}") + + async def cache_bot_info(instance: BotInstance): """Fetch and cache bot info. Call after create_bot.""" if not instance.bot_info: diff --git a/plugins/_telegram_integration/helpers/command_ui.py b/plugins/_telegram_integration/helpers/command_ui.py new file mode 100644 index 000000000..0ce9a5ae6 --- /dev/null +++ b/plugins/_telegram_integration/helpers/command_ui.py @@ -0,0 +1,401 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from agent import AgentContext +from helpers import projects, subagents +from helpers import integration_commands +from helpers.persist_chat import save_tmp_chat +from helpers.state_monitor_integration import mark_dirty_for_context +from plugins._model_config.helpers import model_config +from plugins._telegram_integration.helpers import telegram_client as tc +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_STREAM_ENABLED, + CTX_TG_TOOLS_ENABLED, +) + +PAGE_SIZE = 8 +CALLBACK_PREFIX = "tg" + + +@dataclass(frozen=True) +class PickerItem: + key: str + label: str + + +async def handle_command( + context: AgentContext, + token: str, + chat_id: int, + reply_to_message_id: int | None, + text: str, +) -> bool: + parsed = integration_commands.parse_command(text or "") + if not parsed: + return False + command, args = parsed + if command in {"/model", "/config", "/preset"} and not args: + await send_model_picker(context, token, chat_id, reply_to_message_id, 0) + return True + if command == "/project" and not args: + await send_project_picker(context, token, chat_id, reply_to_message_id, 0) + return True + if command in {"/agent", "/profile"} and not args: + await send_agent_picker(context, token, chat_id, reply_to_message_id, 0) + return True + if command == "/stream": + await send_toggle_picker( + context, + token, + chat_id, + reply_to_message_id, + CTX_TG_STREAM_ENABLED, + "Response streaming", + args, + ) + return True + if command == "/tools": + await send_toggle_picker( + context, + token, + chat_id, + reply_to_message_id, + CTX_TG_TOOLS_ENABLED, + "Tool progress", + args, + ) + return True + return False + + +async def handle_callback( + context: AgentContext, + token: str, + chat_id: int, + message_id: int, + data: str, +) -> bool: + parts = (data or "").split(":") + if len(parts) < 3 or parts[0] != CALLBACK_PREFIX: + return False + kind, action = parts[1], parts[2] + value = parts[3] if len(parts) > 3 else "" + if action == "noop": + return True + if action == "page": + page = _safe_int(value) + if kind == "model": + await edit_model_picker(context, token, chat_id, message_id, page) + elif kind == "project": + await edit_project_picker(context, token, chat_id, message_id, page) + elif kind == "agent": + await edit_agent_picker(context, token, chat_id, message_id, page) + return True + if kind == "model" and action in {"set", "clear"}: + await _select_model(context, _safe_int(value), clear=(action == "clear")) + await edit_model_picker(context, token, chat_id, message_id, 0, selected=True) + return True + if kind == "project" and action in {"set", "clear"}: + await _select_project(context, _safe_int(value), clear=(action == "clear")) + await edit_project_picker(context, token, chat_id, message_id, 0, selected=True) + return True + if kind == "agent" and action == "set": + await _select_agent(context, _safe_int(value)) + await edit_agent_picker(context, token, chat_id, message_id, 0, selected=True) + return True + if kind in {"stream", "tools"} and action in {"on", "off"}: + key = CTX_TG_STREAM_ENABLED if kind == "stream" else CTX_TG_TOOLS_ENABLED + label = "Response streaming" if kind == "stream" else "Tool progress" + context.set_data(key, action == "on") + save_tmp_chat(context) + mark_dirty_for_context(context.id, reason=f"telegram.{kind}_toggle") + await edit_toggle_picker(context, token, chat_id, message_id, key, label) + return True + return True + + +async def send_model_picker( + context: AgentContext, + token: str, + chat_id: int, + reply_to_message_id: int | None, + page: int, +) -> None: + text, markup = _model_view(context, page) + await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup) + + +async def edit_model_picker( + context: AgentContext, + token: str, + chat_id: int, + message_id: int, + page: int, + *, + selected: bool = False, +) -> None: + text, markup = _model_view(context, page, selected=selected) + await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup) + + +async def send_project_picker( + context: AgentContext, + token: str, + chat_id: int, + reply_to_message_id: int | None, + page: int, +) -> None: + text, markup = _project_view(context, page) + await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup) + + +async def edit_project_picker( + context: AgentContext, + token: str, + chat_id: int, + message_id: int, + page: int, + *, + selected: bool = False, +) -> None: + text, markup = _project_view(context, page, selected=selected) + await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup) + + +async def send_agent_picker( + context: AgentContext, + token: str, + chat_id: int, + reply_to_message_id: int | None, + page: int, +) -> None: + text, markup = _agent_view(context, page) + await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup) + + +async def edit_agent_picker( + context: AgentContext, + token: str, + chat_id: int, + message_id: int, + page: int, + *, + selected: bool = False, +) -> None: + text, markup = _agent_view(context, page, selected=selected) + await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup) + + +async def send_toggle_picker( + context: AgentContext, + token: str, + chat_id: int, + reply_to_message_id: int | None, + key: str, + label: str, + args: str = "", +) -> None: + desired = _parse_toggle(args) + if desired is not None: + context.set_data(key, desired) + save_tmp_chat(context) + mark_dirty_for_context(context.id, reason=f"telegram.{key}") + text, markup = _toggle_view(context, key, label) + await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup) + + +async def edit_toggle_picker( + context: AgentContext, + token: str, + chat_id: int, + message_id: int, + key: str, + label: str, +) -> None: + text, markup = _toggle_view(context, key, label) + await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup) + + +def _model_view( + context: AgentContext, + page: int, + *, + selected: bool = False, +) -> tuple[str, dict | None]: + presets = [ + PickerItem(str(preset.get("name", "")), str(preset.get("name", ""))) + for preset in model_config.get_presets() + if isinstance(preset, dict) and preset.get("name") + ] + current = context.get_data("chat_model_override") + current_name = current.get("preset_name") if isinstance(current, dict) else "" + status = f"Current model: {_html(current_name or 'Default')}" + if not model_config.is_chat_override_allowed(context.agent0): + return status + "\nPer-chat model switching is disabled.", None + if selected: + status = "Model updated.\n" + status + rows = _paged_buttons("model", presets, page, current_name) + rows.append([{"text": "Default", "callback_data": "tg:model:clear"}]) + return status, {"inline_keyboard": rows} + + +def _project_view( + context: AgentContext, + page: int, + *, + selected: bool = False, +) -> tuple[str, dict | None]: + items = [ + PickerItem(str(item.get("name", "")), str(item.get("title") or item.get("name") or "")) + for item in projects.get_active_projects_list() or [] + if item.get("name") + ] + current = context.get_data("project") or "" + status = f"Current project: {_html(_label_for(items, current) or 'none')}" + if selected: + status = "Project updated.\n" + status + rows = _paged_buttons("project", items, page, current) + rows.append([{"text": "No project", "callback_data": "tg:project:clear"}]) + return status, {"inline_keyboard": rows} + + +def _agent_view( + context: AgentContext, + page: int, + *, + selected: bool = False, +) -> tuple[str, dict | None]: + items = [ + PickerItem(str(item.get("key", "")), str(item.get("label") or item.get("key") or "")) + for item in subagents.get_all_agents_list() + if item.get("key") + ] + current = getattr(context.agent0.config, "profile", "") or "agent0" + status = f"Current agent: {_html(_label_for(items, current) or current)}" + if context.is_running(): + status += "\nAgent profile can be changed after the current run finishes." + elif selected: + status = "Agent updated.\n" + status + rows = _paged_buttons("agent", items, page, current, disabled=context.is_running()) + if not rows: + return status + "\nNo agent profiles were found.", None + return status, {"inline_keyboard": rows} + + +def _toggle_view(context: AgentContext, key: str, label: str) -> tuple[str, dict]: + enabled = _toggle_enabled(context, key) + kind = "stream" if key == CTX_TG_STREAM_ENABLED else "tools" + state = "enabled" if enabled else "disabled" + text = f"{_html(label)}: {state}" + rows = [[ + {"text": ("On" if enabled else "Turn on"), "callback_data": f"tg:{kind}:on"}, + {"text": ("Off" if not enabled else "Turn off"), "callback_data": f"tg:{kind}:off"}, + ]] + return text, {"inline_keyboard": rows} + + +async def _select_model(context: AgentContext, index: int, *, clear: bool = False) -> None: + if not model_config.is_chat_override_allowed(context.agent0): + return + if clear: + context.set_data("chat_model_override", None) + else: + presets = [preset for preset in model_config.get_presets() if preset.get("name")] + if index < 0 or index >= len(presets): + return + context.set_data("chat_model_override", {"preset_name": presets[index]["name"]}) + save_tmp_chat(context) + mark_dirty_for_context(context.id, reason="telegram.model_select") + + +async def _select_project(context: AgentContext, index: int, *, clear: bool = False) -> None: + if clear: + projects.deactivate_project(context.id) + return + items = [item for item in projects.get_active_projects_list() or [] if item.get("name")] + if index < 0 or index >= len(items): + return + projects.activate_project(context.id, str(items[index]["name"])) + + +async def _select_agent(context: AgentContext, index: int) -> None: + if context.is_running(): + return + from agent import Agent + from initialize import initialize_agent + + items = [item for item in subagents.get_all_agents_list() if item.get("key")] + if index < 0 or index >= len(items): + return + profile = str(items[index]["key"]) + config = initialize_agent(override_settings={"agent_profile": profile}) + context.config = config + agent = context.agent0 + while agent: + agent.config = config + agent = agent.get_data(Agent.DATA_NAME_SUBORDINATE) + save_tmp_chat(context) + mark_dirty_for_context(context.id, reason="telegram.agent_select") + + +def _paged_buttons( + kind: str, + items: list[PickerItem], + page: int, + current: str, + *, + disabled: bool = False, +) -> list[list[dict[str, str]]]: + page = max(0, page) + total_pages = max(1, (len(items) + PAGE_SIZE - 1) // PAGE_SIZE) + page = min(page, total_pages - 1) + start = page * PAGE_SIZE + rows: list[list[dict[str, str]]] = [] + for offset, item in enumerate(items[start:start + PAGE_SIZE], start=start): + marker = "• " if item.key == current else "" + action = "noop" if disabled else "set" + rows.append([{ + "text": f"{marker}{item.label}"[:64], + "callback_data": f"tg:{kind}:{action}:{offset}", + }]) + nav: list[dict[str, str]] = [] + if page > 0: + nav.append({"text": "Prev", "callback_data": f"tg:{kind}:page:{page - 1}"}) + if page < total_pages - 1: + nav.append({"text": "Next", "callback_data": f"tg:{kind}:page:{page + 1}"}) + if nav: + rows.append(nav) + return rows + + +def _toggle_enabled(context: AgentContext, key: str) -> bool: + value = context.get_data(key) + return True if value is None else bool(value) + + +def _parse_toggle(args: str) -> bool | None: + value = (args or "").strip().lower() + if value in {"on", "enable", "enabled", "yes", "true", "1"}: + return True + if value in {"off", "disable", "disabled", "no", "false", "0"}: + return False + return None + + +def _safe_int(value: str) -> int: + try: + return int(value) + except (TypeError, ValueError): + return 0 + + +def _label_for(items: list[PickerItem], key: str) -> str: + for item in items: + if item.key == key: + return item.label + return key + + +def _html(value: str) -> str: + return str(value).replace("&", "&").replace("<", "<").replace(">", ">") diff --git a/plugins/_telegram_integration/helpers/constants.py b/plugins/_telegram_integration/helpers/constants.py index 830e41547..abab54ed2 100644 --- a/plugins/_telegram_integration/helpers/constants.py +++ b/plugins/_telegram_integration/helpers/constants.py @@ -6,11 +6,19 @@ STATE_FILE = "usr/plugins/_telegram_integration/state.json" CTX_TG_BOT = "telegram_bot" CTX_TG_BOT_CFG = "telegram_bot_cfg" CTX_TG_CHAT_ID = "telegram_chat_id" +CTX_TG_CHAT_TYPE = "telegram_chat_type" CTX_TG_USER_ID = "telegram_user_id" CTX_TG_USERNAME = "telegram_username" CTX_TG_TYPING_STOP = "_telegram_typing_stop" CTX_TG_REPLY_TO = "_telegram_reply_to_message_id" +CTX_TG_STREAM_ENABLED = "telegram_stream_enabled" +CTX_TG_TOOLS_ENABLED = "telegram_tools_enabled" # Transient CTX_TG_ATTACHMENTS = "_telegram_response_attachments" CTX_TG_KEYBOARD = "_telegram_response_keyboard" +CTX_TG_PROGRESS_MESSAGE_ID = "_telegram_progress_message_id" +CTX_TG_PROGRESS_LINES = "_telegram_progress_lines" +CTX_TG_RESPONSE_MESSAGE_ID = "_telegram_response_message_id" +CTX_TG_RESPONSE_TEXT = "_telegram_response_text" +CTX_TG_RESPONSE_LAST_UPDATE = "_telegram_response_last_update" diff --git a/plugins/_telegram_integration/helpers/draft_stream.py b/plugins/_telegram_integration/helpers/draft_stream.py new file mode 100644 index 000000000..491c7d2a2 --- /dev/null +++ b/plugins/_telegram_integration/helpers/draft_stream.py @@ -0,0 +1,361 @@ +from __future__ import annotations + +import re +import time + +from agent import AgentContext +from helpers.print_style import PrintStyle +from plugins._telegram_integration.helpers import telegram_client as tc +from plugins._telegram_integration.helpers.bot_manager import get_bot +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_PROGRESS_LINES, + CTX_TG_PROGRESS_MESSAGE_ID, + CTX_TG_REPLY_TO, + CTX_TG_RESPONSE_LAST_UPDATE, + CTX_TG_RESPONSE_MESSAGE_ID, + CTX_TG_RESPONSE_TEXT, + CTX_TG_STREAM_ENABLED, + CTX_TG_TOOLS_ENABLED, +) + +MAX_STREAM_CHARS: int = 3900 +MIN_RESPONSE_UPDATE_SECONDS: float = 1.0 +MAX_PROGRESS_LINES: int = 12 + +TOOL_EMOJIS: dict[str, str] = { + "browser": "🌐", + "code": "⌨️", + "code_execution_tool": "⌨️", + "duckduckgo_search": "🔎", + "file": "📄", + "knowledge_tool": "📚", + "memory": "🧠", + "read_file": "📖", + "search": "🔎", + "search_engine": "🔎", + "search_files": "🔎", + "skill": "📚", + "skill_view": "📚", +} + + +async def start(context: AgentContext) -> None: + # Do not pre-create a placeholder Telegram message. Wait until we have + # real assistant text so the stream starts with meaningful content. + if not _stream_enabled(context): + return + + +async def add_tool_start( + context: AgentContext, + tool_name: str, + args: dict | None = None, +) -> None: + if not _tools_enabled(context): + return + label = _tool_progress_label(tool_name, args or {}) + _append_progress_line(context, f"{_tool_emoji(tool_name)} {label}") + await _send_progress(context) + + +async def add_tool_done(context: AgentContext, tool_name: str, ok: bool = True) -> None: + if not _tools_enabled(context): + return + if ok: + return + _mark_tool_failed(context, tool_name) + await _send_progress(context) + + +async def update_response(context: AgentContext, response_text: str) -> None: + if not _stream_enabled(context): + return + cleaned = _visible_response_text(response_text) + if not cleaned: + return + context.data[CTX_TG_RESPONSE_TEXT] = response_text or "" + now = time.time() + last = float(context.data.get(CTX_TG_RESPONSE_LAST_UPDATE) or 0.0) + if now - last < MIN_RESPONSE_UPDATE_SECONDS: + return + await _update_response_message(context, response_text or "") + context.data[CTX_TG_RESPONSE_LAST_UPDATE] = now + + +async def finalize_response( + context: AgentContext, + response_text: str, + keyboard: list[list[dict]] | None = None, +) -> bool: + message_id = context.data.get(CTX_TG_RESPONSE_MESSAGE_ID) + if not message_id: + return False + ok = await _update_response_message( + context, + response_text or context.data.get(CTX_TG_RESPONSE_TEXT) or "", + keyboard=keyboard, + force=True, + ) + if ok: + context.data.pop(CTX_TG_RESPONSE_MESSAGE_ID, None) + context.data.pop(CTX_TG_RESPONSE_TEXT, None) + context.data.pop(CTX_TG_RESPONSE_LAST_UPDATE, None) + return ok + + +def clear(context: AgentContext) -> None: + for key in ( + CTX_TG_PROGRESS_LINES, + CTX_TG_PROGRESS_MESSAGE_ID, + CTX_TG_RESPONSE_MESSAGE_ID, + CTX_TG_RESPONSE_TEXT, + CTX_TG_RESPONSE_LAST_UPDATE, + ): + context.data.pop(key, None) + + +def _stream_enabled(context: AgentContext) -> bool: + value = context.get_data(CTX_TG_STREAM_ENABLED) + return True if value is None else bool(value) + + +def _tools_enabled(context: AgentContext) -> bool: + value = context.get_data(CTX_TG_TOOLS_ENABLED) + return True if value is None else bool(value) + + +async def _send_progress(context: AgentContext) -> None: + bot = _bot_instance(context) + chat_id = context.data.get(CTX_TG_CHAT_ID) + if not bot or not chat_id: + return + text = "\n".join(context.data.get(CTX_TG_PROGRESS_LINES) or []) + if not text: + return + message_id = context.data.get(CTX_TG_PROGRESS_MESSAGE_ID) + try: + if message_id: + await tc.raw_edit_text(bot.bot.token, int(chat_id), int(message_id), text, parse_mode=None) + return + sent_id = await tc.raw_send_text( + bot.bot.token, + int(chat_id), + text, + parse_mode=None, + ) + if sent_id: + context.data[CTX_TG_PROGRESS_MESSAGE_ID] = sent_id + except Exception as e: + PrintStyle.debug(f"Telegram progress update failed: {e}") + + +async def _ensure_response_message(context: AgentContext, text: str) -> int | None: + message_id = context.data.get(CTX_TG_RESPONSE_MESSAGE_ID) + if message_id: + return int(message_id) + html = _format_response(text) + if not html: + return None + bot = _bot_instance(context) + chat_id = context.data.get(CTX_TG_CHAT_ID) + if not bot or not chat_id: + return None + sent_id = await tc.raw_send_text( + bot.bot.token, + int(chat_id), + html, + reply_to_message_id=_reply_to(context), + parse_mode="HTML", + ) + if sent_id: + context.data[CTX_TG_RESPONSE_MESSAGE_ID] = sent_id + context.data[CTX_TG_RESPONSE_LAST_UPDATE] = time.time() + return sent_id + + +async def _update_response_message( + context: AgentContext, + text: str, + *, + keyboard: list[list[dict]] | None = None, + force: bool = False, +) -> bool: + message_id = await _ensure_response_message(context, text) + bot = _bot_instance(context) + chat_id = context.data.get(CTX_TG_CHAT_ID) + if not message_id or not bot or not chat_id: + return False + markup = _keyboard_markup(keyboard) + html = _format_response(text) + try: + ok = await tc.raw_edit_text( + bot.bot.token, + int(chat_id), + int(message_id), + html, + parse_mode="HTML", + reply_markup=markup, + ) + if ok or force: + return ok + return False + except Exception as e: + PrintStyle.debug(f"Telegram response update failed: {e}") + return False + + +def _append_progress_line(context: AgentContext, line: str) -> None: + lines = list(context.data.get(CTX_TG_PROGRESS_LINES) or []) + if lines and lines[-1] == line: + return + lines.append(line) + context.data[CTX_TG_PROGRESS_LINES] = lines[-MAX_PROGRESS_LINES:] + + +def _mark_tool_failed(context: AgentContext, tool_name: str) -> None: + lines = list(context.data.get(CTX_TG_PROGRESS_LINES) or []) + if not lines: + return + + label = _tool_label(tool_name) + for index in range(len(lines) - 1, -1, -1): + line = lines[index] + if _line_matches_tool(line, label): + _, _, suffix = line.partition(" ") + lines[index] = f"❌ {suffix or label}" + context.data[CTX_TG_PROGRESS_LINES] = lines[-MAX_PROGRESS_LINES:] + return + + lines.append(f"❌ {label}") + context.data[CTX_TG_PROGRESS_LINES] = lines[-MAX_PROGRESS_LINES:] + + +def _bot_instance(context: AgentContext): + bot_name = context.data.get(CTX_TG_BOT) + return get_bot(bot_name) if bot_name else None + + +def _reply_to(context: AgentContext) -> int | None: + value = context.data.get(CTX_TG_REPLY_TO) + try: + return int(value) if value else None + except (TypeError, ValueError): + return None + + +def _tool_label(tool_name: str) -> str: + value = (tool_name or "tool").replace("_", " ").replace("-", " ").strip() + return value or "tool" + + +def _line_matches_tool(line: str, label: str) -> bool: + _, _, suffix = line.partition(" ") + normalized_suffix = suffix.strip().lower() + normalized_label = label.strip().lower() + return normalized_suffix == normalized_label or normalized_suffix.startswith(f"{normalized_label}:") + + +def _tool_progress_label(tool_name: str, args: dict) -> str: + label = _tool_label(tool_name) + detail = _tool_detail(tool_name, args) + return f"{label}: {detail}" if detail else label + + +def _tool_detail(tool_name: str, args: dict) -> str: + if not isinstance(args, dict) or not args: + return "" + + normalized = (tool_name or "").strip().lower() + if "search" in normalized: + return _first_arg(args, ("query", "q", "search", "term", "keywords", "pattern")) + if "browser" in normalized: + return _first_arg(args, ("url", "link", "query", "action")) + if "file" in normalized: + return _first_arg(args, ("path", "file", "filename", "query", "pattern")) + if "skill" in normalized: + return _first_arg(args, ("skill", "name", "query", "path")) + + return _first_arg(args, ("action", "method", "operation")) + + +def _first_arg(args: dict, keys: tuple[str, ...]) -> str: + for key in keys: + value = args.get(key) + if value is None: + continue + text = _compact_detail(value) + if text: + return text + return "" + + +def _compact_detail(value: object) -> str: + if isinstance(value, (list, tuple)): + parts = [_compact_detail(item) for item in value[:2]] + text = ", ".join(part for part in parts if part) + if len(value) > 2: + text = f"{text}, ..." + elif isinstance(value, dict): + return "" + else: + text = str(value).strip() + + text = re.sub(r"\s+", " ", text).strip().strip("\"'") + if len(text) > 80: + text = f"{text[:77].rstrip()}..." + return text + + +def _tool_emoji(tool_name: str) -> str: + normalized = (tool_name or "").strip().lower() + for key, emoji in TOOL_EMOJIS.items(): + if key in normalized: + return emoji + return "🛠️" + + +def _format_response(text: str) -> str: + value = _visible_response_text(text) + if not value: + return "" + return tc.md_to_telegram_html(value) + + +def _strip_incomplete_tool_markup(text: str) -> str: + value = text.lstrip() + value = re.sub(r"^<[^>\n]{0,80}$", "", value) + value = re.sub(r"^```(?:json|xml)?\s*$", "", value, flags=re.IGNORECASE) + return value + + +def _visible_response_text(text: str) -> str: + value = _trim(text or "") + return _strip_incomplete_tool_markup(value) + + +def _trim(text: str) -> str: + if len(text) <= MAX_STREAM_CHARS: + return text + return text[-MAX_STREAM_CHARS:] + + +def _keyboard_markup(keyboard: list[list[dict]] | None) -> dict | None: + if not keyboard: + return None + rows: list[list[dict[str, str]]] = [] + for row in keyboard: + out_row: list[dict[str, str]] = [] + for button in row: + text = str(button.get("text") or "")[:64] + if not text: + continue + if button.get("url"): + out_row.append({"text": text, "url": str(button["url"])}) + else: + data = str(button.get("callback_data", text)) + out_row.append({"text": text, "callback_data": data[:64]}) + if out_row: + rows.append(out_row) + return {"inline_keyboard": rows} if rows else None diff --git a/plugins/_telegram_integration/helpers/handler.py b/plugins/_telegram_integration/helpers/handler.py index ea2dbd055..07c241478 100644 --- a/plugins/_telegram_integration/helpers/handler.py +++ b/plugins/_telegram_integration/helpers/handler.py @@ -21,6 +21,7 @@ from helpers.errors import format_error from initialize import initialize_agent from plugins._telegram_integration.helpers import telegram_client as tc +from plugins._telegram_integration.helpers import command_ui from plugins._telegram_integration.helpers.bot_manager import get_bot from plugins._telegram_integration.helpers.constants import ( PLUGIN_NAME, @@ -29,6 +30,7 @@ from plugins._telegram_integration.helpers.constants import ( CTX_TG_BOT, CTX_TG_BOT_CFG, CTX_TG_CHAT_ID, + CTX_TG_CHAT_TYPE, CTX_TG_USER_ID, CTX_TG_USERNAME, CTX_TG_TYPING_STOP, @@ -141,55 +143,15 @@ async def handle_start(message: TgMessage, bot_name: str, bot_cfg: dict): f"\U0001f44b Hello {user.first_name}! I'm connected to Agent Zero.\n\n" "Send me a message and I'll process it.\n" "Use /clear to reset the conversation.\n" - "Use /project, /config, or /send to control the current chat.", + "Use /project, /model, /agent, or /send to control the current chat.", parse_mode=None, + reply_to_message_id=message.message_id, ) # Ensure a chat context exists await _get_or_create_context(bot_name, bot_cfg, message) -async def handle_clear(message: TgMessage, bot_name: str, bot_cfg: dict): - """Handle /clear command — reset user's chat context.""" - user = message.from_user - if not user: - return - - if not _is_allowed(bot_cfg, user.id, user.username): - return - - key = _map_key(bot_name, user.id, message.chat.id) - - with _chat_map_lock: - state = _load_state() - ctx_id = state.get("chats", {}).get(key) - if ctx_id: - ctx = AgentContext.get(ctx_id) - if ctx: - ctx.reset() - PrintStyle.info(f"Telegram ({bot_name}): cleared chat for user {user.id}") - - instance = get_bot(bot_name) - if instance: - await _send_with_temp_bot( - instance.bot.token, message.chat.id, - "Chat cleared. Send a new message to start fresh.", - parse_mode=None, - ) - - # Send notification - if bot_cfg.get("notify_messages", False): - username_str = f"@{user.username}" if user.username else str(user.id) - NotificationManager.send_notification( - type=NotificationType.INFO, - priority=NotificationPriority.NORMAL, - title="Telegram: chat cleared", - message=f"{username_str} cleared their chat via /clear", - display_time=5, - group="telegram", - ) - - async def handle_message(message: TgMessage, bot_name: str, bot_cfg: dict): """Handle incoming user message.""" user = message.from_user @@ -212,26 +174,38 @@ async def handle_message(message: TgMessage, bot_name: str, bot_cfg: dict): parse_mode=None, ) return + context.data[CTX_TG_CHAT_TYPE] = str(message.chat.type or "") + context.data[CTX_TG_REPLY_TO] = message.message_id + + if await command_ui.handle_command( + context, + instance.bot.token, + message.chat.id, + message.message_id, + text, + ): + return command_reply = integration_commands.try_handle_command(context, text) if command_reply is not None: - await _send_with_temp_bot(instance.bot.token, message.chat.id, command_reply, parse_mode=None) + await _send_with_temp_bot( + instance.bot.token, + message.chat.id, + command_reply, + parse_mode=None, + reply_to_message_id=message.message_id, + ) + return + if integration_commands.extract_command_line(text).startswith("/"): + command = integration_commands.extract_command_line(text).split(" ", 1)[0] + await _send_with_temp_bot( + instance.bot.token, + message.chat.id, + integration_commands.unknown_command_text(command), + parse_mode=None, + reply_to_message_id=message.message_id, + ) return - - # Start persistent typing indicator (thread-based, works across event loops) - typing_stop = _start_typing(instance.bot.token, message.chat.id) - - # Store stop event so send_telegram_reply can cancel typing - context.data[CTX_TG_TYPING_STOP] = typing_stop - - # In group chats, if user replied to the bot's message, reply to the user's message - reply_to_id = None - if message.chat.type != "private" and instance.bot_info: - if (message.reply_to_message - and message.reply_to_message.from_user - and message.reply_to_message.from_user.id == instance.bot_info.id): - reply_to_id = message.message_id - context.data[CTX_TG_REPLY_TO] = reply_to_id # Use temp bot for downloads (cross-event-loop safe) async with _temp_bot(instance.bot.token) as dl_bot: @@ -245,6 +219,24 @@ async def handle_message(message: TgMessage, bot_name: str, bot_cfg: dict): body=text, ) + if context.is_running(): + item = mq.add(context, user_msg, attachments) + save_tmp_chat(context) + await _send_with_temp_bot( + instance.bot.token, + message.chat.id, + f"Queued message #{item.get('seq', len(mq.get_queue(context)))}. Use /send to flush queued work, or /steer to interrupt the active run.", + parse_mode=None, + reply_to_message_id=message.message_id, + ) + return + + # Start persistent typing indicator (thread-based, works across event loops) + typing_stop = _start_typing(instance.bot.token, message.chat.id) + + # Store stop event so send_telegram_reply can cancel typing + context.data[CTX_TG_TYPING_STOP] = typing_stop + msg_id = str(uuid.uuid4()) mq.log_user_message(context, user_msg, attachments, message_id=msg_id, source=" (telegram)") context.communicate(UserMessage( @@ -287,16 +279,40 @@ async def handle_callback_query(query: CallbackQuery, bot_name: str, bot_cfg: di return context = await _get_or_create_context_from_user( - bot_name, bot_cfg, user.id, user.username, query.message.chat.id, + bot_name, bot_cfg, user.id, user.username, query.message.chat.id, str(query.message.chat.type or ""), ) if not context: return + context.data[CTX_TG_REPLY_TO] = query.message.message_id + + instance = get_bot(bot_name) + if instance: + try: + if await command_ui.handle_callback( + context, + instance.bot.token, + query.message.chat.id, + query.message.message_id, + text, + ): + return + except Exception as e: + PrintStyle.error(f"Telegram callback failed: {format_error(e)}") + if text.startswith("tg:"): + return + if text.startswith("tg:"): + return command_reply = integration_commands.try_handle_command(context, text) if command_reply is not None: - instance = get_bot(bot_name) if instance: - await _send_with_temp_bot(instance.bot.token, query.message.chat.id, command_reply, parse_mode=None) + await _send_with_temp_bot( + instance.bot.token, + query.message.chat.id, + command_reply, + parse_mode=None, + reply_to_message_id=query.message.message_id, + ) return agent = context.agent0 @@ -347,7 +363,7 @@ async def _get_or_create_context( if not user: return None return await _get_or_create_context_from_user( - bot_name, bot_cfg, user.id, user.username, message.chat.id, + bot_name, bot_cfg, user.id, user.username, message.chat.id, str(message.chat.type or ""), ) @@ -357,6 +373,7 @@ async def _get_or_create_context_from_user( user_id: int, username: str | None, chat_id: int, + chat_type: str = "", ) -> AgentContext | None: key = _map_key(bot_name, user_id, chat_id) @@ -369,6 +386,7 @@ async def _get_or_create_context_from_user( if ctx_id: ctx = AgentContext.get(ctx_id) if ctx: + ctx.data[CTX_TG_CHAT_TYPE] = chat_type or ctx.data.get(CTX_TG_CHAT_TYPE, "") return ctx # Context was garbage collected, remove stale mapping chats.pop(key, None) @@ -382,6 +400,7 @@ async def _get_or_create_context_from_user( ctx.data[CTX_TG_BOT] = bot_name ctx.data[CTX_TG_BOT_CFG] = bot_cfg ctx.data[CTX_TG_CHAT_ID] = chat_id + ctx.data[CTX_TG_CHAT_TYPE] = chat_type ctx.data[CTX_TG_USER_ID] = user_id ctx.data[CTX_TG_USERNAME] = username or "" @@ -512,7 +531,11 @@ async def send_telegram_reply( if response_text: html_text = tc.md_to_telegram_html(response_text) - if keyboard: + from plugins._telegram_integration.helpers import draft_stream + + if await draft_stream.finalize_response(context, response_text, keyboard): + pass + elif keyboard: await tc.send_text_with_keyboard(reply_bot, chat_id, html_text, keyboard, reply_to_message_id=reply_to) else: await tc.send_text(reply_bot, chat_id, html_text, reply_to_message_id=reply_to) @@ -537,10 +560,22 @@ async def _temp_bot(token: str, **kwargs): await bot.session.close() -async def _send_with_temp_bot(token: str, chat_id: int, text: str, parse_mode: str | None = None): +async def _send_with_temp_bot( + token: str, + chat_id: int, + text: str, + parse_mode: str | None = None, + reply_to_message_id: int | None = None, +): """Send text using a temporary Bot to avoid cross-event-loop session issues.""" async with _temp_bot(token) as bot: - await tc.send_text(bot, chat_id, text, parse_mode=parse_mode) + await tc.send_text( + bot, + chat_id, + text, + reply_to_message_id=reply_to_message_id, + parse_mode=parse_mode, + ) def _start_typing(token: str, chat_id: int) -> threading.Event: diff --git a/plugins/_telegram_integration/helpers/telegram_client.py b/plugins/_telegram_integration/helpers/telegram_client.py index 4ea104123..45db75977 100644 --- a/plugins/_telegram_integration/helpers/telegram_client.py +++ b/plugins/_telegram_integration/helpers/telegram_client.py @@ -1,6 +1,7 @@ import os import re +import aiohttp from aiogram import Bot from aiogram.exceptions import TelegramBadRequest from aiogram.types import ( @@ -17,6 +18,7 @@ _UNSET = object() # sentinel: "not provided" (lets Bot default apply) # Text messages MAX_MESSAGE_LENGTH: int = 4096 # Telegram message length limit +TELEGRAM_API_BASE: str = "https://api.telegram.org" async def send_text( @@ -171,6 +173,88 @@ async def send_typing(bot: Bot, chat_id: int): except Exception: pass + +async def raw_send_text( + token: str, + chat_id: int, + text: str, + reply_to_message_id: int | None = None, + parse_mode: str | None = "HTML", + reply_markup: dict | None = None, +) -> int | None: + payload: dict[str, object] = { + "chat_id": chat_id, + "text": text[:MAX_MESSAGE_LENGTH], + } + if parse_mode: + payload["parse_mode"] = parse_mode + if reply_to_message_id: + payload["reply_parameters"] = {"message_id": int(reply_to_message_id)} + if reply_markup: + payload["reply_markup"] = reply_markup + data = await _raw_post(token, "sendMessage", payload) + result = data.get("result") if isinstance(data, dict) else None + if isinstance(result, dict): + return result.get("message_id") + return None + + +async def raw_edit_text( + token: str, + chat_id: int, + message_id: int, + text: str, + parse_mode: str | None = "HTML", + reply_markup: dict | None = None, +) -> bool: + payload: dict[str, object] = { + "chat_id": chat_id, + "message_id": message_id, + "text": text[:MAX_MESSAGE_LENGTH], + } + if parse_mode: + payload["parse_mode"] = parse_mode + if reply_markup: + payload["reply_markup"] = reply_markup + data = await _raw_post(token, "editMessageText", payload) + if not isinstance(data, dict): + return False + if data.get("ok"): + return True + description = str(data.get("description") or "").lower() + return "message is not modified" in description + + +async def raw_edit_reply_markup( + token: str, + chat_id: int, + message_id: int, + reply_markup: dict | None = None, +) -> bool: + payload: dict[str, object] = { + "chat_id": chat_id, + "message_id": message_id, + } + if reply_markup: + payload["reply_markup"] = reply_markup + data = await _raw_post(token, "editMessageReplyMarkup", payload) + return bool(isinstance(data, dict) and data.get("ok")) + + +async def _raw_post(token: str, method: str, payload: dict[str, object]) -> dict: + url = f"{TELEGRAM_API_BASE}/bot{token}/{method}" + try: + timeout = aiohttp.ClientTimeout(total=10) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.post(url, json=payload) as response: + data = await response.json(content_type=None) + if response.status != 200 or not data.get("ok"): + PrintStyle.debug(f"Telegram {method} failed: {data}") + return data if isinstance(data, dict) else {} + except Exception as e: + PrintStyle.debug(f"Telegram {method} failed: {format_error(e)}") + return {} + # File download async def download_file( diff --git a/plugins/_whatsapp_integration/helpers/handler.py b/plugins/_whatsapp_integration/helpers/handler.py index e1efe415c..9c9e19eb6 100644 --- a/plugins/_whatsapp_integration/helpers/handler.py +++ b/plugins/_whatsapp_integration/helpers/handler.py @@ -207,6 +207,24 @@ async def _route_to_chat( msg_id = str(uuid.uuid4()) media_urls = msg.get("mediaUrls", []) attachments = await _save_incoming_media(media_urls) if media_urls else [] + + if context.is_running(): + item = mq.add(context, user_msg, attachments) + save_tmp_chat(context) + port = int((plugins.get_plugin_config(PLUGIN_NAME) or {}).get("bridge_port", 3100)) + base_url = bridge_manager.get_bridge_url(port) + chat_id = context.data.get(CTX_WA_CHAT_ID, "") or msg.get("chatId", "") + reply_to = msg.get("messageId", "") if context.data.get(CTX_WA_IS_GROUP) else "" + await wa_client.send_message( + base_url, + chat_id, + f"Queued message #{item.get('seq', len(mq.get_queue(context)))}. Use /send to flush queued work, or /steer to interrupt the active run.", + reply_to=reply_to, + ) + await wa_client.send_typing(base_url, chat_id, paused=True) + context.data[CTX_WA_TYPING_ACTIVE] = False + return + mq.log_user_message( context, user_msg, attachments, message_id=msg_id, source=" (whatsapp)", ) From 8bc81c7a6a5030f4899f3e4810480915b0950c3c Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 21:24:28 +0530 Subject: [PATCH 004/231] Send Telegram intermediate responses separately --- plugins/_telegram_integration/README.md | 2 +- .../_50_telegram_response.py | 8 ++- .../helpers/draft_stream.py | 26 +++++++ tests/test_telegram_intermediate_response.py | 70 +++++++++++++++++++ 4 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 tests/test_telegram_intermediate_response.py diff --git a/plugins/_telegram_integration/README.md b/plugins/_telegram_integration/README.md index bf486184e..a9793fa35 100644 --- a/plugins/_telegram_integration/README.md +++ b/plugins/_telegram_integration/README.md @@ -35,7 +35,7 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - **Reply delivery** - Streams tool progress and response text through real Telegram messages updated with `editMessageText`. - Tool progress and the AI response are separate messages; only the AI response replies to the user message. - - `tool_execute_after` intercepts the `response` tool — sends inline progress for `break_loop=false`. + - `tool_execute_after` intercepts the `response` tool — sends `break_loop=false` updates as separate intermediate Telegram messages. - `process_chain_end` auto-sends the final response, with retry logic on failure. - **Formatting** - Converts Markdown output to Telegram-compatible HTML (bold, italic, strikethrough, code, links, blockquotes, lists). diff --git a/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py b/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py index 43fcbae24..c959eaf52 100644 --- a/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py +++ b/plugins/_telegram_integration/extensions/python/tool_execute_after/_50_telegram_response.py @@ -47,6 +47,7 @@ class TelegramResponseIntercept(Extension): async def _send_inline(self, context, tool, response: Response): ensure_dependencies() from plugins._telegram_integration.helpers.handler import send_telegram_reply + from plugins._telegram_integration.helpers import draft_stream agent = self.agent assert agent is not None @@ -55,7 +56,12 @@ class TelegramResponseIntercept(Extension): attachments = context.data.pop(CTX_TG_ATTACHMENTS, []) keyboard = context.data.pop(CTX_TG_KEYBOARD, None) - error = await send_telegram_reply(context, text, attachments or None, keyboard) + if attachments: + error = await send_telegram_reply(context, text, attachments or None, keyboard) + elif await draft_stream.send_intermediate_response(context, text, keyboard): + error = None + else: + error = "Telegram intermediate update was not sent" if error: result = agent.read_prompt("fw.telegram.update_error.md", error=error) diff --git a/plugins/_telegram_integration/helpers/draft_stream.py b/plugins/_telegram_integration/helpers/draft_stream.py index 491c7d2a2..b246f8ab3 100644 --- a/plugins/_telegram_integration/helpers/draft_stream.py +++ b/plugins/_telegram_integration/helpers/draft_stream.py @@ -84,6 +84,32 @@ async def update_response(context: AgentContext, response_text: str) -> None: context.data[CTX_TG_RESPONSE_LAST_UPDATE] = now +async def send_intermediate_response( + context: AgentContext, + response_text: str, + keyboard: list[list[dict]] | None = None, +) -> bool: + html = _format_response(response_text) + if not html: + return False + bot = _bot_instance(context) + chat_id = context.data.get(CTX_TG_CHAT_ID) + if not bot or not chat_id: + return False + try: + sent_id = await tc.raw_send_text( + bot.bot.token, + int(chat_id), + html, + parse_mode="HTML", + reply_markup=_keyboard_markup(keyboard), + ) + return bool(sent_id) + except Exception as e: + PrintStyle.debug(f"Telegram intermediate response failed: {e}") + return False + + async def finalize_response( context: AgentContext, response_text: str, diff --git a/tests/test_telegram_intermediate_response.py b/tests/test_telegram_intermediate_response.py new file mode 100644 index 000000000..cb56946ed --- /dev/null +++ b/tests/test_telegram_intermediate_response.py @@ -0,0 +1,70 @@ +import asyncio + +from plugins._telegram_integration.helpers import draft_stream +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_REPLY_TO, + CTX_TG_RESPONSE_MESSAGE_ID, +) + + +class FakeBot: + class Bot: + token = "token" + + bot = Bot() + + +class FakeContext: + def __init__(self): + self.data = { + CTX_TG_BOT: "main", + CTX_TG_CHAT_ID: 123, + CTX_TG_REPLY_TO: 456, + } + + def get_data(self, key): + return self.data.get(key) + + +def test_intermediate_response_sends_separate_non_reply_message(monkeypatch): + calls = [] + + async def fake_send(token, chat_id, text, reply_to_message_id=None, parse_mode="HTML", reply_markup=None): + calls.append( + { + "token": token, + "chat_id": chat_id, + "text": text, + "reply_to_message_id": reply_to_message_id, + "parse_mode": parse_mode, + "reply_markup": reply_markup, + } + ) + return 789 + + context = FakeContext() + monkeypatch.setattr(draft_stream, "_bot_instance", lambda ctx: FakeBot()) + monkeypatch.setattr(draft_stream.tc, "raw_send_text", fake_send) + + sent = asyncio.run( + draft_stream.send_intermediate_response( + context, + "**Working** on the brief.", + keyboard=[[{"text": "Open", "callback_data": "open"}]], + ) + ) + + assert sent is True + assert calls == [ + { + "token": "token", + "chat_id": 123, + "text": "Working on the brief.", + "reply_to_message_id": None, + "parse_mode": "HTML", + "reply_markup": {"inline_keyboard": [[{"text": "Open", "callback_data": "open"}]]}, + } + ] + assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data From a17cef5b03f500dee1e9e9c3fa0bcdb0c03f3eee Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:00:44 +0530 Subject: [PATCH 005/231] Group Telegram tool streams by response --- .../helpers/draft_stream.py | 21 ++++- tests/test_telegram_intermediate_response.py | 82 +++++++++++++++++++ 2 files changed, 101 insertions(+), 2 deletions(-) diff --git a/plugins/_telegram_integration/helpers/draft_stream.py b/plugins/_telegram_integration/helpers/draft_stream.py index b246f8ab3..a64cb47a6 100644 --- a/plugins/_telegram_integration/helpers/draft_stream.py +++ b/plugins/_telegram_integration/helpers/draft_stream.py @@ -29,10 +29,10 @@ TOOL_EMOJIS: dict[str, str] = { "code": "⌨️", "code_execution_tool": "⌨️", "duckduckgo_search": "🔎", + "read_file": "📖", "file": "📄", "knowledge_tool": "📚", "memory": "🧠", - "read_file": "📖", "search": "🔎", "search_engine": "🔎", "search_files": "🔎", @@ -89,6 +89,12 @@ async def send_intermediate_response( response_text: str, keyboard: list[list[dict]] | None = None, ) -> bool: + if context.data.get(CTX_TG_RESPONSE_MESSAGE_ID): + sent = await finalize_response(context, response_text, keyboard) + if sent: + _reset_progress_group(context) + return sent + html = _format_response(response_text) if not html: return False @@ -104,7 +110,10 @@ async def send_intermediate_response( parse_mode="HTML", reply_markup=_keyboard_markup(keyboard), ) - return bool(sent_id) + sent = bool(sent_id) + if sent: + _reset_progress_group(context) + return sent except Exception as e: PrintStyle.debug(f"Telegram intermediate response failed: {e}") return False @@ -142,6 +151,11 @@ def clear(context: AgentContext) -> None: context.data.pop(key, None) +def _reset_progress_group(context: AgentContext) -> None: + context.data.pop(CTX_TG_PROGRESS_LINES, None) + context.data.pop(CTX_TG_PROGRESS_MESSAGE_ID, None) + + def _stream_enabled(context: AgentContext) -> bool: value = context.get_data(CTX_TG_STREAM_ENABLED) return True if value is None else bool(value) @@ -208,6 +222,7 @@ async def _update_response_message( keyboard: list[list[dict]] | None = None, force: bool = False, ) -> bool: + had_message = bool(context.data.get(CTX_TG_RESPONSE_MESSAGE_ID)) message_id = await _ensure_response_message(context, text) bot = _bot_instance(context) chat_id = context.data.get(CTX_TG_CHAT_ID) @@ -215,6 +230,8 @@ async def _update_response_message( return False markup = _keyboard_markup(keyboard) html = _format_response(text) + if not had_message and not markup and not force: + return True try: ok = await tc.raw_edit_text( bot.bot.token, diff --git a/tests/test_telegram_intermediate_response.py b/tests/test_telegram_intermediate_response.py index cb56946ed..c8c572557 100644 --- a/tests/test_telegram_intermediate_response.py +++ b/tests/test_telegram_intermediate_response.py @@ -4,6 +4,8 @@ from plugins._telegram_integration.helpers import draft_stream from plugins._telegram_integration.helpers.constants import ( CTX_TG_BOT, CTX_TG_CHAT_ID, + CTX_TG_PROGRESS_LINES, + CTX_TG_PROGRESS_MESSAGE_ID, CTX_TG_REPLY_TO, CTX_TG_RESPONSE_MESSAGE_ID, ) @@ -68,3 +70,83 @@ def test_intermediate_response_sends_separate_non_reply_message(monkeypatch): } ] assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data + + +def test_intermediate_response_finalizes_active_stream_and_starts_next_tool_group(monkeypatch): + calls = [] + next_message_id = 100 + + async def fake_send(token, chat_id, text, reply_to_message_id=None, parse_mode="HTML", reply_markup=None): + nonlocal next_message_id + calls.append( + { + "method": "send", + "text": text, + "reply_to_message_id": reply_to_message_id, + "parse_mode": parse_mode, + "reply_markup": reply_markup, + "message_id": next_message_id, + } + ) + next_message_id += 1 + return next_message_id - 1 + + async def fake_edit(token, chat_id, message_id, text, parse_mode="HTML", reply_markup=None): + calls.append( + { + "method": "edit", + "message_id": message_id, + "text": text, + "parse_mode": parse_mode, + "reply_markup": reply_markup, + } + ) + return True + + context = FakeContext() + monkeypatch.setattr(draft_stream, "_bot_instance", lambda ctx: FakeBot()) + monkeypatch.setattr(draft_stream.tc, "raw_send_text", fake_send) + monkeypatch.setattr(draft_stream.tc, "raw_edit_text", fake_edit) + + asyncio.run(draft_stream.add_tool_start(context, "search_engine", {"query": "telegram bot api"})) + asyncio.run(draft_stream.update_response(context, "Found the docs.")) + sent = asyncio.run(draft_stream.send_intermediate_response(context, "Found the docs.")) + asyncio.run(draft_stream.add_tool_start(context, "read_file", {"path": "notes.md"})) + + assert sent is True + assert calls == [ + { + "method": "send", + "text": "🔎 search engine: telegram bot api", + "reply_to_message_id": None, + "parse_mode": None, + "reply_markup": None, + "message_id": 100, + }, + { + "method": "send", + "text": "Found the docs.", + "reply_to_message_id": 456, + "parse_mode": "HTML", + "reply_markup": None, + "message_id": 101, + }, + { + "method": "edit", + "message_id": 101, + "text": "Found the docs.", + "parse_mode": "HTML", + "reply_markup": None, + }, + { + "method": "send", + "text": "📖 read file: notes.md", + "reply_to_message_id": None, + "parse_mode": None, + "reply_markup": None, + "message_id": 102, + }, + ] + assert context.data[CTX_TG_PROGRESS_MESSAGE_ID] == 102 + assert context.data[CTX_TG_PROGRESS_LINES] == ["📖 read file: notes.md"] + assert CTX_TG_RESPONSE_MESSAGE_ID not in context.data From 97ef4797e0ecd0d3e53843f094787bff9a5680d0 Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:07:44 +0530 Subject: [PATCH 006/231] Send Telegram media attachments natively --- .../_telegram_integration/helpers/handler.py | 6 + .../helpers/telegram_client.py | 93 +++++++++++++++ tests/test_telegram_media_delivery.py | 108 ++++++++++++++++++ 3 files changed, 207 insertions(+) create mode 100644 tests/test_telegram_media_delivery.py diff --git a/plugins/_telegram_integration/helpers/handler.py b/plugins/_telegram_integration/helpers/handler.py index 07c241478..3fd0f8c82 100644 --- a/plugins/_telegram_integration/helpers/handler.py +++ b/plugins/_telegram_integration/helpers/handler.py @@ -526,6 +526,12 @@ async def send_telegram_reply( local_path = files.fix_dev_path(path) if tc.is_image_file(local_path): await tc.send_photo(reply_bot, chat_id, local_path, reply_to_message_id=reply_to) + elif tc.is_voice_file(local_path): + await tc.send_voice(reply_bot, chat_id, local_path, reply_to_message_id=reply_to) + elif tc.is_audio_file(local_path): + await tc.send_audio(reply_bot, chat_id, local_path, reply_to_message_id=reply_to) + elif tc.is_video_file(local_path): + await tc.send_video(reply_bot, chat_id, local_path, reply_to_message_id=reply_to) else: await tc.send_file(reply_bot, chat_id, local_path, reply_to_message_id=reply_to) diff --git a/plugins/_telegram_integration/helpers/telegram_client.py b/plugins/_telegram_integration/helpers/telegram_client.py index 45db75977..2fc9f69f1 100644 --- a/plugins/_telegram_integration/helpers/telegram_client.py +++ b/plugins/_telegram_integration/helpers/telegram_client.py @@ -115,6 +115,81 @@ async def send_photo( return None +async def send_voice( + bot: Bot, + chat_id: int, + voice_path: str, + caption: str = "", + reply_to_message_id: int | None = None, +) -> int | None: + """Send a voice message from local path. Returns message_id or None on error.""" + try: + if not os.path.isfile(voice_path): + PrintStyle.error(f"Telegram: voice file not found: {voice_path}") + return None + input_file = FSInputFile(voice_path) + msg = await bot.send_voice( + chat_id=chat_id, + voice=input_file, + caption=caption[:1024] if caption else None, + reply_to_message_id=reply_to_message_id, + ) + return msg.message_id + except Exception as e: + PrintStyle.error(f"Telegram send_voice failed: {format_error(e)}") + return None + + +async def send_audio( + bot: Bot, + chat_id: int, + audio_path: str, + caption: str = "", + reply_to_message_id: int | None = None, +) -> int | None: + """Send an audio message from local path. Returns message_id or None on error.""" + try: + if not os.path.isfile(audio_path): + PrintStyle.error(f"Telegram: audio file not found: {audio_path}") + return None + input_file = FSInputFile(audio_path) + msg = await bot.send_audio( + chat_id=chat_id, + audio=input_file, + caption=caption[:1024] if caption else None, + reply_to_message_id=reply_to_message_id, + ) + return msg.message_id + except Exception as e: + PrintStyle.error(f"Telegram send_audio failed: {format_error(e)}") + return None + + +async def send_video( + bot: Bot, + chat_id: int, + video_path: str, + caption: str = "", + reply_to_message_id: int | None = None, +) -> int | None: + """Send a video message from local path. Returns message_id or None on error.""" + try: + if not os.path.isfile(video_path): + PrintStyle.error(f"Telegram: video file not found: {video_path}") + return None + input_file = FSInputFile(video_path) + msg = await bot.send_video( + chat_id=chat_id, + video=input_file, + caption=caption[:1024] if caption else None, + reply_to_message_id=reply_to_message_id, + ) + return msg.message_id + except Exception as e: + PrintStyle.error(f"Telegram send_video failed: {format_error(e)}") + return None + + # Inline keyboards def build_inline_keyboard( @@ -294,6 +369,9 @@ def _split_text(text: str, max_len: int) -> list[str]: _IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"} +_VOICE_EXTENSIONS = {".ogg", ".oga", ".opus"} +_AUDIO_EXTENSIONS = {".mp3", ".m4a", ".aac", ".wav", ".flac"} +_VIDEO_EXTENSIONS = {".mp4", ".m4v", ".mov", ".webm"} def is_image_file(path: str) -> bool: @@ -301,6 +379,21 @@ def is_image_file(path: str) -> bool: return ext in _IMAGE_EXTENSIONS +def is_voice_file(path: str) -> bool: + _, ext = os.path.splitext(path.lower()) + return ext in _VOICE_EXTENSIONS + + +def is_audio_file(path: str) -> bool: + _, ext = os.path.splitext(path.lower()) + return ext in _AUDIO_EXTENSIONS + + +def is_video_file(path: str) -> bool: + _, ext = os.path.splitext(path.lower()) + return ext in _VIDEO_EXTENSIONS + + def md_to_telegram_html(text: str) -> str: """Convert Markdown to Telegram-compatible HTML.""" stash: list[str] = [] diff --git a/tests/test_telegram_media_delivery.py b/tests/test_telegram_media_delivery.py new file mode 100644 index 000000000..bbdb2a2bc --- /dev/null +++ b/tests/test_telegram_media_delivery.py @@ -0,0 +1,108 @@ +import asyncio + +from plugins._telegram_integration.helpers import handler +from plugins._telegram_integration.helpers import telegram_client as tc +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_REPLY_TO, +) + + +class FakeBotInstance: + class Bot: + token = "token" + + bot = Bot() + + +class FakeContext: + def __init__(self): + self.data = { + CTX_TG_BOT: "main", + CTX_TG_CHAT_ID: 123, + CTX_TG_REPLY_TO: 456, + } + + +class FakeTempBot: + async def __aenter__(self): + return object() + + async def __aexit__(self, exc_type, exc, tb): + return False + + +class FakeMediaBot: + def __init__(self): + self.calls = [] + + async def send_voice(self, **kwargs): + self.calls.append(("voice", kwargs)) + return type("Message", (), {"message_id": 11})() + + async def send_audio(self, **kwargs): + self.calls.append(("audio", kwargs)) + return type("Message", (), {"message_id": 12})() + + async def send_video(self, **kwargs): + self.calls.append(("video", kwargs)) + return type("Message", (), {"message_id": 13})() + + +def test_telegram_client_native_media_helpers(monkeypatch): + bot = FakeMediaBot() + monkeypatch.setattr(tc.os.path, "isfile", lambda path: True) + + voice_id = asyncio.run(tc.send_voice(bot, 123, "voice.ogg", reply_to_message_id=456)) + audio_id = asyncio.run(tc.send_audio(bot, 123, "song.mp3", reply_to_message_id=456)) + video_id = asyncio.run(tc.send_video(bot, 123, "clip.mp4", reply_to_message_id=456)) + + assert (voice_id, audio_id, video_id) == (11, 12, 13) + assert [kind for kind, _ in bot.calls] == ["voice", "audio", "video"] + assert bot.calls[0][1]["chat_id"] == 123 + assert bot.calls[0][1]["reply_to_message_id"] == 456 + + +def test_telegram_reply_routes_native_media_attachments(monkeypatch): + calls = [] + + def fake_temp_bot(*args, **kwargs): + return FakeTempBot() + + async def record(kind, bot, chat_id, path, reply_to_message_id=None, **kwargs): + calls.append( + { + "kind": kind, + "chat_id": chat_id, + "path": path, + "reply_to_message_id": reply_to_message_id, + } + ) + return len(calls) + + monkeypatch.setattr(handler, "get_bot", lambda name: FakeBotInstance()) + monkeypatch.setattr(handler, "_temp_bot", fake_temp_bot) + monkeypatch.setattr(handler.files, "fix_dev_path", lambda path: path) + monkeypatch.setattr(handler.tc, "send_photo", lambda *a, **kw: record("photo", *a, **kw)) + monkeypatch.setattr(handler.tc, "send_voice", lambda *a, **kw: record("voice", *a, **kw)) + monkeypatch.setattr(handler.tc, "send_audio", lambda *a, **kw: record("audio", *a, **kw)) + monkeypatch.setattr(handler.tc, "send_video", lambda *a, **kw: record("video", *a, **kw)) + monkeypatch.setattr(handler.tc, "send_file", lambda *a, **kw: record("document", *a, **kw)) + + error = asyncio.run( + handler.send_telegram_reply( + FakeContext(), + "", + ["image.png", "voice.ogg", "song.mp3", "clip.mp4", "archive.zip"], + ) + ) + + assert error is None + assert calls == [ + {"kind": "photo", "chat_id": 123, "path": "image.png", "reply_to_message_id": 456}, + {"kind": "voice", "chat_id": 123, "path": "voice.ogg", "reply_to_message_id": 456}, + {"kind": "audio", "chat_id": 123, "path": "song.mp3", "reply_to_message_id": 456}, + {"kind": "video", "chat_id": 123, "path": "clip.mp4", "reply_to_message_id": 456}, + {"kind": "document", "chat_id": 123, "path": "archive.zip", "reply_to_message_id": 456}, + ] From 92b728c3d770b5e79b54d836c49771a45ad210d3 Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:11:55 +0530 Subject: [PATCH 007/231] Show friendly Telegram errors --- .../end/_85_telegram_error.py | 42 ++++++++ .../helpers/constants.py | 1 + .../helpers/draft_stream.py | 2 + .../_telegram_integration/helpers/error_ui.py | 100 ++++++++++++++++++ tests/test_telegram_error_ui.py | 96 +++++++++++++++++ 5 files changed, 241 insertions(+) create mode 100644 plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py create mode 100644 plugins/_telegram_integration/helpers/error_ui.py create mode 100644 tests/test_telegram_error_ui.py diff --git a/plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py b/plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py new file mode 100644 index 000000000..443df0d74 --- /dev/null +++ b/plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py @@ -0,0 +1,42 @@ +from helpers.errors import HandledException +from helpers.extension import Extension +from helpers.print_style import PrintStyle +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_ERROR_SENT, + CTX_TG_REPLY_TO, + CTX_TG_TYPING_STOP, +) + + +class TelegramFriendlyError(Extension): + async def execute(self, data: dict = {}, **kwargs): + if not self.agent or self.agent.number != 0: + return + + context = self.agent.context + if not context.data.get(CTX_TG_BOT): + return + + exception = data.get("exception") + if not exception or isinstance(exception, HandledException): + return + + if context.data.get(CTX_TG_ERROR_SENT): + return + + context.data[CTX_TG_ERROR_SENT] = True + + from plugins._telegram_integration.helpers import draft_stream, error_ui + from plugins._telegram_integration.helpers.handler import send_telegram_reply + + text = error_ui.friendly_error_message(exception) + error = await send_telegram_reply(context, text) + if error: + PrintStyle.debug(f"Telegram error reply failed: {error}") + + typing_stop = context.data.pop(CTX_TG_TYPING_STOP, None) + if typing_stop: + typing_stop.set() + draft_stream.clear(context) + context.data.pop(CTX_TG_REPLY_TO, None) diff --git a/plugins/_telegram_integration/helpers/constants.py b/plugins/_telegram_integration/helpers/constants.py index abab54ed2..1ef80c81a 100644 --- a/plugins/_telegram_integration/helpers/constants.py +++ b/plugins/_telegram_integration/helpers/constants.py @@ -22,3 +22,4 @@ CTX_TG_PROGRESS_LINES = "_telegram_progress_lines" CTX_TG_RESPONSE_MESSAGE_ID = "_telegram_response_message_id" CTX_TG_RESPONSE_TEXT = "_telegram_response_text" CTX_TG_RESPONSE_LAST_UPDATE = "_telegram_response_last_update" +CTX_TG_ERROR_SENT = "_telegram_error_sent" diff --git a/plugins/_telegram_integration/helpers/draft_stream.py b/plugins/_telegram_integration/helpers/draft_stream.py index a64cb47a6..22390fdff 100644 --- a/plugins/_telegram_integration/helpers/draft_stream.py +++ b/plugins/_telegram_integration/helpers/draft_stream.py @@ -18,6 +18,7 @@ from plugins._telegram_integration.helpers.constants import ( CTX_TG_RESPONSE_TEXT, CTX_TG_STREAM_ENABLED, CTX_TG_TOOLS_ENABLED, + CTX_TG_ERROR_SENT, ) MAX_STREAM_CHARS: int = 3900 @@ -42,6 +43,7 @@ TOOL_EMOJIS: dict[str, str] = { async def start(context: AgentContext) -> None: + context.data.pop(CTX_TG_ERROR_SENT, None) # Do not pre-create a placeholder Telegram message. Wait until we have # real assistant text so the stream starts with meaningful content. if not _stream_enabled(context): diff --git a/plugins/_telegram_integration/helpers/error_ui.py b/plugins/_telegram_integration/helpers/error_ui.py new file mode 100644 index 000000000..b9db82ebb --- /dev/null +++ b/plugins/_telegram_integration/helpers/error_ui.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import re + + +MAX_DETAIL_CHARS = 500 + + +def friendly_error_message(exception: BaseException) -> str: + summary = _error_summary(exception) + category = _classify_error(exception, summary) + + if category == "auth": + return ( + "**Agent Zero hit a provider setup issue.**\n\n" + "The model provider rejected the request because an API key or credential is missing, invalid, or unauthorized.\n\n" + f"Details: `{summary}`\n\n" + "Please check the model/API key settings, then send the message again." + ) + + if category == "rate_limit": + return ( + "**Agent Zero was rate limited by the model provider.**\n\n" + "The provider is asking us to slow down before trying again.\n\n" + f"Details: `{summary}`" + ) + + if category == "timeout": + return ( + "**Agent Zero did not get a response in time.**\n\n" + "The provider or tool call timed out before the agent could finish this request.\n\n" + f"Details: `{summary}`" + ) + + if category == "provider": + return ( + "**Agent Zero could not complete the model request.**\n\n" + "The model provider returned an error before the agent could finish.\n\n" + f"Details: `{summary}`" + ) + + return ( + "**Agent Zero ran into an error while working on this.**\n\n" + f"Details: `{summary}`" + ) + + +def _classify_error(exception: BaseException, summary: str) -> str: + text = f"{type(exception).__name__} {summary}".lower() + + if any( + marker in text + for marker in ( + "api key", + "api_key", + "apikey", + "auth", + "credential", + "unauthorized", + "forbidden", + "invalid key", + "missing key", + "permission denied", + "401", + "403", + ) + ): + return "auth" + + if any(marker in text for marker in ("rate limit", "ratelimit", "too many requests", "429")): + return "rate_limit" + + if any(marker in text for marker in ("timeout", "timed out", "deadline", "read timed out")): + return "timeout" + + if any( + marker in text + for marker in ( + "litellm", + "openai", + "anthropic", + "model", + "provider", + "badrequest", + "service unavailable", + "overloaded", + "503", + ) + ): + return "provider" + + return "generic" + + +def _error_summary(exception: BaseException) -> str: + text = str(exception).strip() or type(exception).__name__ + text = re.sub(r"\s+", " ", text) + if len(text) > MAX_DETAIL_CHARS: + text = f"{text[: MAX_DETAIL_CHARS - 3].rstrip()}..." + return text diff --git a/tests/test_telegram_error_ui.py b/tests/test_telegram_error_ui.py new file mode 100644 index 000000000..3b73fdf5b --- /dev/null +++ b/tests/test_telegram_error_ui.py @@ -0,0 +1,96 @@ +import asyncio + +from plugins._telegram_integration.extensions.python._functions.agent.Agent.handle_exception.end import ( + _85_telegram_error, +) +from plugins._telegram_integration.helpers import error_ui +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_ERROR_SENT, + CTX_TG_REPLY_TO, + CTX_TG_TYPING_STOP, +) + + +class FakeContext: + def __init__(self): + self.data = { + CTX_TG_BOT: "main", + CTX_TG_REPLY_TO: 456, + } + + +class FakeAgent: + number = 0 + + def __init__(self): + self.context = FakeContext() + + +def test_friendly_error_message_for_missing_api_key(): + message = error_ui.friendly_error_message( + RuntimeError("OPENAI_API_KEY is missing or invalid") + ) + + assert "provider setup issue" in message + assert "API key" in message + assert "OPENAI_API_KEY is missing or invalid" in message + + +def test_friendly_error_message_for_rate_limit(): + message = error_ui.friendly_error_message( + RuntimeError("Rate limit exceeded: too many requests") + ) + + assert "rate limited" in message + assert "too many requests" in message + + +def test_telegram_exception_hook_sends_once_and_cleans_stream_state(monkeypatch): + sends = [] + cleared = [] + typing_stopped = [] + + async def fake_send(context, text, attachments=None, keyboard=None): + sends.append( + { + "text": text, + "reply_to": context.data.get(CTX_TG_REPLY_TO), + "attachments": attachments, + "keyboard": keyboard, + } + ) + return None + + def fake_clear(context): + cleared.append(True) + + class FakeStop: + def set(self): + typing_stopped.append(True) + + agent = FakeAgent() + agent.context.data[CTX_TG_TYPING_STOP] = FakeStop() + + monkeypatch.setattr( + "plugins._telegram_integration.helpers.handler.send_telegram_reply", + fake_send, + ) + monkeypatch.setattr( + "plugins._telegram_integration.helpers.draft_stream.clear", + fake_clear, + ) + + extension = _85_telegram_error.TelegramFriendlyError(agent=agent) + data = {"exception": RuntimeError("provider returned 503 service unavailable")} + + asyncio.run(extension.execute(data=data)) + asyncio.run(extension.execute(data=data)) + + assert len(sends) == 1 + assert "could not complete the model request" in sends[0]["text"] + assert sends[0]["reply_to"] == 456 + assert cleared == [True] + assert typing_stopped == [True] + assert agent.context.data[CTX_TG_ERROR_SENT] is True + assert CTX_TG_REPLY_TO not in agent.context.data From 99207d60a2ce3a086f124ad9d27d096b09971657 Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:14:23 +0530 Subject: [PATCH 008/231] Add Telegram long-run status updates --- .../end/_85_telegram_error.py | 3 +- .../_45_telegram_draft_start.py | 3 +- .../process_chain_end/_55_telegram_reply.py | 8 +- .../helpers/constants.py | 2 + .../helpers/heartbeat.py | 101 ++++++++++++++++++ tests/test_telegram_heartbeat.py | 71 ++++++++++++ 6 files changed, 182 insertions(+), 6 deletions(-) create mode 100644 plugins/_telegram_integration/helpers/heartbeat.py create mode 100644 tests/test_telegram_heartbeat.py diff --git a/plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py b/plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py index 443df0d74..b711ba350 100644 --- a/plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py +++ b/plugins/_telegram_integration/extensions/python/_functions/agent/Agent/handle_exception/end/_85_telegram_error.py @@ -27,7 +27,7 @@ class TelegramFriendlyError(Extension): context.data[CTX_TG_ERROR_SENT] = True - from plugins._telegram_integration.helpers import draft_stream, error_ui + from plugins._telegram_integration.helpers import draft_stream, error_ui, heartbeat from plugins._telegram_integration.helpers.handler import send_telegram_reply text = error_ui.friendly_error_message(exception) @@ -38,5 +38,6 @@ class TelegramFriendlyError(Extension): typing_stop = context.data.pop(CTX_TG_TYPING_STOP, None) if typing_stop: typing_stop.set() + await heartbeat.stop(context) draft_stream.clear(context) context.data.pop(CTX_TG_REPLY_TO, None) diff --git a/plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py b/plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py index af773207d..6fc4fe89c 100644 --- a/plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py +++ b/plugins/_telegram_integration/extensions/python/message_loop_start/_45_telegram_draft_start.py @@ -12,6 +12,7 @@ class TelegramDraftStart(Extension): if not context.data.get(CTX_TG_BOT): return - from plugins._telegram_integration.helpers import draft_stream + from plugins._telegram_integration.helpers import draft_stream, heartbeat await draft_stream.start(context) + await heartbeat.start(context) diff --git a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py index dcd28665e..508c3acf4 100644 --- a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py +++ b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py @@ -26,14 +26,13 @@ class TelegramAutoReply(Extension): return response_text = _extract_last_response(context) - if not response_text: - return attachments = context.data.pop(CTX_TG_ATTACHMENTS, []) keyboard = context.data.pop(CTX_TG_KEYBOARD, None) try: - await self._send_reply(context, response_text, attachments, keyboard) + if response_text: + await self._send_reply(context, response_text, attachments, keyboard) except Exception as e: PrintStyle.error(f"Telegram auto-reply error: {format_error(e)}") finally: @@ -41,8 +40,9 @@ class TelegramAutoReply(Extension): typing_stop = context.data.pop(CTX_TG_TYPING_STOP, None) if typing_stop: typing_stop.set() - from plugins._telegram_integration.helpers import draft_stream + from plugins._telegram_integration.helpers import draft_stream, heartbeat + await heartbeat.stop(context) draft_stream.clear(context) context.data.pop(CTX_TG_REPLY_TO, None) diff --git a/plugins/_telegram_integration/helpers/constants.py b/plugins/_telegram_integration/helpers/constants.py index 1ef80c81a..190ee7fa4 100644 --- a/plugins/_telegram_integration/helpers/constants.py +++ b/plugins/_telegram_integration/helpers/constants.py @@ -23,3 +23,5 @@ CTX_TG_RESPONSE_MESSAGE_ID = "_telegram_response_message_id" CTX_TG_RESPONSE_TEXT = "_telegram_response_text" CTX_TG_RESPONSE_LAST_UPDATE = "_telegram_response_last_update" CTX_TG_ERROR_SENT = "_telegram_error_sent" +CTX_TG_HEARTBEAT_TASK = "_telegram_heartbeat_task" +CTX_TG_HEARTBEAT_STOP = "_telegram_heartbeat_stop" diff --git a/plugins/_telegram_integration/helpers/heartbeat.py b/plugins/_telegram_integration/helpers/heartbeat.py new file mode 100644 index 000000000..188127e41 --- /dev/null +++ b/plugins/_telegram_integration/helpers/heartbeat.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import asyncio +import re +import time +from contextlib import suppress + +from agent import AgentContext +from helpers.print_style import PrintStyle +from plugins._telegram_integration.helpers import telegram_client as tc +from plugins._telegram_integration.helpers.bot_manager import get_bot +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_HEARTBEAT_STOP, + CTX_TG_HEARTBEAT_TASK, +) + + +HEARTBEAT_INTERVAL_SECONDS = 180 + + +async def start(context: AgentContext) -> None: + task = context.data.get(CTX_TG_HEARTBEAT_TASK) + if task and not task.done(): + return + + stop_event = asyncio.Event() + context.data[CTX_TG_HEARTBEAT_STOP] = stop_event + context.data[CTX_TG_HEARTBEAT_TASK] = asyncio.create_task( + _heartbeat_loop(context, stop_event, time.monotonic()) + ) + + +async def stop(context: AgentContext) -> None: + stop_event = context.data.pop(CTX_TG_HEARTBEAT_STOP, None) + if stop_event: + stop_event.set() + + task = context.data.pop(CTX_TG_HEARTBEAT_TASK, None) + if task and not task.done(): + task.cancel() + with suppress(asyncio.CancelledError): + await task + + +async def _heartbeat_loop( + context: AgentContext, + stop_event: asyncio.Event, + started_at: float, +) -> None: + try: + while True: + try: + await asyncio.wait_for(stop_event.wait(), timeout=HEARTBEAT_INTERVAL_SECONDS) + return + except asyncio.TimeoutError: + await _send_heartbeat(context, time.monotonic() - started_at) + except asyncio.CancelledError: + raise + except Exception as exc: + PrintStyle.debug(f"Telegram heartbeat stopped: {exc}") + + +async def _send_heartbeat(context: AgentContext, elapsed_seconds: float) -> bool: + bot_name = context.data.get(CTX_TG_BOT) + chat_id = context.data.get(CTX_TG_CHAT_ID) + bot = get_bot(bot_name) if bot_name else None + if not bot or not chat_id: + return False + + text = heartbeat_text(elapsed_seconds, _current_reason(context)) + sent_id = await tc.raw_send_text( + bot.bot.token, + int(chat_id), + text, + parse_mode=None, + ) + return bool(sent_id) + + +def heartbeat_text(elapsed_seconds: float, reason: str = "") -> str: + minutes = max(1, int(round(elapsed_seconds / 60))) + detail = reason or "working on your request" + return f"Still working... ({minutes} min elapsed - {detail})" + + +def _current_reason(context: AgentContext) -> str: + progress = str(getattr(getattr(context, "log", None), "progress", "") or "") + progress = _clean_progress(progress) + if not progress or progress.lower() == "waiting for input": + return "working on your request" + return progress + + +def _clean_progress(text: str) -> str: + value = re.sub(r"icon://[a-zA-Z0-9_]+(?:\[[^\]]*\])?\s*", "", text) + value = re.sub(r"\s+", " ", value).strip() + if len(value) > 80: + value = f"{value[:77].rstrip()}..." + return value diff --git a/tests/test_telegram_heartbeat.py b/tests/test_telegram_heartbeat.py new file mode 100644 index 000000000..e9626ed42 --- /dev/null +++ b/tests/test_telegram_heartbeat.py @@ -0,0 +1,71 @@ +import asyncio + +from plugins._telegram_integration.helpers import heartbeat +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_HEARTBEAT_STOP, + CTX_TG_HEARTBEAT_TASK, +) + + +class FakeBot: + class Bot: + token = "token" + + bot = Bot() + + +class FakeLog: + progress = "icon://search[Search] A0: Searching official docs" + + +class FakeContext: + def __init__(self): + self.data = { + CTX_TG_BOT: "main", + CTX_TG_CHAT_ID: 123, + } + self.log = FakeLog() + + +def test_heartbeat_text_omits_iteration_count(): + text = heartbeat.heartbeat_text(180, "waiting for provider response") + + assert text == "Still working... (3 min elapsed - waiting for provider response)" + assert "iteration" not in text + + +def test_heartbeat_sends_periodic_status_and_stops(monkeypatch): + calls = [] + + async def fake_send(token, chat_id, text, parse_mode=None, **kwargs): + calls.append( + { + "token": token, + "chat_id": chat_id, + "text": text, + "parse_mode": parse_mode, + } + ) + return len(calls) + + context = FakeContext() + monkeypatch.setattr(heartbeat, "HEARTBEAT_INTERVAL_SECONDS", 0.01) + monkeypatch.setattr(heartbeat, "get_bot", lambda name: FakeBot()) + monkeypatch.setattr(heartbeat.tc, "raw_send_text", fake_send) + + async def run(): + await heartbeat.start(context) + await asyncio.sleep(0.025) + await heartbeat.stop(context) + + asyncio.run(run()) + + assert calls + assert calls[0]["chat_id"] == 123 + assert calls[0]["parse_mode"] is None + assert "Still working..." in calls[0]["text"] + assert "A0: Searching official docs" in calls[0]["text"] + assert CTX_TG_HEARTBEAT_TASK not in context.data + assert CTX_TG_HEARTBEAT_STOP not in context.data From 1ff45e71c24e369beaf8c01bbe8b41c8a80e14b0 Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:17:46 +0530 Subject: [PATCH 009/231] Add Telegram session picker --- helpers/integration_commands.py | 26 +++ .../helpers/command_ui.py | 178 +++++++++++++++++- tests/test_telegram_sessions_picker.py | 78 ++++++++ 3 files changed, 281 insertions(+), 1 deletion(-) create mode 100644 tests/test_telegram_sessions_picker.py diff --git a/helpers/integration_commands.py b/helpers/integration_commands.py index 655a3cc87..e4a09b32b 100644 --- a/helpers/integration_commands.py +++ b/helpers/integration_commands.py @@ -35,6 +35,12 @@ COMMAND_REGISTRY: tuple[IntegrationCommandDef, ...] = ( "Info", ), IntegrationCommandDef("new", "Start a fresh chat context.", "Session"), + IntegrationCommandDef( + "sessions", + "Show or switch recent chat sessions.", + "Session", + aliases=("session",), + ), IntegrationCommandDef("clear", "Reset the current chat context.", "Session", aliases=("reset",)), IntegrationCommandDef( "queue", @@ -168,6 +174,8 @@ def try_handle_command(context: "AgentContext", text: str) -> str | None: return help_text(full=True) if command == "/status": return _handle_status(context) + if command == "/sessions": + return _handle_sessions(context) if command in {"/new", "/clear"}: return _handle_clear(context, new_chat=(command == "/new")) if command == "/send": @@ -251,6 +259,24 @@ def _handle_status(context: "AgentContext") -> str: ) +def _handle_sessions(context: "AgentContext") -> str: + from agent import AgentContext + + contexts = sorted( + AgentContext.all(), + key=lambda item: str(item.output().get("last_message") or ""), + reverse=True, + ) + lines = ["Recent sessions:"] + for item in contexts[:4]: + marker = " (current)" if item.id == context.id else "" + running = " - running" if item.is_running() else "" + lines.append(f"- {item.name or item.id}{marker}{running}") + if len(contexts) > 4: + lines.append(f"And {len(contexts) - 4} more. Use Telegram buttons to page through them.") + return "\n".join(lines) + + def _handle_clear(context: "AgentContext", *, new_chat: bool) -> str: context.reset() mq.remove(context) diff --git a/plugins/_telegram_integration/helpers/command_ui.py b/plugins/_telegram_integration/helpers/command_ui.py index 0ce9a5ae6..a0f108d0f 100644 --- a/plugins/_telegram_integration/helpers/command_ui.py +++ b/plugins/_telegram_integration/helpers/command_ui.py @@ -1,9 +1,10 @@ from __future__ import annotations +import json from dataclasses import dataclass from agent import AgentContext -from helpers import projects, subagents +from helpers import files, projects, subagents from helpers import integration_commands from helpers.persist_chat import save_tmp_chat from helpers.state_monitor_integration import mark_dirty_for_context @@ -12,9 +13,16 @@ from plugins._telegram_integration.helpers import telegram_client as tc from plugins._telegram_integration.helpers.constants import ( CTX_TG_STREAM_ENABLED, CTX_TG_TOOLS_ENABLED, + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_CHAT_TYPE, + CTX_TG_USER_ID, + CTX_TG_USERNAME, + STATE_FILE, ) PAGE_SIZE = 8 +SESSION_PAGE_SIZE = 4 CALLBACK_PREFIX = "tg" @@ -24,6 +32,14 @@ class PickerItem: label: str +@dataclass(frozen=True) +class SessionItem: + context: AgentContext + label: str + last_message: str + running: bool + + async def handle_command( context: AgentContext, token: str, @@ -44,6 +60,9 @@ async def handle_command( if command in {"/agent", "/profile"} and not args: await send_agent_picker(context, token, chat_id, reply_to_message_id, 0) return True + if command in {"/sessions", "/session"} and not args: + await send_session_picker(context, token, chat_id, reply_to_message_id, 0) + return True if command == "/stream": await send_toggle_picker( context, @@ -91,6 +110,8 @@ async def handle_callback( await edit_project_picker(context, token, chat_id, message_id, page) elif kind == "agent": await edit_agent_picker(context, token, chat_id, message_id, page) + elif kind == "session": + await edit_session_picker(context, token, chat_id, message_id, page) return True if kind == "model" and action in {"set", "clear"}: await _select_model(context, _safe_int(value), clear=(action == "clear")) @@ -104,6 +125,10 @@ async def handle_callback( await _select_agent(context, _safe_int(value)) await edit_agent_picker(context, token, chat_id, message_id, 0, selected=True) return True + if kind == "session" and action == "set": + selected_context = await _select_session(context, _safe_int(value)) + await edit_session_picker(selected_context or context, token, chat_id, message_id, 0, selected=bool(selected_context)) + return True if kind in {"stream", "tools"} and action in {"on", "off"}: key = CTX_TG_STREAM_ENABLED if kind == "stream" else CTX_TG_TOOLS_ENABLED label = "Response streaming" if kind == "stream" else "Tool progress" @@ -217,6 +242,30 @@ async def edit_toggle_picker( await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup) +async def send_session_picker( + context: AgentContext, + token: str, + chat_id: int, + reply_to_message_id: int | None, + page: int, +) -> None: + text, markup = _session_view(context, page) + await tc.raw_send_text(token, chat_id, text, reply_to_message_id, "HTML", markup) + + +async def edit_session_picker( + context: AgentContext, + token: str, + chat_id: int, + message_id: int, + page: int, + *, + selected: bool = False, +) -> None: + text, markup = _session_view(context, page, selected=selected) + await tc.raw_edit_text(token, chat_id, message_id, text, "HTML", markup) + + def _model_view( context: AgentContext, page: int, @@ -295,6 +344,48 @@ def _toggle_view(context: AgentContext, key: str, label: str) -> tuple[str, dict return text, {"inline_keyboard": rows} +def _session_view( + context: AgentContext, + page: int, + *, + selected: bool = False, +) -> tuple[str, dict | None]: + items = _session_items() + current_label = _session_label(context) + status = f"Current session: {_html(current_label)}" + if context.is_running(): + status += "\nSession switching is available after the current run finishes." + elif selected: + status = "Session switched.\n" + status + if not items: + return status + "\nNo sessions were found.", None + + total = len(items) + page = _clamp_page(page, total, SESSION_PAGE_SIZE) + start = page * SESSION_PAGE_SIZE + end = min(start + SESSION_PAGE_SIZE, total) + rows: list[list[dict[str, str]]] = [] + for index, item in enumerate(items[start:end], start=start): + marker = "• " if item.context.id == context.id else "" + suffix = " (running)" if item.running else "" + action = "noop" if context.is_running() or item.running else "set" + rows.append([{ + "text": f"{marker}{item.label}{suffix}"[:64], + "callback_data": f"tg:session:{action}:{index}", + }]) + + nav: list[dict[str, str]] = [] + if page > 0: + nav.append({"text": "Prev", "callback_data": f"tg:session:page:{page - 1}"}) + if end < total: + nav.append({"text": "Next", "callback_data": f"tg:session:page:{page + 1}"}) + if nav: + rows.append(nav) + + range_text = f"\nShowing {start + 1}-{end} of {total}." + return status + range_text, {"inline_keyboard": rows} + + async def _select_model(context: AgentContext, index: int, *, clear: bool = False) -> None: if not model_config.is_chat_override_allowed(context.agent0): return @@ -339,6 +430,86 @@ async def _select_agent(context: AgentContext, index: int) -> None: mark_dirty_for_context(context.id, reason="telegram.agent_select") +async def _select_session(context: AgentContext, index: int) -> AgentContext | None: + if context.is_running(): + return None + items = _session_items() + if index < 0 or index >= len(items): + return None + target = items[index].context + if target.is_running(): + return None + + _copy_telegram_binding(context, target) + _set_session_mapping(target) + save_tmp_chat(target) + mark_dirty_for_context(context.id, reason="telegram.session_unselect") + mark_dirty_for_context(target.id, reason="telegram.session_select") + return target + + +def _session_items() -> list[SessionItem]: + contexts = sorted( + AgentContext.all(), + key=lambda item: str(item.output().get("last_message") or ""), + reverse=True, + ) + return [ + SessionItem( + context=item, + label=_session_label(item), + last_message=str(item.output().get("last_message") or ""), + running=item.is_running(), + ) + for item in contexts + ] + + +def _session_label(context: AgentContext) -> str: + return str(context.name or context.id or "Session") + + +def _copy_telegram_binding(source: AgentContext, target: AgentContext) -> None: + for key in ( + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_CHAT_TYPE, + CTX_TG_USER_ID, + CTX_TG_USERNAME, + ): + if key in source.data: + target.data[key] = source.data[key] + + +def _set_session_mapping(context: AgentContext) -> None: + bot_name = str(context.data.get(CTX_TG_BOT) or "") + user_id = context.data.get(CTX_TG_USER_ID) + chat_id = context.data.get(CTX_TG_CHAT_ID) + if not bot_name or user_id is None or chat_id is None: + return + key = f"{bot_name}:{int(user_id)}:{int(chat_id)}" + state = _load_telegram_state() + chats = state.setdefault("chats", {}) + chats[key] = context.id + _save_telegram_state(state) + + +def _load_telegram_state() -> dict: + path = files.get_abs_path(STATE_FILE) + if not files.exists(path): + return {} + try: + return json.loads(files.read_file(path)) + except Exception: + return {} + + +def _save_telegram_state(state: dict) -> None: + path = files.get_abs_path(STATE_FILE) + files.make_dirs(path) + files.write_file(path, json.dumps(state)) + + def _paged_buttons( kind: str, items: list[PickerItem], @@ -390,6 +561,11 @@ def _safe_int(value: str) -> int: return 0 +def _clamp_page(page: int, total_items: int, page_size: int) -> int: + total_pages = max(1, (total_items + page_size - 1) // page_size) + return min(max(0, page), total_pages - 1) + + def _label_for(items: list[PickerItem], key: str) -> str: for item in items: if item.key == key: diff --git a/tests/test_telegram_sessions_picker.py b/tests/test_telegram_sessions_picker.py new file mode 100644 index 000000000..588ecbc7e --- /dev/null +++ b/tests/test_telegram_sessions_picker.py @@ -0,0 +1,78 @@ +import asyncio +from types import SimpleNamespace + +from plugins._telegram_integration.helpers import command_ui +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_CHAT_ID, + CTX_TG_USER_ID, + CTX_TG_USERNAME, +) + + +class FakeContext: + def __init__(self, id, name, last_message, *, running=False): + self.id = id + self.name = name + self.data = {} + self._last_message = last_message + self._running = running + + def output(self): + return {"last_message": self._last_message} + + def is_running(self): + return self._running + + +def test_session_picker_shows_four_recent_sessions(monkeypatch): + contexts = [ + FakeContext("ctx-1", "One", "2026-05-01T10:00:00"), + FakeContext("ctx-2", "Two", "2026-05-02T10:00:00"), + FakeContext("ctx-3", "Three", "2026-05-03T10:00:00"), + FakeContext("ctx-4", "Four", "2026-05-04T10:00:00"), + FakeContext("ctx-5", "Five", "2026-05-05T10:00:00"), + ] + current = contexts[4] + monkeypatch.setattr(command_ui, "AgentContext", SimpleNamespace(all=lambda: contexts)) + + text, markup = command_ui._session_view(current, 0) + + assert "Current session: Five" in text + assert "Showing 1-4 of 5." in text + rows = markup["inline_keyboard"] + assert [row[0]["text"] for row in rows[:4]] == ["• Five", "Four", "Three", "Two"] + assert rows[-1] == [{"text": "Next", "callback_data": "tg:session:page:1"}] + + +def test_select_session_remaps_telegram_chat(monkeypatch): + current = FakeContext("ctx-current", "Current", "2026-05-03T10:00:00") + target = FakeContext("ctx-target", "Target", "2026-05-04T10:00:00") + current.data.update( + { + CTX_TG_BOT: "main", + CTX_TG_CHAT_ID: 123, + CTX_TG_USER_ID: 456, + CTX_TG_USERNAME: "anmol", + } + ) + saved = [] + dirty = [] + state_out = {} + + monkeypatch.setattr(command_ui, "AgentContext", SimpleNamespace(all=lambda: [current, target])) + monkeypatch.setattr(command_ui, "save_tmp_chat", lambda context: saved.append(context.id)) + monkeypatch.setattr(command_ui, "mark_dirty_for_context", lambda context_id, *, reason=None: dirty.append((context_id, reason))) + monkeypatch.setattr(command_ui, "_load_telegram_state", lambda: {"chats": {}}) + monkeypatch.setattr(command_ui, "_save_telegram_state", lambda state: state_out.update(state)) + + selected = asyncio.run(command_ui._select_session(current, 0)) + + assert selected is target + assert target.data[CTX_TG_BOT] == "main" + assert target.data[CTX_TG_CHAT_ID] == 123 + assert target.data[CTX_TG_USER_ID] == 456 + assert target.data[CTX_TG_USERNAME] == "anmol" + assert state_out["chats"]["main:456:123"] == "ctx-target" + assert saved == ["ctx-target"] + assert ("ctx-target", "telegram.session_select") in dirty From 709b8653996508eeda42ebfbdc54a282d5d3c92e Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:35:44 +0530 Subject: [PATCH 010/231] Add Telegram speech mode --- helpers/integration_commands.py | 8 ++ plugins/_telegram_integration/README.md | 3 + .../process_chain_end/_55_telegram_reply.py | 10 ++ .../helpers/command_ui.py | 37 +++++- .../helpers/constants.py | 2 + .../_telegram_integration/helpers/speech.py | 80 +++++++++++++ tests/test_telegram_speech_mode.py | 112 ++++++++++++++++++ 7 files changed, 248 insertions(+), 4 deletions(-) create mode 100644 plugins/_telegram_integration/helpers/speech.py create mode 100644 tests/test_telegram_speech_mode.py diff --git a/helpers/integration_commands.py b/helpers/integration_commands.py index e4a09b32b..b555d2af0 100644 --- a/helpers/integration_commands.py +++ b/helpers/integration_commands.py @@ -70,6 +70,12 @@ COMMAND_REGISTRY: tuple[IntegrationCommandDef, ...] = ( "Configuration", args_hint="[on|off]", ), + IntegrationCommandDef( + "speech", + "Enable or disable Telegram spoken replies.", + "Configuration", + args_hint="[on|off]", + ), IntegrationCommandDef( "project", "Show or switch the active project.", @@ -194,6 +200,8 @@ def try_handle_command(context: "AgentContext", text: str) -> str | None: return _handle_toggle(context, args, "telegram_stream_enabled", "Response streaming") if command == "/tools": return _handle_toggle(context, args, "telegram_tools_enabled", "Tool progress") + if command == "/speech": + return _handle_toggle(context, args, "telegram_speech_enabled", "Speech replies") if command == "/project": return _handle_project(context, args) if command == "/model": diff --git a/plugins/_telegram_integration/README.md b/plugins/_telegram_integration/README.md index a9793fa35..c3c13aae2 100644 --- a/plugins/_telegram_integration/README.md +++ b/plugins/_telegram_integration/README.md @@ -24,6 +24,7 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - `/agent ` switches the active agent profile when the current run is idle. - `/model`, `/project`, and `/agent` show Telegram inline keyboard pickers when used without arguments. - `/stream` toggles live response streaming; `/tools` toggles tool progress messages. + - `/speech` toggles spoken replies. When enabled, final text replies also include Kokoro-generated Telegram audio when Kokoro TTS is available. - `/send` or `/queue send` flushes queued messages for the current chat. - `/steer ` sends an intervention to the active run. - **Group support** @@ -35,6 +36,7 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - **Reply delivery** - Streams tool progress and response text through real Telegram messages updated with `editMessageText`. - Tool progress and the AI response are separate messages; only the AI response replies to the user message. + - Sends native Telegram media: photos, voice/audio files, videos, and document fallback. - `tool_execute_after` intercepts the `response` tool — sends `break_loop=false` updates as separate intermediate Telegram messages. - `process_chain_end` auto-sends the final response, with retry logic on failure. - **Formatting** @@ -63,6 +65,7 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - `helpers/telegram_client.py` — Low-level Telegram API wrapper: send text/file/photo, Markdown→HTML converter, keyboard builder, message splitting. - `helpers/command_ui.py` — Telegram inline keyboard command pickers and callback handling. - `helpers/draft_stream.py` — Editable Telegram message streaming for tool progress and response previews. + - `helpers/speech.py` — Telegram speech-mode TTS synthesis and audio attachment generation. - **Extensions** - `extensions/python/job_loop/_10_telegram_bot.py` — Bot lifecycle manager, starts/stops bots on each tick. - `extensions/python/message_loop_start/_45_telegram_draft_start.py` — Initializes Telegram response streaming. diff --git a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py index 508c3acf4..582d444b2 100644 --- a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py +++ b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py @@ -32,7 +32,17 @@ class TelegramAutoReply(Extension): try: if response_text: + from plugins._telegram_integration.helpers import speech + + speech_attachment, speech_warning = await speech.synthesize_attachment( + context, response_text + ) + if speech_attachment: + attachments = [*attachments, speech_attachment] await self._send_reply(context, response_text, attachments, keyboard) + warning = speech.consume_warning(context, speech_warning) + if warning: + await self._send_reply(context, warning, [], None) except Exception as e: PrintStyle.error(f"Telegram auto-reply error: {format_error(e)}") finally: diff --git a/plugins/_telegram_integration/helpers/command_ui.py b/plugins/_telegram_integration/helpers/command_ui.py index a0f108d0f..6ceb3fe6e 100644 --- a/plugins/_telegram_integration/helpers/command_ui.py +++ b/plugins/_telegram_integration/helpers/command_ui.py @@ -13,6 +13,7 @@ from plugins._telegram_integration.helpers import telegram_client as tc from plugins._telegram_integration.helpers.constants import ( CTX_TG_STREAM_ENABLED, CTX_TG_TOOLS_ENABLED, + CTX_TG_SPEECH_ENABLED, CTX_TG_BOT, CTX_TG_CHAT_ID, CTX_TG_CHAT_TYPE, @@ -85,6 +86,17 @@ async def handle_command( args, ) return True + if command == "/speech": + await send_toggle_picker( + context, + token, + chat_id, + reply_to_message_id, + CTX_TG_SPEECH_ENABLED, + "Speech replies", + args, + ) + return True return False @@ -129,9 +141,8 @@ async def handle_callback( selected_context = await _select_session(context, _safe_int(value)) await edit_session_picker(selected_context or context, token, chat_id, message_id, 0, selected=bool(selected_context)) return True - if kind in {"stream", "tools"} and action in {"on", "off"}: - key = CTX_TG_STREAM_ENABLED if kind == "stream" else CTX_TG_TOOLS_ENABLED - label = "Response streaming" if kind == "stream" else "Tool progress" + if kind in {"stream", "tools", "speech"} and action in {"on", "off"}: + key, label = _toggle_key_label(kind) context.set_data(key, action == "on") save_tmp_chat(context) mark_dirty_for_context(context.id, reason=f"telegram.{kind}_toggle") @@ -334,7 +345,7 @@ def _agent_view( def _toggle_view(context: AgentContext, key: str, label: str) -> tuple[str, dict]: enabled = _toggle_enabled(context, key) - kind = "stream" if key == CTX_TG_STREAM_ENABLED else "tools" + kind = _toggle_kind(key) state = "enabled" if enabled else "disabled" text = f"{_html(label)}: {state}" rows = [[ @@ -542,9 +553,27 @@ def _paged_buttons( def _toggle_enabled(context: AgentContext, key: str) -> bool: value = context.get_data(key) + if key == CTX_TG_SPEECH_ENABLED: + return bool(value) return True if value is None else bool(value) +def _toggle_kind(key: str) -> str: + if key == CTX_TG_STREAM_ENABLED: + return "stream" + if key == CTX_TG_TOOLS_ENABLED: + return "tools" + return "speech" + + +def _toggle_key_label(kind: str) -> tuple[str, str]: + if kind == "stream": + return CTX_TG_STREAM_ENABLED, "Response streaming" + if kind == "tools": + return CTX_TG_TOOLS_ENABLED, "Tool progress" + return CTX_TG_SPEECH_ENABLED, "Speech replies" + + def _parse_toggle(args: str) -> bool | None: value = (args or "").strip().lower() if value in {"on", "enable", "enabled", "yes", "true", "1"}: diff --git a/plugins/_telegram_integration/helpers/constants.py b/plugins/_telegram_integration/helpers/constants.py index 190ee7fa4..9696f2661 100644 --- a/plugins/_telegram_integration/helpers/constants.py +++ b/plugins/_telegram_integration/helpers/constants.py @@ -13,6 +13,7 @@ CTX_TG_TYPING_STOP = "_telegram_typing_stop" CTX_TG_REPLY_TO = "_telegram_reply_to_message_id" CTX_TG_STREAM_ENABLED = "telegram_stream_enabled" CTX_TG_TOOLS_ENABLED = "telegram_tools_enabled" +CTX_TG_SPEECH_ENABLED = "telegram_speech_enabled" # Transient CTX_TG_ATTACHMENTS = "_telegram_response_attachments" @@ -25,3 +26,4 @@ CTX_TG_RESPONSE_LAST_UPDATE = "_telegram_response_last_update" CTX_TG_ERROR_SENT = "_telegram_error_sent" CTX_TG_HEARTBEAT_TASK = "_telegram_heartbeat_task" CTX_TG_HEARTBEAT_STOP = "_telegram_heartbeat_stop" +CTX_TG_SPEECH_WARNING_SENT = "_telegram_speech_warning_sent" diff --git a/plugins/_telegram_integration/helpers/speech.py b/plugins/_telegram_integration/helpers/speech.py new file mode 100644 index 000000000..eb74a5e30 --- /dev/null +++ b/plugins/_telegram_integration/helpers/speech.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import re +import uuid + +from agent import AgentContext +from helpers import files +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_SPEECH_ENABLED, + CTX_TG_SPEECH_WARNING_SENT, + DOWNLOAD_FOLDER, +) + + +MAX_SPEECH_CHARS = 1800 + + +def is_enabled(context: AgentContext) -> bool: + return bool(context.get_data(CTX_TG_SPEECH_ENABLED)) + + +async def synthesize_attachment(context: AgentContext, text: str) -> tuple[str | None, str | None]: + if not is_enabled(context): + return None, None + + speech_text = _speech_text(text) + if not speech_text: + return None, "Speech mode is enabled, but there was no readable text to speak." + + try: + from plugins._kokoro_tts.helpers import runtime + except Exception: + return None, _setup_message() + + if not runtime.is_globally_enabled(): + return None, _setup_message() + + try: + audio = await runtime.synthesize_sentences([speech_text]) + except Exception as exc: + return None, f"Speech mode is enabled, but text-to-speech failed: {exc}" + + if not audio: + return None, "Speech mode is enabled, but text-to-speech returned no audio." + + name = f"telegram_speech_{context.id}_{uuid.uuid4().hex[:10]}.wav" + relative_path = f"{DOWNLOAD_FOLDER}/{name}" + files.write_file_base64(relative_path, audio) + return files.get_abs_path_dockerized(relative_path), None + + +def consume_warning(context: AgentContext, warning: str | None) -> str | None: + if not warning: + context.data.pop(CTX_TG_SPEECH_WARNING_SENT, None) + return None + if context.data.get(CTX_TG_SPEECH_WARNING_SENT): + return None + context.data[CTX_TG_SPEECH_WARNING_SENT] = True + return warning + + +def _setup_message() -> str: + return ( + "Speech mode is enabled, but Kokoro TTS is not available. " + "Enable the Kokoro TTS plugin from Agent Zero settings, then try again." + ) + + +def _speech_text(text: str) -> str: + value = str(text or "") + value = re.sub(r"```.*?```", " ", value, flags=re.DOTALL) + value = re.sub(r"`([^`]+)`", r"\1", value) + value = re.sub(r"!\[([^\]]*)\]\([^)]+\)", r"\1", value) + value = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", value) + value = re.sub(r"<[^>]+>", " ", value) + value = re.sub(r"[*_#>\-]+", " ", value) + value = re.sub(r"\s+", " ", value).strip() + if len(value) > MAX_SPEECH_CHARS: + value = value[:MAX_SPEECH_CHARS].rstrip() + return value diff --git a/tests/test_telegram_speech_mode.py b/tests/test_telegram_speech_mode.py new file mode 100644 index 000000000..575cc2777 --- /dev/null +++ b/tests/test_telegram_speech_mode.py @@ -0,0 +1,112 @@ +import asyncio + +from plugins._telegram_integration.extensions.python.process_chain_end import ( + _55_telegram_reply, +) +from plugins._telegram_integration.helpers import command_ui, speech +from plugins._telegram_integration.helpers.constants import ( + CTX_TG_BOT, + CTX_TG_REPLY_TO, + CTX_TG_SPEECH_ENABLED, +) + + +class FakeContext: + id = "ctx-speech" + + def __init__(self): + self.data = {} + + def get_data(self, key): + return self.data.get(key) + + def set_data(self, key, value): + self.data[key] = value + + +class FakeAgent: + number = 0 + + def __init__(self): + self.context = FakeContext() + self.context.data[CTX_TG_BOT] = "main" + self.context.data[CTX_TG_REPLY_TO] = 456 + + +def test_speech_toggle_defaults_off(): + context = FakeContext() + + text, markup = command_ui._toggle_view(context, CTX_TG_SPEECH_ENABLED, "Speech replies") + + assert "Speech replies: disabled" == text + assert markup["inline_keyboard"][0][0]["callback_data"] == "tg:speech:on" + assert markup["inline_keyboard"][0][1]["callback_data"] == "tg:speech:off" + + +def test_speech_text_is_cleaned_for_tts(): + cleaned = speech._speech_text( + "**Hello** [`Agent Zero`](https://example.com)\n\n```python\nprint('skip')\n```" + ) + + assert cleaned == "Hello Agent Zero" + + +def test_speech_synthesis_writes_audio_attachment(monkeypatch): + context = FakeContext() + context.data[CTX_TG_SPEECH_ENABLED] = True + writes = [] + + class FakeRuntime: + @staticmethod + def is_globally_enabled(): + return True + + @staticmethod + async def synthesize_sentences(sentences): + assert sentences == ["Hello from Agent Zero."] + return "UklGRg==" + + monkeypatch.setattr(speech.files, "write_file_base64", lambda path, audio: writes.append((path, audio))) + monkeypatch.setattr(speech.files, "get_abs_path_dockerized", lambda path: f"/a0/{path}") + monkeypatch.setitem(__import__("sys").modules, "plugins._kokoro_tts.helpers.runtime", FakeRuntime) + + path, warning = asyncio.run(speech.synthesize_attachment(context, "Hello from Agent Zero.")) + + assert warning is None + assert path.startswith("/a0/usr/uploads/telegram_speech_ctx-speech_") + assert path.endswith(".wav") + assert writes[0][1] == "UklGRg==" + + +def test_final_reply_includes_speech_attachment(monkeypatch): + agent = FakeAgent() + sent = [] + + async def fake_synthesize(context, text): + return "/a0/usr/uploads/reply.wav", None + + async def fake_send_reply(context, response_text, attachments, keyboard): + sent.append( + { + "text": response_text, + "attachments": attachments, + "keyboard": keyboard, + } + ) + + monkeypatch.setattr(_55_telegram_reply, "_extract_last_response", lambda context: "Final reply.") + monkeypatch.setattr(speech, "synthesize_attachment", fake_synthesize) + + extension = _55_telegram_reply.TelegramAutoReply(agent=agent) + extension._send_reply = fake_send_reply + + asyncio.run(extension.execute()) + + assert sent == [ + { + "text": "Final reply.", + "attachments": ["/a0/usr/uploads/reply.wav"], + "keyboard": None, + } + ] + assert CTX_TG_REPLY_TO not in agent.context.data From 013a374daa312d57e5d5b0341d008a1b40752f2b Mon Sep 17 00:00:00 2001 From: Anmol Malik Date: Thu, 28 May 2026 23:43:49 +0530 Subject: [PATCH 011/231] Revert "Add Telegram speech mode" This reverts commit c51c235faacbbed44994635d4743b97f7a143a17. --- helpers/integration_commands.py | 8 -- plugins/_telegram_integration/README.md | 3 - .../process_chain_end/_55_telegram_reply.py | 10 -- .../helpers/command_ui.py | 37 +----- .../helpers/constants.py | 2 - .../_telegram_integration/helpers/speech.py | 80 ------------- tests/test_telegram_speech_mode.py | 112 ------------------ 7 files changed, 4 insertions(+), 248 deletions(-) delete mode 100644 plugins/_telegram_integration/helpers/speech.py delete mode 100644 tests/test_telegram_speech_mode.py diff --git a/helpers/integration_commands.py b/helpers/integration_commands.py index b555d2af0..e4a09b32b 100644 --- a/helpers/integration_commands.py +++ b/helpers/integration_commands.py @@ -70,12 +70,6 @@ COMMAND_REGISTRY: tuple[IntegrationCommandDef, ...] = ( "Configuration", args_hint="[on|off]", ), - IntegrationCommandDef( - "speech", - "Enable or disable Telegram spoken replies.", - "Configuration", - args_hint="[on|off]", - ), IntegrationCommandDef( "project", "Show or switch the active project.", @@ -200,8 +194,6 @@ def try_handle_command(context: "AgentContext", text: str) -> str | None: return _handle_toggle(context, args, "telegram_stream_enabled", "Response streaming") if command == "/tools": return _handle_toggle(context, args, "telegram_tools_enabled", "Tool progress") - if command == "/speech": - return _handle_toggle(context, args, "telegram_speech_enabled", "Speech replies") if command == "/project": return _handle_project(context, args) if command == "/model": diff --git a/plugins/_telegram_integration/README.md b/plugins/_telegram_integration/README.md index c3c13aae2..a9793fa35 100644 --- a/plugins/_telegram_integration/README.md +++ b/plugins/_telegram_integration/README.md @@ -24,7 +24,6 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - `/agent ` switches the active agent profile when the current run is idle. - `/model`, `/project`, and `/agent` show Telegram inline keyboard pickers when used without arguments. - `/stream` toggles live response streaming; `/tools` toggles tool progress messages. - - `/speech` toggles spoken replies. When enabled, final text replies also include Kokoro-generated Telegram audio when Kokoro TTS is available. - `/send` or `/queue send` flushes queued messages for the current chat. - `/steer ` sends an intervention to the active run. - **Group support** @@ -36,7 +35,6 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - **Reply delivery** - Streams tool progress and response text through real Telegram messages updated with `editMessageText`. - Tool progress and the AI response are separate messages; only the AI response replies to the user message. - - Sends native Telegram media: photos, voice/audio files, videos, and document fallback. - `tool_execute_after` intercepts the `response` tool — sends `break_loop=false` updates as separate intermediate Telegram messages. - `process_chain_end` auto-sends the final response, with retry logic on failure. - **Formatting** @@ -65,7 +63,6 @@ This plugin connects one or more Telegram bots to Agent Zero. Each bot runs inde - `helpers/telegram_client.py` — Low-level Telegram API wrapper: send text/file/photo, Markdown→HTML converter, keyboard builder, message splitting. - `helpers/command_ui.py` — Telegram inline keyboard command pickers and callback handling. - `helpers/draft_stream.py` — Editable Telegram message streaming for tool progress and response previews. - - `helpers/speech.py` — Telegram speech-mode TTS synthesis and audio attachment generation. - **Extensions** - `extensions/python/job_loop/_10_telegram_bot.py` — Bot lifecycle manager, starts/stops bots on each tick. - `extensions/python/message_loop_start/_45_telegram_draft_start.py` — Initializes Telegram response streaming. diff --git a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py index 582d444b2..508c3acf4 100644 --- a/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py +++ b/plugins/_telegram_integration/extensions/python/process_chain_end/_55_telegram_reply.py @@ -32,17 +32,7 @@ class TelegramAutoReply(Extension): try: if response_text: - from plugins._telegram_integration.helpers import speech - - speech_attachment, speech_warning = await speech.synthesize_attachment( - context, response_text - ) - if speech_attachment: - attachments = [*attachments, speech_attachment] await self._send_reply(context, response_text, attachments, keyboard) - warning = speech.consume_warning(context, speech_warning) - if warning: - await self._send_reply(context, warning, [], None) except Exception as e: PrintStyle.error(f"Telegram auto-reply error: {format_error(e)}") finally: diff --git a/plugins/_telegram_integration/helpers/command_ui.py b/plugins/_telegram_integration/helpers/command_ui.py index 6ceb3fe6e..a0f108d0f 100644 --- a/plugins/_telegram_integration/helpers/command_ui.py +++ b/plugins/_telegram_integration/helpers/command_ui.py @@ -13,7 +13,6 @@ from plugins._telegram_integration.helpers import telegram_client as tc from plugins._telegram_integration.helpers.constants import ( CTX_TG_STREAM_ENABLED, CTX_TG_TOOLS_ENABLED, - CTX_TG_SPEECH_ENABLED, CTX_TG_BOT, CTX_TG_CHAT_ID, CTX_TG_CHAT_TYPE, @@ -86,17 +85,6 @@ async def handle_command( args, ) return True - if command == "/speech": - await send_toggle_picker( - context, - token, - chat_id, - reply_to_message_id, - CTX_TG_SPEECH_ENABLED, - "Speech replies", - args, - ) - return True return False @@ -141,8 +129,9 @@ async def handle_callback( selected_context = await _select_session(context, _safe_int(value)) await edit_session_picker(selected_context or context, token, chat_id, message_id, 0, selected=bool(selected_context)) return True - if kind in {"stream", "tools", "speech"} and action in {"on", "off"}: - key, label = _toggle_key_label(kind) + if kind in {"stream", "tools"} and action in {"on", "off"}: + key = CTX_TG_STREAM_ENABLED if kind == "stream" else CTX_TG_TOOLS_ENABLED + label = "Response streaming" if kind == "stream" else "Tool progress" context.set_data(key, action == "on") save_tmp_chat(context) mark_dirty_for_context(context.id, reason=f"telegram.{kind}_toggle") @@ -345,7 +334,7 @@ def _agent_view( def _toggle_view(context: AgentContext, key: str, label: str) -> tuple[str, dict]: enabled = _toggle_enabled(context, key) - kind = _toggle_kind(key) + kind = "stream" if key == CTX_TG_STREAM_ENABLED else "tools" state = "enabled" if enabled else "disabled" text = f"{_html(label)}: {state}" rows = [[ @@ -553,27 +542,9 @@ def _paged_buttons( def _toggle_enabled(context: AgentContext, key: str) -> bool: value = context.get_data(key) - if key == CTX_TG_SPEECH_ENABLED: - return bool(value) return True if value is None else bool(value) -def _toggle_kind(key: str) -> str: - if key == CTX_TG_STREAM_ENABLED: - return "stream" - if key == CTX_TG_TOOLS_ENABLED: - return "tools" - return "speech" - - -def _toggle_key_label(kind: str) -> tuple[str, str]: - if kind == "stream": - return CTX_TG_STREAM_ENABLED, "Response streaming" - if kind == "tools": - return CTX_TG_TOOLS_ENABLED, "Tool progress" - return CTX_TG_SPEECH_ENABLED, "Speech replies" - - def _parse_toggle(args: str) -> bool | None: value = (args or "").strip().lower() if value in {"on", "enable", "enabled", "yes", "true", "1"}: diff --git a/plugins/_telegram_integration/helpers/constants.py b/plugins/_telegram_integration/helpers/constants.py index 9696f2661..190ee7fa4 100644 --- a/plugins/_telegram_integration/helpers/constants.py +++ b/plugins/_telegram_integration/helpers/constants.py @@ -13,7 +13,6 @@ CTX_TG_TYPING_STOP = "_telegram_typing_stop" CTX_TG_REPLY_TO = "_telegram_reply_to_message_id" CTX_TG_STREAM_ENABLED = "telegram_stream_enabled" CTX_TG_TOOLS_ENABLED = "telegram_tools_enabled" -CTX_TG_SPEECH_ENABLED = "telegram_speech_enabled" # Transient CTX_TG_ATTACHMENTS = "_telegram_response_attachments" @@ -26,4 +25,3 @@ CTX_TG_RESPONSE_LAST_UPDATE = "_telegram_response_last_update" CTX_TG_ERROR_SENT = "_telegram_error_sent" CTX_TG_HEARTBEAT_TASK = "_telegram_heartbeat_task" CTX_TG_HEARTBEAT_STOP = "_telegram_heartbeat_stop" -CTX_TG_SPEECH_WARNING_SENT = "_telegram_speech_warning_sent" diff --git a/plugins/_telegram_integration/helpers/speech.py b/plugins/_telegram_integration/helpers/speech.py deleted file mode 100644 index eb74a5e30..000000000 --- a/plugins/_telegram_integration/helpers/speech.py +++ /dev/null @@ -1,80 +0,0 @@ -from __future__ import annotations - -import re -import uuid - -from agent import AgentContext -from helpers import files -from plugins._telegram_integration.helpers.constants import ( - CTX_TG_SPEECH_ENABLED, - CTX_TG_SPEECH_WARNING_SENT, - DOWNLOAD_FOLDER, -) - - -MAX_SPEECH_CHARS = 1800 - - -def is_enabled(context: AgentContext) -> bool: - return bool(context.get_data(CTX_TG_SPEECH_ENABLED)) - - -async def synthesize_attachment(context: AgentContext, text: str) -> tuple[str | None, str | None]: - if not is_enabled(context): - return None, None - - speech_text = _speech_text(text) - if not speech_text: - return None, "Speech mode is enabled, but there was no readable text to speak." - - try: - from plugins._kokoro_tts.helpers import runtime - except Exception: - return None, _setup_message() - - if not runtime.is_globally_enabled(): - return None, _setup_message() - - try: - audio = await runtime.synthesize_sentences([speech_text]) - except Exception as exc: - return None, f"Speech mode is enabled, but text-to-speech failed: {exc}" - - if not audio: - return None, "Speech mode is enabled, but text-to-speech returned no audio." - - name = f"telegram_speech_{context.id}_{uuid.uuid4().hex[:10]}.wav" - relative_path = f"{DOWNLOAD_FOLDER}/{name}" - files.write_file_base64(relative_path, audio) - return files.get_abs_path_dockerized(relative_path), None - - -def consume_warning(context: AgentContext, warning: str | None) -> str | None: - if not warning: - context.data.pop(CTX_TG_SPEECH_WARNING_SENT, None) - return None - if context.data.get(CTX_TG_SPEECH_WARNING_SENT): - return None - context.data[CTX_TG_SPEECH_WARNING_SENT] = True - return warning - - -def _setup_message() -> str: - return ( - "Speech mode is enabled, but Kokoro TTS is not available. " - "Enable the Kokoro TTS plugin from Agent Zero settings, then try again." - ) - - -def _speech_text(text: str) -> str: - value = str(text or "") - value = re.sub(r"```.*?```", " ", value, flags=re.DOTALL) - value = re.sub(r"`([^`]+)`", r"\1", value) - value = re.sub(r"!\[([^\]]*)\]\([^)]+\)", r"\1", value) - value = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", value) - value = re.sub(r"<[^>]+>", " ", value) - value = re.sub(r"[*_#>\-]+", " ", value) - value = re.sub(r"\s+", " ", value).strip() - if len(value) > MAX_SPEECH_CHARS: - value = value[:MAX_SPEECH_CHARS].rstrip() - return value diff --git a/tests/test_telegram_speech_mode.py b/tests/test_telegram_speech_mode.py deleted file mode 100644 index 575cc2777..000000000 --- a/tests/test_telegram_speech_mode.py +++ /dev/null @@ -1,112 +0,0 @@ -import asyncio - -from plugins._telegram_integration.extensions.python.process_chain_end import ( - _55_telegram_reply, -) -from plugins._telegram_integration.helpers import command_ui, speech -from plugins._telegram_integration.helpers.constants import ( - CTX_TG_BOT, - CTX_TG_REPLY_TO, - CTX_TG_SPEECH_ENABLED, -) - - -class FakeContext: - id = "ctx-speech" - - def __init__(self): - self.data = {} - - def get_data(self, key): - return self.data.get(key) - - def set_data(self, key, value): - self.data[key] = value - - -class FakeAgent: - number = 0 - - def __init__(self): - self.context = FakeContext() - self.context.data[CTX_TG_BOT] = "main" - self.context.data[CTX_TG_REPLY_TO] = 456 - - -def test_speech_toggle_defaults_off(): - context = FakeContext() - - text, markup = command_ui._toggle_view(context, CTX_TG_SPEECH_ENABLED, "Speech replies") - - assert "Speech replies: disabled" == text - assert markup["inline_keyboard"][0][0]["callback_data"] == "tg:speech:on" - assert markup["inline_keyboard"][0][1]["callback_data"] == "tg:speech:off" - - -def test_speech_text_is_cleaned_for_tts(): - cleaned = speech._speech_text( - "**Hello** [`Agent Zero`](https://example.com)\n\n```python\nprint('skip')\n```" - ) - - assert cleaned == "Hello Agent Zero" - - -def test_speech_synthesis_writes_audio_attachment(monkeypatch): - context = FakeContext() - context.data[CTX_TG_SPEECH_ENABLED] = True - writes = [] - - class FakeRuntime: - @staticmethod - def is_globally_enabled(): - return True - - @staticmethod - async def synthesize_sentences(sentences): - assert sentences == ["Hello from Agent Zero."] - return "UklGRg==" - - monkeypatch.setattr(speech.files, "write_file_base64", lambda path, audio: writes.append((path, audio))) - monkeypatch.setattr(speech.files, "get_abs_path_dockerized", lambda path: f"/a0/{path}") - monkeypatch.setitem(__import__("sys").modules, "plugins._kokoro_tts.helpers.runtime", FakeRuntime) - - path, warning = asyncio.run(speech.synthesize_attachment(context, "Hello from Agent Zero.")) - - assert warning is None - assert path.startswith("/a0/usr/uploads/telegram_speech_ctx-speech_") - assert path.endswith(".wav") - assert writes[0][1] == "UklGRg==" - - -def test_final_reply_includes_speech_attachment(monkeypatch): - agent = FakeAgent() - sent = [] - - async def fake_synthesize(context, text): - return "/a0/usr/uploads/reply.wav", None - - async def fake_send_reply(context, response_text, attachments, keyboard): - sent.append( - { - "text": response_text, - "attachments": attachments, - "keyboard": keyboard, - } - ) - - monkeypatch.setattr(_55_telegram_reply, "_extract_last_response", lambda context: "Final reply.") - monkeypatch.setattr(speech, "synthesize_attachment", fake_synthesize) - - extension = _55_telegram_reply.TelegramAutoReply(agent=agent) - extension._send_reply = fake_send_reply - - asyncio.run(extension.execute()) - - assert sent == [ - { - "text": "Final reply.", - "attachments": ["/a0/usr/uploads/reply.wav"], - "keyboard": None, - } - ] - assert CTX_TG_REPLY_TO not in agent.context.data From 67224672e89dbd330aab0745230c49b0f6d6fa3a Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Thu, 28 May 2026 22:14:36 +0200 Subject: [PATCH 012/231] Add Remote Control tunnel providers Rename the Remote Link UI and user-facing messages to Remote Control. Refactor tunnel startup into provider helpers, remove ngrok support, and keep Cloudflare Tunnel, Microsoft Dev Tunnels, Serveo, and Tailscale wired through the Remote Control selector. Add provider-time binary preparation, Tailscale userspace tailscaled startup, Tailscale login URL display, Microsoft Dev Tunnel progress feedback, and focused regression coverage for the Remote Control provider flows. --- .../register-remote-link.js | 2 +- helpers/cli_tunnel.py | 262 ++++++ helpers/cloudflare_tunnel._py | 10 +- helpers/cloudflare_tunnel.py | 11 + helpers/microsoft_tunnel.py | 144 ++++ helpers/runtime.py | 2 +- helpers/serveo_tunnel.py | 11 + helpers/tailscale_tunnel.py | 433 ++++++++++ helpers/tunnel_common.py | 56 ++ helpers/tunnel_manager.py | 113 ++- tests/test_tunnel_remote_link.py | 757 ++++++++++++++++++ .../settings/external/external-settings.html | 4 +- webui/components/settings/settings-store.js | 2 +- .../settings/tunnel/remote-link.html | 2 +- .../settings/tunnel/tunnel-section.html | 37 +- .../settings/tunnel/tunnel-store.js | 58 +- .../sidebar/top-section/header-icons.html | 2 +- 17 files changed, 1818 insertions(+), 88 deletions(-) create mode 100644 helpers/cli_tunnel.py create mode 100644 helpers/cloudflare_tunnel.py create mode 100644 helpers/microsoft_tunnel.py create mode 100644 helpers/serveo_tunnel.py create mode 100644 helpers/tailscale_tunnel.py create mode 100644 helpers/tunnel_common.py create mode 100644 tests/test_tunnel_remote_link.py diff --git a/extensions/webui/right_canvas_register_surfaces/register-remote-link.js b/extensions/webui/right_canvas_register_surfaces/register-remote-link.js index ca8dc6dfb..dd0d95d4d 100644 --- a/extensions/webui/right_canvas_register_surfaces/register-remote-link.js +++ b/extensions/webui/right_canvas_register_surfaces/register-remote-link.js @@ -1,3 +1,3 @@ export default async function registerRemoteLinkAction() { - // Remote Link is opened from the sidebar dropdown, not the right canvas rail. + // Remote Control is opened from the sidebar dropdown, not the right canvas rail. } diff --git a/helpers/cli_tunnel.py b/helpers/cli_tunnel.py new file mode 100644 index 000000000..05f702fc9 --- /dev/null +++ b/helpers/cli_tunnel.py @@ -0,0 +1,262 @@ +import os +import platform +import queue +import shutil +import subprocess +import tarfile +import tempfile +import threading +import time +import urllib.request +import zipfile +from collections import deque +from pathlib import Path + +from flaredantic import NotifyEvent + +from helpers import files +from helpers.tunnel_common import TunnelHelper + + +RUNTIME_BIN_DIR = Path(files.get_abs_path("tmp", "bin")) + + +def executable_name(name): + return f"{name}.exe" if platform.system().lower() == "windows" else name + + +def chmod_executable(path): + if platform.system().lower() != "windows": + path.chmod(0o755) + + +def notify_download(notify, message, data=None): + if callable(notify): + notify(NotifyEvent.DOWNLOADING, message, data) + + +def notify_download_complete(notify, message, data=None): + if callable(notify): + notify(NotifyEvent.DOWNLOAD_COMPLETE, message, data) + + +def download_file(url, destination, notify=None): + destination.parent.mkdir(parents=True, exist_ok=True) + notify_download(notify, f"Downloading {url}") + temp_fd, temp_name = tempfile.mkstemp( + prefix=f"{destination.name}.", + suffix=".download", + dir=str(destination.parent), + ) + os.close(temp_fd) + temp_path = Path(temp_name) + try: + with urllib.request.urlopen(url, timeout=60) as response: + with temp_path.open("wb") as handle: + shutil.copyfileobj(response, handle) + temp_path.replace(destination) + notify_download_complete(notify, f"Downloaded {destination.name}") + return destination + except Exception: + temp_path.unlink(missing_ok=True) + raise + + +def platform_parts(): + system = platform.system().lower() + machine = platform.machine().lower() + if machine in {"x86_64", "amd64"}: + arch = "amd64" + elif machine in {"aarch64", "arm64"}: + arch = "arm64" + elif machine in {"i386", "i686", "x86"}: + arch = "386" + elif machine.startswith("arm"): + arch = "arm" + else: + raise RuntimeError(f"Unsupported CPU architecture for tunnel binary: {machine}") + + if system not in {"linux", "darwin", "windows"}: + raise RuntimeError(f"Unsupported operating system for tunnel binary: {system}") + return system, arch + + +def extract_named_members_from_tar(archive_path, destination_dir, member_names): + destination_dir.mkdir(parents=True, exist_ok=True) + extracted = {} + wanted = set(member_names) + with tarfile.open(archive_path, "r:*") as archive: + for member in archive.getmembers(): + member_name = Path(member.name).name + if member_name not in wanted or not member.isfile(): + continue + target = destination_dir / executable_name(member_name) + source = archive.extractfile(member) + if source is None: + continue + with source, target.open("wb") as handle: + shutil.copyfileobj(source, handle) + chmod_executable(target) + extracted[member_name] = target + missing = wanted - set(extracted) + if missing: + raise RuntimeError( + f"Archive {archive_path.name} did not contain expected binaries: " + f"{', '.join(sorted(missing))}." + ) + return extracted + + +def extract_named_members_from_zip(archive_path, destination_dir, member_names): + destination_dir.mkdir(parents=True, exist_ok=True) + extracted = {} + wanted = set(member_names) + with zipfile.ZipFile(archive_path) as archive: + for member in archive.infolist(): + member_name = Path(member.filename).name + normalized_name = member_name.removesuffix(".exe") + if normalized_name not in wanted or member.is_dir(): + continue + target = destination_dir / executable_name(normalized_name) + with archive.open(member) as source, target.open("wb") as handle: + shutil.copyfileobj(source, handle) + chmod_executable(target) + extracted[normalized_name] = target + missing = wanted - set(extracted) + if missing: + raise RuntimeError( + f"Archive {archive_path.name} did not contain expected binaries: " + f"{', '.join(sorted(missing))}." + ) + return extracted + + +class CliTunnelHelper(TunnelHelper): + label = "CLI tunnel" + + def __init__( + self, + *, + label, + binary, + port, + command, + url_pattern, + missing_binary_message, + shutdown_command=None, + timeout=30, + notify=None, + binary_resolver=None, + preflight=None, + ): + super().__init__(port, notify=notify) + self.label = label + self.binary = binary + self.command = command + self.url_pattern = url_pattern + self.missing_binary_message = missing_binary_message + self.shutdown_command = shutdown_command + self.timeout = timeout + self.tunnel_process = None + self.binary_path = None + self.binary_resolver = binary_resolver + self.preflight = preflight + self.command_prefix = [] + + def _extract_url(self, line): + match = self.url_pattern.search(line) + if not match: + return None + return match.group(1 if match.lastindex else 0).rstrip(".,)") + + def start(self): + binary_path = ( + self.binary_resolver(self._notify) + if callable(self.binary_resolver) + else shutil.which(self.binary) + ) + if not binary_path: + raise RuntimeError(self.missing_binary_message) + self.binary_path = binary_path + if callable(self.preflight): + preflight_result = self.preflight(binary_path, notify=self._notify) + if isinstance(preflight_result, dict): + self.command_prefix = list(preflight_result.get("command_prefix") or []) + + command = [binary_path, *self.command_prefix, *self.command[1:]] + self.notify_starting(self.label) + self.tunnel_process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + + output_queue = queue.Queue() + + def read_output(): + if self.tunnel_process is None or self.tunnel_process.stdout is None: + return + for line in self.tunnel_process.stdout: + output_queue.put(line) + + threading.Thread(target=read_output, daemon=True).start() + + deadline = time.monotonic() + self.timeout + recent_output = deque(maxlen=8) + while time.monotonic() < deadline: + try: + line = output_queue.get(timeout=0.1) + except queue.Empty: + if self.tunnel_process.poll() is not None: + break + continue + + cleaned_line = line.strip() + if cleaned_line: + recent_output.append(cleaned_line) + url = self._extract_url(cleaned_line) + if url: + self.tunnel_url = url + self.notify_url_ready(self.label, url) + return self.tunnel_url + + details = " ".join(recent_output) + if self.tunnel_process.poll() is not None: + raise RuntimeError( + f"{self.label} exited before it reported a remote URL." + + (f" Output: {details}" if details else "") + ) + self._terminate_process() + raise RuntimeError( + f"{self.label} did not report a remote URL within {self.timeout} seconds." + + (f" Output: {details}" if details else "") + ) + + def _terminate_process(self): + if not self.tunnel_process or self.tunnel_process.poll() is not None: + return + self.tunnel_process.terminate() + try: + self.tunnel_process.wait(timeout=8) + except subprocess.TimeoutExpired: + self.tunnel_process.kill() + self.tunnel_process.wait(timeout=3) + + def stop(self): + self._terminate_process() + + if self.shutdown_command: + binary_path = self.binary_path or shutil.which(self.binary) + if binary_path: + subprocess.run( + [binary_path, *self.command_prefix, *self.shutdown_command[1:]], + check=False, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + timeout=10, + ) + self.tunnel_url = None + self.notify_stopped(self.label) + return True diff --git a/helpers/cloudflare_tunnel._py b/helpers/cloudflare_tunnel._py index 5ef12de8e..3d0e825f2 100644 --- a/helpers/cloudflare_tunnel._py +++ b/helpers/cloudflare_tunnel._py @@ -120,11 +120,11 @@ class CloudflareTunnel: return def start(self): - """Starts the cloudflare tunnel""" + """Starts Cloudflare Tunnel""" if not self.cloudflared_path: self.download_cloudflared() - PrintStyle().print("\nStarting Cloudflare tunnel...") + PrintStyle().print("\nStarting Cloudflare Tunnel...") # Start tunnel process self.tunnel_process = subprocess.Popen( [ @@ -147,11 +147,11 @@ class CloudflareTunnel: ).start() def stop(self): - """Stops the cloudflare tunnel""" + """Stops Cloudflare Tunnel""" self._stop_event.set() if self.tunnel_process: - PrintStyle().print("\nStopping Cloudflare tunnel...") + PrintStyle().print("\nStopping Cloudflare Tunnel...") self.tunnel_process.terminate() self.tunnel_process.wait() self.tunnel_process = None - self.tunnel_url = None \ No newline at end of file + self.tunnel_url = None diff --git a/helpers/cloudflare_tunnel.py b/helpers/cloudflare_tunnel.py new file mode 100644 index 000000000..ba446f32a --- /dev/null +++ b/helpers/cloudflare_tunnel.py @@ -0,0 +1,11 @@ +from flaredantic import FlareConfig, FlareTunnel + +from helpers.tunnel_common import FlaredanticTunnelHelper + + +class CloudflareTunnel(FlaredanticTunnelHelper): + label = "Cloudflare Tunnel" + + def build_tunnel(self): + config = FlareConfig(port=self.port, verbose=True) + return FlareTunnel(config) diff --git a/helpers/microsoft_tunnel.py b/helpers/microsoft_tunnel.py new file mode 100644 index 000000000..40173e376 --- /dev/null +++ b/helpers/microsoft_tunnel.py @@ -0,0 +1,144 @@ +import getpass +import hashlib +import os +import socket + +from flaredantic import MicrosoftConfig, MicrosoftTunnel, NotifyEvent + +try: + from flaredantic.core.exceptions import MicrosoftTunnelError +except Exception: # pragma: no cover - keeps tests independent from package internals + MicrosoftTunnelError = RuntimeError + +from helpers import files +from helpers.tunnel_common import FlaredanticTunnelHelper + + +MICROSOFT_TUNNEL_ID_ENV_KEYS = ( + "A0_MICROSOFT_DEV_TUNNEL_ID", + "MICROSOFT_DEV_TUNNEL_ID", +) +MICROSOFT_TUNNEL_TIMEOUT = 120 + + +def default_microsoft_tunnel_id(): + for env_key in MICROSOFT_TUNNEL_ID_ENV_KEYS: + configured = (os.environ.get(env_key) or "").strip() + if configured: + return configured + + seed = "|".join([ + getpass.getuser(), + socket.gethostname(), + files.get_abs_path("usr"), + ]) + digest = hashlib.sha256(seed.encode("utf-8")).hexdigest()[:10] + return f"agent-zero-{digest}" + + +class AgentZeroMicrosoftTunnel(MicrosoftTunnel): + def notify(self, event, message, data=None): + try: + return super().notify(event, message, data) + except AttributeError: + self.agent_zero_notifications.append({ + "event": event.value if hasattr(event, "value") else event, + "message": message, + "data": data, + }) + return None + + @property + def agent_zero_notifications(self): + if not hasattr(self, "_agent_zero_notifications"): + self._agent_zero_notifications = [] + return self._agent_zero_notifications + + def _notify_progress(self, message, data=None): + self.notify(NotifyEvent.INFO, message, data) + + def _ensure_logged_in(self): + parent = getattr(super(), "_ensure_logged_in", None) + if callable(parent): + parent() + self._notify_progress( + "Microsoft Dev Tunnels login confirmed. Preparing your tunnel..." + ) + + def _ensure_tunnel(self): + tunnel_id = self.config.tunnel_id + port = str(self.config.port) + + self._notify_progress( + f"Checking Microsoft Dev Tunnel `{tunnel_id}`...", + {"tunnel_id": tunnel_id}, + ) + show = self._run_cmd(["show", tunnel_id]) + if show.returncode != 0: + self._notify_progress( + f"Creating Microsoft Dev Tunnel `{tunnel_id}`...", + {"tunnel_id": tunnel_id}, + ) + create = self._run_cmd(["create", tunnel_id]) + if create.returncode != 0: + raise MicrosoftTunnelError(f"Failed to create tunnel: {create.stdout}") + else: + self._notify_progress( + f"Microsoft Dev Tunnel `{tunnel_id}` already exists. Checking port {port}...", + {"tunnel_id": tunnel_id, "port": port}, + ) + + self._notify_progress( + f"Checking Microsoft Dev Tunnel port {port}...", + {"tunnel_id": tunnel_id, "port": port}, + ) + port_show = self._run_cmd(["port", "show", tunnel_id, "-p", port]) + if port_show.returncode != 0: + self._notify_progress( + f"Creating Microsoft Dev Tunnel port {port}...", + {"tunnel_id": tunnel_id, "port": port}, + ) + port_create = self._run_cmd([ + "port", + "create", + tunnel_id, + "-p", + port, + "--protocol", + "http", + ]) + if port_create.returncode != 0: + raise MicrosoftTunnelError( + f"Failed to create port: {port_create.stdout}" + ) + + self._notify_progress( + "Microsoft Dev Tunnel setup is ready. Starting the secure host..." + ) + + +class MicrosoftDevTunnel(FlaredanticTunnelHelper): + label = "Microsoft Dev Tunnels" + + def build_tunnel(self): + config = MicrosoftConfig( + port=self.port, + verbose=True, + timeout=MICROSOFT_TUNNEL_TIMEOUT, + tunnel_id=default_microsoft_tunnel_id(), + ) + return AgentZeroMicrosoftTunnel(config) + + def start(self): + try: + return super().start() + except Exception as e: + if "Timeout waiting for Microsoft Dev Tunnels URL" not in str(e): + raise + tunnel_id = default_microsoft_tunnel_id() + raise RuntimeError( + "Microsoft Dev Tunnels did not return a URL. Agent Zero uses " + f"the tunnel id `{tunnel_id}` to avoid flaredantic's global " + "`flaredantic` tunnel-id collision. If this still fails, set " + "`A0_MICROSOFT_DEV_TUNNEL_ID` to a fresh unique value and try again." + ) from e diff --git a/helpers/runtime.py b/helpers/runtime.py index 1d4b084e1..5751de5f7 100644 --- a/helpers/runtime.py +++ b/helpers/runtime.py @@ -31,7 +31,7 @@ def initialize(): "--cloudflare_tunnel", type=bool, default=False, - help="Use cloudflare tunnel for public URL", + help="Use Cloudflare Tunnel for public URL", ) parser.add_argument( "--development", type=bool, default=False, help="Development mode" diff --git a/helpers/serveo_tunnel.py b/helpers/serveo_tunnel.py new file mode 100644 index 000000000..5c5e36878 --- /dev/null +++ b/helpers/serveo_tunnel.py @@ -0,0 +1,11 @@ +from flaredantic import ServeoConfig, ServeoTunnel + +from helpers.tunnel_common import FlaredanticTunnelHelper + + +class ServeoTunnelHelper(FlaredanticTunnelHelper): + label = "Serveo" + + def build_tunnel(self): + config = ServeoConfig(port=self.port) + return ServeoTunnel(config) diff --git a/helpers/tailscale_tunnel.py b/helpers/tailscale_tunnel.py new file mode 100644 index 000000000..25e13361e --- /dev/null +++ b/helpers/tailscale_tunnel.py @@ -0,0 +1,433 @@ +from collections import deque +import json +import os +import queue +import re +import shutil +import subprocess +import threading +import time +import urllib.parse +import urllib.request +from pathlib import Path + +from flaredantic import NotifyEvent + +from helpers import cli_tunnel, files +from helpers.cli_tunnel import CliTunnelHelper + + +TAILSCALE_URL_RE = re.compile(r"https://[a-zA-Z0-9.-]+\.ts\.net[^\s\"']*") +TAILSCALE_LOGIN_URL_RE = re.compile(r"https://login\.tailscale\.com/[^\s\"']+") +TAILSCALE_STABLE_PACKAGES_URL = "https://pkgs.tailscale.com/stable/?v=latest" +TAILSCALE_UP_TIMEOUT = 180 +TAILSCALE_DAEMON_START_TIMEOUT = 12 +TAILSCALE_RUNTIME_DIR = Path(files.get_abs_path("tmp", "tailscale")) +TAILSCALE_STATE_DIR = Path(files.get_abs_path("usr", "tailscale")) +TAILSCALE_SOCKET_PATH = TAILSCALE_RUNTIME_DIR / "tailscaled.sock" +TAILSCALE_DAEMON_LOG_PATH = TAILSCALE_RUNTIME_DIR / "tailscaled.log" +TAILSCALE_DAEMON_PID_PATH = TAILSCALE_RUNTIME_DIR / "tailscaled.pid" + + +def tailscale_arch(): + _, arch = cli_tunnel.platform_parts() + return arch + + +def tailscale_archive_url(): + arch = tailscale_arch() + with urllib.request.urlopen(TAILSCALE_STABLE_PACKAGES_URL, timeout=30) as response: + html = response.read().decode("utf-8", errors="replace") + pattern = re.compile(rf'href="([^"]*tailscale_[^"]+_{re.escape(arch)}\.tgz)"') + match = pattern.search(html) + if not match: + raise RuntimeError(f"Could not find a Tailscale static binary for {arch}.") + return urllib.parse.urljoin(TAILSCALE_STABLE_PACKAGES_URL, match.group(1)) + + +def install_tailscale(notify=None): + existing = shutil.which("tailscale") + if existing and resolve_tailscaled_binary(existing): + return existing + + install_path = cli_tunnel.RUNTIME_BIN_DIR / cli_tunnel.executable_name("tailscale") + daemon_path = cli_tunnel.RUNTIME_BIN_DIR / cli_tunnel.executable_name("tailscaled") + if install_path.exists() and daemon_path.exists(): + return str(install_path) + + download_url = tailscale_archive_url() + archive_path = ( + cli_tunnel.RUNTIME_BIN_DIR + / urllib.parse.urlparse(download_url).path.rsplit("/", 1)[-1] + ) + cli_tunnel.download_file(download_url, archive_path, notify=notify) + try: + extracted = cli_tunnel.extract_named_members_from_tar( + archive_path, + cli_tunnel.RUNTIME_BIN_DIR, + {"tailscale", "tailscaled"}, + ) + return str(extracted["tailscale"]) + finally: + archive_path.unlink(missing_ok=True) + + +def resolve_tailscaled_binary(binary_path): + sibling = Path(binary_path).with_name(cli_tunnel.executable_name("tailscaled")) + if sibling.exists(): + return str(sibling) + return shutil.which("tailscaled") + + +def notify_info(notify, message, data=None): + if callable(notify): + notify(NotifyEvent.INFO, message, data) + + +def tailscale_socket_args(socket_path=None): + return ["--socket", str(socket_path)] if socket_path else [] + + +def tailscale_command(binary_path, args, socket_path=None): + return [binary_path, *tailscale_socket_args(socket_path), *args] + + +def tailscale_status(binary_path, socket_path=None): + return subprocess.run( + tailscale_command(binary_path, ["status", "--json"], socket_path), + check=False, + text=True, + capture_output=True, + timeout=12, + ) + + +def compact_output(lines): + return " ".join(line.strip() for line in lines if line and line.strip()) + + +def tailscale_daemon_hint(output): + lowered = output.lower() + return any( + marker in lowered + for marker in ( + "failed to connect to local tailscaled", + "tailscaled.sock", + "no such file or directory", + "connection refused", + "is tailscaled running", + ) + ) + + +def tailscale_daemon_ready(binary_path, socket_path): + completed = tailscale_status(binary_path, socket_path=socket_path) + output = compact_output([completed.stderr, completed.stdout]) + return not tailscale_daemon_hint(output) + + +def read_recent_tailscaled_log(): + if not TAILSCALE_DAEMON_LOG_PATH.exists(): + return "" + try: + lines = TAILSCALE_DAEMON_LOG_PATH.read_text( + encoding="utf-8", + errors="replace", + ).splitlines() + except OSError: + return "" + return compact_output(lines[-12:]) + + +def start_tailscaled(binary_path, notify=None): + daemon_path = resolve_tailscaled_binary(binary_path) + if not daemon_path: + raise RuntimeError( + "Tailscale was prepared, but Agent Zero could not find the " + "`tailscaled` daemon binary. Try Tailscale Remote Control again so " + "Agent Zero can re-download the static Tailscale package." + ) + + if tailscale_daemon_ready(binary_path, TAILSCALE_SOCKET_PATH): + return TAILSCALE_SOCKET_PATH + + TAILSCALE_RUNTIME_DIR.mkdir(parents=True, exist_ok=True) + TAILSCALE_STATE_DIR.mkdir(parents=True, exist_ok=True) + TAILSCALE_SOCKET_PATH.unlink(missing_ok=True) + + notify_info( + notify, + "Starting Tailscale's background service inside this container...", + ) + + log_handle = TAILSCALE_DAEMON_LOG_PATH.open("ab") + process = subprocess.Popen( + [ + daemon_path, + "--tun=userspace-networking", + "--socket", + str(TAILSCALE_SOCKET_PATH), + "--statedir", + str(TAILSCALE_STATE_DIR), + "--state", + str(TAILSCALE_STATE_DIR / "tailscaled.state"), + ], + stdout=log_handle, + stderr=subprocess.STDOUT, + stdin=subprocess.DEVNULL, + start_new_session=True, + close_fds=True, + ) + log_handle.close() + TAILSCALE_DAEMON_PID_PATH.write_text(str(process.pid), encoding="utf-8") + + deadline = time.monotonic() + TAILSCALE_DAEMON_START_TIMEOUT + while time.monotonic() < deadline: + if process.poll() is not None: + break + if tailscale_daemon_ready(binary_path, TAILSCALE_SOCKET_PATH): + notify_info(notify, "Tailscale background service is running.") + return TAILSCALE_SOCKET_PATH + time.sleep(0.25) + + details = read_recent_tailscaled_log() + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=3) + TAILSCALE_DAEMON_PID_PATH.unlink(missing_ok=True) + message = ( + "Agent Zero downloaded Tailscale, but could not start the `tailscaled` " + "background service in this container." + ) + if details: + message = f"{message} Details: {details}" + raise RuntimeError(message) + + +def stop_managed_tailscaled(): + if not TAILSCALE_DAEMON_PID_PATH.exists(): + return + try: + pid = int(TAILSCALE_DAEMON_PID_PATH.read_text(encoding="utf-8").strip()) + except Exception: + TAILSCALE_DAEMON_PID_PATH.unlink(missing_ok=True) + return + try: + cmdline = Path(f"/proc/{pid}/cmdline").read_bytes().decode( + "utf-8", + errors="replace", + ) + except Exception: + cmdline = "" + if "tailscaled" in cmdline: + try: + os.kill(pid, 15) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + try: + os.kill(pid, 0) + except ProcessLookupError: + break + time.sleep(0.1) + else: + os.kill(pid, 9) + except ProcessLookupError: + pass + except PermissionError: + pass + TAILSCALE_DAEMON_PID_PATH.unlink(missing_ok=True) + TAILSCALE_SOCKET_PATH.unlink(missing_ok=True) + + +def tailscale_up_failure_message(output, *, timed_out=False): + details = compact_output(output) + if tailscale_daemon_hint(details): + message = ( + "Agent Zero started Tailscale and ran `tailscale up`, but the " + "Tailscale background service stopped responding." + ) + elif timed_out: + message = ( + "Agent Zero ran `tailscale up`, but it did not finish before the " + "setup timeout. If a Tailscale login link was shown, approve this " + "container in your browser, then start Remote Control again." + ) + else: + message = ( + "Agent Zero ran `tailscale up`, but Tailscale did not finish joining " + "this container to your tailnet. Complete any Tailscale login or " + "admin approval it requested, then try Tailscale Remote Control again." + ) + if details: + message = f"{message} Details: {details}" + return message + + +def run_tailscale_up( + binary_path, + notify=None, + timeout=TAILSCALE_UP_TIMEOUT, + socket_path=None, +): + notify_info( + notify, + "Tailscale needs this container to join your tailnet. Running `tailscale up` now...", + ) + process = subprocess.Popen( + tailscale_command(binary_path, ["up"], socket_path), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + output_queue = queue.Queue() + + def read_output(): + if process.stdout is None: + return + for line in process.stdout: + output_queue.put(line) + + reader = threading.Thread(target=read_output, daemon=True) + reader.start() + deadline = time.monotonic() + timeout + recent_output = deque(maxlen=10) + login_announced = False + + def record_line(line): + nonlocal login_announced + cleaned = line.strip() + if not cleaned: + return + recent_output.append(cleaned) + login_match = TAILSCALE_LOGIN_URL_RE.search(cleaned) + if login_match and not login_announced: + login_url = login_match.group(0).rstrip(".,)") + notify_info( + notify, + "Open the Tailscale login link and approve this container. " + "Agent Zero will continue when Tailscale finishes setup.", + {"provider": "tailscale", "url": login_url}, + ) + login_announced = True + + while True: + try: + record_line(output_queue.get(timeout=0.1)) + except queue.Empty: + if process.poll() is not None: + reader.join(timeout=1) + while True: + try: + record_line(output_queue.get_nowait()) + except queue.Empty: + break + break + if time.monotonic() >= deadline: + process.terminate() + try: + process.wait(timeout=8) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=3) + reader.join(timeout=1) + while True: + try: + record_line(output_queue.get_nowait()) + except queue.Empty: + break + raise RuntimeError( + tailscale_up_failure_message(recent_output, timed_out=True) + ) + continue + + returncode = process.wait(timeout=3) + if returncode != 0: + raise RuntimeError(tailscale_up_failure_message(recent_output)) + + notify_info( + notify, + "Tailscale setup completed. Checking the tailnet connection...", + ) + + +def ensure_tailscale_ready(binary_path, notify=None): + socket_path = None + completed = tailscale_status(binary_path) + if completed.returncode != 0 and tailscale_daemon_hint( + compact_output([completed.stderr, completed.stdout]) + ): + socket_path = start_tailscaled(binary_path, notify=notify) + completed = tailscale_status(binary_path, socket_path=socket_path) + + if completed.returncode != 0: + run_tailscale_up(binary_path, notify=notify, socket_path=socket_path) + completed = tailscale_status(binary_path, socket_path=socket_path) + if completed.returncode != 0: + details = compact_output([completed.stderr, completed.stdout]) + message = ( + "Agent Zero ran `tailscale up`, but Tailscale status is still " + "not available. Tailscale may still need browser approval, admin " + "approval, or a running `tailscaled` service." + ) + if details: + message = f"{message} Details: {details}" + raise RuntimeError(message) + + try: + payload = json.loads(completed.stdout or "{}") + except json.JSONDecodeError: + return + + backend_state = str(payload.get("BackendState") or "").lower() + if backend_state and backend_state != "running": + run_tailscale_up(binary_path, notify=notify, socket_path=socket_path) + completed = tailscale_status(binary_path, socket_path=socket_path) + try: + payload = json.loads(completed.stdout or "{}") + except json.JSONDecodeError: + return {"command_prefix": tailscale_socket_args(socket_path)} + backend_state = str(payload.get("BackendState") or "").lower() + if backend_state and backend_state != "running": + raise RuntimeError( + "Agent Zero ran `tailscale up`, but this node is still not ready " + "for Tailscale Funnel " + f"(state: {backend_state}). Complete Tailscale sign-in or admin " + "approval, then try again." + ) + + return {"command_prefix": tailscale_socket_args(socket_path)} + + +class TailscaleTunnel(CliTunnelHelper): + def __init__(self, port, notify=None): + target = f"http://127.0.0.1:{port}" + super().__init__( + label="Tailscale Funnel", + binary="tailscale", + port=port, + command=["tailscale", "funnel", "--yes", target], + url_pattern=TAILSCALE_URL_RE, + missing_binary_message=( + "Tailscale could not be prepared in this environment. Install " + "Tailscale, make sure the `tailscaled` service is available to " + "this container, enable Funnel for the tailnet, then try again." + ), + shutdown_command=["tailscale", "funnel", "--yes", target, "off"], + binary_resolver=lambda notify_callback: install_tailscale( + notify=notify_callback + ), + preflight=ensure_tailscale_ready, + notify=notify, + ) + + def stop(self): + try: + return super().stop() + finally: + if self.command_prefix: + stop_managed_tailscaled() diff --git a/helpers/tunnel_common.py b/helpers/tunnel_common.py new file mode 100644 index 000000000..9daa0a972 --- /dev/null +++ b/helpers/tunnel_common.py @@ -0,0 +1,56 @@ +from flaredantic import NotifyEvent + + +def event_value(event): + return event.value if hasattr(event, "value") else event + + +class TunnelHelper: + def __init__(self, port, notify=None): + self.port = port + self.notify_callback = notify + self.tunnel = None + self.tunnel_url = None + + def _notify(self, event, message, data=None): + if callable(self.notify_callback): + self.notify_callback(event, message, data) + + def notify_starting(self, label): + self._notify( + NotifyEvent.CREATING_TUNNEL, + f"Starting {label} on port {self.port}...", + ) + + def notify_url_ready(self, label, url): + self._notify( + NotifyEvent.TUNNEL_URL, + f"{label} URL is ready", + {"url": url}, + ) + + def notify_stopped(self, label): + self._notify(NotifyEvent.TUNNEL_STOPPED, f"{label} stopped") + + +class FlaredanticTunnelHelper(TunnelHelper): + label = "Remote Control" + + def build_tunnel(self): + raise NotImplementedError + + def start(self): + self.notify_starting(self.label) + self.tunnel = self.build_tunnel() + self.tunnel.start() + self.tunnel_url = getattr(self.tunnel, "tunnel_url", None) + if self.tunnel_url: + self.notify_url_ready(self.label, self.tunnel_url) + return self.tunnel_url + + def stop(self): + if self.tunnel: + self.tunnel.stop() + self.tunnel_url = None + self.notify_stopped(self.label) + return True diff --git a/helpers/tunnel_manager.py b/helpers/tunnel_manager.py index 4a2ddbeb2..984fd841a 100644 --- a/helpers/tunnel_manager.py +++ b/helpers/tunnel_manager.py @@ -1,13 +1,42 @@ -from flaredantic import ( - FlareTunnel, FlareConfig, - ServeoConfig, ServeoTunnel, - MicrosoftTunnel, MicrosoftConfig, - notifier, NotifyData, NotifyEvent -) import threading +import time from collections import deque +from flaredantic import NotifyData, NotifyEvent, notifier + +from helpers.cloudflare_tunnel import CloudflareTunnel +from helpers.microsoft_tunnel import MicrosoftDevTunnel from helpers.print_style import PrintStyle +from helpers.serveo_tunnel import ServeoTunnelHelper +from helpers.tailscale_tunnel import TailscaleTunnel +from helpers.tunnel_common import event_value + + +SUPPORTED_TUNNEL_PROVIDERS = { + "cloudflared", + "microsoft", + "serveo", + "tailscale", +} +TUNNEL_PROVIDER_ALIASES = { + "cloudflare": "cloudflared", + "cloudflare_tunnel": "cloudflared", + "cloudflare-tunnel": "cloudflared", + "tailscale_funnel": "tailscale", + "tailscale-funnel": "tailscale", +} + + +def normalize_provider(provider): + normalized = (provider or "serveo").strip().lower() + normalized = TUNNEL_PROVIDER_ALIASES.get(normalized, normalized) + if normalized not in SUPPORTED_TUNNEL_PROVIDERS: + supported = ", ".join(sorted(SUPPORTED_TUNNEL_PROVIDERS)) + raise ValueError( + f"Unsupported remote control provider '{provider}'. Choose one of: {supported}." + ) + return normalized + # Singleton to manage the tunnel instance class TunnelManager: @@ -30,80 +59,86 @@ class TunnelManager: self._subscribed = False def _on_notify(self, data: NotifyData): - """Handle notifications from flaredantic""" + """Handle notifications from flaredantic.""" self.notifications.append({ "event": data.event.value, "message": data.message, - "data": data.data + "data": data.data, }) def _ensure_subscribed(self): - """Subscribe to flaredantic notifications if not already""" + """Subscribe to flaredantic notifications if not already.""" if not self._subscribed: notifier.subscribe(self._on_notify) self._subscribed = True def get_notifications(self): - """Get and clear pending notifications""" + """Get and clear pending notifications.""" notifications = list(self.notifications) self.notifications.clear() return notifications def get_last_error(self): - """Check for recent error in notifications without clearing""" - for n in reversed(list(self.notifications)): - if n['event'] == NotifyEvent.ERROR.value: - return n['message'] + """Check for recent error in notifications without clearing.""" + for notification in reversed(list(self.notifications)): + if notification["event"] == NotifyEvent.ERROR.value: + return notification["message"] return None + def _append_notification(self, event, message, data=None): + self.notifications.append({ + "event": event_value(event), + "message": message, + "data": data, + }) + + def _create_tunnel(self, port, provider): + if provider == "cloudflared": + return CloudflareTunnel(port, notify=self._append_notification) + if provider == "microsoft": + return MicrosoftDevTunnel(port, notify=self._append_notification) + if provider == "tailscale": + return TailscaleTunnel(port, notify=self._append_notification) + return ServeoTunnelHelper(port, notify=self._append_notification) + def start_tunnel(self, port=80, provider="serveo"): - """Start a new tunnel or return the existing one's URL""" + """Start a new tunnel or return the existing one's URL.""" if self.is_running and self.tunnel_url: return self.tunnel_url - self.provider = provider self._ensure_subscribed() self.notifications.clear() + try: + self.provider = normalize_provider(provider) + except Exception as e: + error_msg = str(e) + PrintStyle.error(f"Error starting tunnel: {error_msg}") + self._append_notification(NotifyEvent.ERROR, error_msg) + return None try: - # Start tunnel in a separate thread to avoid blocking + # Start tunnel in a separate thread to avoid blocking. def run_tunnel(): try: - if self.provider == "cloudflared": - config = FlareConfig(port=port, verbose=True) - self.tunnel = FlareTunnel(config) - elif self.provider == "microsoft": - config = MicrosoftConfig(port=port, verbose=True) # type: ignore - self.tunnel = MicrosoftTunnel(config) - else: # Default to serveo - config = ServeoConfig(port=port) # type: ignore - self.tunnel = ServeoTunnel(config) - + self.tunnel = self._create_tunnel(port, self.provider) self.tunnel.start() self.tunnel_url = self.tunnel.tunnel_url self.is_running = True except Exception as e: error_msg = str(e) PrintStyle.error(f"Error in tunnel thread: {error_msg}") - self.notifications.append({ - "event": NotifyEvent.ERROR.value, - "message": error_msg, - "data": None - }) + self._append_notification(NotifyEvent.ERROR, error_msg) tunnel_thread = threading.Thread(target=run_tunnel) tunnel_thread.daemon = True tunnel_thread.start() - # Wait for tunnel to start (no timeout - user may need time for login) - import time + # No timeout: Microsoft login can legitimately require user interaction. while True: if self.tunnel_url: break - # Check if we have errors - if any(n['event'] == NotifyEvent.ERROR.value for n in self.notifications): + if any(n["event"] == NotifyEvent.ERROR.value for n in self.notifications): break - # Check if thread died without producing URL if not tunnel_thread.is_alive(): break time.sleep(0.1) @@ -114,7 +149,7 @@ class TunnelManager: return None def stop_tunnel(self): - """Stop the running tunnel""" + """Stop the running tunnel.""" if self.tunnel and self.is_running: try: self.tunnel.stop() @@ -127,5 +162,5 @@ class TunnelManager: return False def get_tunnel_url(self): - """Get the current tunnel URL if available""" + """Get the current tunnel URL if available.""" return self.tunnel_url if self.is_running else None diff --git a/tests/test_tunnel_remote_link.py b/tests/test_tunnel_remote_link.py new file mode 100644 index 000000000..0e8fdbdf4 --- /dev/null +++ b/tests/test_tunnel_remote_link.py @@ -0,0 +1,757 @@ +import io +import enum +import importlib +import os +import sys +import tarfile +import types +import zipfile +from pathlib import Path + +import pytest + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + + +class FakeNotifyEvent(enum.Enum): + DOWNLOADING = "downloading" + DOWNLOAD_PROGRESS = "download_progress" + DOWNLOAD_COMPLETE = "download_complete" + CREATING_TUNNEL = "creating_tunnel" + TUNNEL_URL = "tunnel_url" + TUNNEL_STOPPED = "tunnel_stopped" + ERROR = "error" + INFO = "info" + + +class FakeConfig: + def __init__(self, **kwargs): + self.kwargs = kwargs + for key, value in kwargs.items(): + setattr(self, key, value) + + +class FakeTunnel: + def __init__(self, config): + self.config = config + self.tunnel_url = "" + self.stopped = False + self.notifications = [] + + def start(self): + self.tunnel_url = "https://example.test" + return self.tunnel_url + + def stop(self): + self.stopped = True + return True + + def notify(self, event, message, data=None): + self.notifications.append({ + "event": event.value if hasattr(event, "value") else event, + "message": message, + "data": data, + }) + + +class FakeNotifier: + def subscribe(self, callback): + self.callback = callback + + +def write_tar_archive(path, members): + path.parent.mkdir(parents=True, exist_ok=True) + with tarfile.open(path, "w:gz") as archive: + for name, content in members.items(): + payload = content.encode("utf-8") + info = tarfile.TarInfo(name) + info.size = len(payload) + info.mode = 0o755 + archive.addfile(info, io.BytesIO(payload)) + + +def write_zip_archive(path, members): + path.parent.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(path, "w") as archive: + for name, content in members.items(): + archive.writestr(name, content) + + +HELPER_MODULES = [ + "helpers.cli_tunnel", + "helpers.cloudflare_tunnel", + "helpers.microsoft_tunnel", + "helpers.serveo_tunnel", + "helpers.tailscale_tunnel", + "helpers.tunnel_common", + "helpers.tunnel_manager", +] + + +def remote_link_modules(): + return types.SimpleNamespace( + cli=importlib.import_module("helpers.cli_tunnel"), + microsoft=importlib.import_module("helpers.microsoft_tunnel"), + tailscale=importlib.import_module("helpers.tailscale_tunnel"), + ) + + +@pytest.fixture() +def tunnel_manager_module(monkeypatch): + fake_flaredantic = types.SimpleNamespace( + FlareConfig=FakeConfig, + FlareTunnel=FakeTunnel, + MicrosoftConfig=FakeConfig, + MicrosoftTunnel=FakeTunnel, + NotifyData=object, + NotifyEvent=FakeNotifyEvent, + ServeoConfig=FakeConfig, + ServeoTunnel=FakeTunnel, + notifier=FakeNotifier(), + ) + monkeypatch.setitem(sys.modules, "flaredantic", fake_flaredantic) + helpers_package = sys.modules.get("helpers") + for module_name in HELPER_MODULES: + sys.modules.pop(module_name, None) + if helpers_package and hasattr(helpers_package, module_name.rsplit(".", 1)[-1]): + delattr(helpers_package, module_name.rsplit(".", 1)[-1]) + module = importlib.import_module("helpers.tunnel_manager") + yield module + for module_name in HELPER_MODULES: + sys.modules.pop(module_name, None) + if helpers_package and hasattr(helpers_package, module_name.rsplit(".", 1)[-1]): + delattr(helpers_package, module_name.rsplit(".", 1)[-1]) + + +def test_remote_link_provider_options_match_supported_remote_link_providers(): + html = ( + PROJECT_ROOT / "webui/components/settings/tunnel/tunnel-section.html" + ).read_text(encoding="utf-8") + + assert html.count("' in html + assert '' in html + assert '' in html + assert '' in html + assert '' not in html + assert "Cloudflare Tunnel is the quickest shareable URL" in html + assert "Agent Zero will start its sign-in flow" in html + + +def test_tunnel_provider_normalization_preserves_aliases(tunnel_manager_module): + manager = tunnel_manager_module + + assert manager.normalize_provider("cloudflare") == "cloudflared" + assert manager.normalize_provider("Cloudflare-Tunnel") == "cloudflared" + assert manager.normalize_provider("tailscale-funnel") == "tailscale" + + with pytest.raises(ValueError, match="Unsupported remote control provider"): + manager.normalize_provider("lantern") + + +def test_tailscale_cli_commands_are_wired(tunnel_manager_module): + manager = tunnel_manager_module.TunnelManager() + + tailscale = manager._create_tunnel(50001, "tailscale") + + assert tailscale.command == [ + "tailscale", + "funnel", + "--yes", + "http://127.0.0.1:50001", + ] + assert tailscale.shutdown_command == [ + "tailscale", + "funnel", + "--yes", + "http://127.0.0.1:50001", + "off", + ] + + +def test_remote_link_providers_have_dedicated_helper_modules(): + helper_files = { + "cloudflared": PROJECT_ROOT / "helpers/cloudflare_tunnel.py", + "microsoft": PROJECT_ROOT / "helpers/microsoft_tunnel.py", + "serveo": PROJECT_ROOT / "helpers/serveo_tunnel.py", + "tailscale": PROJECT_ROOT / "helpers/tailscale_tunnel.py", + } + + assert all(path.exists() for path in helper_files.values()) + + +def test_microsoft_dev_tunnel_uses_unique_a0_tunnel_id(tunnel_manager_module): + manager = tunnel_manager_module.TunnelManager() + tunnel = manager._create_tunnel(50001, "microsoft") + + assert tunnel.start() == "https://example.test" + + config = tunnel.tunnel.config.kwargs + assert config["tunnel_id"].startswith("agent-zero-") + assert config["tunnel_id"] != "flaredantic" + assert config["timeout"] == 120 + + +def test_microsoft_dev_tunnel_id_can_be_overridden( + tunnel_manager_module, + monkeypatch, +): + modules = remote_link_modules() + monkeypatch.setenv("A0_MICROSOFT_DEV_TUNNEL_ID", "agent-zero-custom") + + assert modules.microsoft.default_microsoft_tunnel_id() == "agent-zero-custom" + + +def test_microsoft_dev_tunnel_timeout_error_is_enriched( + tunnel_manager_module, +): + modules = remote_link_modules() + + class FailingMicrosoftTunnel: + def __init__(self, config): + self.config = config + + def start(self): + raise RuntimeError("Timeout waiting for Microsoft Dev Tunnels URL") + + def stop(self): + return None + + modules.microsoft.AgentZeroMicrosoftTunnel = FailingMicrosoftTunnel + tunnel = modules.microsoft.MicrosoftDevTunnel(50001) + + with pytest.raises(RuntimeError, match="global `flaredantic` tunnel-id collision"): + tunnel.start() + + +def test_microsoft_dev_tunnel_emits_setup_progress_notifications( + tunnel_manager_module, +): + modules = remote_link_modules() + config = FakeConfig(port=80, tunnel_id="agent-zero-test") + tunnel = modules.microsoft.AgentZeroMicrosoftTunnel(config) + commands = [] + + def fake_run_cmd(args): + commands.append(args) + if args[0] == "show" or args[:2] == ["port", "show"]: + return types.SimpleNamespace(returncode=1, stdout="missing") + return types.SimpleNamespace(returncode=0, stdout="ok") + + tunnel._run_cmd = fake_run_cmd + + tunnel._ensure_tunnel() + + messages = [notification["message"] for notification in tunnel.notifications] + assert messages == [ + "Checking Microsoft Dev Tunnel `agent-zero-test`...", + "Creating Microsoft Dev Tunnel `agent-zero-test`...", + "Checking Microsoft Dev Tunnel port 80...", + "Creating Microsoft Dev Tunnel port 80...", + "Microsoft Dev Tunnel setup is ready. Starting the secure host...", + ] + assert commands == [ + ["show", "agent-zero-test"], + ["create", "agent-zero-test"], + ["port", "show", "agent-zero-test", "-p", "80"], + ["port", "create", "agent-zero-test", "-p", "80", "--protocol", "http"], + ] + + +@pytest.mark.parametrize( + ("provider", "expected_label"), + [ + ("cloudflared", "Cloudflare Tunnel"), + ("microsoft", "Microsoft Dev Tunnels"), + ("serveo", "Serveo"), + ], +) +def test_flaredantic_provider_helpers_emit_manager_notifications( + tunnel_manager_module, + provider, + expected_label, +): + manager = tunnel_manager_module.TunnelManager() + tunnel = manager._create_tunnel(50001, provider) + + assert tunnel.start() == "https://example.test" + + assert manager.notifications[0] == { + "event": "creating_tunnel", + "message": f"Starting {expected_label} on port 50001...", + "data": None, + } + assert manager.notifications[-1] == { + "event": "tunnel_url", + "message": f"{expected_label} URL is ready", + "data": {"url": "https://example.test"}, + } + + +def test_zip_extraction_accepts_windows_exe_members( + tunnel_manager_module, + monkeypatch, + tmp_path, +): + modules = remote_link_modules() + archive_path = tmp_path / "tailscale.zip" + destination = tmp_path / "bin" + write_zip_archive(archive_path, {"tailscale.exe": "binary"}) + monkeypatch.setattr(modules.cli.platform, "system", lambda: "Windows") + + extracted = modules.cli.extract_named_members_from_zip( + archive_path, + destination, + {"tailscale"}, + ) + + assert extracted == {"tailscale": destination / "tailscale.exe"} + assert (destination / "tailscale.exe").read_text(encoding="utf-8") == "binary" + + +def test_tailscale_installs_runtime_binaries_from_static_archive( + tunnel_manager_module, + monkeypatch, + tmp_path, +): + modules = remote_link_modules() + monkeypatch.setattr(modules.tailscale.shutil, "which", lambda binary: None) + monkeypatch.setattr(modules.cli, "RUNTIME_BIN_DIR", tmp_path / "bin") + monkeypatch.setattr( + modules.tailscale, + "tailscale_archive_url", + lambda: "https://pkgs.tailscale.com/stable/tailscale_1.84.0_amd64.tgz", + ) + + def fake_download(url, destination, notify=None): + write_tar_archive( + destination, + { + "tailscale_1.84.0_amd64/tailscale": "#!/bin/sh\n", + "tailscale_1.84.0_amd64/tailscaled": "#!/bin/sh\n", + }, + ) + return destination + + monkeypatch.setattr(modules.cli, "download_file", fake_download) + + installed = Path(modules.tailscale.install_tailscale()) + + assert installed == tmp_path / "bin" / "tailscale" + assert (tmp_path / "bin" / "tailscaled").exists() + assert os.access(tmp_path / "bin" / "tailscale", os.X_OK) + assert os.access(tmp_path / "bin" / "tailscaled", os.X_OK) + assert not (tmp_path / "bin" / "tailscale_1.84.0_amd64.tgz").exists() + + +def test_tailscale_installer_downloads_static_pair_when_system_daemon_is_missing( + tunnel_manager_module, + monkeypatch, + tmp_path, +): + modules = remote_link_modules() + monkeypatch.setattr( + modules.tailscale.shutil, + "which", + lambda binary: "/usr/bin/tailscale" if binary == "tailscale" else None, + ) + monkeypatch.setattr(modules.cli, "RUNTIME_BIN_DIR", tmp_path / "bin") + monkeypatch.setattr( + modules.tailscale, + "tailscale_archive_url", + lambda: "https://pkgs.tailscale.com/stable/tailscale_1.84.0_amd64.tgz", + ) + + def fake_download(url, destination, notify=None): + write_tar_archive( + destination, + { + "tailscale_1.84.0_amd64/tailscale": "#!/bin/sh\n", + "tailscale_1.84.0_amd64/tailscaled": "#!/bin/sh\n", + }, + ) + return destination + + monkeypatch.setattr(modules.cli, "download_file", fake_download) + + assert modules.tailscale.install_tailscale() == str(tmp_path / "bin" / "tailscale") + assert (tmp_path / "bin" / "tailscaled").exists() + + +def test_tailscale_static_package_url_is_discovered_from_official_listing( + tunnel_manager_module, + monkeypatch, +): + modules = remote_link_modules() + html = ( + 'arm' + 'amd64' + ) + + class FakeResponse: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback): + return None + + def read(self): + return html.encode("utf-8") + + monkeypatch.setattr(modules.tailscale, "tailscale_arch", lambda: "amd64") + monkeypatch.setattr( + modules.tailscale.urllib.request, + "urlopen", + lambda url, timeout=30: FakeResponse(), + ) + + assert ( + modules.tailscale.tailscale_archive_url() + == "https://pkgs.tailscale.com/stable/tailscale_1.84.0_amd64.tgz" + ) + + +def test_tar_extraction_sanitizes_member_paths( + tunnel_manager_module, + tmp_path, +): + modules = remote_link_modules() + archive_path = tmp_path / "tailscale.tgz" + destination = tmp_path / "safe" + write_tar_archive(archive_path, {"../tailscale": "binary"}) + + extracted = modules.cli.extract_named_members_from_tar( + archive_path, + destination, + {"tailscale"}, + ) + + assert extracted == {"tailscale": destination / "tailscale"} + assert (destination / "tailscale").read_text(encoding="utf-8") == "binary" + assert not (tmp_path / "tailscale").exists() + + +@pytest.mark.parametrize( + ("provider", "module_name", "installer_name", "expected_message"), + [ + ("tailscale", "tailscale", "install_tailscale", "tailscale download failed"), + ], +) +def test_cli_provider_installer_failures_return_actionable_error( + tunnel_manager_module, + monkeypatch, + provider, + module_name, + installer_name, + expected_message, +): + manager_module = tunnel_manager_module + modules = remote_link_modules() + manager = manager_module.TunnelManager() + + def fail_install(notify=None): + raise RuntimeError(expected_message) + + monkeypatch.setattr(getattr(modules, module_name), installer_name, fail_install) + + assert manager.start_tunnel(port=50001, provider=provider) is None + assert expected_message in manager.get_last_error() + + +def test_tailscale_preflight_starts_managed_daemon_then_runs_up_with_socket( + tunnel_manager_module, + monkeypatch, + tmp_path, +): + modules = remote_link_modules() + runtime_dir = tmp_path / "runtime" + state_dir = tmp_path / "state" + socket_path = runtime_dir / "tailscaled.sock" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + tailscale = bin_dir / "tailscale" + tailscaled = bin_dir / "tailscaled" + tailscale.write_text("#!/bin/sh\n", encoding="utf-8") + tailscaled.write_text("#!/bin/sh\n", encoding="utf-8") + monkeypatch.setattr(modules.tailscale, "TAILSCALE_RUNTIME_DIR", runtime_dir) + monkeypatch.setattr(modules.tailscale, "TAILSCALE_STATE_DIR", state_dir) + monkeypatch.setattr(modules.tailscale, "TAILSCALE_SOCKET_PATH", socket_path) + monkeypatch.setattr( + modules.tailscale, + "TAILSCALE_DAEMON_LOG_PATH", + runtime_dir / "tailscaled.log", + ) + monkeypatch.setattr( + modules.tailscale, + "TAILSCALE_DAEMON_PID_PATH", + runtime_dir / "tailscaled.pid", + ) + status_results = iter([ + types.SimpleNamespace( + returncode=1, + stdout="", + stderr="failed to connect to local tailscaled", + ), + types.SimpleNamespace( + returncode=1, + stdout="", + stderr="failed to connect to local tailscaled", + ), + types.SimpleNamespace(returncode=1, stdout="Logged out.", stderr=""), + types.SimpleNamespace(returncode=1, stdout="Logged out.", stderr=""), + types.SimpleNamespace( + returncode=0, + stdout='{"BackendState": "Running"}', + stderr="", + ), + ]) + status_calls = [] + + def fake_status(binary_path, socket_path=None): + status_calls.append((binary_path, socket_path)) + return next(status_results) + + monkeypatch.setattr(modules.tailscale, "tailscale_status", fake_status) + popen_commands = [] + + class FakeProcess: + def __init__(self, command, **kwargs): + popen_commands.append(command) + self.pid = 12345 + self.is_tailscale_up = command[0] == str(tailscale) + self.stdout = iter(["Success.\n"]) if command[0] == str(tailscale) else None + + def poll(self): + return 0 if self.is_tailscale_up else None + + def terminate(self): + return None + + def wait(self, timeout=None): + return 0 + + monkeypatch.setattr(modules.tailscale.subprocess, "Popen", FakeProcess) + + result = modules.tailscale.ensure_tailscale_ready(str(tailscale)) + + assert result == {"command_prefix": ["--socket", str(socket_path)]} + assert status_calls == [ + (str(tailscale), None), + (str(tailscale), socket_path), + (str(tailscale), socket_path), + (str(tailscale), socket_path), + (str(tailscale), socket_path), + ] + assert popen_commands == [ + [ + str(tailscaled), + "--tun=userspace-networking", + "--socket", + str(socket_path), + "--statedir", + str(state_dir), + "--state", + str(state_dir / "tailscaled.state"), + ], + [str(tailscale), "--socket", str(socket_path), "up"], + ] + + +def test_tailscale_preflight_emits_login_url_from_tailscale_up( + tunnel_manager_module, + monkeypatch, +): + modules = remote_link_modules() + monkeypatch.setattr( + modules.tailscale, + "tailscale_status", + lambda binary_path, socket_path=None: types.SimpleNamespace( + returncode=1, + stdout="", + stderr="", + ), + ) + + class FakeProcess: + stdout = iter( + [ + "To authenticate, visit:\n", + "https://login.tailscale.com/a/abcdef\n", + ] + ) + + def poll(self): + return 1 + + def terminate(self): + return None + + def wait(self, timeout=None): + return 1 + + notifications = [] + monkeypatch.setattr( + modules.tailscale.subprocess, + "Popen", + lambda *args, **kwargs: FakeProcess(), + ) + + with pytest.raises(RuntimeError, match="joining this container to your tailnet"): + modules.tailscale.ensure_tailscale_ready( + "/tmp/tailscale", + notify=lambda event, message, data=None: notifications.append( + {"event": event.value, "message": message, "data": data} + ), + ) + + assert notifications[0]["message"] == ( + "Tailscale needs this container to join your tailnet. Running `tailscale up` now..." + ) + assert notifications[1] == { + "event": "info", + "message": ( + "Open the Tailscale login link and approve this container. " + "Agent Zero will continue when Tailscale finishes setup." + ), + "data": { + "provider": "tailscale", + "url": "https://login.tailscale.com/a/abcdef", + }, + } + + +def test_tailscale_preflight_runs_up_then_accepts_running_status( + tunnel_manager_module, + monkeypatch, +): + modules = remote_link_modules() + status_results = iter([ + types.SimpleNamespace(returncode=1, stdout="", stderr="not logged in"), + types.SimpleNamespace( + returncode=0, + stdout='{"BackendState": "Running"}', + stderr="", + ), + ]) + notifications = [] + + monkeypatch.setattr( + modules.tailscale, + "tailscale_status", + lambda binary_path, socket_path=None: next(status_results), + ) + + class FakeProcess: + stdout = iter(["Success.\n"]) + + def poll(self): + return 0 + + def wait(self, timeout=None): + return 0 + + monkeypatch.setattr( + modules.tailscale.subprocess, + "Popen", + lambda *args, **kwargs: FakeProcess(), + ) + + modules.tailscale.ensure_tailscale_ready( + "/tmp/tailscale", + notify=lambda event, message, data=None: notifications.append( + {"event": event.value, "message": message, "data": data} + ), + ) + + assert notifications[-1] == { + "event": "info", + "message": "Tailscale setup completed. Checking the tailnet connection...", + "data": None, + } + + +def test_cli_tunnel_preflight_prefix_is_used_for_start_and_shutdown( + tunnel_manager_module, + monkeypatch, +): + modules = remote_link_modules() + popen_commands = [] + run_commands = [] + + class FakeProcess: + def __init__(self, command, **kwargs): + popen_commands.append(command) + self.stdout = iter(["https://agent-zero.ts.net\n"]) + + def poll(self): + return None + + def terminate(self): + return None + + def wait(self, timeout=None): + return 0 + + monkeypatch.setattr(modules.cli.subprocess, "Popen", FakeProcess) + monkeypatch.setattr( + modules.cli.subprocess, + "run", + lambda command, **kwargs: run_commands.append(command) + or types.SimpleNamespace(returncode=0), + ) + + tunnel = modules.cli.CliTunnelHelper( + label="Tailscale Funnel", + binary="tailscale", + port=50001, + command=["tailscale", "funnel", "--yes", "http://127.0.0.1:50001"], + shutdown_command=[ + "tailscale", + "funnel", + "--yes", + "http://127.0.0.1:50001", + "off", + ], + url_pattern=modules.tailscale.TAILSCALE_URL_RE, + missing_binary_message="missing", + binary_resolver=lambda notify=None: "/tmp/tailscale", + preflight=lambda binary_path, notify=None: { + "command_prefix": ["--socket", "/tmp/tailscaled.sock"] + }, + ) + + assert tunnel.start() == "https://agent-zero.ts.net" + tunnel.stop() + + assert popen_commands == [ + [ + "/tmp/tailscale", + "--socket", + "/tmp/tailscaled.sock", + "funnel", + "--yes", + "http://127.0.0.1:50001", + ] + ] + assert run_commands == [ + [ + "/tmp/tailscale", + "--socket", + "/tmp/tailscaled.sock", + "funnel", + "--yes", + "http://127.0.0.1:50001", + "off", + ] + ] diff --git a/webui/components/settings/external/external-settings.html b/webui/components/settings/external/external-settings.html index 067fc1674..a9d7416dd 100644 --- a/webui/components/settings/external/external-settings.html +++ b/webui/components/settings/external/external-settings.html @@ -41,8 +41,8 @@
  • - Remote Link - Remote Link + Remote Control + Remote Control
    • diff --git a/webui/components/settings/settings-store.js b/webui/components/settings/settings-store.js index bd45b35c4..1b74622a4 100644 --- a/webui/components/settings/settings-store.js +++ b/webui/components/settings/settings-store.js @@ -47,7 +47,7 @@ const TAB_ITEMS = Object.freeze([ { id: "section-secrets", label: "Secrets", icon: "lock" }, { id: "section-auth", label: "Authentication", icon: "passkey" }, { id: "section-external-api", label: "External API", icon: "api" }, - { id: "section-tunnel", label: "Remote Link", icon: "share" }, + { id: "section-tunnel", label: "Remote Control", icon: "share" }, ], }, { diff --git a/webui/components/settings/tunnel/remote-link.html b/webui/components/settings/tunnel/remote-link.html index 65f42ae05..d0f21dc80 100644 --- a/webui/components/settings/tunnel/remote-link.html +++ b/webui/components/settings/tunnel/remote-link.html @@ -1,6 +1,6 @@ - Remote Link + Remote Control diff --git a/webui/components/settings/tunnel/tunnel-section.html b/webui/components/settings/tunnel/tunnel-section.html index 4c13d1743..6dc0b3e5f 100644 --- a/webui/components/settings/tunnel/tunnel-section.html +++ b/webui/components/settings/tunnel/tunnel-section.html @@ -1,6 +1,6 @@ - Remote Link + Remote Control @@ -20,23 +20,23 @@ share
    -
    Remote Link
    +
    Remote Control
    - Create a temporary HTTPS link for this Agent Zero instance, then open it from another browser or scan it from your phone. + Create temporary HTTPS access for this Agent Zero instance, then open it from another browser or scan it from your phone.
    - Use sign-in before sharing a remote link. The URL can reach your instance while the link is active. + Use sign-in before sharing remote control access. The URL can reach your instance while access is active.
    Link provider
    - Cloudflare is the quickest path for a shareable URL. Microsoft Dev Tunnels may ask you to sign in. + Cloudflare Tunnel is the quickest shareable URL. Tailscale is private to your tailnet and Agent Zero will start its sign-in flow when needed. Microsoft Dev Tunnels may ask you to approve a GitHub device login.
    @@ -45,7 +45,8 @@ x-model="$store.tunnelStore.provider" :disabled="$store.tunnelStore.isLoading" > - + + @@ -54,22 +55,22 @@ -
    @@ -430,20 +429,6 @@
    -
    -
    -
    Run LiteParse in subprocess
    -
    - Isolate LiteParse native runtime failures from the Web UI process. -
    -
    -
    - -
    -
    diff --git a/tests/test_document_query_plugin.py b/tests/test_document_query_plugin.py index c1d3e9484..c03475b47 100644 --- a/tests/test_document_query_plugin.py +++ b/tests/test_document_query_plugin.py @@ -10,6 +10,7 @@ from plugins._document_query.helpers.fetch import FetchedDocument, fetch_public_ from plugins._document_query.helpers.document_query import DocumentQueryHelper from plugins._document_query.helpers.parsers.base import BaseParser from plugins._document_query.helpers.parsers import get_parsers_for_mimetype +from plugins._document_query.helpers.parsers import liteparse as liteparse_module from plugins._document_query.helpers.parsers.liteparse import LiteParseParser from plugins._document_query.helpers.parsers.text import TextParser @@ -119,8 +120,9 @@ def test_default_config_bounds_liteparse_runtime_concurrency(): assert "parser_concurrency: 1" in default_config assert "context_intro_chunks: 2" in default_config - assert "liteparse_num_workers: 1" in default_config - assert "liteparse_subprocess: true" in default_config + assert "liteparse_num_workers: 2" in default_config + assert "liteparse_ocr_auto_disable_pages: 30" in default_config + assert "liteparse_subprocess" not in default_config def test_config_panel_exposes_document_query_settings(): @@ -153,11 +155,11 @@ def test_config_panel_exposes_document_query_settings(): "liteparse_preserve_very_small_text", "liteparse_output_format", "liteparse_num_workers", - "liteparse_subprocess", "pdf_ocr_fallback", "thread_offload", ]: assert f"config.{setting}" in config_html + assert "liteparse_subprocess" not in config_html def test_document_query_thumbnail_matches_plugin_hub_limits(): @@ -173,9 +175,169 @@ def test_document_query_thumbnail_matches_plugin_hub_limits(): def test_liteparse_parser_caps_workers_by_default(): parser = LiteParseParser() - assert parser._liteparse_kwargs({})["num_workers"] == 1 + assert parser._liteparse_kwargs({})["num_workers"] == 2 assert parser._liteparse_kwargs({"liteparse_num_workers": "3"})["num_workers"] == 3 - assert parser._liteparse_kwargs({"liteparse_num_workers": ""})["num_workers"] == 1 + assert parser._liteparse_kwargs({"liteparse_num_workers": ""})["num_workers"] == 2 + + +def test_liteparse_parser_always_uses_subprocess(monkeypatch): + fetched = FetchedDocument( + uri="/tmp/report.pdf", + source_uri="/tmp/report.pdf", + scheme="file", + mimetype="application/pdf", + content=b"", + local_path="/tmp/report.pdf", + ) + parser = LiteParseParser() + + monkeypatch.setattr(parser, "_parse_subprocess", lambda _document, _config: "ok") + + def fail_in_process(_document, _config): + raise AssertionError("LiteParse must stay isolated from the Web UI process") + + monkeypatch.setattr(parser, "_parse_in_process", fail_in_process) + + assert parser._parse_sync(fetched, {"liteparse_subprocess": False}) == "ok" + + +def test_liteparse_auto_disables_ocr_for_large_text_pdf(monkeypatch): + parser = LiteParseParser() + fetched = FetchedDocument( + uri="/tmp/report.pdf", + source_uri="/tmp/report.pdf", + scheme="file", + mimetype="application/pdf", + content=b"", + local_path="/tmp/report.pdf", + ) + monkeypatch.setattr( + liteparse_module, + "_pdf_text_profile", + lambda _file_path, _config: liteparse_module._PdfTextProfile( + page_count=277, + sampled_pages=5, + text_chars=2500, + ), + ) + + kwargs = parser._liteparse_kwargs({}, fetched, "/tmp/report.pdf") + + assert kwargs["ocr_enabled"] is False + + +def test_liteparse_keeps_ocr_for_small_pdf(monkeypatch): + parser = LiteParseParser() + fetched = FetchedDocument( + uri="/tmp/bill.pdf", + source_uri="/tmp/bill.pdf", + scheme="file", + mimetype="application/pdf", + content=b"", + local_path="/tmp/bill.pdf", + ) + monkeypatch.setattr( + liteparse_module, + "_pdf_text_profile", + lambda _file_path, _config: liteparse_module._PdfTextProfile( + page_count=10, + sampled_pages=5, + text_chars=2500, + ), + ) + + kwargs = parser._liteparse_kwargs({}, fetched, "/tmp/bill.pdf") + + assert kwargs["ocr_enabled"] is True + + +def test_liteparse_keeps_ocr_for_large_text_sparse_pdf(monkeypatch): + parser = LiteParseParser() + fetched = FetchedDocument( + uri="/tmp/scan.pdf", + source_uri="/tmp/scan.pdf", + scheme="file", + mimetype="application/pdf", + content=b"", + local_path="/tmp/scan.pdf", + ) + monkeypatch.setattr( + liteparse_module, + "_pdf_text_profile", + lambda _file_path, _config: liteparse_module._PdfTextProfile( + page_count=277, + sampled_pages=5, + text_chars=20, + ), + ) + + kwargs = parser._liteparse_kwargs({}, fetched, "/tmp/scan.pdf") + + assert kwargs["ocr_enabled"] is True + + +def test_liteparse_respects_explicit_ocr_disabled(monkeypatch): + parser = LiteParseParser() + fetched = FetchedDocument( + uri="/tmp/bill.pdf", + source_uri="/tmp/bill.pdf", + scheme="file", + mimetype="application/pdf", + content=b"", + local_path="/tmp/bill.pdf", + ) + monkeypatch.setattr( + liteparse_module, + "_pdf_text_profile", + lambda _file_path, _config: liteparse_module._PdfTextProfile( + page_count=10, + sampled_pages=5, + text_chars=0, + ), + ) + + kwargs = parser._liteparse_kwargs( + {"liteparse_ocr_enabled": False}, + fetched, + "/tmp/bill.pdf", + ) + + assert kwargs["ocr_enabled"] is False + + +def test_liteparse_target_pages_can_keep_ocr_enabled_for_large_pdf(monkeypatch): + parser = LiteParseParser() + fetched = FetchedDocument( + uri="/tmp/report.pdf", + source_uri="/tmp/report.pdf", + scheme="file", + mimetype="application/pdf", + content=b"", + local_path="/tmp/report.pdf", + ) + monkeypatch.setattr( + liteparse_module, + "_pdf_text_profile", + lambda _file_path, _config: liteparse_module._PdfTextProfile( + page_count=277, + sampled_pages=5, + text_chars=2500, + ), + ) + + small_range = parser._liteparse_kwargs( + {"liteparse_target_pages": "1-10"}, + fetched, + "/tmp/report.pdf", + ) + large_range = parser._liteparse_kwargs( + {"liteparse_target_pages": "1-40"}, + fetched, + "/tmp/report.pdf", + ) + + assert small_range["ocr_enabled"] is True + assert large_range["ocr_enabled"] is False def test_query_optimize_prompt_filename_is_spelled_correctly(): From 3c1bdfa5f91e43f80a65f45aef492ea04a7aac4a Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Sat, 30 May 2026 21:48:43 +0200 Subject: [PATCH 029/231] Add background computer-use tool contract Expose native window listing, indexed window state, and element_action dispatch modes through the Agent Zero connector tool. Update host computer-use prompts and platform skills so agents prefer background structural targeting and only rely on screenshots for foreground or uncertain actions. --- .../agent.system.tool.computer_use_remote.md | 12 +- .../skills/host-computer-use-linux/SKILL.md | 4 + .../skills/host-computer-use-macos/SKILL.md | 4 + .../skills/host-computer-use-windows/SKILL.md | 32 +++- .../skills/host-computer-use/SKILL.md | 21 ++- .../tools/computer_use_remote.py | 140 +++++++++++++++++- 6 files changed, 199 insertions(+), 14 deletions(-) diff --git a/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md b/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md index b8e0beae0..c0003090d 100644 --- a/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md +++ b/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md @@ -2,17 +2,17 @@ Runtime-gated beta desktop control through a connected A0 CLI on the user's host machine. The callable contract is available in the tool prompt. Availability, backend support, and trust mode are checked when the tool runs, together with CLI presence, local enablement, and re-arm state. Computer Use enablement is scoped to the current CLI session, not scoped to a single chat context. -Use this for native host desktop UI inspection, screenshots, clicking, scrolling, typing, key presses, and status checks. Do not use it for ordinary web-page navigation or host-browser control; use the browser tool for web pages unless browser automation cannot express the task. For complex desktop workflows, load and follow skill `host-computer-use` before proceeding. +Use this for native host desktop UI inspection, screenshots, background-safe window/element actions when supported, clicking, scrolling, typing, key presses, and status checks. Do not use it for ordinary web-page navigation or host-browser control; use the browser tool for web pages unless browser automation cannot express the task. For complex desktop workflows, load and follow skill `host-computer-use` before proceeding. This is the only desktop-control path for the user's connected host/local computer. Do not substitute the `linux-desktop` skill, the Agent Zero Desktop/Xpra surface, `desktopctl.sh`, `code_execution_tool`, or Docker/server shell commands for host screen actions; those target the internal Agent Zero runtime and cannot see or control the user's host screen. If the tool reports no CLI, disabled computer use, or `COMPUTER_USE_REARM_REQUIRED`, stop and tell the user to run `/computer-use on` in A0 CLI and approve any host permission prompt. -Call `start_session` before screen-driven tasks. Use `status` for state only, `capture` for screenshots without an action, and `stop_session` when the desktop task is complete. Interactive coordinate actions should use normalized global-screen coordinates from the most recent capture. +Call `start_session` before screen-driven tasks. Use `status` for state only, `capture` for screenshots without an action, and `stop_session` when the desktop task is complete. When the backend advertises native window and element-index features, prefer `list_windows` -> `get_window_state` -> `element_action` with `dispatch: "background"` before using global coordinates. Interactive coordinate actions should use normalized global-screen coordinates from the most recent capture. Some actions are backend-specific and intentionally documented only in backend skills. If `status` or `start_session` reports backend-specific features or tells you to load a backend skill, load and follow that skill before using those backend-only actions. For structural targeting details, load and follow the backend-specific skill such as `host-computer-use-macos` or `host-computer-use-windows`; do not apply one backend's guidance to another backend. -State-changing actions automatically attach a fresh screen after they run. Treat key presses, clicks, scrolling, and typing as attempts, not success: inspect the latest attached screen, or one explicit `capture` if it is unclear or unchanged, before saying the requested outcome happened. If the tool says a screen was attached but you cannot actually inspect the image, stop and report that visual verification is unavailable; do not continue by assuming the host state. A `type` result only proves keystrokes were sent; it does not prove that text landed in the intended place. +State-changing actions automatically attach a fresh screen after they run unless the backend returns a definitive structural background result. Treat key presses, clicks, scrolling, and typing as attempts, not success; treat foreground fallbacks the same way. Inspect the latest attached screen, or one explicit `capture` if it is unclear or unchanged, before saying the requested outcome happened. If the tool says a screen was attached but you cannot actually inspect the image, stop and report that visual verification is unavailable; do not continue by assuming the host state. A `type` result only proves keystrokes were sent; it does not prove that text landed in the intended place. ```json { @@ -24,10 +24,14 @@ State-changing actions automatically attach a fresh screen after they run. Treat ``` Required argument: -- `action`: one of `start_session`, `status`, `capture`, `move`, `click`, `scroll`, `key`, `type`, `stop_session`; backend skills may document additional backend-only action values +- `action`: one of `start_session`, `status`, `capture`, `list_windows`, `get_window_state`, `element_action`, `move`, `click`, `scroll`, `key`, `type`, `stop_session`; backend skills may document additional backend-only action values Optional arguments by action: - `session_id`: session returned by `start_session` +- `pid`, `window_id`: target a native app/window for `get_window_state` and `element_action` +- `element_index`: target an element from the latest `get_window_state` +- `operation`: action such as `invoke`, `press`, `set_value`, `focus`, or backend-specific operations +- `dispatch`: `background`, `auto`, or `foreground`; prefer `background` for `element_action` - `x`, `y`: normalized `[0,1]` global-screen coordinates for `move` and `click` - `button`: `left`, `right`, or `middle` for `click` - `count`: click count for `click` diff --git a/plugins/_a0_connector/skills/host-computer-use-linux/SKILL.md b/plugins/_a0_connector/skills/host-computer-use-linux/SKILL.md index df0fca704..a9d38adc7 100644 --- a/plugins/_a0_connector/skills/host-computer-use-linux/SKILL.md +++ b/plugins/_a0_connector/skills/host-computer-use-linux/SKILL.md @@ -20,6 +20,8 @@ Linux backends can advertise structural AT-SPI features: When these features are present, prefer structural targeting over pixel clicks for named buttons, menu items, text fields, dialogs, toolbar items, tab strips, and application windows. +If the backend also advertises `native-window-list`, `window-state`, `element-index-targeting`, or `background-dispatch`, prefer the generic background loop from `host-computer-use`: `list_windows` -> `get_window_state` -> `element_action`. If those features are absent, use the AT-SPI snapshot/action flow below. + Use `ax_snapshot` to inspect the Linux AT-SPI tree: ```json @@ -68,6 +70,8 @@ Targeting options: Use screenshots for proof after every state-changing action. AT-SPI actions and keyboard events are attempts, not proof, and Wayland focus can reject or redirect input when the active window changes. +True background dispatch on Linux is compositor, toolkit, and app dependent. Do not claim a Linux action was background-safe unless the tool result explicitly says `actual_dispatch=background`. + On GNOME/Wayland, useful shortcuts include: - `Super+H`: hide the active window diff --git a/plugins/_a0_connector/skills/host-computer-use-macos/SKILL.md b/plugins/_a0_connector/skills/host-computer-use-macos/SKILL.md index fdd04a835..e01695e7c 100644 --- a/plugins/_a0_connector/skills/host-computer-use-macos/SKILL.md +++ b/plugins/_a0_connector/skills/host-computer-use-macos/SKILL.md @@ -19,6 +19,8 @@ macOS backends can advertise structural Accessibility features: When these features are present, prefer structural targeting over pixel clicks for named controls such as buttons, menu items, text fields, sheets, alerts, toolbar items, and sidebar rows. +If the backend also advertises `native-window-list`, `window-state`, `element-index-targeting`, or `background-dispatch`, prefer the generic background loop from `host-computer-use`: `list_windows` -> `get_window_state` -> `element_action`. If those features are absent, use the AX snapshot/action flow below. + Use `ax_snapshot` to inspect the frontmost app's bounded Accessibility tree: ```json @@ -65,6 +67,8 @@ Targeting options: AX actions are attempts, not proof. They attach a fresh screenshot after state-changing actions; inspect that image before saying the requested outcome happened. +Do not assume macOS work happened in the background unless the tool result explicitly says `actual_dispatch=background`. If the result says `background_unavailable`, use foreground dispatch only when foreground control is acceptable. + ## macOS Window Actions For active-app window tasks, macOS shortcuts are usually: diff --git a/plugins/_a0_connector/skills/host-computer-use-windows/SKILL.md b/plugins/_a0_connector/skills/host-computer-use-windows/SKILL.md index 5b34b2c45..c3392a71a 100644 --- a/plugins/_a0_connector/skills/host-computer-use-windows/SKILL.md +++ b/plugins/_a0_connector/skills/host-computer-use-windows/SKILL.md @@ -17,10 +17,37 @@ Windows backends can advertise structural UI Automation features: - `uia-structural-targeting` - `uia-element-action` - `uia-window-management` +- `native-window-list` +- `window-state` +- `element-index-targeting` +- `background-dispatch` +- `foreground-dispatch-fallback` -When these features are present, prefer structural targeting over pixel clicks for named controls such as buttons, menu items, text fields, dialogs, toolbar items, browser address bars, composer fields, and list rows. +When these features are present, prefer background structural targeting over pixel clicks for named controls such as buttons, menu items, text fields, dialogs, toolbar items, browser address bars, composer fields, and list rows. -Use `uia_snapshot` to inspect the bounded Windows UI Automation tree: +Preferred loop: + +1. Use `list_windows` to find the target native window. +2. Use `get_window_state` with the target `pid` and/or `window_id`. +3. Use `element_action` with the returned `element_index` and `dispatch: "background"`. +4. If the backend reports `background_unavailable`, switch to `dispatch: "auto"` or `dispatch: "foreground"` only when foreground control is acceptable. + +Example: + +```json +{ + "tool_name": "computer_use_remote", + "tool_args": { + "action": "element_action", + "window_id": "uia-hwnd:123456", + "element_index": 7, + "operation": "invoke", + "dispatch": "background" + } +} +``` + +Use `uia_snapshot` to inspect the bounded Windows UI Automation tree when the newer window-state loop is not available: ```json { @@ -72,6 +99,7 @@ Targeting options: Action selection: - Prefer the actions listed on the target node. If a node offers `invoke`, use `invoke`, not `click`. +- Prefer `element_action` over `uia_action` when the backend advertises `element-index-targeting`; it preserves the background-first dispatch contract. - For window focus, hiding, restoring, or maximizing, use `focus_window`, `minimize`, `restore`, or `maximize`; do not click titlebar buttons. - For typing into an app, first structurally focus the app/window if needed, then `set_value` on the target field. A global `type` result only proves keys were sent, not that they landed in the intended control. - After a window operation, navigation, menu open/close, dialog transition, or other layout change, take a fresh `uia_snapshot` before reusing a path. diff --git a/plugins/_a0_connector/skills/host-computer-use/SKILL.md b/plugins/_a0_connector/skills/host-computer-use/SKILL.md index d7d6f8b9d..b8815c5d9 100644 --- a/plugins/_a0_connector/skills/host-computer-use/SKILL.md +++ b/plugins/_a0_connector/skills/host-computer-use/SKILL.md @@ -17,7 +17,7 @@ triggers: # Host Computer Use -This skill unlocks the beta `computer_use_remote` tool for connected local desktop control through A0 CLI. +This skill unlocks the beta `computer_use_remote` tool for connected local desktop control through A0 CLI. Prefer native background-safe computer use when the connected backend advertises window and element-index features. ## When to Use @@ -55,9 +55,13 @@ Use: Arguments: -- `action`: `start_session`, `status`, `capture`, `move`, `click`, `scroll`, `key`, `type`, `stop_session` +- `action`: `start_session`, `status`, `capture`, `list_windows`, `get_window_state`, `element_action`, `move`, `click`, `scroll`, `key`, `type`, `stop_session` - `session_id`: optional after `start_session` - backend skills may document additional backend-only action values; use them only when backend metadata advertises matching support and after loading the backend-specific skill +- `list_windows`: returns native top-level window records when the backend supports them +- `get_window_state`: pass `pid` and/or `window_id`; returns a target-window accessibility tree with stable `element_index` values for the current state +- `element_action`: pass `element_index` from the latest `get_window_state`; optional `operation`, `value`/`text`, and `dispatch` +- `dispatch`: `background`, `auto`, or `foreground`; default to `background` for element actions - `move`: `x`, `y` normalized to `[0,1]` - `click`: optional `x`, `y`, optional `button` (`left`, `right`, `middle`), optional `count` - `scroll`: `dx`, `dy` @@ -72,10 +76,13 @@ If any tool result contains `COMPUTER_USE_REARM_REQUIRED` or `status=rearm requi 1. Call `start_session` first. 2. Read the returned `backend_id`, `backend_family`, and `features`; load a backend-specific Computer Use skill when the task needs backend-only affordances. -3. Decide final success from the latest screenshot, not from memory. -4. Interactive actions already attach a fresh screenshot after they run; inspect it before claiming the requested outcome succeeded. -5. Use `status` for state without starting a session. -6. Use `capture` only when you need another screenshot without taking an action. +3. If the backend advertises `native-window-list`, call `list_windows` before using coordinates. +4. If the backend advertises `window-state` and `element-index-targeting`, call `get_window_state` for the target `pid`/`window_id`, then use `element_action` with `dispatch: "background"` by default. +5. If `element_action` reports `background_unavailable`, use `dispatch: "auto"` or `dispatch: "foreground"` only when foreground control is acceptable for the user/task. +6. Decide final success from the latest screenshot or a definitive structural result, not from memory. +7. Interactive actions already attach a fresh screenshot after they run; inspect it before claiming the requested outcome succeeded. +8. Use `status` for state without starting a session. +9. Use `capture` only when you need another screenshot without taking an action. ## Backend Skills @@ -88,7 +95,9 @@ If any tool result contains `COMPUTER_USE_REARM_REQUIRED` or `status=rearm requi - Only the latest screenshot or a definitive tool result counts as evidence. - If a tool result says a screenshot was attached but you cannot actually see the image, stop and report that visual verification is unavailable. Do not continue with another action from an assumed host state. +- Prefer the background-safe native loop when advertised: `list_windows` -> `get_window_state` -> `element_action`. - Outside advertised structural accessibility support, use normalized global screen coordinates; do not assume window ids, element indexes, background-safe input, or semantic click targets unless the runtime explicitly advertises them. +- Treat `dispatch: "foreground"` as intentional control of the user's visible desktop. Use it only after deciding that a background action is unavailable or unsuitable. - On Linux, AT-SPI structural targeting uses backend-specific actions documented in `host-computer-use-linux`; do not apply macOS AX-specific assumptions unless the backend is macOS. - Prefer accessibility and semantic UI paths first: shortcuts, command palettes, menu accelerators, address/search bars, focus traversal, and other keyboard-accessible controls. - Prefer `key` and `type` over pointer actions whenever a reliable keyboard path exists. diff --git a/plugins/_a0_connector/tools/computer_use_remote.py b/plugins/_a0_connector/tools/computer_use_remote.py index 32f73e9d6..e8a03168c 100644 --- a/plugins/_a0_connector/tools/computer_use_remote.py +++ b/plugins/_a0_connector/tools/computer_use_remote.py @@ -34,6 +34,7 @@ REARM_REQUIRED_DEFAULT_MESSAGE = ( ) _AUTO_CAPTURE_ACTIONS = { "start_session", + "element_action", "ax_action", "uia_action", "move", @@ -58,6 +59,9 @@ _SUPPORTED_ACTIONS = { "start_session", "status", "capture", + "list_windows", + "get_window_state", + "element_action", "ax_snapshot", "ax_action", "uia_snapshot", @@ -80,7 +84,8 @@ class ComputerUseRemote(Tool): return Response( message=( "action is required and must be one of: " - "start_session, status, capture, ax_snapshot, ax_action, " + "start_session, status, capture, list_windows, get_window_state, " + "element_action, ax_snapshot, ax_action, " "uia_snapshot, uia_action, " "move, click, scroll, key, type, stop_session" ), @@ -226,6 +231,11 @@ class ComputerUseRemote(Tool): data = result.get("result") result_data = dict(data) if isinstance(data, dict) else {} + if action == "element_action": + actual_dispatch = str(result_data.get("actual_dispatch") or "").strip().lower() + if actual_dispatch in {"background", "none"}: + return "" + session_id = str(result_data.get("session_id") or self.args.get("session_id") or "").strip() if not session_id: return "" @@ -324,6 +334,34 @@ class ComputerUseRemote(Tool): payload["text"] = self.args.get("text", "") if self._coerce_bool(self.args.get("submit")): payload["submit"] = True + elif action == "list_windows": + for key in ("include_hidden", "include_offscreen", "max_windows"): + if key in self.args: + payload[key] = self.args.get(key) + elif action == "get_window_state": + for key in ("pid", "window_id", "mode", "max_depth", "max_nodes"): + if key in self.args: + payload[key] = self.args.get(key) + elif action == "element_action": + for key in ( + "pid", + "window_id", + "element_index", + "path", + "operation", + "name", + "dispatch", + "value", + "text", + "submit", + ): + if key in self.args: + payload[key] = self.args.get(key) + target = self.args.get("target") + if isinstance(target, dict): + payload["target"] = dict(target) + if "selector" in self.args: + payload["selector"] = self.args.get("selector") elif action == "ax_snapshot": if "max_depth" in self.args: payload["max_depth"] = self._coerce_int(self.args.get("max_depth"), name="max_depth") @@ -386,6 +424,12 @@ class ComputerUseRemote(Tool): if action == "capture": summary = self._record_capture(data) return f"Current screen attached: {summary} {CAPTURE_VERIFICATION_NOTE}" + if action == "list_windows": + return self._format_window_list(data) + if action == "get_window_state": + return self._format_window_state(data) + if action == "element_action": + return self._format_element_action(data) if action == "ax_snapshot": return self._format_ax_snapshot(data) if action == "ax_action": @@ -552,6 +596,94 @@ class ComputerUseRemote(Tool): f"active_contexts={active_text}.{rearm_guidance}" ) + def _format_window_list(self, data: dict[str, Any]) -> str: + windows = data.get("windows") if isinstance(data.get("windows"), list) else [] + count = data.get("count", len(windows)) + if not windows: + return ( + "No native windows were returned by the computer-use backend. " + "Use capture or backend-specific snapshots as a fallback." + ) + lines = [ + ( + f"Computer-use native window list: {count} window(s). " + "Prefer get_window_state(pid/window_id) before element_action." + ) + ] + for item in windows[:40]: + if not isinstance(item, dict): + continue + title = str(item.get("title") or item.get("name") or "").strip() + app_name = str(item.get("app_name") or item.get("app") or "").strip() + pid = item.get("pid") + window_id = str(item.get("window_id") or "").strip() + role = str(item.get("role") or "window").strip() + parts = [f"- window_id={window_id or '?'}"] + if pid is not None: + parts.append(f"pid={pid}") + if app_name: + parts.append(f"app={app_name!r}") + if title: + parts.append(f"title={title!r}") + parts.append(f"role={role}") + frame = item.get("frame") + if isinstance(frame, dict): + parts.append( + f"frame=({frame.get('x', '?')},{frame.get('y', '?')} " + f"{frame.get('width', '?')}x{frame.get('height', '?')})" + ) + flags: list[str] = [] + for flag in ("is_on_screen", "on_current_space", "focused", "visible"): + if flag in item: + flags.append(f"{flag}={item.get(flag)}") + if flags: + parts.append(" ".join(flags)) + lines.append(" ".join(parts)) + if len(windows) > 40: + lines.append("... window list truncated; request a narrower app/window target.") + return "\n".join(lines) + + def _format_window_state(self, data: dict[str, Any]) -> str: + tree = data.get("tree") if isinstance(data.get("tree"), dict) else {} + window = data.get("window") if isinstance(data.get("window"), dict) else {} + app = data.get("app") if isinstance(data.get("app"), dict) else {} + title = str(window.get("title") or app.get("name") or "target window").strip() + window_id = str(window.get("window_id") or data.get("window_id") or "").strip() + node_count = data.get("node_count", "?") + truncated = " truncated" if data.get("truncated") else "" + mode = str(data.get("mode") or "auto").strip() + return ( + f"Window state for {title!r}" + f"{f' window_id={window_id}' if window_id else ''}: " + f"{node_count} element(s){truncated}, mode={mode}. " + "Use element_action with element_index; dispatch defaults to background." + f"{self._structural_tree_outline(tree)}" + ) + + def _format_element_action(self, data: dict[str, Any]) -> str: + target = data.get("target") if isinstance(data.get("target"), dict) else {} + operation = str(data.get("operation") or "?") + requested_dispatch = str(data.get("requested_dispatch") or data.get("dispatch") or "background") + actual_dispatch = str(data.get("actual_dispatch") or data.get("dispatch") or requested_dispatch) + fallback = bool(data.get("foreground_fallback_used") or data.get("fallback_used")) + index = target.get("element_index", data.get("element_index", "?")) + label = ( + self._uia_target_label(target) + if target.get("selector") or target.get("automation_id") + else self._ax_target_label(target) + ) + if data.get("background_unavailable"): + reason = str(data.get("reason") or "background dispatch was unavailable") + return ( + f"Background element action unavailable for element_index={index}: {reason}. " + "Use dispatch='auto' or dispatch='foreground' only if foreground control is acceptable." + ) + dispatch_text = ( + f"requested_dispatch={requested_dispatch}, actual_dispatch={actual_dispatch}" + f"{', foreground_fallback_used=true' if fallback else ''}" + ) + return f"Performed {operation} on element_index={index} {label}; {dispatch_text}." + def _format_ax_snapshot(self, data: dict[str, Any]) -> str: app = data.get("app") if isinstance(data.get("app"), dict) else {} tree = data.get("tree") if isinstance(data.get("tree"), dict) else {} @@ -611,7 +743,11 @@ class ComputerUseRemote(Tool): indent = " " * max(0, depth) role = str(node.get("role") or "element") path = node.get("path", []) - parts = [f"{indent}- path={path} role={role}"] + element_index = node.get("element_index") + prefix = f"{indent}-" + if element_index is not None: + prefix = f"{prefix} element_index={element_index}" + parts = [f"{prefix} path={path} role={role}"] for key in ("title", "name", "description", "automation_id", "class_name", "selector"): value = node.get(key) if isinstance(value, str) and value.strip(): From ce4156a1109892a884ed65f7e591508020c9ea05 Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Sat, 30 May 2026 22:22:15 +0200 Subject: [PATCH 030/231] Harden internal Xpra desktop control Move the internal Linux Desktop skill toward a structured-first workflow: state/window commands before screenshots, batched desktopctl sequences in one process, and final screenshots only when pixels matter. Normalize internal system Desktop resize requests so narrow portrait canvas sizes do not shrink the real X11 root display and distort screenshots. Keep the live Docker runtime in sync with the changed Desktop files. --- helpers/virtual_desktop.py | 33 ++++++++++++ plugins/_desktop/helpers/desktop_session.py | 2 + plugins/_desktop/helpers/desktop_state.py | 5 +- .../_desktop/skills/linux-desktop/SKILL.md | 52 +++++++++++++------ .../linux-desktop/scripts/desktopctl.sh | 18 +++++-- plugins/_desktop/webui/desktop-store.js | 2 +- tests/test_office_canvas_setup.py | 5 ++ tests/test_office_desktop_state.py | 28 ++++++++++ 8 files changed, 121 insertions(+), 24 deletions(-) diff --git a/helpers/virtual_desktop.py b/helpers/virtual_desktop.py index cd1020492..d417cf4ad 100644 --- a/helpers/virtual_desktop.py +++ b/helpers/virtual_desktop.py @@ -23,6 +23,7 @@ MAX_WIDTH = 1920 MAX_HEIGHT = 1080 MIN_WIDTH = 360 MIN_HEIGHT = 240 +MIN_DESKTOP_ASPECT_RATIO = 4 / 3 SESSION_PATH = "/desktop/session" XPRA_HTML_ROOT_CANDIDATES = ( Path("/usr/share/xpra/www"), @@ -220,6 +221,38 @@ def normalize_size( ) +def normalize_desktop_display_size( + width: int | float | str, + height: int | float | str, + *, + max_width: int = MAX_WIDTH, + max_height: int = MAX_HEIGHT, + min_width: int = MIN_WIDTH, + min_height: int = MIN_HEIGHT, + min_aspect_ratio: float = MIN_DESKTOP_ASPECT_RATIO, +) -> tuple[int, int]: + normalized_width, normalized_height = normalize_size( + width, + height, + max_width=max_width, + max_height=max_height, + min_width=min_width, + min_height=min_height, + ) + if normalized_height <= 0: + return normalized_width, normalized_height + if normalized_width / normalized_height >= min_aspect_ratio: + return normalized_width, normalized_height + return normalize_size( + DEFAULT_WIDTH, + DEFAULT_HEIGHT, + max_width=max_width, + max_height=max_height, + min_width=min_width, + min_height=min_height, + ) + + def resize_display( *, display: int, diff --git a/plugins/_desktop/helpers/desktop_session.py b/plugins/_desktop/helpers/desktop_session.py index c2a7e270b..0545a5684 100644 --- a/plugins/_desktop/helpers/desktop_session.py +++ b/plugins/_desktop/helpers/desktop_session.py @@ -399,6 +399,8 @@ class DesktopSessionManager: if not session: return {"ok": False, "error": "LibreOffice desktop session not found."} is_system_desktop = session.session_id == SYSTEM_SESSION_ID and session.extension == "desktop" + if is_system_desktop: + width, height = virtual_desktop.normalize_desktop_display_size(width, height) result = virtual_desktop.resize_display( display=session.display, width=width, diff --git a/plugins/_desktop/helpers/desktop_state.py b/plugins/_desktop/helpers/desktop_state.py index 825a7d6cb..cbfbb7a55 100644 --- a/plugins/_desktop/helpers/desktop_state.py +++ b/plugins/_desktop/helpers/desktop_state.py @@ -721,9 +721,10 @@ def compact_prompt_context(state: dict[str, Any] | None = None) -> str: lines.append(f"- screenshot_context={context_id}") context_arg = f" --context-id {context_id}" if context_id else "" lines.append( - "- next=plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh observe --json --screenshot" + "- next=plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh state --json" + f"{context_arg} for structured checks; use observe --json --screenshot" f"{context_arg} " - "before any coordinate action; prefer focus/key/paste/save/app-native helpers first." + "before coordinate or visual-OCR actions; prefer sequence/focus/key/paste/save/app-native helpers first." ) lines.append( "- verify=for terminal/CLI-agent output, use the screenshot path from a fresh final " diff --git a/plugins/_desktop/skills/linux-desktop/SKILL.md b/plugins/_desktop/skills/linux-desktop/SKILL.md index 54931bae4..9ede17908 100644 --- a/plugins/_desktop/skills/linux-desktop/SKILL.md +++ b/plugins/_desktop/skills/linux-desktop/SKILL.md @@ -31,14 +31,16 @@ Use the Desktop as a full Linux GUI when the user explicitly needs a visual work ## Operating Model -The Desktop is an observe-act-verify control surface. Use this decision hierarchy: +The Desktop is a structured-first X11 control surface. Use this decision hierarchy: 1. Prefer structured tools such as `office_artifact` for deterministic Office file creation, reads, and edits. -2. Prefer app-native helpers for visible live edits, such as `desktopctl.sh calc-set-cell` for Calc/UNO spreadsheet changes. -3. Prefer launcher commands, window focus, keyboard shortcuts, menus, paste, and save commands. -4. Use coordinate clicks only as a last resort, and only after a fresh Desktop observation. -5. After any GUI action, verify through Desktop state, active window titles, screenshots, saved file state, or exported output. -6. For terminal or CLI-agent work, verify against a fresh final `observe --json --screenshot` captured after the command has finished or visibly returned to an input prompt. Agent-facing Desktop screenshots are ephemeral refs; `desktopctl` shell observations with `--context-id` return chat-scoped screenshot paths. Do not report from an earlier screenshot path. +2. Prefer structured Desktop state and window commands: `check`, `state --json`, `windows`, `active-window`, `geometry`, and `wait-window`. +3. Prefer app-native helpers for visible live edits, such as `desktopctl.sh calc-set-cell` for Calc/UNO spreadsheet changes. +4. Prefer launcher commands, window focus, keyboard shortcuts, menus, paste, and save commands. Batch multi-step actions with `sequence`/`batch` so one helper process drives the flow. +5. Use screenshots for visual inspection, OCR, coordinate work, and final evidence when the task depends on pixels. Do not take screenshots merely to learn window titles or readiness that structured state already reports. +6. Use coordinate clicks only as a last resort, and only after a fresh screenshot observation. +7. After any GUI action, verify through Desktop state, active window titles, screenshots when visually necessary, saved file state, or exported output. +8. For terminal or CLI-agent work, prefer deterministic command output or saved transcripts. When exact visible terminal text matters, verify against a fresh final `observe --json --screenshot` captured after the command has finished or visibly returned to an input prompt. Agent-facing Desktop screenshots are ephemeral refs; `desktopctl` shell observations with `--context-id` return chat-scoped screenshot paths. Do not report from an earlier screenshot path. Keep these standing rules: @@ -58,7 +60,7 @@ Use the helper script when the Desktop is already open and you need reliable app DESKTOP=/a0/plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh $DESKTOP check $DESKTOP state --json -$DESKTOP observe --json --screenshot +$DESKTOP observe --json $DESKTOP launch calc $DESKTOP wait-window LibreOffice $DESKTOP windows LibreOffice @@ -68,7 +70,7 @@ $DESKTOP key ctrl+s The script targets the persistent `agent-zero-desktop` X display, sets `DISPLAY`, `XAUTHORITY`, and `HOME` to the XFCE profile, then uses `xdotool` for input. Startup normally prepares this session. If `check` fails during explicit Desktop work, report that the Desktop runtime is not ready instead of installing packages ad hoc. -If `observe --json --screenshot` shows a reachable display, visible Desktop/window entries, and a fresh screenshot, the Desktop is usable even when `active_window` is `null`; a bare XFCE desktop can have no active application window. Treat missing screenshots, missing display, or unavailable `xdotool`/`xwd` as blockers and stop with the specific readiness message instead of repeating clicks or inventing a fallback. Shell screenshots captured with `--context-id` live in the owning chat's screenshot folder; screenshots without a chat context remain temporary. +If `state --json` or `observe --json` shows a reachable display and visible Desktop/window entries, the Desktop is usable even when `active_window` is `null`; a bare XFCE desktop can have no active application window. Treat missing display or unavailable `xdotool` as blockers and stop with the specific readiness message instead of repeating clicks or inventing a fallback. Use `observe --json --screenshot` only when you need pixels, and treat unavailable `xwd` as a screenshot blocker rather than a general Desktop blocker. Shell screenshots captured with `--context-id` live in the owning chat's screenshot folder; screenshots without a chat context remain temporary. For direct app launches without coordinates: @@ -85,6 +87,20 @@ $DESKTOP paste-text "Text to insert" $DESKTOP key ctrl+s ``` +For multi-step window or terminal actions, batch commands through one helper process: + +```bash +DESKTOP=/a0/plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh +$DESKTOP sequence - <<'EOF' +focus Terminal +paste-text echo ready +key Return +state --json +EOF +``` + +Use one command per line. `paste-text` joins the remaining words on that line, so it is suitable for ordinary command text; for complex multiline payloads, write the payload to a file or invoke `paste-text` directly. + For live spreadsheet coworking, use the Calc helper instead of hand-written UNO snippets: ```bash @@ -120,10 +136,10 @@ When browser automation is available, the higher-level QA flow is: Terminal apps are visual state, not structured logs. When the task depends on exact terminal output, follow this stricter loop: -1. Run `desktopctl.sh observe --json --screenshot` immediately before acting to record the starting window and screenshot path. -2. Use `focus`, `paste-text` or `type`, and `key Return` to drive the terminal. Prefer CLI-native commands and keyboard input over clicks. +1. Run `desktopctl.sh state --json` or `desktopctl.sh windows Terminal` before acting to confirm the target window. Add `observe --json --screenshot` only if the current prompt or screen contents must be read visually. +2. Use `sequence`, `focus`, `paste-text` or `type`, and `key Return` to drive the terminal. Prefer CLI-native commands and keyboard input over clicks. 3. Wait until the CLI has visibly produced a response or returned to an input prompt. -4. Run a new final `desktopctl.sh observe --json --screenshot`. +4. If exact visible text matters and no deterministic transcript was saved, run a new final `desktopctl.sh observe --json --screenshot`. 5. Verify exact text only from the screenshot path returned by that final observation, or from a newer screenshot. Never use an earlier screenshot path as final evidence. 6. If the final screenshot is cropped, stale, or unreadable, capture another screenshot or report the result as unverified with that specific reason. @@ -140,13 +156,17 @@ Example for a nested CLI-agent smoke test: ```bash DESKTOP=/a0/plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh -$DESKTOP focus "Terminal" -$DESKTOP paste-text 'TARGET_CLI="example-cli-agent"; FALLBACK_CMD=""; if command -v "$TARGET_CLI" >/dev/null 2>&1; then "$TARGET_CLI"; elif [ -n "$FALLBACK_CMD" ]; then sh -lc "$FALLBACK_CMD"; else echo "CLI agent not found: $TARGET_CLI"; fi' -$DESKTOP key Return +$DESKTOP sequence - <<'EOF' +focus Terminal +paste-text TARGET_CLI="example-cli-agent"; FALLBACK_CMD=""; if command -v "$TARGET_CLI" >/dev/null 2>&1; then "$TARGET_CLI"; elif [ -n "$FALLBACK_CMD" ]; then sh -lc "$FALLBACK_CMD"; else echo "CLI agent not found: $TARGET_CLI"; fi +key Return +EOF $DESKTOP observe --json --screenshot # Verify the screenshot shows the target CLI prompt, not a shell prompt, before sending natural language: -$DESKTOP paste-text 'Reply with exactly the requested smoke-test token.' -$DESKTOP key Return +$DESKTOP sequence - <<'EOF' +paste-text Reply with exactly the requested smoke-test token. +key Return +EOF $DESKTOP observe --json --screenshot ``` diff --git a/plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh b/plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh index 0abbf7e64..555ae0d0c 100755 --- a/plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh +++ b/plugins/_desktop/skills/linux-desktop/scripts/desktopctl.sh @@ -71,7 +71,8 @@ Commands: drag X1 Y1 X2 Y2 Drag from X1,Y1 to X2,Y2 in Desktop coordinates. right-click X Y Move and right-click at X,Y in Desktop coordinates. paste-text TEXT Put TEXT on the Desktop clipboard and paste it with an app-native shortcut. - sequence FILE|- Run a newline-delimited command sequence. + sequence FILE|- Run a newline-delimited command sequence in this process. + batch FILE|- Alias for sequence. key KEY... Send one or more xdotool key names. type TEXT Type text into the focused window. click X Y Move and click at X,Y in Desktop coordinates. @@ -228,7 +229,7 @@ run_sequence_line() { \#*) return 0 ;; esac # shellcheck disable=SC2086 - "$0" $line + dispatch_command $line } run_sequence() { @@ -279,7 +280,11 @@ launch_app() { esac } -case "$command_name" in +dispatch_command() { + local command_name="${1:-help}" + shift || true + + case "$command_name" in help|-h|--help) usage ;; @@ -387,7 +392,7 @@ case "$command_name" in fi paste_text "$@" ;; - sequence) + sequence|batch) source_file="${1:?sequence requires FILE or -}" run_sequence "$source_file" ;; @@ -448,4 +453,7 @@ case "$command_name" in usage >&2 exit 2 ;; -esac + esac +} + +dispatch_command "$command_name" "$@" diff --git a/plugins/_desktop/webui/desktop-store.js b/plugins/_desktop/webui/desktop-store.js index 75f9d4371..1408ecad6 100644 --- a/plugins/_desktop/webui/desktop-store.js +++ b/plugins/_desktop/webui/desktop-store.js @@ -1818,7 +1818,7 @@ const model = { html, body, #screen { width: 100% !important; height: 100% !important; - overflow: hidden !important; + overflow: auto !important; } #float_menu, .windowhead, diff --git a/tests/test_office_canvas_setup.py b/tests/test_office_canvas_setup.py index 535237acf..28935ba67 100644 --- a/tests/test_office_canvas_setup.py +++ b/tests/test_office_canvas_setup.py @@ -272,7 +272,9 @@ def test_desktop_plugin_owns_routes_runtime_surface_and_state_paths(): assert "canvas.width = normalizedWidth" in desktop_store assert "canvas.height = normalizedHeight" in desktop_store assert "canvas?.clientWidth || canvas?.width" in desktop_store + assert "overflow: auto !important;" in desktop_store assert "Installing Agent Zero Desktop runtime dependencies" in desktop_session + assert "normalize_desktop_display_size" in desktop_session assert "__a0XpraOffsetWarnPatched" in desktop_store assert "window does not fit in canvas, offsets" in desktop_store assert "decode error packet" in desktop_store @@ -529,6 +531,9 @@ def test_office_and_desktop_skills_are_rehomed_and_renamed(): assert "Open in Desktop action" in desktop_skill assert "$BASE_DIR/usr/plugins/_desktop/profiles/$SESSION" in desktopctl assert "$BASE_DIR/usr/plugins/_desktop/sessions/$SESSION.json" in desktopctl + assert "sequence|batch)" in desktopctl + assert "dispatch_command $line" in desktopctl + assert '"$0" $line' not in desktopctl def test_skill_catalog_and_connector_boundaries_are_static_guarded(): diff --git a/tests/test_office_desktop_state.py b/tests/test_office_desktop_state.py index bf4f18bea..12ae906a9 100644 --- a/tests/test_office_desktop_state.py +++ b/tests/test_office_desktop_state.py @@ -10,6 +10,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) +from helpers import virtual_desktop from plugins._desktop.helpers import desktop_state @@ -296,6 +297,33 @@ def test_desktop_state_default_screenshot_returns_ephemeral_ref(tmp_path, monkey assert not (tmp_path / "ctx_id").exists() +def test_desktop_prompt_context_recommends_structured_state_before_screenshots(): + context = desktop_state.compact_prompt_context( + { + "display": ":120", + "size": {"width": 1280, "height": 720}, + "pointer": {"x": 10, "y": 20}, + "active_window": {"title": "Terminal", "class": "Xfce4-terminal"}, + "windows": [{"title": "Terminal", "class": "Xfce4-terminal"}], + "screenshot": {}, + "context_id": "ctx_id", + "errors": [], + } + ) + + assert "state --json --context-id ctx_id for structured checks" in context + assert "observe --json --screenshot --context-id ctx_id before coordinate or visual-OCR actions" in context + assert "before any coordinate action" not in context + + +def test_virtual_desktop_system_display_normalization_rejects_portrait_viewports(): + assert virtual_desktop.normalize_desktop_display_size(395, 1080) == ( + virtual_desktop.DEFAULT_WIDTH, + virtual_desktop.DEFAULT_HEIGHT, + ) + assert virtual_desktop.normalize_desktop_display_size(1600, 900) == (1600, 900) + + def test_xwd_fallback_parser_handles_truecolor_pixels(tmp_path, monkeypatch): raw_path = tmp_path / "shot.xwd" target = tmp_path / "shot.png" From d8710c3f8145731d8bfa0900f0c3b2b2707d0e20 Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Sat, 30 May 2026 22:39:44 +0200 Subject: [PATCH 031/231] Preserve computer-use capability metadata Store and return the connector computer-use contract_version and nested capabilities payload in Agent Zero, surface the contract version in tool status text, and teach the remote computer-use prompt/skill to prefer the structured background dispatch contract over OS-name assumptions. --- plugins/_a0_connector/helpers/ws_runtime.py | 13 +++++ .../agent.system.tool.computer_use_remote.md | 4 +- .../skills/host-computer-use/SKILL.md | 17 +++--- .../tools/computer_use_remote.py | 5 ++ ...test_a0_connector_computer_use_metadata.py | 52 +++++++++++++++++++ 5 files changed, 82 insertions(+), 9 deletions(-) create mode 100644 tests/test_a0_connector_computer_use_metadata.py diff --git a/plugins/_a0_connector/helpers/ws_runtime.py b/plugins/_a0_connector/helpers/ws_runtime.py index 94711a2ba..1987e5350 100644 --- a/plugins/_a0_connector/helpers/ws_runtime.py +++ b/plugins/_a0_connector/helpers/ws_runtime.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import copy import threading import time from dataclasses import dataclass @@ -58,6 +59,8 @@ class ComputerUseMetadata: backend_id: str backend_family: str features: tuple[str, ...] + contract_version: int + capabilities: dict[str, Any] support_reason: str updated_at: float @@ -337,6 +340,12 @@ def store_sid_computer_use_metadata(sid: str, payload: dict[str, Any]) -> Comput 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")), @@ -348,6 +357,8 @@ def store_sid_computer_use_metadata(sid: str, payload: dict[str, Any]) -> Comput 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(), ) @@ -377,6 +388,8 @@ def computer_use_metadata_for_sid(sid: str) -> dict[str, Any] | None: "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, } diff --git a/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md b/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md index c0003090d..163f23db4 100644 --- a/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md +++ b/plugins/_a0_connector/prompts/agent.system.tool.computer_use_remote.md @@ -8,7 +8,7 @@ This is the only desktop-control path for the user's connected host/local comput If the tool reports no CLI, disabled computer use, or `COMPUTER_USE_REARM_REQUIRED`, stop and tell the user to run `/computer-use on` in A0 CLI and approve any host permission prompt. -Call `start_session` before screen-driven tasks. Use `status` for state only, `capture` for screenshots without an action, and `stop_session` when the desktop task is complete. When the backend advertises native window and element-index features, prefer `list_windows` -> `get_window_state` -> `element_action` with `dispatch: "background"` before using global coordinates. Interactive coordinate actions should use normalized global-screen coordinates from the most recent capture. +Call `start_session` before screen-driven tasks. Use `status` for state only, `capture` for screenshots without an action, and `stop_session` when the desktop task is complete. Read `backend_id`, `backend_family`, `features`, and the structured `capabilities` object in status/session results. When capabilities report native windows, window state, element indexes, and background dispatch, prefer `list_windows` -> `get_window_state` -> `element_action` with `dispatch: "background"` before using global coordinates. Interactive coordinate actions should use normalized global-screen coordinates from the most recent capture. Some actions are backend-specific and intentionally documented only in backend skills. If `status` or `start_session` reports backend-specific features or tells you to load a backend skill, load and follow that skill before using those backend-only actions. For structural targeting details, load and follow the backend-specific skill such as `host-computer-use-macos` or `host-computer-use-windows`; do not apply one backend's guidance to another backend. @@ -39,3 +39,5 @@ Optional arguments by action: - `key` or `keys`: key press value for `key` - `text`: text to type for `type` - `submit`: boolean Enter-after-type flag for `type` + +Status/session results may include `contract_version` and `capabilities`. Treat `capabilities.identity.pid`, `capabilities.identity.window_id`, `capabilities.identity.element_index`, and `capabilities.dispatch.background` as the authoritative cross-platform contract for whether the native background loop is available. Use `features` for backend-specific refinements and skill selection. diff --git a/plugins/_a0_connector/skills/host-computer-use/SKILL.md b/plugins/_a0_connector/skills/host-computer-use/SKILL.md index b8815c5d9..3d4cc543b 100644 --- a/plugins/_a0_connector/skills/host-computer-use/SKILL.md +++ b/plugins/_a0_connector/skills/host-computer-use/SKILL.md @@ -75,14 +75,15 @@ If any tool result contains `COMPUTER_USE_REARM_REQUIRED` or `status=rearm requi ## Core Loop 1. Call `start_session` first. -2. Read the returned `backend_id`, `backend_family`, and `features`; load a backend-specific Computer Use skill when the task needs backend-only affordances. -3. If the backend advertises `native-window-list`, call `list_windows` before using coordinates. -4. If the backend advertises `window-state` and `element-index-targeting`, call `get_window_state` for the target `pid`/`window_id`, then use `element_action` with `dispatch: "background"` by default. -5. If `element_action` reports `background_unavailable`, use `dispatch: "auto"` or `dispatch: "foreground"` only when foreground control is acceptable for the user/task. -6. Decide final success from the latest screenshot or a definitive structural result, not from memory. -7. Interactive actions already attach a fresh screenshot after they run; inspect it before claiming the requested outcome succeeded. -8. Use `status` for state without starting a session. -9. Use `capture` only when you need another screenshot without taking an action. +2. Read the returned `backend_id`, `backend_family`, `features`, `contract_version`, and `capabilities`; load a backend-specific Computer Use skill when the task needs backend-only affordances. +3. Prefer the structured `capabilities` object over guessing from OS names. Use `capabilities.identity.pid`, `capabilities.identity.window_id`, `capabilities.identity.element_index`, and `capabilities.dispatch.background` as the portable contract for the native background loop. +4. If the backend advertises native window listing through capabilities or `native-window-list`, call `list_windows` before using coordinates. +5. If the backend advertises window state and element-index targeting through capabilities or features, call `get_window_state` for the target `pid`/`window_id`, then use `element_action` with `dispatch: "background"` by default. +6. If `element_action` reports `background_unavailable`, use `dispatch: "auto"` or `dispatch: "foreground"` only when foreground control is acceptable for the user/task. +7. Decide final success from the latest screenshot or a definitive structural result, not from memory. +8. Interactive actions already attach a fresh screenshot after they run; inspect it before claiming the requested outcome succeeded. +9. Use `status` for state without starting a session. +10. Use `capture` only when you need another screenshot without taking an action. ## Backend Skills diff --git a/plugins/_a0_connector/tools/computer_use_remote.py b/plugins/_a0_connector/tools/computer_use_remote.py index e8a03168c..f1d5d2c5a 100644 --- a/plugins/_a0_connector/tools/computer_use_remote.py +++ b/plugins/_a0_connector/tools/computer_use_remote.py @@ -505,6 +505,11 @@ class ComputerUseRemote(Tool): if backend_family: backend_text = f"{backend_text}/{backend_family}" parts.append(f"backend={backend_text}") + contract_version = data.get("contract_version") + if not contract_version and isinstance(data.get("capabilities"), dict): + contract_version = data["capabilities"].get("contract_version") + if contract_version: + parts.append(f"contract=v{contract_version}") if features: parts.append(f"features={', '.join(features)}") return ", ".join(parts) diff --git a/tests/test_a0_connector_computer_use_metadata.py b/tests/test_a0_connector_computer_use_metadata.py new file mode 100644 index 000000000..f803b939e --- /dev/null +++ b/tests/test_a0_connector_computer_use_metadata.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from plugins._a0_connector.helpers import ws_runtime + + +def test_computer_use_metadata_preserves_structured_capabilities() -> None: + sid = "sid-capabilities" + ws_runtime.clear_sid_computer_use_metadata(sid) + + ws_runtime.store_sid_computer_use_metadata( + sid, + { + "supported": True, + "enabled": True, + "trust_mode": "allow", + "status": "active", + "last_error": "", + "restore_token_present": True, + "artifact_root": "/a0/tmp/_a0_connector/computer_use", + "backend_id": "windows", + "backend_family": "windows", + "features": ["native-window-list", "element-index-targeting"], + "contract_version": 1, + "capabilities": { + "contract_version": 1, + "identity": { + "pid": True, + "window_id": True, + "element_index": True, + }, + "dispatch": { + "default": "background", + "background": True, + }, + }, + "support_reason": "Windows UIA backend is available.", + }, + ) + + metadata = ws_runtime.computer_use_metadata_for_sid(sid) + + assert metadata is not None + assert metadata["contract_version"] == 1 + assert metadata["capabilities"]["identity"]["element_index"] is True + assert metadata["capabilities"]["dispatch"]["default"] == "background" + + metadata["capabilities"]["dispatch"]["default"] = "foreground" + fresh_metadata = ws_runtime.computer_use_metadata_for_sid(sid) + assert fresh_metadata is not None + assert fresh_metadata["capabilities"]["dispatch"]["default"] == "background" + + ws_runtime.clear_sid_computer_use_metadata(sid) From 45a223dae0895376c470926d722ea31fd0f22c43 Mon Sep 17 00:00:00 2001 From: Louis Giliberto Date: Sat, 30 May 2026 20:12:50 -0700 Subject: [PATCH 032/231] FIX: Fix per-subagent model configs --- plugins/_model_config/webui/config.html | 14 +++++++++++--- plugins/_model_config/webui/model-config-store.js | 2 +- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/plugins/_model_config/webui/config.html b/plugins/_model_config/webui/config.html index ebfde4021..31e477ade 100644 --- a/plugins/_model_config/webui/config.html +++ b/plugins/_model_config/webui/config.html @@ -7,12 +7,20 @@ -