# Phase 2 Critical Path Implementation Plan

> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

**Goal:** Enable the revenue path: user registers → subscribes to Pro → gets API key → uses CLI proxy with plan-based limits.

**Architecture:** The Python proxy is the single authority for API keys. It exposes key management endpoints (`/api/keys/*`) that the Next.js dashboard calls. Each key stores the user's plan in its metadata. The proxy's existing `RateLimitMiddleware` is extended to enforce plan-based RPM limits. Stripe checkout is handled entirely in Next.js (already has webhook handler).

**Tech Stack:** Python/FastAPI (proxy), Next.js/Auth.js (dashboard), Prisma/SQLite (user DB), Stripe (billing)

---

## Task 1: Key Management API — Tests

**Files:**
- Create: `tests/test_key_management_api.py`

**Step 1: Write failing tests for key management endpoints**

```python
"""Tests for the /api/keys/* management endpoints."""
from __future__ import annotations

import pytest
from httpx import ASGITransport, AsyncClient

from llmhost.proxy.app import create_app


@pytest.fixture
def app():
    return create_app()


@pytest.fixture
async def client(app):
    transport = ASGITransport(app=app)
    async with AsyncClient(transport=transport, base_url="http://test") as c:
        yield c


@pytest.mark.asyncio
async def test_create_key(client: AsyncClient) -> None:
    resp = await client.post("/api/keys", json={"label": "my-laptop", "plan": "free"})
    assert resp.status_code == 201
    data = resp.json()
    assert "key_id" in data
    assert "raw_key" in data
    assert data["label"] == "my-laptop"
    assert data["plan"] == "free"
    assert data["raw_key"].startswith("lmh_")  # prefixed keys


@pytest.mark.asyncio
async def test_list_keys(client: AsyncClient) -> None:
    # Create a key first
    await client.post("/api/keys", json={"label": "test-key", "plan": "pro"})
    resp = await client.get("/api/keys")
    assert resp.status_code == 200
    keys = resp.json()["keys"]
    assert len(keys) >= 1
    assert "raw_key" not in str(keys)  # raw key never exposed in list
    assert keys[-1]["label"] == "test-key"
    assert keys[-1]["plan"] == "pro"


@pytest.mark.asyncio
async def test_revoke_key(client: AsyncClient) -> None:
    create_resp = await client.post("/api/keys", json={"label": "to-revoke"})
    key_id = create_resp.json()["key_id"]
    resp = await client.delete(f"/api/keys/{key_id}")
    assert resp.status_code == 200
    # Verify key no longer works
    list_resp = await client.get("/api/keys")
    ids = [k["id"] for k in list_resp.json()["keys"]]
    assert key_id not in ids


@pytest.mark.asyncio
async def test_revoke_nonexistent_key(client: AsyncClient) -> None:
    resp = await client.delete("/api/keys/nonexistent")
    assert resp.status_code == 404


@pytest.mark.asyncio
async def test_create_key_missing_label(client: AsyncClient) -> None:
    resp = await client.post("/api/keys", json={})
    assert resp.status_code == 422


@pytest.mark.asyncio
async def test_created_key_authenticates(client: AsyncClient) -> None:
    create_resp = await client.post("/api/keys", json={"label": "auth-test"})
    raw_key = create_resp.json()["raw_key"]
    # Use the key to hit an authenticated endpoint
    resp = await client.get("/v1/models", headers={"Authorization": f"Bearer {raw_key}"})
    assert resp.status_code == 200
```

**Step 2: Run tests to verify they fail**

Run: `pytest tests/test_key_management_api.py -v`
Expected: FAIL — `/api/keys` routes don't exist yet

---

## Task 2: Key Management API — Implementation

**Files:**
- Create: `llmhosts/proxy/key_routes.py`
- Modify: `llmhosts/proxy/app.py` (add router include)
- Modify: `llmhosts/proxy/auth.py` (add plan field + key prefix)

**Step 1: Add plan field and key prefix to ProxyKeyStore**

In `llmhosts/proxy/auth.py`, update `create_key` to accept a `plan` parameter and prefix keys with `lmh_`:

```python
def create_key(self, label: str, plan: str = "free") -> tuple[str, str]:
    key_id = uuid.uuid4().hex[:12]
    raw_key = "lmh_" + secrets.token_urlsafe(32)
    key_hash = _hash_key(raw_key)
    self._keys[key_id] = {
        "hash": key_hash,
        "label": label,
        "plan": plan,
        "created_at": datetime.now(tz=timezone.utc).isoformat(),
        "last_used": None,
    }
    self._save()
    logger.info("Created proxy key %s (%s) plan=%s", key_id, label, plan)
    return key_id, raw_key
```

Update `list_keys` to include plan:

```python
def list_keys(self) -> list[dict[str, Any]]:
    return [
        {
            "id": key_id,
            "label": meta["label"],
            "plan": meta.get("plan", "free"),
            "created_at": meta["created_at"],
            "last_used": meta["last_used"],
        }
        for key_id, meta in self._keys.items()
    ]
```

Add `get_plan` method:

```python
def get_plan(self, raw_key: str) -> str:
    incoming_hash = _hash_key(raw_key)
    for meta in self._keys.values():
        if meta["hash"] == incoming_hash:
            return meta.get("plan", "free")
    return "free"
```

**Step 2: Create key management routes**

Create `llmhosts/proxy/key_routes.py`:

```python
"""Key management API routes: /api/keys/*"""
from __future__ import annotations

from typing import Any

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel

router = APIRouter(prefix="/api/keys", tags=["keys"])


class CreateKeyRequest(BaseModel):
    label: str
    plan: str = "free"
    user_id: str | None = None


@router.post("", status_code=201)
async def create_key(request: Request, body: CreateKeyRequest) -> JSONResponse:
    key_store = request.app.state.key_store
    key_id, raw_key = key_store.create_key(label=body.label, plan=body.plan)
    return JSONResponse(
        status_code=201,
        content={
            "key_id": key_id,
            "raw_key": raw_key,
            "label": body.label,
            "plan": body.plan,
        },
    )


@router.get("")
async def list_keys(request: Request) -> dict[str, Any]:
    key_store = request.app.state.key_store
    return {"keys": key_store.list_keys()}


@router.delete("/{key_id}")
async def revoke_key(request: Request, key_id: str) -> JSONResponse:
    key_store = request.app.state.key_store
    if key_store.revoke_key(key_id):
        return JSONResponse(content={"status": "revoked", "key_id": key_id})
    return JSONResponse(status_code=404, content={"error": "Key not found"})
```

**Step 3: Register the router in app.py**

In `llmhosts/proxy/app.py`, add:

```python
from llmhost.proxy.key_routes import router as key_routes_router
# ... in create_app():
app.include_router(key_routes_router)
```

**Step 4: Run tests to verify they pass**

Run: `pytest tests/test_key_management_api.py -v`
Expected: ALL PASS

**Step 5: Run full suite to verify no regressions**

Run: `pytest tests/ -v`
Expected: 282+ passed, 0 failed

**Step 6: Commit**

```bash
git add llmhosts/proxy/key_routes.py llmhosts/proxy/auth.py llmhosts/proxy/app.py tests/test_key_management_api.py
git commit -m "feat(api): key management endpoints with plan support"
```

---

## Task 3: Plan-Based Rate Limiting — Tests

**Files:**
- Create: `tests/test_plan_ratelimit.py`

**Step 1: Write failing tests for plan-based rate limits**

```python
"""Tests for plan-based rate limiting."""
from __future__ import annotations

import pytest
from httpx import ASGITransport, AsyncClient

from llmhost.config import LLMHostsConfig, SecurityConfig
from llmhost.proxy.app import create_app


@pytest.fixture
def auth_app():
    config = LLMHostsConfig()
    config.security = SecurityConfig(require_auth=True, rate_limit_rpm=60)
    return create_app(config)


@pytest.fixture
async def auth_client(auth_app):
    transport = ASGITransport(app=auth_app)
    async with AsyncClient(transport=transport, base_url="http://test") as c:
        yield c


@pytest.mark.asyncio
async def test_free_plan_gets_60_rpm(auth_client: AsyncClient) -> None:
    # Create a free key
    create_resp = await auth_client.post("/api/keys", json={"label": "free-key", "plan": "free"})
    raw_key = create_resp.json()["raw_key"]
    resp = await auth_client.get(
        "/v1/models",
        headers={"Authorization": f"Bearer {raw_key}"},
    )
    assert resp.status_code == 200
    assert resp.headers.get("X-RateLimit-Limit") == "60"


@pytest.mark.asyncio
async def test_pro_plan_gets_300_rpm(auth_client: AsyncClient) -> None:
    create_resp = await auth_client.post("/api/keys", json={"label": "pro-key", "plan": "pro"})
    raw_key = create_resp.json()["raw_key"]
    resp = await auth_client.get(
        "/v1/models",
        headers={"Authorization": f"Bearer {raw_key}"},
    )
    assert resp.status_code == 200
    assert resp.headers.get("X-RateLimit-Limit") == "300"


@pytest.mark.asyncio
async def test_team_plan_gets_600_rpm(auth_client: AsyncClient) -> None:
    create_resp = await auth_client.post("/api/keys", json={"label": "team-key", "plan": "team"})
    raw_key = create_resp.json()["raw_key"]
    resp = await auth_client.get(
        "/v1/models",
        headers={"Authorization": f"Bearer {raw_key}"},
    )
    assert resp.status_code == 200
    assert resp.headers.get("X-RateLimit-Limit") == "600"
```

**Step 2: Run tests to verify they fail**

Run: `pytest tests/test_plan_ratelimit.py -v`
Expected: FAIL — rate limiter doesn't know about plans yet

---

## Task 4: Plan-Based Rate Limiting — Implementation

**Files:**
- Modify: `llmhosts/proxy/ratelimit.py` (plan-aware buckets)
- Modify: `llmhosts/proxy/auth.py` (attach plan to request state)

**Step 1: Update AuthMiddleware to attach plan to request state**

In `llmhosts/proxy/auth.py`, in the `dispatch` method, after successful key validation:

```python
if key_id is not None:
    request.state.authenticated = True
    request.state.api_key_id = key_id
    # Attach plan for rate limiter
    request.state.plan = self._key_store.get_plan_by_id(key_id)
    return await call_next(request)
```

Add `get_plan_by_id` to `ProxyKeyStore`:

```python
def get_plan_by_id(self, key_id: str) -> str:
    meta = self._keys.get(key_id)
    return meta.get("plan", "free") if meta else "free"
```

**Step 2: Update RateLimitMiddleware for plan-based limits**

In `llmhosts/proxy/ratelimit.py`, add plan RPM mapping and plan-aware bucket selection:

```python
# Plan-based RPM limits
PLAN_RPM: dict[str, int] = {
    "free": 60,
    "pro": 300,
    "team": 600,
    "enterprise": 0,  # 0 = unlimited
}
```

Update `_get_bucket` and `dispatch` to use plan-specific rate:

```python
def _get_bucket(self, key: str, plan: str = "free") -> TokenBucket:
    bucket = self._buckets.get(key)
    plan_rpm = PLAN_RPM.get(plan, self.rpm)
    if plan_rpm == 0:  # unlimited
        plan_rpm = 999_999
    plan_rate = plan_rpm / 60.0
    plan_burst = max(10, plan_rpm // 6)  # ~10s burst
    if bucket is None or bucket.rate != plan_rate:
        bucket = TokenBucket(rate=plan_rate, burst=plan_burst)
        self._buckets[key] = bucket
    return bucket
```

Update `dispatch` to read plan from request state:

```python
plan = getattr(request.state, "plan", "free") if hasattr(request, "state") else "free"
bucket = self._get_bucket(key, plan)
# Update response header to show plan-specific limit
plan_rpm = PLAN_RPM.get(plan, self.rpm)
if plan_rpm == 0:
    plan_rpm = 999_999
# Use plan_rpm instead of self.rpm in headers
```

**Step 3: Run tests**

Run: `pytest tests/test_plan_ratelimit.py tests/test_ratelimit.py -v`
Expected: ALL PASS (new tests + existing tests)

**Step 4: Run full suite**

Run: `pytest tests/ -v`
Expected: 288+ passed (6 new tests added)

**Step 5: Commit**

```bash
git add llmhosts/proxy/ratelimit.py llmhosts/proxy/auth.py tests/test_plan_ratelimit.py
git commit -m "feat(proxy): plan-based rate limiting (free=60, pro=300, team=600 RPM)"
```

---

## Task 5: Stripe Checkout API — Tests

**Files:**
- Create: `frontend/src/app/api/stripe/checkout/route.test.ts` (or manual test doc)

Since the Stripe checkout is entirely in Next.js and requires Stripe SDK mocking, we document the manual test plan and implement directly.

**Step 1: Create checkout session API route**

Create `frontend/src/app/api/stripe/checkout/route.ts`:

```typescript
import { NextRequest, NextResponse } from "next/server";
import { auth } from "@/lib/auth";
import { stripe } from "@/lib/stripe";
import { db } from "@/lib/db";

export async function POST(req: NextRequest) {
  const session = await auth();
  if (!session?.user?.id) {
    return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
  }

  if (!stripe) {
    return NextResponse.json(
      { error: "Stripe is not configured" },
      { status: 503 }
    );
  }

  const { priceId } = await req.json();
  if (!priceId) {
    return NextResponse.json(
      { error: "priceId is required" },
      { status: 400 }
    );
  }

  // Get or create Stripe customer
  const user = await db.user.findUnique({
    where: { id: session.user.id },
    select: { stripeCustomerId: true, email: true },
  });

  let customerId = user?.stripeCustomerId;
  if (!customerId) {
    const customer = await stripe.customers.create({
      email: user?.email ?? undefined,
      metadata: { userId: session.user.id },
    });
    customerId = customer.id;
    await db.user.update({
      where: { id: session.user.id },
      data: { stripeCustomerId: customerId },
    });
  }

  const checkoutSession = await stripe.checkout.sessions.create({
    customer: customerId,
    mode: "subscription",
    line_items: [{ price: priceId, quantity: 1 }],
    success_url: `${req.nextUrl.origin}/dashboard/settings?billing=success`,
    cancel_url: `${req.nextUrl.origin}/dashboard/settings?billing=cancelled`,
    metadata: { userId: session.user.id },
  });

  return NextResponse.json({ url: checkoutSession.url });
}
```

**Step 2: Create billing portal API route**

Create `frontend/src/app/api/stripe/portal/route.ts`:

```typescript
import { NextRequest, NextResponse } from "next/server";
import { auth } from "@/lib/auth";
import { stripe } from "@/lib/stripe";
import { db } from "@/lib/db";

export async function POST(req: NextRequest) {
  const session = await auth();
  if (!session?.user?.id) {
    return NextResponse.json({ error: "Unauthorized" }, { status: 401 });
  }

  if (!stripe) {
    return NextResponse.json(
      { error: "Stripe is not configured" },
      { status: 503 }
    );
  }

  const user = await db.user.findUnique({
    where: { id: session.user.id },
    select: { stripeCustomerId: true },
  });

  if (!user?.stripeCustomerId) {
    return NextResponse.json(
      { error: "No billing account found" },
      { status: 404 }
    );
  }

  const portalSession = await stripe.billingPortal.sessions.create({
    customer: user.stripeCustomerId,
    return_url: `${req.nextUrl.origin}/dashboard/settings`,
  });

  return NextResponse.json({ url: portalSession.url });
}
```

**Step 3: Commit**

```bash
git add frontend/src/app/api/stripe/checkout/route.ts frontend/src/app/api/stripe/portal/route.ts
git commit -m "feat(stripe): checkout session + billing portal API routes"
```

---

## Task 6: Plans Library + Pricing Constants

**Files:**
- Create: `frontend/src/lib/plans.ts` (if not exists, or update)
- Create: `llmhosts/plans.py` (shared plan constants for Python side)

**Step 1: Create Python plan constants**

Create `llmhosts/plans.py`:

```python
"""Plan definitions shared across the LLMHosts proxy."""
from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class PlanLimits:
    rpm: int           # requests per minute (0 = unlimited)
    devices: int       # max registered devices (0 = unlimited)
    tunnel_public: bool
    hive_mode: bool
    rbac: bool
    cloud_dashboard: bool


PLANS: dict[str, PlanLimits] = {
    "free": PlanLimits(rpm=60, devices=1, tunnel_public=False, hive_mode=False, rbac=False, cloud_dashboard=False),
    "pro": PlanLimits(rpm=300, devices=5, tunnel_public=True, hive_mode=False, rbac=False, cloud_dashboard=True),
    "team": PlanLimits(rpm=600, devices=25, tunnel_public=True, hive_mode=True, rbac=True, cloud_dashboard=True),
    "enterprise": PlanLimits(rpm=0, devices=0, tunnel_public=True, hive_mode=True, rbac=True, cloud_dashboard=True),
}


def get_plan_limits(plan: str) -> PlanLimits:
    return PLANS.get(plan, PLANS["free"])
```

**Step 2: Write test for plans**

```python
# In tests/test_plans.py
from llmhost.plans import PLANS, get_plan_limits


def test_free_plan_defaults():
    p = get_plan_limits("free")
    assert p.rpm == 60
    assert p.devices == 1
    assert not p.tunnel_public


def test_pro_plan():
    p = get_plan_limits("pro")
    assert p.rpm == 300
    assert p.devices == 5
    assert p.tunnel_public


def test_unknown_plan_falls_back_to_free():
    p = get_plan_limits("nonexistent")
    assert p.rpm == 60
```

**Step 3: Run and commit**

Run: `pytest tests/test_plans.py -v`
Expected: PASS

```bash
git add llmhosts/plans.py tests/test_plans.py
git commit -m "feat: plan definitions with limits (free/pro/team/enterprise)"
```

---

## Task 7: Wire Plan Limits Into Rate Limiter

**Files:**
- Modify: `llmhosts/proxy/ratelimit.py` (use plans.py instead of hardcoded dict)

**Step 1: Replace hardcoded PLAN_RPM with plans.py**

```python
from llmhost.plans import get_plan_limits

# Remove PLAN_RPM dict, use:
plan_limits = get_plan_limits(plan)
plan_rpm = plan_limits.rpm if plan_limits.rpm > 0 else 999_999
```

**Step 2: Run all tests**

Run: `pytest tests/ -v`
Expected: ALL PASS

**Step 3: Commit**

```bash
git add llmhosts/proxy/ratelimit.py
git commit -m "refactor: wire plan limits from plans.py into rate limiter"
```

---

## Task 8: Integration Test — Full Auth + Plan Flow

**Files:**
- Create: `tests/test_auth_plan_flow.py`

**Step 1: Write integration test for the full flow**

```python
"""Integration test: create key with plan → authenticate → verify plan limits."""
from __future__ import annotations

import pytest
from httpx import ASGITransport, AsyncClient

from llmhost.config import LLMHostsConfig, SecurityConfig
from llmhost.proxy.app import create_app


@pytest.fixture
def configured_app():
    config = LLMHostsConfig()
    config.security = SecurityConfig(require_auth=True)
    return create_app(config)


@pytest.fixture
async def client(configured_app):
    transport = ASGITransport(app=configured_app)
    async with AsyncClient(transport=transport, base_url="http://test") as c:
        yield c


@pytest.mark.asyncio
async def test_full_flow_create_key_and_use(client: AsyncClient) -> None:
    # 1. Create a pro API key
    resp = await client.post("/api/keys", json={"label": "integration-test", "plan": "pro"})
    assert resp.status_code == 201
    raw_key = resp.json()["raw_key"]
    assert raw_key.startswith("lmh_")

    # 2. Use the key to list models
    resp = await client.get("/v1/models", headers={"Authorization": f"Bearer {raw_key}"})
    assert resp.status_code == 200
    assert resp.headers.get("X-RateLimit-Limit") == "300"  # pro plan

    # 3. List keys (should see our key)
    resp = await client.get("/api/keys")
    keys = resp.json()["keys"]
    assert any(k["label"] == "integration-test" and k["plan"] == "pro" for k in keys)

    # 4. Revoke the key
    key_id = next(k["id"] for k in keys if k["label"] == "integration-test")
    resp = await client.delete(f"/api/keys/{key_id}")
    assert resp.status_code == 200

    # 5. Key should no longer work
    resp = await client.get("/v1/models", headers={"Authorization": f"Bearer {raw_key}"})
    assert resp.status_code == 401


@pytest.mark.asyncio
async def test_anonymous_rejected_when_auth_required(client: AsyncClient) -> None:
    resp = await client.get("/v1/models")
    assert resp.status_code == 401


@pytest.mark.asyncio
async def test_health_always_works(client: AsyncClient) -> None:
    resp = await client.get("/health")
    assert resp.status_code == 200
```

**Step 2: Run and commit**

Run: `pytest tests/test_auth_plan_flow.py -v`
Expected: ALL PASS

```bash
git add tests/test_auth_plan_flow.py
git commit -m "test: integration test for full auth + plan + revoke flow"
```

---

## Task 9: Final Verification + Push

**Step 1: Run full test suite with coverage**

Run: `pytest tests/ -v --tb=short`
Expected: 295+ passed, 0 failed

**Step 2: Run linter**

Run: `ruff check llmhost/ tests/`
Expected: No errors

**Step 3: Push all commits**

```bash
git push origin main
```

**Step 4: Sync develop**

```bash
git checkout develop && git merge --ff-only main && git push origin develop && git checkout main
```

**Step 5: Update session tracker**

Update `docs/tracker/2026-02-16_1130_session-016.md` with completed work.

---

## Summary

| Task | What | Tests Added | Status |
|------|------|-------------|--------|
| 1-2 | Key Management API (`/api/keys/*`) | 6 | TODO |
| 3-4 | Plan-based rate limiting | 3 | TODO |
| 5 | Stripe checkout + portal routes | 0 (Next.js) | TODO |
| 6 | Plan definitions (`plans.py`) | 3 | TODO |
| 7 | Wire plans into rate limiter | 0 (refactor) | TODO |
| 8 | Integration test (full flow) | 3 | TODO |
| 9 | Final verify + push | 0 | TODO |

**Total new tests: ~15**
**Expected final count: ~297 tests**
