thinkstack-core 4.0.0__py3-none-any.whl

thinkstack_core/codex/proxy.py

@@ -0,0 +1,412 @@
+"""
+thinkstack_core.codex.proxy
+=========================
+Lightweight HTTP proxy server for intercepting OpenAI Codex CLI traffic.
+
+Uses only stdlib (http.server, urllib.request, threading) — no FastAPI or
+httpx dependency required.
+
+Default port: 8766  (main ThinkStack proxy runs on 8765)
+Upstream:     https://api.openai.com
+
+Routes:
+  POST /v1/*         → inject RACP, forward, capture response
+  GET  /health       → 200 {"status": "ok"}
+  *    (everything else) → transparent forward
+
+Errors always return 502 JSON — never crash the proxy.
+"""
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import urllib.error
+import urllib.request
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger("thinkstack.codex.proxy")
+
+DEFAULT_PORT = 8766
+DEFAULT_UPSTREAM = "https://api.openai.com"
+
+# Headers that must NOT be forwarded to the upstream (they are hop-by-hop or
+# set by the proxy itself when re-encoding the body).
+_HOP_BY_HOP = frozenset(
+    [
+        "connection",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "te",
+        "trailers",
+        "transfer-encoding",
+        "upgrade",
+        "host",
+        "content-length",  # recalculated after body modification
+    ]
+)
+
+
+# ---------------------------------------------------------------------------
+# Request handler
+# ---------------------------------------------------------------------------
+
+class _CodexRequestHandler(BaseHTTPRequestHandler):
+    """One instance per incoming HTTP connection."""
+
+    # Injected by CodexProxy on the server object
+    upstream: str = DEFAULT_UPSTREAM
+
+    # Silence the default access log — we write our own on error
+    def log_message(self, fmt: str, *args) -> None:  # type: ignore[override]
+        pass
+
+    # ------------------------------------------------------------------
+    # Routing
+    # ------------------------------------------------------------------
+
+    def do_GET(self) -> None:
+        if self.path == "/health":
+            self._send_json(200, {"status": "ok", "proxy": "thinkstack-codex"})
+            return
+        self._forward_transparent()
+
+    def do_POST(self) -> None:
+        if self.path.startswith("/v1/"):
+            self._handle_v1_post()
+        else:
+            self._forward_transparent()
+
+    # Support additional methods by transparent proxy
+    def do_PUT(self) -> None:
+        self._forward_transparent()
+
+    def do_DELETE(self) -> None:
+        self._forward_transparent()
+
+    def do_PATCH(self) -> None:
+        self._forward_transparent()
+
+    def do_HEAD(self) -> None:
+        self._forward_transparent()
+
+    def do_OPTIONS(self) -> None:
+        self._forward_transparent()
+
+    # ------------------------------------------------------------------
+    # v1 POST handler — RACP injection + capture
+    # ------------------------------------------------------------------
+
+    def _handle_v1_post(self) -> None:
+        try:
+            body_bytes = self._read_body()
+        except Exception as exc:
+            self._send_error_502(f"failed to read request body: {exc}")
+            return
+
+        # Parse JSON — if not parseable, forward raw
+        try:
+            request_body = json.loads(body_bytes) if body_bytes else {}
+        except Exception:
+            request_body = None
+
+        # Inject RACP prefix
+        modified_body_bytes = body_bytes
+        if request_body is not None:
+            try:
+                from thinkstack_core.codex.capture import CodexCaptureHandler
+                import uuid
+
+                session_id = str(uuid.uuid4())
+                handler = CodexCaptureHandler()
+                modified_body = handler.pre_execute(request_body, session_id)
+                modified_body_bytes = json.dumps(modified_body).encode("utf-8")
+            except Exception as exc:
+                logger.warning("thinkstack codex RACP injection failed — %s", exc)
+                session_id = ""
+                handler = None
+        else:
+            session_id = ""
+            handler = None
+
+        # Forward to upstream
+        upstream_url = self.upstream.rstrip("/") + self.path
+        fwd_headers = self._build_forward_headers(len(modified_body_bytes))
+
+        try:
+            req = urllib.request.Request(
+                upstream_url,
+                data=modified_body_bytes,
+                headers=fwd_headers,
+                method="POST",
+            )
+            with urllib.request.urlopen(req) as resp:
+                resp_bytes = resp.read()
+                resp_status = resp.status
+                resp_headers = dict(resp.headers)
+        except urllib.error.HTTPError as exc:
+            resp_bytes = exc.read()
+            resp_status = exc.code
+            resp_headers = dict(exc.headers)
+        except Exception as exc:
+            self._send_error_502(f"upstream request failed: {exc}")
+            return
+
+        # Send response back to Codex CLI
+        self._relay_response(resp_status, resp_headers, resp_bytes)
+
+        # Fire post_execute in background — never blocks the response
+        if handler is not None and request_body is not None:
+            try:
+                resp_body = json.loads(resp_bytes) if resp_bytes else {}
+            except Exception:
+                resp_body = {}
+            t = threading.Thread(
+                target=_safe_post_execute,
+                args=(handler, request_body, resp_body, session_id),
+                daemon=True,
+            )
+            t.start()
+
+    # ------------------------------------------------------------------
+    # Transparent forward
+    # ------------------------------------------------------------------
+
+    def _forward_transparent(self) -> None:
+        try:
+            body_bytes = self._read_body()
+        except Exception as exc:
+            self._send_error_502(f"failed to read request body: {exc}")
+            return
+
+        upstream_url = self.upstream.rstrip("/") + self.path
+        fwd_headers = self._build_forward_headers(len(body_bytes) if body_bytes else 0)
+
+        try:
+            req = urllib.request.Request(
+                upstream_url,
+                data=body_bytes if body_bytes else None,
+                headers=fwd_headers,
+                method=self.command,
+            )
+            with urllib.request.urlopen(req) as resp:
+                resp_bytes = resp.read()
+                resp_status = resp.status
+                resp_headers = dict(resp.headers)
+        except urllib.error.HTTPError as exc:
+            resp_bytes = exc.read()
+            resp_status = exc.code
+            resp_headers = dict(exc.headers)
+        except Exception as exc:
+            self._send_error_502(f"upstream request failed: {exc}")
+            return
+
+        self._relay_response(resp_status, resp_headers, resp_bytes)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _read_body(self) -> bytes:
+        length = int(self.headers.get("Content-Length", 0))
+        if length > 0:
+            return self.rfile.read(length)
+        return b""
+
+    def _build_forward_headers(self, content_length: int) -> dict:
+        """Build headers to send upstream, stripping hop-by-hop headers."""
+        headers: dict = {}
+        for key, val in self.headers.items():
+            if key.lower() not in _HOP_BY_HOP:
+                headers[key] = val
+        if content_length > 0:
+            headers["Content-Length"] = str(content_length)
+        # Ensure Host is set to the upstream host
+        from urllib.parse import urlparse
+        parsed = urlparse(self.upstream)
+        headers["Host"] = parsed.netloc
+        return headers
+
+    def _relay_response(self, status: int, headers: dict, body: bytes) -> None:
+        """Write upstream response back to the Codex CLI connection."""
+        self.send_response(status)
+        for key, val in headers.items():
+            if key.lower() not in ("transfer-encoding", "connection"):
+                self.send_header(key, val)
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        if body:
+            self.wfile.write(body)
+
+    def _send_json(self, status: int, payload: dict) -> None:
+        body = json.dumps(payload).encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(body)))
+        self.end_headers()
+        self.wfile.write(body)
+
+    def _send_error_502(self, detail: str) -> None:
+        logger.warning("thinkstack codex proxy 502 — %s", detail)
+        self._send_json(502, {"error": "bad_gateway", "detail": detail})
+
+
+# ---------------------------------------------------------------------------
+# Background post_execute helper
+# ---------------------------------------------------------------------------
+
+def _safe_post_execute(
+    handler: object,
+    request_body: dict,
+    response_body: dict,
+    session_id: str,
+) -> None:
+    try:
+        handler.post_execute(request_body, response_body, session_id)  # type: ignore[attr-defined]
+    except Exception as exc:
+        logger.warning("thinkstack codex post_execute background error — %s", exc)
+
+
+# ---------------------------------------------------------------------------
+# CodexProxy
+# ---------------------------------------------------------------------------
+
+class CodexProxy:
+    """
+    Lightweight stdlib-only HTTP proxy for intercepting Codex CLI traffic.
+
+    Parameters
+    ----------
+    port:
+        Local port to bind on (default 8766).
+    upstream:
+        Base URL of the real OpenAI API (default https://api.openai.com).
+    """
+
+    def __init__(
+        self,
+        port: int = DEFAULT_PORT,
+        upstream: str = DEFAULT_UPSTREAM,
+    ) -> None:
+        self.port = port
+        self.upstream = upstream
+        self._server: Optional[HTTPServer] = None
+        self._thread: Optional[threading.Thread] = None
+
+    def start(self) -> None:
+        """Start the proxy server (blocking — press Ctrl+C to stop)."""
+        self._server = self._make_server()
+        try:
+            self._server.serve_forever()
+        except KeyboardInterrupt:
+            pass
+        finally:
+            self._server.server_close()
+
+    def start_background(self) -> threading.Thread:
+        """
+        Start the proxy in a daemon thread.
+
+        Returns the thread (already started).  Call stop() to shut down.
+        """
+        self._server = self._make_server()
+        t = threading.Thread(target=self._server.serve_forever, daemon=True)
+        t.start()
+        self._thread = t
+        return t
+
+    def stop(self) -> None:
+        """Shut down the proxy server gracefully."""
+        if self._server is not None:
+            self._server.shutdown()
+            self._server.server_close()
+            self._server = None
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _make_server(self) -> HTTPServer:
+        upstream = self.upstream
+
+        class _Handler(_CodexRequestHandler):
+            pass
+
+        _Handler.upstream = upstream  # type: ignore[assignment]
+        server = HTTPServer(("127.0.0.1", self.port), _Handler)
+        server.timeout = 30
+        return server
+
+
+# ---------------------------------------------------------------------------
+# install_codex_hook
+# ---------------------------------------------------------------------------
+
+def install_codex_hook(project_root: str) -> tuple[bool, str]:
+    """
+    Create `.codex/config.json` under *project_root* pointing at the local
+    ThinkStack Codex proxy.
+
+    Idempotent — does not overwrite an existing config.
+
+    Returns (True, config_path) on success, (False, error_message) on failure.
+    """
+    try:
+        codex_dir = Path(project_root) / ".codex"
+        codex_dir.mkdir(parents=True, exist_ok=True)
+
+        config_path = codex_dir / "config.json"
+
+        # Idempotent — do not overwrite if already present
+        if config_path.exists():
+            return True, str(config_path)
+
+        config = {
+            "apiBase": f"http://localhost:{DEFAULT_PORT}/v1",
+            "provider": "openai",
+        }
+        config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
+        return True, str(config_path)
+
+    except Exception as exc:
+        return False, str(exc)
+
+
+# ---------------------------------------------------------------------------
+# get_codex_hook_status
+# ---------------------------------------------------------------------------
+
+def get_codex_hook_status(project_root: str) -> dict:
+    """
+    Return the install and runtime status of the ThinkStack Codex proxy.
+
+    Returns::
+
+        {
+            "codex_config": bool,       # .codex/config.json exists
+            "codex_config_path": str,   # absolute path to config file
+            "proxy_running": bool,      # True if GET /health returns 200
+        }
+    """
+    config_path = Path(project_root) / ".codex" / "config.json"
+    codex_config = config_path.exists()
+
+    proxy_running = False
+    try:
+        req = urllib.request.Request(
+            f"http://localhost:{DEFAULT_PORT}/health",
+            method="GET",
+        )
+        with urllib.request.urlopen(req, timeout=2) as resp:
+            proxy_running = resp.status == 200
+    except Exception:
+        proxy_running = False
+
+    return {
+        "codex_config": codex_config,
+        "codex_config_path": str(config_path),
+        "proxy_running": proxy_running,
+    }

thinkstack_core/compat.py

@@ -0,0 +1,103 @@
+"""
+Backward-compatibility shims for the thinkstack → ThinkStack rename.
+
+Provides fallbacks for deprecated environment variables, MCP tool names,
+and auth headers so existing users can migrate without breaking.
+Shims emit deprecation warnings and will be removed in v5.0.0.
+"""
+from __future__ import annotations
+
+import os
+import warnings
+from pathlib import Path
+
+
+# ---------------------------------------------------------------------------
+# Environment variable shims
+# ---------------------------------------------------------------------------
+
+_thinkstack_TO_THINKSTACK = {
+    "THINKSTACK_DISABLE": "thinkstack_DISABLE",
+    "THINKSTACK_MCP_URL": "thinkstack_MCP_URL",
+    "THINKSTACK_API_KEY": "thinkstack_API_KEY",
+    "THINKSTACK_MCP_DOMAIN": "thinkstack_MCP_DOMAIN",
+    "THINKSTACK_ORG_ID": "thinkstack_ORG_ID",
+    "THINKSTACK_GCC_PATH": "thinkstack_GCC_PATH",
+    "THINKSTACK_REGISTRY_PATH": "thinkstack_REGISTRY_PATH",
+    "THINKSTACK_SIGN_LOG": "thinkstack_SIGN_LOG",
+    "THINKSTACK_STORAGE": "thinkstack_STORAGE",
+    "THINKSTACK_BUCKET": "thinkstack_BUCKET",
+    "THINKSTACK_ACCESS_KEY_ID": "thinkstack_ACCESS_KEY_ID",
+    "THINKSTACK_SECRET_ACCESS_KEY": "thinkstack_SECRET_ACCESS_KEY",
+    "THINKSTACK_AZURE_CONTAINER": "thinkstack_AZURE_CONTAINER",
+    "THINKSTACK_SESSION_CLOUD_URL": "thinkstack_SESSION_CLOUD_URL",
+    "THINKSTACK_AGGPHI": "thinkstack_AGGPHI",
+    "THINKSTACK_AGGPHI_LLM": "thinkstack_AGGPHI_LLM",
+    "THINKSTACK_AGGPHI_MODEL": "thinkstack_AGGPHI_MODEL",
+    "THINKSTACK_JWKS_URLS": "thinkstack_JWKS_URLS",
+    "THINKSTACK_METRICS_WEBHOOK_URL": "thinkstack_METRICS_WEBHOOK_URL",
+    "THINKSTACK_DATADOG_API_KEY": "thinkstack_DATADOG_API_KEY",
+    "THINKSTACK_SPLUNK_HEC_URL": "thinkstack_SPLUNK_HEC_URL",
+    "THINKSTACK_SERVICENOW_INSTANCE": "thinkstack_SERVICENOW_INSTANCE",
+    "THINKSTACK_GITHUB_WEBHOOK_SECRET": "thinkstack_GITHUB_WEBHOOK_SECRET",
+    "THINKSTACK_DAEMON_ROOT": "thinkstack_DAEMON_ROOT",
+    "THINKSTACK_AGENT_ID": "thinkstack_AGENT_ID",
+    "THINKSTACK_AGENT_NAME": "thinkstack_AGENT_NAME",
+    "THINKSTACK_AGENT_SCOPE": "thinkstack_AGENT_SCOPE",
+    "THINKSTACK_PAT_ENCRYPTION_KEY": "thinkstack_PAT_ENCRYPTION_KEY",
+}
+
+
+def getenv(
+    name: str, default: str | None = None, *, _warned: set[str] | None = None
+) -> str | None:
+    """Read an environment variable with backward-compat fallback.
+
+    Tries *name* first (e.g. ``THINKSTACK_MCP_URL``). If unset, falls
+    back to the old ``thinkstack_*`` name and issues a deprecation warning.
+    """
+    value = os.environ.get(name)
+    if value is not None:
+        return value
+
+    old_name = _thinkstack_TO_THINKSTACK.get(name)
+    if old_name is not None:
+        old_value = os.environ.get(old_name)
+        if old_value is not None:
+            if _warned is None or old_name not in _warned:
+                warnings.warn(
+                    f"{old_name} is deprecated, use {name} instead",
+                    DeprecationWarning,
+                    stacklevel=2,
+                )
+                if _warned is not None:
+                    _warned.add(old_name)
+            return old_value
+
+    return default
+
+
+# ---------------------------------------------------------------------------
+# Local runtime directory auto-migration
+# ---------------------------------------------------------------------------
+
+_OLD_thinkstack_DIR = Path.home() / ".thinkstack"
+_NEW_THINKSTACK_DIR = Path.home() / ".thinkstack"
+
+
+def migrate_runtime_dir() -> None:
+    """Move ``~/.thinkstack/`` to ``~/.thinkstack/`` if the old dir exists
+    and the new one doesn't. Called once at import time."""
+    if _OLD_thinkstack_DIR.exists() and not _NEW_THINKSTACK_DIR.exists():
+        import shutil
+        shutil.move(str(_OLD_thinkstack_DIR), str(_NEW_THINKSTACK_DIR))
+        warnings.warn(
+            f"Moved {_OLD_thinkstack_DIR} → {_NEW_THINKSTACK_DIR} "
+            "(thinkstack is now thinkstack)",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+
+# Run auto-migration on import
+migrate_runtime_dir()