commandcode-usage/mock_server.py at main · MAXeaglet/commandcode-usage · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#!/usr/bin/env python3
"""Mock of the commandcode.ai /alpha endpoints for offline testing.

Serves different situations per path prefix:
  /ok            — full normal shape (whoami/credits/subscriptions/usage)
  /exhausted     — five-hour window exceeded
  /partial       — credits ok, subscriptions & usage fail (resilience path)
  /snake         — snake_case field variants
  /garbage       — HTML body (bad JSON path) on credits
  (any request without Authorization) — 401

Run: python3 mock_server.py [port]   (default 18090)
"""
import json
import sys
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

NOW_MS = int(time.time() * 1000)

WHOAMI = {
    "user": {"id": "u_123", "name": "Max", "userName": "maxeagle"},
    "org": {"id": "org_9"},
}

def credits(shape="ok"):
    if shape == "exhausted":
        five = {"used": 5.0, "cap": 5.0, "exceeded": True,
                "resetAt": NOW_MS + 42 * 60 * 1000}
        limited, exceeded = True, "fiveHour"
    elif shape == "snake":
        five = {"used": 1.2, "cap": 5.0, "exceeded": False,
                "reset_at": NOW_MS + 3600 * 1000}
        limited, exceeded = False, None
    else:
        five = {"used": 1.23, "cap": 5.0, "exceeded": False,
                "resetAt": NOW_MS + 3600 * 1000}
        limited, exceeded = False, None
    weekly = {"used": 12.5, "cap": 40.0, "exceeded": False,
              "resetAt": NOW_MS + 3 * 86400 * 1000}
    monthly, below = (0.5, True) if shape == "lowbal" else (18.5, False)
    return {
        "credits": {"monthlyCredits": monthly, "purchasedCredits": 0,
                    "freeCredits": 2.0, "planId": "individual-goat",
                    "belowThreshold": below, "creditThreshold": 1.0},
        "windowLimits": {"limited": limited, "exceeded": exceeded,
                         "fiveHour": five, "weekly": weekly},
    }


def subscriptions(shape="ok"):
    return {
        "success": True,
        "data": {
            "planId": "individual-goat",
            "status": "active",
            "currentPeriodStart": NOW_MS - 18 * 86400 * 1000,
            "currentPeriodEnd": NOW_MS + 12 * 86400 * 1000,
            "cancelAtPeriodEnd": shape == "cancel",
            "pendingPhase": None,
        },
    }


USAGE = {
    "totalCount": 1883, "totalCost": 6.0011675234,
    "averageCost": 0.0031870247070631963, "successRate": 100,
    "completedCount": 1883, "failedCount": 0,
    "totalTokensIn": 487_592_979, "totalTokensOut": 1_210_434,
    "totalTokens": 488_803_413,
    "totalCredits": 6.001167523399999, "totalFreeCredits": 0,
    "totalMonthlyCredits": 6.001167523399999, "totalPurchasedCredits": 0,
    "periodBasis": "billing-period",
}


# 密钥里带 shape 名即选中该形态（方便联调）：sk-exhausted-… / sk-snake-… 等
def shape_for_key(auth):
    key = auth or ""
    for s in ("exhausted", "partial", "snake", "garbage", "badkey", "cancel", "lowbal"):
        if s in key.lower():
            return s
    return "ok"


class Handler(BaseHTTPRequestHandler):
    # 必须 HTTP/1.1 + 连接复用：Worker 的 fetch 会复用 keep-alive 连接，
    # 默认的 HTTP/1.0（每次响应即断）会让第二次请求撞在已关闭的 socket 上而失败。
    protocol_version = "HTTP/1.1"

    def do_GET(self):
        if not self.headers.get("Authorization"):
            self._send(401, {"error": "unauthorized"})
            return
        shape = "badkey" if "badkey" in (self.headers.get("Authorization") or "").lower() \
            else shape_for_key(self.headers.get("Authorization"))
        path = self.path.split("?")[0]

        if path == "/alpha/whoami":
            if shape == "badkey":
                self._send(401, {"error": "unauthorized"})
                return
            self._send(200, WHOAMI)
        elif path == "/alpha/billing/credits":
            if shape == "garbage":
                body = b"<html>login page</html>"
                self.send_response(200)
                self.send_header("Content-Type", "text/html")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
                return
            self._send(200, credits(shape))
        elif path == "/alpha/billing/subscriptions":
            if shape == "partial":
                self._send(500, {"error": "internal"})
                return
            self._send(200, subscriptions(shape))
        elif path == "/alpha/usage/summary":
            if shape == "partial":
                self._send(500, {"error": "internal"})
                return
            self._send(200, USAGE)
        else:
            self._send(404, {"error": "not found"})

    def _send(self, code, obj):
        body = json.dumps(obj).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, *args):
        pass


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 18090
    print(f"mock listening on 127.0.0.1:{port} "
          f"(shape by key: sk-exhausted-… / sk-partial-… / sk-snake-… / sk-garbage-… / "
          f"sk-badkey-… / sk-cancel-… / sk-lowbal-… / 其他=正常)", flush=True)
    # 多线程：keep-alive 下空闲连接会占住单线程服务器
    ThreadingHTTPServer(("127.0.0.1", port), Handler).serve_forever()