# -*- coding: utf-8 -*- # pylint: disable=protected-access """Wrap ``Workspace.stream_query`` to emit pet lifecycle events. Emissions: - ``query.received`` / ``query.running`` — fired up-front - ``query.first_token`` (Replying) — first ``Message`` envelope with ``type=message`` and ``status=in_progress`` - ``tool.detected`` (Using ) — ``Message`` with ``type=plugin_call`` (``DataContent.data["name"]``) - ``tool.result`` — ``Message`` with ``type=plugin_call_output`` and ``status=completed`` - ``query.done`` / ``query.cancelled`` / ``query.error`` — driven by stream exit """ from __future__ import annotations import asyncio import logging from typing import Any from emitter import emit_pet_event, schedule_emit_pet_event logger = logging.getLogger("qwenpaw.pet_desktop") _ORIGINAL_STREAM_QUERY = None _PATCHED = False # Bound methods captured by ``make_process_from_runner`` freeze the # function pointer at access time and bypass class-level patches; we # track our replacements so ``restore_agent_runner`` can undo them. _PATCHED_CHANNEL_PROCESSES: list[tuple[Any, Any]] = [] def _request_meta(runner: Any, request: Any) -> dict[str, Any]: return { "agent_id": getattr(runner, "agent_id", "default"), "agent_name": getattr(runner, "agent_name", "QwenPaw"), "session_id": getattr(request, "session_id", "") if request else "", "user_id": getattr(request, "user_id", "") if request else "", "channel": getattr(request, "channel", "") if request else "", } def _content_text(content_block: Any) -> str: """Pull a text string out of a TextContent-shaped block.""" text = getattr(content_block, "text", None) if text is None and isinstance(content_block, dict): text = content_block.get("text") return str(text) if text else "" def _content_data(content_block: Any) -> dict | None: """Pull ``data`` dict out of a DataContent-shaped block.""" data = getattr(content_block, "data", None) if data is None and isinstance(content_block, dict): data = content_block.get("data") return data if isinstance(data, dict) else None def _classify_event( event: Any, seen_first_text: bool, ) -> tuple[str | None, str | None]: """Map a single envelope event to a pet lifecycle emission. Returns ``(event_name, text)`` or ``(None, None)`` for events the pet doesn't speak about. ``seen_first_text`` guards the ``query.first_token`` emission to fire only once per turn. """ object_kind = getattr(event, "object", None) if object_kind is None and isinstance(event, dict): object_kind = event.get("object") if object_kind != "message": return None, None msg_type = getattr(event, "type", None) if msg_type is None and isinstance(event, dict): msg_type = event.get("type") msg_type_val = ( msg_type.value if hasattr(msg_type, "value") else str(msg_type or "") ) status = getattr(event, "status", None) if status is None and isinstance(event, dict): status = event.get("status") status_val = ( status.value if hasattr(status, "value") else str(status or "") ) content = getattr(event, "content", None) if content is None and isinstance(event, dict): content = event.get("content") content_list = content if isinstance(content, list) else [] if msg_type_val == "plugin_call" and status_val == "completed": # First content block carries {name, call_id, arguments}. data = _content_data(content_list[0]) if content_list else None name = (data or {}).get("name") or "tool" return "tool.detected", f"Using {str(name)[:40]}" if msg_type_val == "plugin_call_output" and status_val == "completed": return "tool.result", "Tool result" # Fire on the empty in_progress envelope, not the completed one — # TextContent deltas arrive as separate ``object="content"`` events # which the classifier ignores, so "Replying" needs to flip when the # reply starts, not after it finishes. if ( msg_type_val == "message" and status_val == "in_progress" and not seen_first_text ): return "query.first_token", "Replying" return None, None def patch_agent_runner() -> None: """Install an observing wrapper around Workspace.stream_query.""" global _ORIGINAL_STREAM_QUERY, _PATCHED if _PATCHED: return from pineagents.app.workspace.workspace import Workspace _ORIGINAL_STREAM_QUERY = Workspace.stream_query async def patched_stream_query(self, request, *args, **kwargs): meta = _request_meta(self, request) seen_first_text = False last_event = None schedule_emit_pet_event( "query.received", text="New message", **meta, ) schedule_emit_pet_event("query.running", text="Thinking", **meta) try: async for envelope in _ORIGINAL_STREAM_QUERY( self, request, *args, **kwargs, ): event, text = _classify_event(envelope, seen_first_text) if event == "query.first_token": seen_first_text = True if event and event != last_event: schedule_emit_pet_event(event, text=text, **meta) last_event = event yield envelope await asyncio.to_thread( emit_pet_event, "query.done", text="Done", **meta, ) except asyncio.CancelledError: schedule_emit_pet_event( "query.cancelled", text="Interrupted", duration_ms=1200, **meta, ) raise except Exception as exc: schedule_emit_pet_event( "query.error", text=type(exc).__name__, duration_ms=2500, **meta, ) raise Workspace.stream_query = patched_stream_query _PATCHED = True logger.info("QwenPaw Pet patched Workspace.stream_query") # Existing channels hold bound-method copies of # ``runner.stream_query`` from before the patch; the class-level # assignment above doesn't reach them. Walk the live channels and # also intercept the factory for any workspace brought up later. _wire_existing_channels(patched_stream_query) _patch_make_process_factory(patched_stream_query) def _wire_existing_channels(patched_fn) -> None: """Replace each live channel's ``_process`` with a proxy that calls ``patched_fn`` against the channel's owning runner. No-ops when the FastAPI app or its ``multi_agent_manager`` isn't importable (CLI-only mode, headless runner). """ try: from pineagents.app._app import app as fastapi_app except Exception: logger.debug("Pet patch: FastAPI app not importable; skip rewire") return manager = getattr( getattr(fastapi_app, "state", None), "multi_agent_manager", None, ) if manager is None: logger.debug("Pet patch: no multi_agent_manager; skip rewire") return workspaces = getattr(manager, "agents", {}) or {} rewired = 0 for ws in workspaces.values(): try: services = ws._service_manager.services except AttributeError: continue cm = services.get("channel_manager") if cm is None: continue runner = services.get("runner") if runner is None: continue for channel in getattr(cm, "channels", []) or []: original = getattr(channel, "_process", None) if original is None: continue def _make_proxy(_runner, _fn): def _proxy(request, *args, **kwargs): return _fn(_runner, request, *args, **kwargs) _proxy.__qualname__ = "qwenpaw_pet.patched_channel_process" return _proxy channel._process = _make_proxy(runner, patched_fn) _PATCHED_CHANNEL_PROCESSES.append((channel, original)) rewired += 1 if rewired: logger.info( "QwenPaw Pet rewired %d channel _process refs to patched " "stream_query", rewired, ) def _patch_make_process_factory(patched_fn) -> None: """Wrap ``make_process_from_runner`` so workspaces created after this point also receive the patched proxy instead of a bound method.""" try: from pineagents.app.channels import utils as ch_utils except Exception: logger.debug( "Pet patch: channels.utils not importable; skip factory patch", ) return original = getattr(ch_utils, "make_process_from_runner", None) if original is None or getattr(original, "_pet_patched", False): return def patched_factory(runner): def _proxy(request, *args, **kwargs): return patched_fn(runner, request, *args, **kwargs) _proxy.__qualname__ = "qwenpaw_pet.patched_channel_process" return _proxy patched_factory._pet_patched = True # type: ignore[attr-defined] patched_factory._original = original # type: ignore[attr-defined] ch_utils.make_process_from_runner = patched_factory def restore_agent_runner() -> None: """Restore the original stream_query method and channel _process refs.""" global _PATCHED if not _PATCHED or _ORIGINAL_STREAM_QUERY is None: return from pineagents.app.workspace.workspace import Workspace Workspace.stream_query = _ORIGINAL_STREAM_QUERY # Undo channel rewiring for channel, original in _PATCHED_CHANNEL_PROCESSES: try: channel._process = original except Exception: pass _PATCHED_CHANNEL_PROCESSES.clear() # Undo factory patch try: from pineagents.app.channels import utils as ch_utils factory = getattr(ch_utils, "make_process_from_runner", None) if factory is not None and getattr(factory, "_pet_patched", False): ch_utils.make_process_from_runner = factory._original except Exception: pass _PATCHED = False logger.info("QwenPaw Pet restored Workspace.stream_query")