Streamlit Master Architect

from __future__ import annotations

import argparse
import ast
import json
import os
import re
import sys
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Literal

try:
    import tomllib  # py311+
except ModuleNotFoundError:  # pragma: no cover
    tomllib = None  # type: ignore[assignment]

JsonDict = dict[str, Any]
Severity = Literal["low", "medium", "high"]


DEFAULT_EXCLUDES = {
    ".git",
    ".hg",
    ".svn",
    ".mypy_cache",
    ".pytest_cache",
    ".ruff_cache",
    "__pycache__",
    "build",
    "dist",
    "node_modules",
    "opensrc",
    "venv",
    ".venv",
    "env",
    ".env",
    "site-packages",
}


WIDGET_FUNCS = {
    # Common input widgets (subset; used for "missing key in loop" heuristic).
    "button",
    "checkbox",
    "color_picker",
    "date_input",
    "datetime_input",
    "file_uploader",
    "camera_input",
    "multiselect",
    "number_input",
    "radio",
    "selectbox",
    "slider",
    "text_input",
    "text_area",
    "time_input",
    "toggle",
}


DEPRECATED_OR_RISKY = {
    "st.cache": ("high", "Deprecated caching API; migrate to st.cache_data / st.cache_resource."),
    "st.experimental_memo": ("high", "Deprecated; migrate to st.cache_data."),
    "st.experimental_singleton": ("high", "Deprecated; migrate to st.cache_resource."),
    "st.experimental_rerun": ("medium", "Prefer st.rerun (stable)."),
    "st.experimental_set_query_params": ("medium", "Prefer st.query_params."),
    "st.experimental_get_query_params": ("medium", "Prefer st.query_params."),
    "st.bokeh_chart": ("high", "Native Bokeh support removed in modern Streamlit; replace with Altair/Plotly/etc."),
}


RISKY_FLAGS = {
    "unsafe_allow_html=True": (
        "high",
        "Potential XSS sink (e.g., st.markdown). Avoid unless content is trusted and sanitized.",
    ),
    "unsafe_allow_javascript=True": (
        "high",
        "High-risk JS execution (st.html). Never use with untrusted input.",
    ),
}


@dataclass(frozen=True)
class Finding:
    severity: Severity
    code: str
    message: str
    locations: list[str]


def _read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8", errors="replace")


def _pypi_latest_streamlit_version(timeout_s: float = 10.0) -> str | None:
    url = "https://pypi.org/pypi/streamlit/json"
    req = urllib.request.Request(url, headers={"User-Agent": "streamlit-master-architect/0.1"})
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            payload = json.loads(resp.read().decode("utf-8", errors="replace"))
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError):
        return None
    info = payload.get("info", {})
    if isinstance(info, dict):
        v = info.get("version")
        return str(v) if v else None
    return None


def _installed_streamlit_version() -> str | None:
    try:
        from importlib.metadata import PackageNotFoundError, version
    except Exception:  # pragma: no cover
        return None
    try:
        return version("streamlit")
    except PackageNotFoundError:
        return None
    except Exception:
        return None


def _locked_streamlit_version(root: Path) -> str | None:
    if tomllib is None:
        return None

    uv_lock = root / "uv.lock"
    if uv_lock.exists():
        try:
            data = tomllib.loads(_read_text(uv_lock))
        except Exception:
            data = None
        if isinstance(data, dict):
            pkgs = data.get("package", [])
            if isinstance(pkgs, list):
                for p in pkgs:
                    if isinstance(p, dict) and p.get("name") == "streamlit":
                        v = p.get("version")
                        return str(v) if v else None

    poetry_lock = root / "poetry.lock"
    if poetry_lock.exists():
        try:
            data = tomllib.loads(_read_text(poetry_lock))
        except Exception:
            data = None
        if isinstance(data, dict):
            pkgs = data.get("package", [])
            if isinstance(pkgs, list):
                for p in pkgs:
                    if isinstance(p, dict) and p.get("name") == "streamlit":
                        v = p.get("version")
                        return str(v) if v else None

    return None


def _iter_python_files(root: Path, *, excludes: set[str]) -> Iterable[Path]:
    for path in root.rglob("*.py"):
        rel_parts = path.relative_to(root).parts
        if any(p in excludes for p in rel_parts):
            continue
        yield path


def _parse_requirements(path: Path) -> list[str]:
    specs: list[str] = []
    for raw in _read_text(path).splitlines():
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        # Remove inline comments.
        line = line.split("#", 1)[0].strip()
        # Ignore pip options and URL installs.
        if line.startswith("-"):
            continue
        if "://" in line:
            continue
        if re.match(r"(?i)^streamlit(\[.*\])?([<>=!~]=?.*)?$", line):
            specs.append(line)
    return specs


def _scan_pyproject(path: Path) -> list[str]:
    if tomllib is None:
        return []
    try:
        data = tomllib.loads(_read_text(path))
    except Exception:
        return []

    found: list[str] = []

    project = data.get("project", {})
    if isinstance(project, dict):
        deps = project.get("dependencies", [])
        if isinstance(deps, list):
            for d in deps:
                if isinstance(d, str) and d.lower().startswith("streamlit"):
                    found.append(d)
        opt = project.get("optional-dependencies", {})
        if isinstance(opt, dict):
            for _group, group_deps in opt.items():
                if isinstance(group_deps, list):
                    for d in group_deps:
                        if isinstance(d, str) and d.lower().startswith("streamlit"):
                            found.append(d)

    tool = data.get("tool", {})
    if isinstance(tool, dict):
        poetry = tool.get("poetry", {})
        if isinstance(poetry, dict):
            deps = poetry.get("dependencies", {})
            if isinstance(deps, dict):
                v = deps.get("streamlit")
                if isinstance(v, str):
                    found.append(f"streamlit {v}")
                elif isinstance(v, dict):
                    found.append(f"streamlit {json.dumps(v, sort_keys=True)}")

            groups = poetry.get("group", {})
            if isinstance(groups, dict):
                for _gname, gdata in groups.items():
                    if not isinstance(gdata, dict):
                        continue
                    gdeps = gdata.get("dependencies", {})
                    if isinstance(gdeps, dict):
                        v = gdeps.get("streamlit")
                        if isinstance(v, str):
                            found.append(f"streamlit {v}")
                        elif isinstance(v, dict):
                            found.append(f"streamlit {json.dumps(v, sort_keys=True)}")

    # Deduplicate while keeping stable order.
    out: list[str] = []
    for x in found:
        if x not in out:
            out.append(x)
    return out


def _collect_dependency_specs(root: Path) -> list[JsonDict]:
    specs: list[JsonDict] = []

    req = root / "requirements.txt"
    if req.exists():
        for s in _parse_requirements(req):
            specs.append({"file": str(req), "spec": s})

    pyproject = root / "pyproject.toml"
    if pyproject.exists():
        for s in _scan_pyproject(pyproject):
            specs.append({"file": str(pyproject), "spec": s})

    return specs


def _get_attr_chain(expr: ast.AST) -> list[str] | None:
    parts: list[str] = []
    cur: ast.AST | None = expr
    while isinstance(cur, ast.Attribute):
        parts.append(cur.attr)
        cur = cur.value
    if isinstance(cur, ast.Name):
        parts.append(cur.id)
        return list(reversed(parts))
    return None


def _scan_streamlit_usage(py_file: Path) -> tuple[dict[str, int], list[Finding]]:
    text = _read_text(py_file)
    try:
        tree = ast.parse(text, filename=str(py_file))
    except SyntaxError:
        return {}, [Finding(severity="low", code="parse_error", message="Failed to parse Python file.", locations=[str(py_file)])]

    module_aliases: set[str] = set()
    imported_names: dict[str, str] = {}

    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                if alias.name == "streamlit":
                    module_aliases.add(alias.asname or "streamlit")
        elif isinstance(node, ast.ImportFrom) and node.module == "streamlit":
            for alias in node.names:
                imported_names[alias.asname or alias.name] = alias.name

    usage: dict[str, int] = {}
    findings: list[Finding] = []

    # Simple text-level risky flag detection (fast, includes non-Streamlit sinks).
    for needle, (sev, msg) in RISKY_FLAGS.items():
        if needle in text:
            findings.append(
                Finding(
                    severity=sev,
                    code="security_flag",
                    message=f"{msg} (found `{needle}`)",
                    locations=[str(py_file)],
                )
            )

    # Deprecated API usage detection via AST calls.
    for node in ast.walk(tree):
        if not isinstance(node, ast.Call):
            continue

        # Case A: st.foo(...)
        chain = _get_attr_chain(node.func)
        call_name: str | None = None
        if chain and chain[0] in module_aliases:
            call_name = ".".join(chain)
        elif isinstance(node.func, ast.Name) and node.func.id in imported_names:
            call_name = f"streamlit.{imported_names[node.func.id]}"

        if not call_name:
            continue

        usage[call_name] = usage.get(call_name, 0) + 1

        # Map aliases like "st.experimental_memo" regardless of alias name.
        canonical = call_name
        for alias in module_aliases:
            if canonical.startswith(f"{alias}."):
                canonical = "st." + canonical[len(alias) + 1 :]
                break

        if canonical in DEPRECATED_OR_RISKY:
            sev, msg = DEPRECATED_OR_RISKY[canonical]
            loc = f"{py_file}:{getattr(node, 'lineno', 1)}"
            findings.append(Finding(severity=sev, code="deprecated_api", message=f"{canonical}: {msg}", locations=[loc]))

        # Heuristic: widget call in loop without key.
        # Only applies to direct st.<widget>(...) or st.sidebar.<widget>(...) forms.
        if chain and chain[0] in module_aliases:
            # last segment could be widget name
            widget_name = chain[-1]
            if widget_name in WIDGET_FUNCS:
                has_key = any(isinstance(k, ast.keyword) and k.arg == "key" for k in node.keywords)
                if not has_key:
                    # Walk parents via a second pass: easiest is to flag only if text contains "for" on same or previous line.
                    # This is conservative and avoids building a full parent map.
                    line_no = getattr(node, "lineno", 0)
                    if line_no:
                        lines = text.splitlines()
                        window = "\n".join(lines[max(0, line_no - 3) : line_no])
                        if re.search(r"\bfor\b", window):
                            loc = f"{py_file}:{line_no}"
                            findings.append(
                                Finding(
                                    severity="medium",
                                    code="missing_key_in_loop",
                                    message=f"Possible widget call inside a loop without `key=`: {call_name}",
                                    locations=[loc],
                                )
                            )

    return usage, findings


def _aggregate_findings(findings: Iterable[Finding]) -> list[JsonDict]:
    # Merge identical (severity, code, message) by accumulating locations.
    merged: dict[tuple[str, str, str], set[str]] = {}
    for f in findings:
        key = (f.severity, f.code, f.message)
        merged.setdefault(key, set()).update(f.locations)
    out: list[JsonDict] = []
    for (sev, code, msg), locs in sorted(merged.items(), key=lambda x: (x[0][0], x[0][1], x[0][2])):
        out.append({"severity": sev, "code": code, "message": msg, "locations": sorted(locs)})
    return out


def _to_markdown(report: JsonDict) -> str:
    lines: list[str] = []
    lines.append(f"# Streamlit project audit\n")
    lines.append(f"- Root: `{report['project_root']}`")

    st_info = report.get("streamlit", {})
    if isinstance(st_info, dict):
        lines.append(f"- Installed Streamlit: `{st_info.get('installed_version')}`")
        lines.append(f"- Locked Streamlit (lockfile): `{st_info.get('locked_version')}`")
        lines.append(f"- Latest Streamlit (PyPI): `{st_info.get('latest_version')}`")

    specs = report.get("dependency_specs", [])
    lines.append("\n## Dependency specs\n")
    if not specs:
        lines.append("- (none found)")
    else:
        for s in specs:
            lines.append(f"- `{s.get('file')}`: `{s.get('spec')}`")

    usage = report.get("top_calls", [])
    lines.append("\n## Streamlit usage (top calls)\n")
    if not usage:
        lines.append("- (no Streamlit calls detected)")
    else:
        for item in usage:
            lines.append(f"- `{item['call']}`: {item['count']}")

    issues = report.get("issues", [])
    lines.append("\n## Findings\n")
    if not issues:
        lines.append("- (no issues detected)")
    else:
        for i in issues:
            sev = i.get("severity", "medium")
            msg = i.get("message", "")
            lines.append(f"- **{sev}**: {msg}")
            locs = i.get("locations", [])
            if isinstance(locs, list) and locs:
                for loc in locs[:20]:
                    lines.append(f"  - `{loc}`")
                if len(locs) > 20:
                    lines.append(f"  - … (+{len(locs) - 20} more)")

    recs = report.get("recommendations", [])
    lines.append("\n## Recommendations\n")
    if not recs:
        lines.append("- (none)")
    else:
        for r in recs:
            lines.append(f"- {r}")

    return "\n".join(lines) + "\n"


def main() -> int:
    parser = argparse.ArgumentParser(description="Audit a Streamlit project for version, deps, and risky/deprecated APIs.")
    parser.add_argument("--root", type=str, default=".", help="Project root to scan.")
    parser.add_argument("--format", type=str, default="json", choices=["json", "md"], help="Output format.")
    parser.add_argument("--output", type=str, default="", help="Write report to a file instead of stdout.")
    parser.add_argument(
        "--check-latest",
        action="store_true",
        help="(deprecated) Enable latest Streamlit version check (default on).",
    )
    parser.add_argument("--no-check-latest", action="store_true", help="Disable PyPI latest-version check.")
    parser.add_argument("--top", type=int, default=30, help="Top N Streamlit calls to include.")
    args = parser.parse_args()

    root = Path(args.root).resolve()
    if not root.exists():
        raise FileNotFoundError(f"Root not found: {root}")

    check_latest = not args.no_check_latest

    installed = _installed_streamlit_version()
    locked = _locked_streamlit_version(root)
    latest = _pypi_latest_streamlit_version() if check_latest else None
    dep_specs = _collect_dependency_specs(root)

    all_usage: dict[str, int] = {}
    findings: list[Finding] = []

    for py in _iter_python_files(root, excludes=DEFAULT_EXCLUDES):
        usage, file_findings = _scan_streamlit_usage(py)
        for k, v in usage.items():
            all_usage[k] = all_usage.get(k, 0) + v
        findings.extend(file_findings)

    # Detect beta APIs by name prefix (works regardless of version).
    beta_calls = [k for k in all_usage.keys() if ".beta_" in k]
    if beta_calls:
        findings.append(
            Finding(
                severity="high",
                code="deprecated_api",
                message="Detected st.beta_* APIs; these are legacy and should be migrated to stable equivalents.",
                locations=beta_calls[:50],
            )
        )

    top_calls = sorted(all_usage.items(), key=lambda kv: kv[1], reverse=True)[: max(1, args.top)]

    recommendations: list[str] = []
    current_for_compare = installed or locked
    if latest and current_for_compare and latest != current_for_compare:
        recommendations.append(
            f"Consider upgrading Streamlit from {current_for_compare} to {latest} after reading release notes and running tests."
        )
    if not dep_specs:
        recommendations.append("No Streamlit dependency spec found (requirements.txt/pyproject.toml). Ensure Streamlit is pinned or constrained for reproducible deploys.")
    if any(f.code == "security_flag" for f in findings):
        recommendations.append("Review all unsafe HTML/JS flags; ensure inputs are trusted/sanitized and usage is isolated.")
    if any(f.code == "deprecated_api" for f in findings):
        recommendations.append("Migrate deprecated Streamlit APIs to their stable equivalents (see Findings).")

    report: JsonDict = {
        "project_root": str(root),
        "streamlit": {"installed_version": installed, "locked_version": locked, "latest_version": latest},
        "dependency_specs": dep_specs,
        "top_calls": [{"call": k, "count": v} for k, v in top_calls],
        "issues": _aggregate_findings(findings),
        "recommendations": recommendations,
    }

    out: str
    if args.format == "md":
        out = _to_markdown(report)
    else:
        out = json.dumps(report, indent=2, sort_keys=True)

    if args.output:
        Path(args.output).write_text(out + ("\n" if not out.endswith("\n") else ""), encoding="utf-8")
    else:
        sys.stdout.write(out + ("\n" if not out.endswith("\n") else ""))

    return 0


if __name__ == "__main__":
    raise SystemExit(main())

from __future__ import annotations

import json
import queue
import subprocess
import threading
import time
from dataclasses import dataclass
from typing import Any, Mapping


JsonValue = Any


class MCPProtocolError(RuntimeError):
    """Raised when MCP JSON-RPC framing or protocol assumptions fail."""


class MCPRequestError(RuntimeError):
    """Raised when an MCP request returns an error response."""


@dataclass(frozen=True)
class MCPTool:
    name: str
    description: str | None
    input_schema: Mapping[str, JsonValue] | None


@dataclass(frozen=True)
class MCPResponse:
    id: int
    result: JsonValue | None
    error: JsonValue | None


def _encode_framed_message(payload: Mapping[str, JsonValue]) -> bytes:
    body = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    header = f"Content-Length: {len(body)}\r\n\r\n".encode("ascii")
    return header + body


def _read_exact(stream: Any, n: int) -> bytes:
    buf = bytearray()
    while len(buf) < n:
        chunk = stream.read(n - len(buf))
        if not chunk:
            raise MCPProtocolError("Unexpected EOF while reading framed message body.")
        buf.extend(chunk)
    return bytes(buf)


def _read_framed_message(stream: Any) -> Mapping[str, JsonValue]:
    headers: dict[str, str] = {}
    while True:
        line = stream.readline()
        if line is None or line == b"":
            raise MCPProtocolError("EOF while reading headers.")
        if line in (b"\r\n", b"\n"):
            break
        decoded = line.decode("ascii", errors="strict").strip()
        if ":" not in decoded:
            continue
        k, v = decoded.split(":", 1)
        headers[k.strip().lower()] = v.strip()

    if "content-length" not in headers:
        raise MCPProtocolError(f"Missing Content-Length header. Headers: {headers}")

    try:
        length = int(headers["content-length"])
    except ValueError as e:
        raise MCPProtocolError(f"Invalid Content-Length: {headers['content-length']!r}") from e

    body = _read_exact(stream, length)
    try:
        msg = json.loads(body.decode("utf-8"))
    except Exception as e:
        raise MCPProtocolError(f"Invalid JSON body: {body[:200]!r}...") from e

    if not isinstance(msg, dict):
        raise MCPProtocolError(f"Expected JSON object message, got: {type(msg)}")

    return msg


class MCPStdioClient:
    """Minimal MCP stdio client (JSON-RPC 2.0 with Content-Length framing)."""

    def __init__(
        self,
        command: list[str],
        *,
        cwd: str | None = None,
        env: Mapping[str, str] | None = None,
        startup_timeout_s: float = 15.0,
        request_timeout_s: float = 60.0,
        log_notifications: bool = True,
    ) -> None:
        self._command = command
        self._cwd = cwd
        self._env = dict(env) if env is not None else None
        self._startup_timeout_s = startup_timeout_s
        self._request_timeout_s = request_timeout_s
        self._log_notifications = log_notifications

        self._proc: subprocess.Popen[bytes] | None = None
        self._rx_thread: threading.Thread | None = None
        self._rx_queue: queue.Queue[Mapping[str, JsonValue]] = queue.Queue()
        self._next_id = 1

    def start(self) -> None:
        if self._proc is not None:
            return
        self._proc = subprocess.Popen(
            self._command,
            cwd=self._cwd,
            env=self._env,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if self._proc.stdin is None or self._proc.stdout is None:
            raise MCPProtocolError("Failed to open stdio pipes to MCP server process.")

        self._rx_thread = threading.Thread(target=self._reader_loop, daemon=True)
        self._rx_thread.start()

        deadline = time.time() + self._startup_timeout_s
        while time.time() < deadline:
            if self._proc.poll() is not None:
                stderr = self._safe_read_stderr()
                raise MCPProtocolError(
                    f"MCP server process exited early (code {self._proc.returncode}).\n{stderr}"
                )
            if not self._rx_queue.empty():
                return
            time.sleep(0.05)

    def terminate(self) -> None:
        if self._proc is None:
            return
        if self._proc.poll() is None:
            self._proc.terminate()
            try:
                self._proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self._proc.kill()
        self._proc = None

    def initialize(
        self,
        *,
        protocol_version: str = "2024-11-05",
        client_name: str = "streamlit-master-architect",
        client_version: str = "0.1.0",
        capabilities: Mapping[str, JsonValue] | None = None,
    ) -> JsonValue:
        caps = dict(capabilities) if capabilities is not None else {}
        result = self.request(
            "initialize",
            {
                "protocolVersion": protocol_version,
                "clientInfo": {"name": client_name, "version": client_version},
                "capabilities": caps,
            },
        )
        try:
            self.notify("initialized", {})
        except Exception:
            pass
        return result

    def list_tools(self) -> list[MCPTool]:
        res = self.request("tools/list", {})
        tools_raw = res.get("tools", []) if isinstance(res, dict) else []
        tools: list[MCPTool] = []
        for t in tools_raw:
            if not isinstance(t, dict):
                continue
            tools.append(
                MCPTool(
                    name=str(t.get("name", "")),
                    description=(str(t["description"]) if t.get("description") is not None else None),
                    input_schema=(t.get("inputSchema") if isinstance(t.get("inputSchema"), dict) else None),
                )
            )
        return tools

    def call_tool(self, name: str, arguments: Mapping[str, JsonValue]) -> JsonValue:
        return self.request("tools/call", {"name": name, "arguments": dict(arguments)})

    def request(self, method: str, params: Mapping[str, JsonValue]) -> JsonValue:
        req_id = self._next_id
        self._next_id += 1

        self._send({"jsonrpc": "2.0", "id": req_id, "method": method, "params": dict(params)})

        deadline = time.time() + self._request_timeout_s
        while time.time() < deadline:
            msg = self._rx_queue.get(timeout=max(0.05, self._request_timeout_s / 200))
            if "id" not in msg:
                self._handle_notification(msg)
                continue
            if msg.get("id") != req_id:
                self._rx_queue.put(msg)
                continue

            resp = MCPResponse(id=req_id, result=msg.get("result"), error=msg.get("error"))
            if resp.error is not None:
                raise MCPRequestError(f"MCP request {method} failed: {resp.error}")
            return resp.result

        raise TimeoutError(f"MCP request timed out: {method}")

    def notify(self, method: str, params: Mapping[str, JsonValue]) -> None:
        self._send({"jsonrpc": "2.0", "method": method, "params": dict(params)})

    def _send(self, payload: Mapping[str, JsonValue]) -> None:
        if self._proc is None or self._proc.stdin is None:
            raise MCPProtocolError("MCP process not started.")
        framed = _encode_framed_message(payload)
        try:
            self._proc.stdin.write(framed)
            self._proc.stdin.flush()
        except BrokenPipeError as e:
            stderr = self._safe_read_stderr()
            raise MCPProtocolError(f"MCP server stdin closed.\n{stderr}") from e

    def _reader_loop(self) -> None:
        assert self._proc is not None
        assert self._proc.stdout is not None
        stdout = self._proc.stdout
        while True:
            if self._proc.poll() is not None:
                return
            try:
                msg = _read_framed_message(stdout)
            except Exception as e:
                if self._proc.poll() is not None:
                    return
                self._rx_queue.put({"jsonrpc": "2.0", "method": "mcp.protocol_error", "params": {"error": str(e)}})
                return
            self._rx_queue.put(msg)

    def _handle_notification(self, msg: Mapping[str, JsonValue]) -> None:
        if not self._log_notifications:
            return
        method = msg.get("method")
        if not method:
            return
        if method in ("notifications/message", "log", "logging/message", "mcp.protocol_error"):
            params = msg.get("params", {})
            print(f"[MCP notification] {method}: {params}")

    def _safe_read_stderr(self) -> str:
        if self._proc is None or self._proc.stderr is None:
            return ""
        try:
            return self._proc.stderr.read().decode("utf-8", errors="replace")
        except Exception:
            return ""

from __future__ import annotations

import argparse
import base64
import json
import os
import shlex
import socket
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any, Mapping

from mcp_stdio_client import MCPStdioClient, MCPTool


def _find_free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return int(s.getsockname()[1])


def _wait_http_ok(url: str, timeout_s: float = 30.0) -> None:
    deadline = time.time() + timeout_s
    last_err: Exception | None = None
    while time.time() < deadline:
        try:
            with urllib.request.urlopen(url, timeout=2) as resp:
                if 200 <= resp.status < 500:
                    return
        except Exception as e:
            last_err = e
        time.sleep(0.25)
    raise TimeoutError(f"App did not become reachable at {url}. Last error: {last_err}")


def _start_streamlit(app_path: Path, *, port: int, streamlit_cmd: str) -> subprocess.Popen[str]:
    cmd = shlex.split(streamlit_cmd) + [
        "run",
        str(app_path),
        "--server.headless=true",
        f"--server.port={port}",
        "--server.address=127.0.0.1",
    ]
    return subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        env=os.environ.copy(),
    )


def _stop_process(proc: subprocess.Popen[Any], name: str) -> None:
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
    try:
        if proc.stdout is not None:
            out = proc.stdout.read()
            if out:
                print(f"\n[{name} output]\n{out}\n")
    except Exception:
        pass


def _tool_props(tool: MCPTool) -> list[str]:
    schema = tool.input_schema or {}
    props = schema.get("properties", {}) if isinstance(schema, dict) else {}
    if isinstance(props, dict):
        return [str(k) for k in props.keys()]
    return []


def _tool_required(tool: MCPTool) -> list[str]:
    schema = tool.input_schema or {}
    required = schema.get("required", []) if isinstance(schema, dict) else []
    if isinstance(required, list):
        return [str(k) for k in required]
    return []


def _pick_tool(tools: list[MCPTool], keywords: list[str]) -> MCPTool | None:
    for kw in keywords:
        for t in tools:
            if kw.lower() in t.name.lower():
                return t
    return None


def _best_effort_args(tool: MCPTool, *, url: str, screenshot_path: Path) -> Mapping[str, Any]:
    props = [p.lower() for p in _tool_props(tool)]
    args: dict[str, Any] = {}

    if "url" in props:
        args["url"] = url
    elif "href" in props:
        args["href"] = url
    elif "target" in props:
        args["target"] = url

    if "path" in props:
        args["path"] = str(screenshot_path)
    if "filename" in props:
        args["filename"] = str(screenshot_path)
    if "fullpage" in props:
        args["fullPage"] = True
    if "full_page" in props:
        args["full_page"] = True

    # Common "console messages" flags
    if "errorsonly" in props:
        args["errorsOnly"] = True
    if "errors_only" in props:
        args["errors_only"] = True

    return args


def main() -> int:
    parser = argparse.ArgumentParser(description="Run Streamlit E2E via Playwright MCP server (stdio).")
    parser.add_argument("--app", type=str, required=True, help="Path to Streamlit app entrypoint (.py).")
    parser.add_argument("--artifacts", type=str, default="artifacts", help="Artifacts output directory.")
    parser.add_argument(
        "--playwright-mcp-cmd",
        type=str,
        default="npx -y @playwright/mcp@latest",
        help="Command to start Playwright MCP server (stdio).",
    )
    parser.add_argument("--streamlit-cmd", type=str, default="streamlit", help="Streamlit CLI command.")
    parser.add_argument("--headless", action="store_true", help="Pass --headless to Playwright MCP server if supported.")
    args = parser.parse_args()

    app_path = Path(args.app).resolve()
    artifacts_dir = Path(args.artifacts).resolve()
    artifacts_dir.mkdir(parents=True, exist_ok=True)

    port = _find_free_port()
    base_url = f"http://127.0.0.1:{port}"
    screenshot_path = artifacts_dir / "smoke.png"
    tools_dump_path = artifacts_dir / "tools.json"
    console_dump_path = artifacts_dir / "console.json"

    streamlit_proc = _start_streamlit(app_path, port=port, streamlit_cmd=args.streamlit_cmd)
    try:
        _wait_http_ok(base_url, timeout_s=45.0)
        print(f"[ok] Streamlit reachable at {base_url}")

        mcp_cmd = shlex.split(args.playwright_mcp_cmd)
        if args.headless:
            mcp_cmd.append("--headless")

        mcp = MCPStdioClient(mcp_cmd, request_timeout_s=120.0)
        mcp.start()
        init_res = mcp.initialize(protocol_version="2024-11-05")
        print(f"[ok] MCP initialize: {json.dumps(init_res)[:200]}...")

        tools = mcp.list_tools()
        tools_dump_path.write_text(
            json.dumps(
                [{"name": t.name, "description": t.description, "inputSchema": t.input_schema} for t in tools],
                indent=2,
            ),
            encoding="utf-8",
        )
        print(f"[ok] Discovered {len(tools)} MCP tools (saved to {tools_dump_path})")

        nav_tool = _pick_tool(tools, ["navigate", "goto", "open", "url"])
        shot_tool = _pick_tool(tools, ["screenshot", "snapshot", "capture"])
        console_tool = _pick_tool(tools, ["console_messages", "console", "console-messages", "logs"])

        if nav_tool is None:
            raise RuntimeError("Could not find a navigation-like tool in tools/list output.")
        if shot_tool is None:
            raise RuntimeError("Could not find a screenshot-like tool in tools/list output.")

        nav_args = _best_effort_args(nav_tool, url=base_url, screenshot_path=screenshot_path)
        print(f"[run] Calling navigate tool: {nav_tool.name} args={nav_args}")
        mcp.call_tool(nav_tool.name, nav_args)

        shot_args = _best_effort_args(shot_tool, url=base_url, screenshot_path=screenshot_path)
        print(f"[run] Calling screenshot tool: {shot_tool.name} args={shot_args}")
        shot_res = mcp.call_tool(shot_tool.name, shot_args)

        if not screenshot_path.exists() and isinstance(shot_res, dict):
            maybe_b64 = shot_res.get("data") or shot_res.get("base64")
            if isinstance(maybe_b64, str) and maybe_b64:
                screenshot_path.write_bytes(base64.b64decode(maybe_b64))

        if screenshot_path.exists():
            print(f"[ok] Screenshot saved to {screenshot_path}")
        else:
            print("[warn] Screenshot not written to disk. Inspect tool response in logs/artifacts.")

        # Best-effort console capture (non-fatal).
        if console_tool is not None:
            try:
                console_args = _best_effort_args(console_tool, url=base_url, screenshot_path=screenshot_path)
                required = {k.lower() for k in _tool_required(console_tool)}
                if required and not required.issubset({k.lower() for k in console_args.keys()}):
                    print(f"[warn] Console tool requires {sorted(required)}; skipping capture.")
                else:
                    print(f"[run] Calling console tool: {console_tool.name} args={console_args}")
                    console_res = mcp.call_tool(console_tool.name, console_args)
                    console_dump_path.write_text(
                        json.dumps(console_res, indent=2, sort_keys=True, default=str),
                        encoding="utf-8",
                    )
                    print(f"[ok] Console messages saved to {console_dump_path}")
            except Exception as e:
                print(f"[warn] Console capture failed: {e}")

        mcp.terminate()
        return 0
    finally:
        _stop_process(streamlit_proc, name="streamlit")


if __name__ == "__main__":
    raise SystemExit(main())

from __future__ import annotations

import argparse
import shutil
from pathlib import Path


def _templates_dir() -> Path:
    return Path(__file__).resolve().parents[1] / "templates"


def list_templates() -> list[str]:
    root = _templates_dir()
    if not root.exists():
        return []
    return sorted([p.name for p in root.iterdir() if p.is_dir() and not p.name.startswith("_")])


def copy_template(*, template: str, dest: Path) -> None:
    src = _templates_dir() / template
    if not src.exists() or not src.is_dir():
        raise FileNotFoundError(f"Template not found: {src}")
    shutil.copytree(src, dest, dirs_exist_ok=False)


def main() -> int:
    parser = argparse.ArgumentParser(description="Scaffold a Streamlit app from bundled templates.")
    parser.add_argument("--list", action="store_true", help="List available templates and exit.")
    parser.add_argument("--template", type=str, help="Template name (see --list).")
    parser.add_argument("--dest", type=str, help="Destination directory (must not exist).")
    args = parser.parse_args()

    if args.list:
        for name in list_templates():
            print(name)
        return 0

    if not args.template or not args.dest:
        parser.error("--template and --dest are required (or use --list).")

    dest = Path(args.dest).resolve()
    if dest.exists():
        raise FileExistsError(f"Destination already exists: {dest}")

    copy_template(template=args.template, dest=dest)
    print(f"[ok] Created {dest} from template '{args.template}'")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

from __future__ import annotations

import argparse
import re
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


DOCS_ROOT = "https://docs.streamlit.io"
LLMS_TXT_URL = f"{DOCS_ROOT}/llms.txt"


_MD_LINK_RE = re.compile(r"\[[^\]]+\]\(([^)]+)\)")


def _fetch_text(url: str, *, timeout_s: float = 30.0) -> str:
    req = urllib.request.Request(url, headers={"User-Agent": "streamlit-master-architect/0.1"})
    with urllib.request.urlopen(req, timeout=timeout_s) as resp:
        return resp.read().decode("utf-8", errors="replace")


def _normalize_url(raw: str) -> str | None:
    raw = raw.strip()
    if not raw:
        return None
    if raw.startswith("#"):
        return None
    if raw.startswith("mailto:"):
        return None
    if raw.startswith("http://") or raw.startswith("https://"):
        return raw
    if raw.startswith("/"):
        return f"{DOCS_ROOT}{raw}"
    return None


def parse_llms_txt(markdown: str) -> list[str]:
    urls: list[str] = []
    for m in _MD_LINK_RE.finditer(markdown):
        u = _normalize_url(m.group(1))
        if u:
            urls.append(u)
    # Always include llms index itself.
    urls.append(LLMS_TXT_URL)
    deduped = sorted(set(urls))
    return deduped


def _safe_filename(url: str) -> str:
    # Keep stable, filesystem-safe names.
    path = url.replace(DOCS_ROOT, "").lstrip("/")
    if not path:
        path = "root"
    path = re.sub(r"[^a-zA-Z0-9._/-]+", "_", path)
    path = path.strip("_").replace("/", "__")
    return f"{path}.html"


@dataclass(frozen=True)
class FetchResult:
    url: str
    ok: bool
    error: str | None = None
    out_path: Path | None = None


def fetch_pages(
    urls: Iterable[str],
    *,
    out_dir: Path,
    max_pages: int,
    concurrency: int,
    sleep_s: float,
) -> list[FetchResult]:
    out_dir.mkdir(parents=True, exist_ok=True)
    urls_list = list(urls) if max_pages <= 0 else list(urls)[:max_pages]

    def _one(url: str) -> FetchResult:
        if sleep_s:
            time.sleep(sleep_s)
        try:
            html = _fetch_text(url)
        except (urllib.error.URLError, TimeoutError) as e:
            return FetchResult(url=url, ok=False, error=str(e))
        out_path = out_dir / _safe_filename(url)
        out_path.write_text(html, encoding="utf-8")
        return FetchResult(url=url, ok=True, out_path=out_path)

    results: list[FetchResult] = []
    with ThreadPoolExecutor(max_workers=max(1, concurrency)) as ex:
        futs = [ex.submit(_one, u) for u in urls_list]
        for fut in as_completed(futs):
            results.append(fut.result())
    return results


def main() -> int:
    parser = argparse.ArgumentParser(description="Sync Streamlit docs starting from llms.txt.")
    parser.add_argument("--out", type=str, default="docs_snapshot", help="Output directory.")
    parser.add_argument("--fetch", action="store_true", help="Fetch pages (HTML) in addition to URL list.")
    parser.add_argument(
        "--max-pages",
        type=int,
        default=50,
        help="Max pages to fetch when --fetch is set. Use 0 to fetch all URLs from llms.txt.",
    )
    parser.add_argument("--concurrency", type=int, default=8, help="Fetch concurrency for --fetch.")
    parser.add_argument("--sleep-s", type=float, default=0.0, help="Optional sleep between fetches (per worker).")
    args = parser.parse_args()

    out_dir = Path(args.out).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    llms = _fetch_text(LLMS_TXT_URL)
    (out_dir / "llms.txt").write_text(llms, encoding="utf-8")

    urls = parse_llms_txt(llms)
    (out_dir / "urls.txt").write_text("\n".join(urls) + "\n", encoding="utf-8")
    print(f"[ok] Parsed {len(urls)} URLs from {LLMS_TXT_URL}")

    if not args.fetch:
        print(f"[ok] Wrote {out_dir / 'urls.txt'}")
        return 0

    pages_dir = out_dir / "pages"
    results = fetch_pages(
        urls,
        out_dir=pages_dir,
        max_pages=args.max_pages,
        concurrency=max(1, args.concurrency),
        sleep_s=max(0.0, args.sleep_s),
    )
    ok = sum(1 for r in results if r.ok)
    bad = len(results) - ok
    print(f"[ok] Fetched {ok} pages ({bad} failed) into {pages_dir}")
    return 0 if bad == 0 else 2


if __name__ == "__main__":
    raise SystemExit(main())

[server]
headless = true
address = "127.0.0.1"
port = 8501

[client]
showErrorDetails = true

[theme]
base = "light"
primaryColor = "#5B8DEF"
backgroundColor = "#FFFFFF"
secondaryBackgroundColor = "#F6F7F9"
textColor = "#111111"
font = "sans serif"

streamlit>=1.52.2,<2.0.0
pandas>=2.2.0

from __future__ import annotations

import datetime as dt
import io

import pandas as pd
import streamlit as st


APP_TITLE = "SMA Demo — Single Page"


def _set_page() -> None:
    st.set_page_config(page_title=APP_TITLE, page_icon="🧱", layout="wide")


@st.cache_data(show_spinner="Loading demo dataset…")
def load_demo_data(n: int = 2_000) -> pd.DataFrame:
    now = dt.datetime.now()
    df = pd.DataFrame(
        {
            "ts": pd.date_range(now - dt.timedelta(days=30), periods=n, freq="h"),
            "value": pd.Series(range(n)).astype(float).rolling(24, min_periods=1).mean(),
            "group": (pd.Series(range(n)) % 5).map(lambda x: f"g{x}"),
        }
    )
    return df


def _make_csv_bytes(df: pd.DataFrame) -> bytes:
    buf = io.StringIO()
    df.to_csv(buf, index=False)
    return buf.getvalue().encode("utf-8")


def main() -> None:
    _set_page()

    st.title("Streamlit patterns demo")
    st.caption(
        "Demonstrates caching, datetime input, deferred download generation, chat input (audio optional), and safe HTML."
    )

    df = load_demo_data()

    with st.sidebar:
        st.header("Controls")
        group = st.selectbox("Group", sorted(df["group"].unique()), key="group")

        def dt_input(label: str, *, key_prefix: str) -> dt.datetime:
            now = dt.datetime.now()
            if hasattr(st, "datetime_input"):
                v: dt.datetime | str = "now"
                out = st.datetime_input(label, value=v, key=key_prefix)
                return out if isinstance(out, dt.datetime) else now
            d = st.date_input(f"{label} (date)", value=now.date(), key=f"{key_prefix}__date")
            t = st.time_input(f"{label} (time)", value=now.time().replace(second=0, microsecond=0), key=f"{key_prefix}__time")
            return dt.datetime.combine(d, t)

        start = dt_input("Start datetime", key_prefix="start_dt")
        end = dt_input("End datetime", key_prefix="end_dt")

    if isinstance(start, dt.datetime) and isinstance(end, dt.datetime) and start > end:
        st.error("Start must be <= end.")
        st.stop()

    filtered = df[df["group"] == group].copy()
    if isinstance(start, dt.datetime):
        filtered = filtered[filtered["ts"] >= start.replace(tzinfo=None)]
    if isinstance(end, dt.datetime):
        filtered = filtered[filtered["ts"] <= end.replace(tzinfo=None)]

    st.subheader("Preview")
    st.dataframe(filtered.head(200), height=320)

    st.subheader("Chart")
    st.line_chart(filtered.set_index("ts")["value"])

    st.subheader("Deferred download (callable)")
    st.write("Use a callable to generate download data only when the user clicks.")

    def generate_csv() -> bytes:
        return _make_csv_bytes(filtered)

    def download_ui() -> None:
        st.download_button(
            label="Download filtered CSV",
            data=generate_csv,
            file_name=f"demo-{group}.csv",
            mime="text/csv",
        )

    # Streamlit docs recommend wrapping downloads in a fragment to prevent full reruns on click.
    if hasattr(st, "fragment"):

        @st.fragment
        def _download_fragment() -> None:
            download_ui()

        _download_fragment()
    else:
        download_ui()

    st.divider()
    st.subheader("Chat input (audio optional in 1.52.x)")

    try:
        user_input = st.chat_input(
            "Ask about the dataset…",
            accept_audio=True,
            audio_sample_rate=16_000,
        )
    except TypeError:
        user_input = st.chat_input("Ask about the dataset…")
    if user_input:
        user_text = user_input if isinstance(user_input, str) else str(user_input)
        with st.chat_message("user"):
            st.write(user_text)

        with st.chat_message("assistant"):
            st.write("This template does not call an external LLM. Wire your provider in your app code.")

    st.divider()
    st.subheader("Safe HTML (CSS-only) example")

    css = """
    <style>
      .sma-note { padding: 0.75rem 1rem; border-radius: 0.75rem; border: 1px solid rgba(127,127,127,0.35); }
      .sma-note b { font-weight: 700; }
    </style>
    <div class="sma-note"><b>Note:</b> Prefer theming/config over custom HTML. Do not execute untrusted JS.</div>
    """
    if hasattr(st, "html"):
        st.html(css)
    else:
        st.info("`st.html` is unavailable in this Streamlit version; using plain text fallback.")
        st.markdown("**Note:** Prefer theming/config over custom HTML. Do not execute untrusted JS.")


if __name__ == "__main__":
    main()

from __future__ import annotations

from pathlib import Path

from streamlit.testing.v1 import AppTest


def test_app_loads() -> None:
    app = Path(__file__).resolve().parents[1] / "streamlit_app.py"
    at = AppTest.from_file(str(app)).run()
    assert any("Streamlit patterns demo" in h.value for h in at.title), "Expected page title not found."

export default function smaCounter(component) {
  const { data, parentElement, setStateValue, setTriggerValue } = component

  // Basic resilient access: Python kwargs are passed as an object.
  const label = (data && data.label) ? String(data.label) : "Click"

  // Avoid clobbering parentElement; render into a container.
  const root = document.createElement("div")
  root.style.display = "flex"
  root.style.gap = "0.5rem"
  root.style.alignItems = "center"
  root.style.fontFamily = "system-ui, -apple-system, Segoe UI, Roboto, sans-serif"

  const button = document.createElement("button")
  button.textContent = label
  button.type = "button"

  const countEl = document.createElement("span")
  countEl.textContent = "0"

  // Mount
  root.appendChild(button)
  root.appendChild(countEl)
  parentElement.appendChild(root)

  let count = 0
  const onClick = () => {
    count += 1
    countEl.textContent = String(count)
    setStateValue("count", count)
    setTriggerValue("clicked", true)
  }
  button.addEventListener("click", onClick)

  // Cleanup on unmount
  return () => {
    button.removeEventListener("click", onClick)
    root.remove()
  }
}

streamlit>=1.52.2,<2.0.0

from __future__ import annotations

from pathlib import Path
from typing import Any

import streamlit as st


def _load_js() -> str:
    return (Path(__file__).parent / "frontend" / "component.js").read_text(encoding="utf-8")


def main() -> None:
    st.set_page_config(page_title="SMA — Component v2", page_icon="🧩", layout="centered")
    st.title("Custom component v2 (minimal)")
    st.caption("Demonstrates the v2 contract: data → JS, state/trigger → Python callbacks.")

    if not hasattr(st, "components") or not hasattr(st.components, "v2"):
        st.error("This template requires Streamlit custom components v2 (`st.components.v2`).")
        st.stop()

    counter = st.components.v2.component(
        name="sma_counter",
        js=_load_js(),
        isolate_styles=True,
    )

    def on_clicked() -> None:
        st.toast("Clicked!", icon="✅")

    result: Any = counter(
        label="Click me",
        on_clicked_change=on_clicked,
        key="counter",
    )

    st.subheader("Component result")
    st.write(result)


if __name__ == "__main__":
    main()

streamlit>=1.52.2,<2.0.0

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Generator, Literal, TypedDict

import streamlit as st


Role = Literal["user", "assistant"]


class ChatMessage(TypedDict):
    role: Role
    content: str


@dataclass(frozen=True)
class ChatConfig:
    title: str = "SMA — Chat App Skeleton"
    icon: str = "💬"


def _set_page() -> None:
    st.set_page_config(page_title=ChatConfig.title, page_icon=ChatConfig.icon, layout="wide")


def _ss_init() -> None:
    if "messages" not in st.session_state:
        st.session_state["messages"] = []


def _fake_streaming_llm(prompt: str) -> Generator[str, None, None]:
    for tok in (prompt.upper().split()[:50] or ["(empty)"]):
        yield tok + " "
        time.sleep(0.03)


def main() -> None:
    _set_page()
    _ss_init()

    st.title("Chat App Skeleton (streaming-ready)")
    st.caption("Wire your LLM provider in a dedicated module; keep rendering separate from logic.")

    for msg in st.session_state["messages"]:
        with st.chat_message(msg["role"]):
            st.write(msg["content"])

    try:
        user_input = st.chat_input(
            "Send a message (audio optional)…",
            accept_audio=True,
            audio_sample_rate=16_000,
        )
    except TypeError:
        user_input = st.chat_input("Send a message…")

    if user_input:
        user_text = user_input if isinstance(user_input, str) else str(user_input)
        st.session_state["messages"].append({"role": "user", "content": user_text})
        with st.chat_message("user"):
            st.write(user_text)

        with st.chat_message("assistant"):
            streamed = st.write_stream(_fake_streaming_llm(user_text))

        st.session_state["messages"].append({"role": "assistant", "content": str(streamed)})


if __name__ == "__main__":
    main()