@qa-gentic/stlc-agents 1.0.16 → 1.0.18

package/src/stlc_agents/shared/install_hook.py

@@ -0,0 +1,154 @@
+"""
+install_hook.py  —  stlc_agents.shared.install_hook
+─────────────────────────────────────────────────────
+Called automatically by postinstall.js after `pip install qa-gentic-stlc-agents`.
+Also callable manually: python -m stlc_agents.shared.install_hook
+
+Applies the cost tracking patch to all 5 MCP server files by importing
+and running the same logic as scripts/apply_cost_tracking.py, but resolved
+relative to the installed package location (works in site-packages, .venv, etc).
+"""
+
+from __future__ import annotations
+import re
+import sys
+from pathlib import Path
+
+
+SERVERS = [
+    ("agent_gherkin_generator",     "qa-gherkin-generator"),
+    ("agent_test_case_manager",     "qa-test-case-manager"),
+    ("agent_playwright_generator",  "qa-playwright-generator"),
+    ("agent_helix_writer",          "qa-helix-writer"),
+    ("agent_jira_manager",          "qa-jira-manager"),
+]
+
+IMPORT_MARKER = "from stlc_agents.shared.cost_tracker import track"
+TIME_IMPORT   = "import time"
+
+OLD_RETURN = (
+    'return [types.TextContent(type="text", '
+    'text=json.dumps(result, indent=2, ensure_ascii=False))]'
+)
+NEW_RETURN = "return track(result, tool_name=name, server={server!r}, t0=t0)"
+
+OLD_ERR_BLOCK = """\
+        return [types.TextContent(
+            type="text",
+            text=json.dumps({"error": str(exc), "tool": name}, indent=2),
+        )]"""
+
+NEW_ERR_BLOCK = """\
+        err_result = {{"error": str(exc), "tool": name}}
+        return track(err_result, tool_name=name, server={server!r}, t0=t0)"""
+
+
+def _root() -> Path:
+    """Resolve the stlc_agents package root regardless of install method."""
+    import stlc_agents
+    return Path(stlc_agents.__file__).parent
+
+
+def patch_server(agent_dir: str, server_name: str, root: Path) -> str:
+    """Patch one server file. Returns 'patched' | 'already_patched' | 'not_found' | 'no_change'."""
+    path = root / agent_dir / "server.py"
+    if not path.exists():
+        return "not_found"
+
+    src = path.read_text(encoding="utf-8")
+    if IMPORT_MARKER in src:
+        return "already_patched"
+
+    original = src
+
+    # 1. import time
+    if TIME_IMPORT not in src:
+        src = src.replace("import sys\n", "import sys\nimport time\n", 1)
+
+    # 2. cost_tracker import — after last `from stlc_agents...` line
+    last_match = None
+    for m in re.finditer(r"^from stlc_agents\..+\n", src, re.MULTILINE):
+        last_match = m
+    if last_match:
+        pos = last_match.end()
+        src = src[:pos] + "from stlc_agents.shared.cost_tracker import track\n" + src[pos:]
+    else:
+        src = src.replace(
+            "from dotenv import load_dotenv\n",
+            "from dotenv import load_dotenv\nfrom stlc_agents.shared.cost_tracker import track\n",
+            1,
+        )
+
+    # 3. t0 = time.monotonic() inside call_tool()
+    src = src.replace(
+        "@app.call_tool()\nasync def call_tool(name: str, arguments: dict)"
+        " -> list[types.TextContent]:\n    try:\n",
+        "@app.call_tool()\nasync def call_tool(name: str, arguments: dict)"
+        " -> list[types.TextContent]:\n    t0 = time.monotonic()\n    try:\n",
+        1,
+    )
+
+    # 4. Replace all result return lines
+    new_ret = NEW_RETURN.format(server=server_name)
+    for indent in ("        ", "            ", "                "):
+        src = src.replace(f"{indent}{OLD_RETURN}", f"{indent}{new_ret}")
+
+    # 5. Replace error-path block
+    src = src.replace(OLD_ERR_BLOCK, NEW_ERR_BLOCK.format(server=server_name))
+
+    if src == original:
+        return "no_change"
+
+    path.write_text(src, encoding="utf-8")
+    return "patched"
+
+
+def apply_cost_tracking() -> None:
+    """Entry point — called by postinstall.js and the console_script."""
+    root = _root()
+
+    ok   = "\x1b[32m✓\x1b[0m"
+    skip = "\x1b[33m–\x1b[0m"
+    err  = "\x1b[31m✗\x1b[0m"
+
+    print("\n  stlc-agents · Activating cost tracking on MCP servers...\n")
+
+    any_patched = False
+    for agent_dir, server_name in SERVERS:
+        status = patch_server(agent_dir, server_name, root)
+        if status == "patched":
+            print(f"  {ok}  {agent_dir}  ({server_name})")
+            any_patched = True
+        elif status == "already_patched":
+            print(f"  {skip}  {agent_dir}  — already active")
+        elif status == "not_found":
+            print(f"  {err}  {agent_dir}/server.py  — not found (skip)")
+        elif status == "no_change":
+            print(f"  {skip}  {agent_dir}  — no matching pattern (manual patch needed)")
+
+    print()
+    if any_patched:
+        print("  Cost tracking is now active.  On every MCP tool call you will see:")
+        print("  [stlc-cost] <server> · <tool>  ~<N>K tokens  $<cost>  (session: $<total>)")
+        print()
+        print("  Session logs:  ~/.qa-stlc/cost-<session-id>.jsonl")
+        print("  View report:   qa-stlc cost")
+        print("  View all:      qa-stlc cost --all")
+        print()
+        print("  Environment variables:")
+        print("    STLC_COST_TRACKING=false          disable output")
+        print("    STLC_CODING_AGENT_MODEL=<model>   set your agent's model for exact pricing")
+        print("      e.g. claude-sonnet-4-6 | claude-opus-4-6 | gpt-4o")
+        print("    STLC_COST_LOG_DIR=<path>          change log directory")
+    else:
+        print("  All servers already have cost tracking active.")
+
+    print()
+
+
+def main() -> None:
+    apply_cost_tracking()
+
+
+if __name__ == "__main__":
+    main()

package/src/stlc_agents/shared/pricing.py

@@ -0,0 +1,72 @@
+"""
+pricing.py  —  Model pricing registry for stlc-agents cost tracking.
+
+Prices: USD per million tokens (MTok).
+Source: Anthropic official docs, April 2026.
+
+Models this repo actually calls:
+  - claude-sonnet-4-20250514  (LocatorHealer AI Vision, default)
+  - gpt-4o                    (LocatorHealer AI Vision, copilot provider)
+  + whatever coding agent the user runs (Claude / Copilot / Cursor / Windsurf)
+    — the user declares this via STLC_CODING_AGENT_MODEL env var.
+"""
+
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class ModelPricing:
+    model_id: str
+    display_name: str
+    provider: str
+    input_per_mtok: float        # USD / 1M input tokens
+    output_per_mtok: float       # USD / 1M output tokens
+    cache_write_per_mtok: float  # USD / 1M cache-write tokens
+    cache_read_per_mtok: float   # USD / 1M cache-read tokens
+
+    def cost(
+        self,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        cache_write_tokens: int = 0,
+        cache_read_tokens: int = 0,
+    ) -> float:
+        return (
+            (input_tokens        / 1_000_000) * self.input_per_mtok
+            + (output_tokens     / 1_000_000) * self.output_per_mtok
+            + (cache_write_tokens/ 1_000_000) * self.cache_write_per_mtok
+            + (cache_read_tokens / 1_000_000) * self.cache_read_per_mtok
+        )
+
+
+_REGISTRY: list[ModelPricing] = [
+    # ── Anthropic ──────────────────────────────────────────────────────────
+    ModelPricing("claude-sonnet-4-20250514", "Claude Sonnet 4",    "anthropic",  3.00, 15.00,  3.75, 0.30),
+    ModelPricing("claude-sonnet-4-6",        "Claude Sonnet 4.6",  "anthropic",  3.00, 15.00,  3.75, 0.30),
+    ModelPricing("claude-haiku-4-5-20251001","Claude Haiku 4.5",   "anthropic",  1.00,  5.00,  1.25, 0.10),
+    ModelPricing("claude-opus-4-6",          "Claude Opus 4.6",    "anthropic",  5.00, 25.00,  6.25, 0.50),
+    ModelPricing("claude-opus-4-7",          "Claude Opus 4.7",    "anthropic",  5.00, 25.00,  6.25, 0.50),
+    # ── OpenAI / Copilot ──────────────────────────────────────────────────
+    ModelPricing("gpt-4o",                   "GPT-4o",             "openai",     2.50, 10.00,  0.00, 0.00),
+    ModelPricing("gpt-4o-mini",              "GPT-4o Mini",        "openai",     0.15,  0.60,  0.00, 0.00),
+]
+
+_by_id: dict[str, ModelPricing] = {p.model_id: p for p in _REGISTRY}
+
+
+def get_pricing(model_id: str) -> Optional[ModelPricing]:
+    """Exact match first, then longest substring match."""
+    key = model_id.lower().strip()
+    if key in _by_id:
+        return _by_id[key]
+    # Substring: "claude-sonnet-4-20250514" ⊇ "sonnet-4"
+    for p in _REGISTRY:
+        if key in p.model_id or p.model_id in key:
+            return p
+    return None
+
+
+def list_models() -> list[ModelPricing]:
+    return list(_REGISTRY)