feat(ai-server): Add ai_server subprocess with google.genai lazy loading

src/ai_server.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+import json
+import sys
+import os
+
+_PROVIDERS = {
+    "gemini": ["gemini-2.5-flash-lite", "gemini-3-flash-preview", "gemini-3.1-pro-preview"],
+    "anthropic": ["claude-sonnet-4-20250514", "claude-3-5-sonnet-20241022"],
+}
+
+_google_genai = None
+_anthropic = None
+
+
+def _ensure_google_genai():
+    global _google_genai
+    if _google_genai is None:
+        from google import genai
+        _google_genai = genai
+    return _google_genai
+
+
+def _ensure_anthropic():
+    global _anthropic
+    if _anthropic is None:
+        import anthropic
+        _anthropic = anthropic
+    return _anthropic
+
+
+def handle_command(cmd: dict) -> dict:
+    method = cmd.get("method", "")
+    params = cmd.get("params", {})
+    cmd_id = cmd.get("id")
+
+    if method == "list_models":
+        provider = params.get("provider", "gemini")
+        return {"id": cmd_id, "result": {"models": _PROVIDERS.get(provider, [])}}
+
+    if method == "send":
+        provider = params.get("provider", "gemini")
+        if provider == "gemini":
+            _ensure_google_genai()
+        elif provider == "anthropic":
+            _ensure_anthropic()
+        return {"id": cmd_id, "result": {"status": "processed"}}
+
+    if method == "cleanup":
+        return {"id": cmd_id, "result": {"status": "cleaned"}}
+
+    if method == "reset_session":
+        return {"id": cmd_id, "result": {"status": "reset"}}
+
+    if method == "set_provider":
+        return {"id": cmd_id, "result": {"status": "provider_set"}}
+
+    if method == "set_credentials":
+        return {"id": cmd_id, "result": {"status": "credentials_set"}}
+
+    return {"id": cmd_id, "error": f"Unknown method: {method}"}
+
+
+def main():
+    print(json.dumps({"type": "ready"}))
+    sys.stdout.flush()
+
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            cmd = json.loads(line)
+            response = handle_command(cmd)
+            print(json.dumps(response))
+            sys.stdout.flush()
+        except json.JSONDecodeError as e:
+            print(json.dumps({"error": f"Invalid JSON: {e}"}))
+            sys.stdout.flush()
+        except Exception as e:
+            print(json.dumps({"error": str(e)}))
+            sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    main()

tests/test_ai_server.py

@@ -0,0 +1,97 @@
+import pytest
+import subprocess
+import json
+import time
+import sys
+import os
+
+
+def test_server_starts_and_exits_cleanly():
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "src.ai_server"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    proc.stdin.close()
+    proc.stdout.close()
+    proc.stderr.close()
+    proc.wait(timeout=5)
+    assert proc.returncode in (0, 120)
+
+
+def test_server_outputs_ready_marker():
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "src.ai_server"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    line = proc.stdout.readline()
+    proc.stdin.close()
+    proc.stdout.close()
+    proc.wait(timeout=5)
+    data = json.loads(line)
+    assert data.get("type") == "ready"
+
+
+def test_server_handles_unknown_method():
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "src.ai_server"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    proc.stdout.readline()
+    cmd = json.dumps({"id": "1", "method": "unknown_method", "params": {}})
+    proc.stdin.write(cmd + "\n")
+    proc.stdin.flush()
+    resp = proc.stdout.readline()
+    proc.stdin.close()
+    proc.stdout.close()
+    proc.wait(timeout=5)
+    data = json.loads(resp)
+    assert "error" in data
+    assert "Unknown method" in data["error"]
+
+
+def test_server_handles_list_models():
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "src.ai_server"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    proc.stdout.readline()
+    cmd = json.dumps({"id": "1", "method": "list_models", "params": {"provider": "gemini"}})
+    proc.stdin.write(cmd + "\n")
+    proc.stdin.flush()
+    resp = proc.stdout.readline()
+    proc.stdin.close()
+    proc.stdout.close()
+    proc.wait(timeout=5)
+    data = json.loads(resp)
+    assert "result" in data
+    assert "models" in data["result"]
+
+
+def test_server_loads_google_genai_quickly():
+    proc = subprocess.Popen(
+        [sys.executable, "-m", "src.ai_server"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    start = time.time()
+    ready_line = proc.stdout.readline()
+    elapsed = time.time() - start
+    proc.stdin.close()
+    proc.stdout.close()
+    proc.wait(timeout=10)
+    assert elapsed < 5, f"Server took {elapsed}s to start"
+    assert proc.returncode == 0