144 lines
6.8 KiB
Python
144 lines
6.8 KiB
Python
"""Repair model-mangled ``computer_use`` screenshot paths in final responses.
|
|
|
|
Some models rewrite the Windows path ``computer_use`` reported into a POSIX-looking
|
|
one (``C:\\Users\\Alice\\...`` -> ``/Users/Alice/...``) inside an explicit ``MEDIA:``
|
|
directive, so delivery-path validation rejects it and drops the attachment.
|
|
Deliberately narrow: only rewrites paths in a response that *already* carries a
|
|
``MEDIA:`` directive whose ``computer_use_<uuid>`` basename exactly matches a
|
|
canonical path returned this turn. Never auto-attaches; validation still runs.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
from typing import Any, Dict, Iterator, List
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Absolute-path prefix for canonical capture paths: Windows drive letter, POSIX
|
|
# root, or UNC share. Shared string so the summary regex stays in sync.
|
|
_ABS_PATH_PREFIX_PATTERN = r"(?:[A-Za-z]:[/\\]|/|\\\\)"
|
|
_ABS_PATH_PREFIX_RE = re.compile(r"^" + _ABS_PATH_PREFIX_PATTERN)
|
|
_CAPTURE_BASENAME_PATTERN = r"computer_use_[0-9a-f]{32}\.(?:png|jpe?g)"
|
|
|
|
_COMPUTER_USE_CAPTURE_BASENAME_RE = re.compile(r"^" + _CAPTURE_BASENAME_PATTERN + r"$", re.IGNORECASE)
|
|
_COMPUTER_USE_CAPTURE_SUMMARY_RE = re.compile(
|
|
r"\(shareable screenshot saved to "
|
|
r"(?P<path>" + _ABS_PATH_PREFIX_PATTERN + r"[^\r\n]*?" + _CAPTURE_BASENAME_PATTERN + r")\)",
|
|
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")
|
|
name = str((call.get("function") or {}).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:
|
|
"""Canonical (lowercased) capture basename for either separator style, or ''."""
|
|
basename = re.split(r"[/\\]", str(path or "").strip().strip("`\"'"))[-1]
|
|
return basename.lower() if _COMPUTER_USE_CAPTURE_BASENAME_RE.fullmatch(basename) else ""
|
|
|
|
|
|
def _iter_computer_use_capture_paths(content: Any) -> Iterator[str]:
|
|
"""Yield persisted screenshot paths from computer_use result content (JSON, a
|
|
multimodal list, or text; the latter two keep the canonical path in the summary
|
|
line since the envelope's ``meta`` is not stored in the tool message)."""
|
|
if isinstance(content, str):
|
|
stripped = content.strip()
|
|
if not stripped.startswith(("{", "[")):
|
|
for match in _COMPUTER_USE_CAPTURE_SUMMARY_RE.finditer(content):
|
|
yield match.group("path").strip()
|
|
return
|
|
# Parse JSON, never regex-scan it: JSON escaping doubles backslashes, so a raw-text
|
|
# hit would yield a path that exists nowhere. Fail closed on truncated JSON.
|
|
try:
|
|
payload = json.loads(stripped)
|
|
except Exception:
|
|
return
|
|
if isinstance(payload, (dict, list)):
|
|
yield from _iter_computer_use_capture_paths(payload)
|
|
elif isinstance(content, list):
|
|
for part in content:
|
|
yield from _iter_computer_use_capture_paths(part)
|
|
elif isinstance(content, dict):
|
|
meta = content.get("meta")
|
|
for holder in (content, meta if isinstance(meta, dict) else {}):
|
|
if isinstance(holder.get("screenshot_path"), str):
|
|
yield holder["screenshot_path"]
|
|
# Producer shapes (computer_use ``_capture_response``): content/text = multimodal
|
|
# parts; text_summary/summary = 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 _current_turn_messages(messages: List[Dict[str, Any]], history_offset: int) -> List[Dict[str, Any]]:
|
|
if not history_offset or len(messages) >= history_offset:
|
|
return messages[history_offset:]
|
|
# Compression can invalidate the slice boundary: recover the turn from its last user
|
|
# message, fail closed if none remains. Narrower than run.py's scan-everything fallback
|
|
# for auto-attach — that decides whether to ATTACH; this only rewrites emitted paths.
|
|
for index in range(len(messages) - 1, -1, -1):
|
|
if messages[index].get("role") == "user":
|
|
return messages[index:]
|
|
return []
|
|
|
|
|
|
def repair_explicit_computer_use_media_paths(response: str, messages: List[Dict[str, Any]], history_offset: int = 0) -> str:
|
|
"""Recover model-mangled paths in explicit ``MEDIA:`` directives whose basename
|
|
matches (case-insensitively) a canonical screenshot path from this turn.
|
|
Fail-open: the repair is cosmetic, so any error returns the response unchanged."""
|
|
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 _canonical_capture_paths(turn_messages: List[Dict[str, Any]]) -> Dict[str, str]:
|
|
"""``{lowercase basename: absolute canonical path}`` from this turn's computer_use results."""
|
|
call_id_names = tool_name_by_call_id(turn_messages)
|
|
canonical: 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 "")
|
|
if str(msg.get("name") or msg.get("tool_name") or call_id_names.get(call_id) or "") != "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[basename] = path
|
|
return canonical
|
|
|
|
|
|
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
|
|
canonical_by_basename = _canonical_capture_paths(_current_turn_messages(messages, history_offset))
|
|
if not canonical_by_basename:
|
|
return response
|
|
|
|
# Lazy: keeps the import cheap for standalone cron processes that may never hit MEDIA:.
|
|
from gateway.platforms.base import BasePlatformAdapter
|
|
|
|
repaired = response
|
|
for emitted_path, _is_voice in BasePlatformAdapter.extract_media(response)[0]:
|
|
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
|