mcp-stata 1.2.2__py3-none-any.whl → 1.6.2__py3-none-any.whl

mcp_stata/ui_http.py

@@ -0,0 +1,540 @@
+import hashlib
+import json
+import secrets
+import threading
+import time
+import uuid
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Any, Callable, Optional
+
+from .stata_client import StataClient
+
+
+def _stable_hash(payload: dict[str, Any]) -> str:
+    return hashlib.sha1(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()
+
+
+@dataclass
+class UIChannelInfo:
+    base_url: str
+    token: str
+    expires_at: int
+
+
+@dataclass
+class ViewHandle:
+    view_id: str
+    dataset_id: str
+    frame: str
+    obs_indices: list[int]
+    filtered_n: int
+    created_at: float
+    last_access: float
+
+
+class UIChannelManager:
+    def __init__(
+        self,
+        client: StataClient,
+        *,
+        host: str = "127.0.0.1",
+        port: int = 0,
+        token_ttl_s: int = 20 * 60,
+        view_ttl_s: int = 30 * 60,
+        max_limit: int = 500,
+        max_vars: int = 200,
+        max_chars: int = 500,
+        max_request_bytes: int = 1_000_000,
+    ):
+        self._client = client
+        self._host = host
+        self._port = port
+        self._token_ttl_s = token_ttl_s
+        self._view_ttl_s = view_ttl_s
+        self._max_limit = max_limit
+        self._max_vars = max_vars
+        self._max_chars = max_chars
+        self._max_request_bytes = max_request_bytes
+
+        self._lock = threading.Lock()
+        self._httpd: ThreadingHTTPServer | None = None
+        self._thread: threading.Thread | None = None
+
+        self._token: str | None = None
+        self._expires_at: int = 0
+
+        self._dataset_version: int = 0
+        self._dataset_id_cache: str | None = None
+        self._dataset_id_cache_at_version: int = -1
+
+        self._views: dict[str, ViewHandle] = {}
+
+    def notify_potential_dataset_change(self) -> None:
+        with self._lock:
+            self._dataset_version += 1
+            self._dataset_id_cache = None
+            self._views.clear()
+
+    def get_channel(self) -> UIChannelInfo:
+        self._ensure_http_server()
+        with self._lock:
+            self._ensure_token()
+            assert self._httpd is not None
+            port = self._httpd.server_address[1]
+            base_url = f"http://{self._host}:{port}"
+            return UIChannelInfo(base_url=base_url, token=self._token or "", expires_at=self._expires_at)
+
+    def capabilities(self) -> dict[str, bool]:
+        return {"dataBrowser": True, "filtering": True}
+
+    def current_dataset_id(self) -> str:
+        with self._lock:
+            if self._dataset_id_cache is not None and self._dataset_id_cache_at_version == self._dataset_version:
+                return self._dataset_id_cache
+
+        state = self._client.get_dataset_state()
+        payload = {
+            "version": self._dataset_version,
+            "frame": state.get("frame"),
+            "n": state.get("n"),
+            "k": state.get("k"),
+            "sortlist": state.get("sortlist"),
+        }
+        digest = _stable_hash(payload)
+
+        with self._lock:
+            self._dataset_id_cache = digest
+            self._dataset_id_cache_at_version = self._dataset_version
+            return digest
+
+    def get_view(self, view_id: str) -> Optional[ViewHandle]:
+        now = time.time()
+        with self._lock:
+            self._evict_expired_locked(now)
+            view = self._views.get(view_id)
+            if view is None:
+                return None
+            view.last_access = now
+            return view
+
+    def create_view(self, *, dataset_id: str, frame: str, filter_expr: str) -> ViewHandle:
+        current_id = self.current_dataset_id()
+        if dataset_id != current_id:
+            raise DatasetChangedError(current_id)
+
+        try:
+            obs_indices = self._client.compute_view_indices(filter_expr)
+        except ValueError as e:
+            raise InvalidFilterError(str(e))
+        except RuntimeError as e:
+            msg = str(e) or "No data in memory"
+            if "no data" in msg.lower():
+                raise NoDataInMemoryError(msg)
+            raise
+        now = time.time()
+        view_id = f"view_{uuid.uuid4().hex}"
+        view = ViewHandle(
+            view_id=view_id,
+            dataset_id=current_id,
+            frame=frame,
+            obs_indices=obs_indices,
+            filtered_n=len(obs_indices),
+            created_at=now,
+            last_access=now,
+        )
+        with self._lock:
+            self._evict_expired_locked(now)
+            self._views[view_id] = view
+        return view
+
+    def delete_view(self, view_id: str) -> bool:
+        with self._lock:
+            return self._views.pop(view_id, None) is not None
+
+    def validate_token(self, header_value: str | None) -> bool:
+        if not header_value:
+            return False
+        if not header_value.startswith("Bearer "):
+            return False
+        token = header_value[len("Bearer ") :].strip()
+        with self._lock:
+            self._ensure_token()
+            if self._token is None:
+                return False
+            if time.time() * 1000 >= self._expires_at:
+                return False
+            return secrets.compare_digest(token, self._token)
+
+    def limits(self) -> tuple[int, int, int, int]:
+        return self._max_limit, self._max_vars, self._max_chars, self._max_request_bytes
+
+    def _ensure_token(self) -> None:
+        now_ms = int(time.time() * 1000)
+        if self._token is None or now_ms >= self._expires_at:
+            self._token = secrets.token_urlsafe(32)
+            self._expires_at = int((time.time() + self._token_ttl_s) * 1000)
+
+    def _evict_expired_locked(self, now: float) -> None:
+        expired: list[str] = []
+        for key, view in self._views.items():
+            if now - view.last_access >= self._view_ttl_s:
+                expired.append(key)
+        for key in expired:
+            self._views.pop(key, None)
+
+    def _ensure_http_server(self) -> None:
+        with self._lock:
+            if self._httpd is not None:
+                return
+
+            manager = self
+
+            class Handler(BaseHTTPRequestHandler):
+                def _send_json(self, status: int, payload: dict[str, Any]) -> None:
+                    data = json.dumps(payload).encode("utf-8")
+                    self.send_response(status)
+                    self.send_header("Content-Type", "application/json")
+                    self.send_header("Content-Length", str(len(data)))
+                    self.end_headers()
+                    self.wfile.write(data)
+
+                def _error(self, status: int, code: str, message: str, *, stata_rc: int | None = None) -> None:
+                    body: dict[str, Any] = {"error": {"code": code, "message": message}}
+                    if stata_rc is not None:
+                        body["error"]["stataRc"] = stata_rc
+                    self._send_json(status, body)
+
+                def _require_auth(self) -> bool:
+                    if manager.validate_token(self.headers.get("Authorization")):
+                        return True
+                    self._error(401, "auth_failed", "Unauthorized")
+                    return False
+
+                def _read_json(self) -> dict[str, Any] | None:
+                    max_limit, max_vars, max_chars, max_bytes = manager.limits()
+                    _ = (max_limit, max_vars, max_chars)
+
+                    length = int(self.headers.get("Content-Length", "0") or "0")
+                    if length <= 0:
+                        return {}
+                    if length > max_bytes:
+                        self._error(400, "request_too_large", "Request too large")
+                        return None
+                    raw = self.rfile.read(length)
+                    try:
+                        parsed = json.loads(raw.decode("utf-8"))
+                    except Exception:
+                        self._error(400, "invalid_request", "Invalid JSON")
+                        return None
+                    if not isinstance(parsed, dict):
+                        self._error(400, "invalid_request", "Expected JSON object")
+                        return None
+                    return parsed
+
+                def do_GET(self) -> None:
+                    if not self._require_auth():
+                        return
+
+                    if self.path == "/v1/dataset":
+                        try:
+                            state = manager._client.get_dataset_state()
+                            dataset_id = manager.current_dataset_id()
+                            self._send_json(
+                                200,
+                                {
+                                    "dataset": {
+                                        "id": dataset_id,
+                                        "frame": state.get("frame"),
+                                        "n": state.get("n"),
+                                        "k": state.get("k"),
+                                        "changed": state.get("changed"),
+                                    }
+                                },
+                            )
+                            return
+                        except NoDataInMemoryError as e:
+                            self._error(400, "no_data_in_memory", str(e), stata_rc=e.stata_rc)
+                            return
+                        except Exception as e:
+                            self._error(500, "internal_error", str(e))
+                            return
+
+                    if self.path == "/v1/vars":
+                        try:
+                            state = manager._client.get_dataset_state()
+                            dataset_id = manager.current_dataset_id()
+                            variables = manager._client.list_variables_rich()
+                            self._send_json(
+                                200,
+                                {
+                                    "dataset": {"id": dataset_id, "frame": state.get("frame")},
+                                    "variables": variables,
+                                },
+                            )
+                            return
+                        except NoDataInMemoryError as e:
+                            self._error(400, "no_data_in_memory", str(e), stata_rc=e.stata_rc)
+                            return
+                        except Exception as e:
+                            self._error(500, "internal_error", str(e))
+                            return
+
+                    self._error(404, "not_found", "Not found")
+
+                def do_POST(self) -> None:
+                    if not self._require_auth():
+                        return
+
+                    if self.path == "/v1/page":
+                        body = self._read_json()
+                        if body is None:
+                            return
+                        try:
+                            resp = handle_page_request(manager, body, view_id=None)
+                            self._send_json(200, resp)
+                            return
+                        except HTTPError as e:
+                            self._error(e.status, e.code, e.message, stata_rc=e.stata_rc)
+                            return
+                        except Exception as e:
+                            self._error(500, "internal_error", str(e))
+                            return
+
+                    if self.path == "/v1/views":
+                        body = self._read_json()
+                        if body is None:
+                            return
+                        dataset_id = str(body.get("datasetId", ""))
+                        frame = str(body.get("frame", "default"))
+                        filter_expr = str(body.get("filterExpr", ""))
+                        if not dataset_id or not filter_expr:
+                            self._error(400, "invalid_request", "datasetId and filterExpr are required")
+                            return
+                        try:
+                            view = manager.create_view(dataset_id=dataset_id, frame=frame, filter_expr=filter_expr)
+                            self._send_json(
+                                200,
+                                {
+                                    "dataset": {"id": view.dataset_id, "frame": view.frame},
+                                    "view": {"id": view.view_id, "filteredN": view.filtered_n},
+                                },
+                            )
+                            return
+                        except DatasetChangedError as e:
+                            self._error(409, "dataset_changed", "Dataset changed")
+                            return
+                        except ValueError as e:
+                            self._error(400, "invalid_filter", str(e))
+                            return
+                        except RuntimeError as e:
+                            msg = str(e) or "No data in memory"
+                            if "no data" in msg.lower():
+                                self._error(400, "no_data_in_memory", msg)
+                                return
+                            self._error(500, "internal_error", msg)
+                            return
+                        except Exception as e:
+                            self._error(500, "internal_error", str(e))
+                            return
+
+                    if self.path.startswith("/v1/views/") and self.path.endswith("/page"):
+                        parts = self.path.split("/")
+                        if len(parts) != 5:
+                            self._error(404, "not_found", "Not found")
+                            return
+                        view_id = parts[3]
+                        body = self._read_json()
+                        if body is None:
+                            return
+                        try:
+                            resp = handle_page_request(manager, body, view_id=view_id)
+                            self._send_json(200, resp)
+                            return
+                        except HTTPError as e:
+                            self._error(e.status, e.code, e.message, stata_rc=e.stata_rc)
+                            return
+                        except Exception as e:
+                            self._error(500, "internal_error", str(e))
+                            return
+
+                    if self.path == "/v1/filters/validate":
+                        body = self._read_json()
+                        if body is None:
+                            return
+                        filter_expr = str(body.get("filterExpr", ""))
+                        if not filter_expr:
+                            self._error(400, "invalid_request", "filterExpr is required")
+                            return
+                        try:
+                            manager._client.validate_filter_expr(filter_expr)
+                            self._send_json(200, {"ok": True})
+                            return
+                        except ValueError as e:
+                            self._error(400, "invalid_filter", str(e))
+                            return
+                        except RuntimeError as e:
+                            msg = str(e) or "No data in memory"
+                            if "no data" in msg.lower():
+                                self._error(400, "no_data_in_memory", msg)
+                                return
+                            self._error(500, "internal_error", msg)
+                            return
+                        except Exception as e:
+                            self._error(500, "internal_error", str(e))
+                            return
+
+                    self._error(404, "not_found", "Not found")
+
+                def do_DELETE(self) -> None:
+                    if not self._require_auth():
+                        return
+
+                    if self.path.startswith("/v1/views/"):
+                        parts = self.path.split("/")
+                        if len(parts) != 4:
+                            self._error(404, "not_found", "Not found")
+                            return
+                        view_id = parts[3]
+                        if manager.delete_view(view_id):
+                            self._send_json(200, {"ok": True})
+                        else:
+                            self._error(404, "not_found", "Not found")
+                        return
+
+                    self._error(404, "not_found", "Not found")
+
+                def log_message(self, format: str, *args: Any) -> None:
+                    return
+
+            httpd = ThreadingHTTPServer((self._host, self._port), Handler)
+            t = threading.Thread(target=httpd.serve_forever, daemon=True)
+            t.start()
+            self._httpd = httpd
+            self._thread = t
+
+
+class HTTPError(Exception):
+    def __init__(self, status: int, code: str, message: str, *, stata_rc: int | None = None):
+        super().__init__(message)
+        self.status = status
+        self.code = code
+        self.message = message
+        self.stata_rc = stata_rc
+
+
+class DatasetChangedError(Exception):
+    def __init__(self, current_dataset_id: str):
+        super().__init__("dataset_changed")
+        self.current_dataset_id = current_dataset_id
+
+
+class NoDataInMemoryError(Exception):
+    def __init__(self, message: str = "No data in memory", *, stata_rc: int | None = None):
+        super().__init__(message)
+        self.stata_rc = stata_rc
+
+
+class InvalidFilterError(Exception):
+    def __init__(self, message: str, *, stata_rc: int | None = None):
+        super().__init__(message)
+        self.message = message
+        self.stata_rc = stata_rc
+
+
+def handle_page_request(manager: UIChannelManager, body: dict[str, Any], *, view_id: str | None) -> dict[str, Any]:
+    max_limit, max_vars, max_chars, _ = manager.limits()
+
+    if view_id is None:
+        dataset_id = str(body.get("datasetId", ""))
+        frame = str(body.get("frame", "default"))
+    else:
+        view = manager.get_view(view_id)
+        if view is None:
+            raise HTTPError(404, "not_found", "View not found")
+        dataset_id = view.dataset_id
+        frame = view.frame
+
+    offset = int(body.get("offset", 0) or 0)
+    limit = int(body.get("limit", 0) or 0)
+    vars_req = body.get("vars", [])
+    include_obs_no = bool(body.get("includeObsNo", False))
+    max_chars_req = int(body.get("maxChars", max_chars) or max_chars)
+
+    if offset < 0:
+        raise HTTPError(400, "invalid_request", "offset must be >= 0")
+    if limit <= 0:
+        raise HTTPError(400, "invalid_request", "limit must be > 0")
+    if limit > max_limit:
+        raise HTTPError(400, "request_too_large", f"limit must be <= {max_limit}")
+    if max_chars_req <= 0:
+        raise HTTPError(400, "invalid_request", "maxChars must be > 0")
+    if max_chars_req > max_chars:
+        raise HTTPError(400, "request_too_large", f"maxChars must be <= {max_chars}")
+
+    if not isinstance(vars_req, list) or not all(isinstance(v, str) for v in vars_req):
+        raise HTTPError(400, "invalid_request", "vars must be a list of strings")
+    if len(vars_req) > max_vars:
+        raise HTTPError(400, "request_too_large", f"vars length must be <= {max_vars}")
+
+    current_id = manager.current_dataset_id()
+    if dataset_id != current_id:
+        raise HTTPError(409, "dataset_changed", "Dataset changed")
+
+    if view_id is None:
+        obs_indices = None
+        filtered_n: int | None = None
+    else:
+        assert view is not None
+        obs_indices = view.obs_indices
+        filtered_n = view.filtered_n
+
+    try:
+        dataset_state = manager._client.get_dataset_state()
+        page = manager._client.get_page(
+            offset=offset,
+            limit=limit,
+            vars=vars_req,
+            include_obs_no=include_obs_no,
+            max_chars=max_chars_req,
+            obs_indices=obs_indices,
+        )
+    except RuntimeError as e:
+        # StataClient uses RuntimeError("No data in memory") for empty dataset.
+        msg = str(e) or "No data in memory"
+        if "no data" in msg.lower():
+            raise HTTPError(400, "no_data_in_memory", msg)
+        raise HTTPError(500, "internal_error", msg)
+    except ValueError as e:
+        msg = str(e)
+        if msg.lower().startswith("invalid variable"):
+            raise HTTPError(400, "invalid_variable", msg)
+        raise HTTPError(400, "invalid_request", msg)
+    except Exception as e:
+        raise HTTPError(500, "internal_error", str(e))
+
+    view_obj: dict[str, Any] = {
+        "offset": offset,
+        "limit": limit,
+        "returned": page["returned"],
+        "filteredN": filtered_n,
+    }
+    if view_id is not None:
+        view_obj["viewId"] = view_id
+
+    return {
+        "dataset": {
+            "id": current_id,
+            "frame": dataset_state.get("frame"),
+            "n": dataset_state.get("n"),
+            "k": dataset_state.get("k"),
+        },
+        "view": view_obj,
+        "vars": page["vars"],
+        "rows": page["rows"],
+        "display": {
+            "maxChars": max_chars_req,
+            "truncatedCells": page["truncated_cells"],
+            "missing": ".",
+        },
+    }