| 1 | from __future__ import annotations |
| 2 | |
| 3 | from helpers.api import ApiHandler, Request |
| 4 | from plugins._desktop.helpers import desktop_session |
| 5 | from plugins._office.helpers import document_store |
| 6 | from plugins._office.helpers import libreoffice |
| 7 | |
| 8 | |
| 9 | class DesktopSession(ApiHandler): |
| 10 | async def process(self, input: dict, request: Request) -> dict: |
| 11 | action = str(input.get("action") or "desktop").lower().strip() |
| 12 | |
| 13 | if action == "status": |
| 14 | return desktop_session.collect_desktop_status() |
| 15 | if action in {"desktop", "open", "session"}: |
| 16 | return self._desktop() |
| 17 | if action in {"open_document", "document"}: |
| 18 | return self._open_document(input, request) |
| 19 | if action in {"save", "desktop_save"}: |
| 20 | return self._save(input) |
| 21 | if action in {"sync", "desktop_sync", "heartbeat"}: |
| 22 | return self._sync(input) |
| 23 | if action in {"state", "desktop_state"}: |
| 24 | return self._state(input) |
| 25 | if action in {"shutdown", "desktop_shutdown"}: |
| 26 | return self._shutdown(input) |
| 27 | return {"ok": False, "error": f"Unsupported desktop session action: {action}"} |
| 28 | |
| 29 | def _desktop(self) -> dict: |
| 30 | desktop = desktop_session.get_manager().ensure_system_desktop() |
| 31 | if not desktop.get("available"): |
| 32 | return { |
| 33 | "ok": False, |
| 34 | "error": desktop.get("error") or "Desktop session is unavailable.", |
| 35 | "status": desktop.get("status") or {}, |
| 36 | "desktop": desktop, |
| 37 | "libreoffice": libreoffice.collect_status(), |
| 38 | } |
| 39 | document = { |
| 40 | "file_id": desktop_session.SYSTEM_FILE_ID, |
| 41 | "path": desktop["path"], |
| 42 | "basename": desktop["title"], |
| 43 | "title": desktop["title"], |
| 44 | "extension": "desktop", |
| 45 | "size": 0, |
| 46 | "version": 0, |
| 47 | } |
| 48 | return { |
| 49 | "ok": True, |
| 50 | "session_id": desktop["session_id"], |
| 51 | "desktop_session_id": desktop["session_id"], |
| 52 | "file_id": desktop_session.SYSTEM_FILE_ID, |
| 53 | "title": desktop["title"], |
| 54 | "extension": "desktop", |
| 55 | "path": desktop["path"], |
| 56 | "text": "", |
| 57 | "document": document, |
| 58 | "version": 0, |
| 59 | "desktop": desktop, |
| 60 | "store_session_id": "", |
| 61 | "mode": "desktop", |
| 62 | } |
| 63 | |
| 64 | def _open_document(self, input: dict, request: Request) -> dict: |
| 65 | context_id = str(input.get("ctxid") or input.get("context_id") or "").strip() |
| 66 | file_id = str(input.get("file_id") or "").strip() |
| 67 | try: |
| 68 | doc = ( |
| 69 | document_store.get_document(file_id) |
| 70 | if file_id |
| 71 | else document_store.register_document(str(input.get("path") or ""), context_id=context_id) |
| 72 | ) |
| 73 | except Exception as exc: |
| 74 | return {"ok": False, "error": str(exc)} |
| 75 | |
| 76 | ext = str(doc.get("extension") or "").lower() |
| 77 | if ext in document_store.EDITOR_TEXT_EXTENSIONS: |
| 78 | return { |
| 79 | "ok": False, |
| 80 | "error": "Text documents use the Editor surface.", |
| 81 | "requires_editor": True, |
| 82 | "document": _public_doc(doc), |
| 83 | } |
| 84 | if ext not in desktop_session.OFFICIAL_EXTENSIONS: |
| 85 | return {"ok": False, "error": f".{ext} documents do not use the Desktop surface."} |
| 86 | |
| 87 | store_session = document_store.create_session( |
| 88 | doc["file_id"], |
| 89 | user_id=str(input.get("user_id") or "agent-zero-user"), |
| 90 | permission="write", |
| 91 | origin=self._origin(request), |
| 92 | ) |
| 93 | desktop = desktop_session.get_manager().open(doc, refresh=input.get("refresh") is True) |
| 94 | if not desktop.get("available"): |
| 95 | document_store.close_session(session_id=store_session["session_id"]) |
| 96 | return { |
| 97 | "ok": False, |
| 98 | "error": desktop.get("error") or desktop.get("reason") or "Desktop session is unavailable.", |
| 99 | "desktop": desktop, |
| 100 | "libreoffice": libreoffice.collect_status(), |
| 101 | } |
| 102 | return { |
| 103 | "ok": True, |
| 104 | "session_id": desktop["session_id"], |
| 105 | "desktop_session_id": desktop["session_id"], |
| 106 | "file_id": doc["file_id"], |
| 107 | "title": doc["basename"], |
| 108 | "extension": doc["extension"], |
| 109 | "path": doc["path"], |
| 110 | "text": "", |
| 111 | "document": _public_doc(doc), |
| 112 | "version": document_store.item_version(doc), |
| 113 | "desktop": desktop, |
| 114 | "store_session_id": store_session["session_id"], |
| 115 | "mode": "edit", |
| 116 | } |
| 117 | |
| 118 | def _save(self, input: dict) -> dict: |
| 119 | session_id = str(input.get("desktop_session_id") or input.get("session_id") or "").strip() |
| 120 | if not session_id: |
| 121 | return {"ok": False, "error": "desktop_session_id is required."} |
| 122 | return desktop_session.get_manager().save( |
| 123 | session_id, |
| 124 | file_id=str(input.get("file_id") or ""), |
| 125 | ) |
| 126 | |
| 127 | def _sync(self, input: dict) -> dict: |
| 128 | return desktop_session.get_manager().sync( |
| 129 | session_id=str(input.get("desktop_session_id") or input.get("session_id") or ""), |
| 130 | file_id=str(input.get("file_id") or ""), |
| 131 | ) |
| 132 | |
| 133 | def _state(self, input: dict) -> dict: |
| 134 | return desktop_session.get_manager().state( |
| 135 | include_screenshot=bool(input.get("include_screenshot") is True), |
| 136 | context_id=str(input.get("ctxid") or input.get("context_id") or ""), |
| 137 | ) |
| 138 | |
| 139 | def _shutdown(self, input: dict) -> dict: |
| 140 | return desktop_session.get_manager().shutdown_system_desktop( |
| 141 | save_first=input.get("save_first") is not False, |
| 142 | source=str(input.get("source") or "api"), |
| 143 | ) |
| 144 | |
| 145 | def _origin(self, request: Request) -> str: |
| 146 | origin = request.headers.get("Origin") or request.host_url.rstrip("/") |
| 147 | return origin.rstrip("/") |
| 148 | |
| 149 | |
| 150 | def _public_doc(doc: dict) -> dict: |
| 151 | return { |
| 152 | "file_id": doc["file_id"], |
| 153 | "path": document_store.display_path(doc["path"]), |
| 154 | "basename": doc["basename"], |
| 155 | "title": doc["basename"], |
| 156 | "extension": doc["extension"], |
| 157 | "size": doc["size"], |
| 158 | "version": document_store.item_version(doc), |
| 159 | "last_modified": doc["last_modified"], |
| 160 | } |