Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

gateway/media_repair.py

@@ -0,0 +1,213 @@
+"""Repair model-mangled ``computer_use`` screenshot paths in final responses.
+
+``computer_use`` persists a bounded screenshot into the Hermes image cache and
+tells the model its absolute path.  Some models rewrite a Windows path into a
+POSIX-looking one (``C:\\Users\\Alice\\...`` -> ``/Users/Alice/...``) when
+emitting an explicit ``MEDIA:`` directive, so delivery-path validation rejects
+the nonexistent path and the attachment is dropped.
+
+The repair is deliberately narrow: it only rewrites paths inside a response
+that *already* carries an explicit ``MEDIA:`` directive, and only when the
+directive's generated ``computer_use_<uuid>`` basename exactly matches a
+canonical screenshot path returned by ``computer_use`` in the current turn.
+It never auto-attaches captures, and normal media path validation still runs
+after the repair.
+
+This lives in its own module (mirroring ``gateway/media_policy.py``) so every
+delivery surface — the messaging gateway's main turn path, gateway background
+tasks, and cron job delivery — shares one implementation.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from typing import Any, Dict, Iterator, List
+
+logger = logging.getLogger(__name__)
+
+# Absolute-path prefix accepted for canonical capture paths: Windows drive
+# letter, POSIX root, or UNC share. Kept as a pattern string so the summary
+# regex below and the compiled prefix check stay in sync.
+_ABS_PATH_PREFIX_PATTERN = r"(?:[A-Za-z]:[/\\]|/|\\\\)"
+_ABS_PATH_PREFIX_RE = re.compile(r"^" + _ABS_PATH_PREFIX_PATTERN)
+
+_COMPUTER_USE_CAPTURE_BASENAME_RE = re.compile(
+    r"^computer_use_[0-9a-f]{32}\.(?:png|jpe?g)$",
+    re.IGNORECASE,
+)
+_COMPUTER_USE_CAPTURE_SUMMARY_RE = re.compile(
+    r"\(shareable screenshot saved to "
+    r"(?P<path>" + _ABS_PATH_PREFIX_PATTERN + r"[^\r\n]*?"
+    r"computer_use_[0-9a-f]{32}\.(?:png|jpe?g))\)",
+    re.IGNORECASE,
+)
+
+
+def tool_name_by_call_id(messages: List[Dict[str, Any]]) -> Dict[str, str]:
+    """Map assistant tool-call ids to tool names for the given messages."""
+    mapping: Dict[str, str] = {}
+    for msg in messages:
+        if msg.get("role") != "assistant":
+            continue
+        for call in msg.get("tool_calls") or []:
+            call_id = call.get("id") or call.get("call_id")
+            fn = call.get("function") or {}
+            name = str(fn.get("name") or call.get("name") or "")
+            if call_id and name:
+                mapping[str(call_id)] = name
+    return mapping
+
+
+def _computer_use_capture_basename(path: Any) -> str:
+    """Return a canonical capture basename for either path separator style."""
+    value = str(path or "").strip().strip("`\"'")
+    basename = re.split(r"[/\\]", value)[-1]
+    if _COMPUTER_USE_CAPTURE_BASENAME_RE.fullmatch(basename):
+        return basename.lower()
+    return ""
+
+
+def _iter_computer_use_capture_paths(content: Any) -> Iterator[str]:
+    """Yield persisted screenshot paths from computer_use result content.
+
+    The tool can return JSON, a multimodal content list, or a text fallback.
+    The latter two retain the canonical path in the human-readable summary
+    even though the multimodal envelope's ``meta`` dictionary is not stored in
+    the tool message.
+    """
+    if isinstance(content, str):
+        stripped = content.strip()
+        if stripped.startswith(("{", "[")):
+            # JSON-looking content: parse first, never regex-scan the raw
+            # text. JSON escaping doubles backslashes, so a summary-regex hit
+            # on the raw string would yield ``C:\\\\Users\\\\...`` — a path
+            # that exists nowhere. Fail closed on unparseable JSON (tool
+            # output truncation) rather than repair to a corrupted path.
+            try:
+                payload = json.loads(stripped)
+            except Exception:
+                return
+            if isinstance(payload, (dict, list)):
+                yield from _iter_computer_use_capture_paths(payload)
+            return
+        for match in _COMPUTER_USE_CAPTURE_SUMMARY_RE.finditer(content):
+            yield match.group("path").strip()
+        return
+
+    if isinstance(content, list):
+        for part in content:
+            yield from _iter_computer_use_capture_paths(part)
+        return
+
+    if not isinstance(content, dict):
+        return
+
+    screenshot_path = content.get("screenshot_path")
+    if isinstance(screenshot_path, str):
+        yield screenshot_path
+    meta = content.get("meta")
+    if isinstance(meta, dict) and isinstance(meta.get("screenshot_path"), str):
+        yield meta["screenshot_path"]
+    # Producer shapes (tools/computer_use/tool.py::_capture_response):
+    # "content"/"text" — multimodal envelope parts; "text_summary"/"summary" —
+    # the human-readable summary carrying the "(shareable screenshot saved
+    # to ...)" line.
+    for field in ("content", "text", "text_summary", "summary"):
+        nested = content.get(field)
+        if isinstance(nested, (str, dict, list)):
+            yield from _iter_computer_use_capture_paths(nested)
+
+
+def repair_explicit_computer_use_media_paths(
+    response: str,
+    messages: List[Dict[str, Any]],
+    history_offset: int = 0,
+) -> str:
+    """Recover model-mangled paths for explicitly requested screenshots.
+
+    Repair only an already-explicit ``MEDIA:`` directive whose unique
+    generated basename case-insensitively matches a canonical screenshot
+    path from this turn. This does not auto-attach ordinary computer-use
+    captures, and normal media path validation still runs after the repair.
+
+    Fail-open: the repair is cosmetic, so an unexpected error returns the
+    response unchanged rather than aborting delivery.
+    """
+    try:
+        return _repair_explicit_computer_use_media_paths_inner(
+            response, messages, history_offset
+        )
+    except Exception:
+        logger.debug("computer_use media path repair failed", exc_info=True)
+        return response
+
+
+def _repair_explicit_computer_use_media_paths_inner(
+    response: str,
+    messages: List[Dict[str, Any]],
+    history_offset: int = 0,
+) -> str:
+    if "MEDIA:" not in response:
+        return response
+
+    if history_offset and len(messages) >= history_offset:
+        turn_messages = messages[history_offset:]
+    elif history_offset:
+        # Compression can invalidate the original slice boundary. Recover the
+        # current turn from its last user message; fail closed if none
+        # remains. (Deliberately narrower than the scan-everything fallback
+        # in gateway/run.py::_collect_auto_append_media_tags — that helper
+        # decides whether to ATTACH, this one only rewrites paths the model
+        # already explicitly emitted.)
+        last_user = next(
+            (
+                index
+                for index in range(len(messages) - 1, -1, -1)
+                if messages[index].get("role") == "user"
+            ),
+            None,
+        )
+        turn_messages = messages[last_user:] if last_user is not None else []
+    else:
+        turn_messages = messages
+
+    call_id_names = tool_name_by_call_id(turn_messages)
+
+    canonical_by_basename: Dict[str, str] = {}
+    for msg in turn_messages:
+        if msg.get("role") not in {"tool", "function"}:
+            continue
+        call_id = str(msg.get("tool_call_id") or msg.get("call_id") or "")
+        tool_name = str(
+            msg.get("name")
+            or msg.get("tool_name")
+            or call_id_names.get(call_id)
+            or ""
+        )
+        if tool_name != "computer_use":
+            continue
+        for path in _iter_computer_use_capture_paths(msg.get("content")):
+            basename = _computer_use_capture_basename(path)
+            if basename and _ABS_PATH_PREFIX_RE.match(path):
+                canonical_by_basename[basename] = path
+
+    if not canonical_by_basename:
+        return response
+
+    # Lazy on purpose: keeps `import gateway.media_repair` cheap for
+    # standalone cron processes that may never hit a MEDIA: response.
+    # No import cycle either way (base.py imports neither this module
+    # nor gateway.run at module level).
+    from gateway.platforms.base import BasePlatformAdapter
+
+    media_files, _ = BasePlatformAdapter.extract_media(response)
+    repaired = response
+    for emitted_path, _is_voice in media_files:
+        canonical = canonical_by_basename.get(
+            _computer_use_capture_basename(emitted_path)
+        )
+        if canonical and emitted_path != canonical:
+            repaired = repaired.replace(emitted_path, canonical)
+    return repaired