diff --git a/CHANGELOG.md b/CHANGELOG.md index 67dd305..29d650d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,20 @@ All notable user-visible changes. One bullet per change, each citing its PR or issue — no citation, no bullet. Dates are the day the change landed on main. +## [Unreleased] + +### Fixed + +- First message of a chat no longer sends `parent_message_id` (upstream began rejecting `0` with biz_code 26 "invalid message id" on 2026-10-31, and edges disagree on whether `null` is accepted). Follow-up turns keep the `parent = previous response id` chain. (PR #92) +- An upstream failure delivered as `application/json` on the SSE endpoint (as that biz_code 26 was) now names its code and message in the raised error and logs the response content-type with the raw non-SSE body, instead of surfacing as an anonymous empty stream. (PR #92) +- Upstream business errors (HTTP 200 with `code=40003 INVALID_TOKEN`, "messages too frequent", a killed web session) surface as a typed 401/429 and rotate the token instead of raising `TypeError` inside the route and pinning the dead account forever. The pre-turn create hops (`create_new_chat`, `create_challange_pow`) go through the same classifier, outside the per-signature lock so the retry cannot deadlock. (PR #92) +- Upstream hang-ups mid-stream log a warning and emit the terminal error frame instead of a traceback and a dead socket. (PR #92) +- A single upstream turn is bounded by `DEEPSEEKER_REQUEST_DEADLINE` (default 600 s); a slowly-dripping upstream no longer holds a request until the client gives up. (PR #92) + +### Added + +- `POST /v1/messages/count_tokens` (also at `/messages/count_tokens`) — Claude Code's context probes returned 404 before. (PR #92) + ## [6.0] — 2026-10-21 ### Fixed diff --git a/README.md b/README.md index 440d407..edaf5a5 100644 --- a/README.md +++ b/README.md @@ -210,6 +210,7 @@ Copy `.env.example` to `.env` and edit your secret values. The table below is th - `GET /v1/files/{file_id}/content` - **Anthropic Base**: `http://localhost:4000` - `POST /v1/messages` (also at `/messages`) + - `POST /v1/messages/count_tokens` (also at `/messages/count_tokens`) - `POST /v1/files/upload` - **Auth Key**: configured in `.env` (`DEEPSEEKER_API_KEY`) diff --git a/app.py b/app.py index afafc1d..975f4db 100644 --- a/app.py +++ b/app.py @@ -36,11 +36,6 @@ from functions import ( - effort_level, - reply_id, - request_effort, - request_tools, - token_busy, CookieGenerationError, StreamToolParser, UpstreamError, @@ -52,6 +47,7 @@ data_dir, delete_sessions_for_chat, delete_token, + effort_level, find_session, get_auth_token, get_file_content, @@ -69,9 +65,13 @@ pick_token, record_file, record_usage, + reply_id, + request_effort, + request_tools, restore_token, save_session, send_message, + token_busy, upload_file, usage_by_day_key, ) @@ -539,6 +539,45 @@ async def _wrapped(): return _wrapped() +# Hard ceiling on one upstream turn. The shared ClientSession uses +# ClientTimeout(total=None, connect=15, sock_read=600), and sock_read resets on +# every keepalive byte — a slowly-dripping upstream never trips it, so a single +# request could hang until the client gave up (the access log has 19.7-minute +# /v1/messages entries). DEADLINE is for the UPSTREAM work only: the SSE +# generators outlive handle_chat's return, so it is applied inside them and +# cancelled on exit, never around the route. +DEADLINE = max(1, int(os.getenv("DEEPSEEKER_REQUEST_DEADLINE", "600"))) + + +def _with_deadline(gen, what="upstream", seconds=None): + """Cancel `gen` and raise UpstreamError(504) once `seconds` elapse.""" + + async def _wrapped(): + limit = DEADLINE if seconds is None else seconds + try: + async with asyncio.timeout(limit): + async for chunk in gen: + yield chunk + except TimeoutError: + logger.warning("%s exceeded %ss deadline; cancelled", what, limit) + raise UpstreamError(504, f"{what} deadline exceeded ({limit}s)") from None + + return _wrapped() + + +def _is_truncated_stream_error(exc): + """True when upstream hung up mid-stream (aiohttp payload/transfer errors). + + These arrive after the SSE headers are already on the wire, so the client + sees a dead socket unless the generator emits a terminal error frame first. + """ + import aiohttp + + return isinstance( + exc, (aiohttp.ClientPayloadError, aiohttp.ClientError) + ) and not isinstance(exc, TimeoutError) + + def get_current_admin(request: Request): sid = request.cookies.get("session_id") if not sid or sid not in SESSIONS or time.time() - SESSIONS[sid] > SESSION_TTL: @@ -807,6 +846,92 @@ async def _preflight_stream(gen): _stream_retry = contextvars.ContextVar("stream_retry", default=None) +def _setup_retry_kwargs( + messages, + model, + thinking, + search, + stream, + tools, + is_anthropic, + req_model, + scope, + _attempt, + _auth_rotated, + _context_retried, + is_responses, + response_opts, +): + """The handle_chat call this attempt would make, as a dict — so a + setup-phase failure can hand it back to handle_chat with the attempt + budget advanced and the dead token excluded.""" + return dict( + messages=messages, + model=model, + thinking=thinking, + search=search, + stream=stream, + tools=tools, + is_anthropic=is_anthropic, + req_model=req_model, + scope=scope, + _attempt=_attempt, + _auth_rotated=_auth_rotated, + _context_retried=_context_retried, + is_responses=is_responses, + response_opts=response_opts, + ) + + +async def _setup_failure_outcome(e, token_id, retry_kwargs): + """Classify a failure that happened while PREPARING a turn, not sending it. + + Preparing means `create_new_chat` / `create_challange_pow` / `upload_file`: + the hops that run before `send_message`. Those calls live outside the try + that owns token rotation, so an auth failure there (upstream answers HTTP + 200 with `code=40003 INVALID_TOKEN`, which `_biz_data` turns into + `UpstreamBizError(401)`) escaped handle_chat entirely and surfaced as a + bare 500 — and, worse, the dead token was never marked, so it kept being + picked and every later turn failed the same way. + + Returns the retry_kwargs dict to re-enter handle_chat with, or the terminal + response when the budget is spent. + + DELIBERATELY does not call handle_chat itself (it did, and that was the + bug): the create-path caller holds `_sig_locks[sig]`, an asyncio.Lock is + not reentrant, so the retry blocked forever on the lock it already owned + and every request hung until the client gave up. Callers must invoke this, + then leave their lock, then retry. + """ + code = _upstream_http_code(e) + auth_class = code in (401, 403, 429) + if token_id is not None and auth_class: + await _db(mark_limited, token_id, code) + logger.warning( + "Upstream HTTP %s while preparing the turn (token #%s): %s", + code, + token_id, + _redact_secret_material(str(e))[:200], + ) + else: + logger.error( + "Upstream failure while preparing the turn (token #%s): %s", + token_id, + _redact_secret_material(str(e))[:300], + exc_info=e, + ) + attempt = retry_kwargs.get("_attempt", 0) + if attempt + 1 >= MAX_UPSTREAM_ATTEMPTS: + return _terminal_error_response(e, retry_kwargs.get("is_anthropic", False)) + retry_kwargs["_attempt"] = attempt + 1 + retry_kwargs["_auth_rotated"] = bool( + retry_kwargs.get("_auth_rotated") + ) or auth_class + if token_id is not None: + retry_kwargs["_exclude_token"] = token_id + return retry_kwargs + + async def handle_chat( messages, model, @@ -936,14 +1061,17 @@ async def handle_chat( for fid in file_ids: await _db(record_file, fid, new_token_id) reply_id.set(None) - gen = send_message( - new_session_id, - new_tok["token"], - prompt, - 0, - thinking, - search, - file_ids, + gen = _with_deadline( + send_message( + new_session_id, + new_tok["token"], + prompt, + 0, + thinking, + search, + file_ids, + ), + "upstream stream (rotation)", ) gen = await _preflight_stream(gen) except Exception as e: @@ -1103,6 +1231,7 @@ async def handle_chat( ) else: create_lock = _take_lock("sig", _sig_locks, sig, SIG_LOCKS_MAX) + setup_retry = None async with create_lock: sess = await _db(find_session, sig) if not sess: @@ -1147,35 +1276,74 @@ async def handle_chat( # that summary via build_prompt(rollover_summary=...). A single large # first exchange is untouched (the first-message path accepts ~1M # tokens); this only fires for accumulated session context. - if _context_retried or needs_rollover(messages, tools): - scratch_chat = await create_new_chat(tok["token"]) - summary_gen = send_message( - scratch_chat, - tok["token"], - build_summary_request_prompt(messages), - 0, - False, - False, - [], - ) - summary = strip_summary_tags(await collect_response(summary_gen)) - rollover_summary = summary[ - : MAX_SUMMARY_TOKENS * 4 - ] # ~4 chars/token cap - logger.info( - "Context rollover: handoff summary of ~%d tokens prepared in scratch chat %s", - count_tok(rollover_summary), - scratch_chat, - ) + # The upstream chat is created HERE, before the try that owns + # token rotation further down — so an auth failure on this hop + # (HTTP 200 + code=40003 INVALID_TOKEN) used to escape as a + # bare 500 with the dead token never marked. Route it through + # the same classifier the send path uses. + try: + if _context_retried or needs_rollover(messages, tools): + scratch_chat = await create_new_chat(tok["token"]) + summary_gen = send_message( + scratch_chat, + tok["token"], + build_summary_request_prompt(messages), + 0, + False, + False, + [], + ) + summary = strip_summary_tags( + await collect_response(summary_gen) + ) + rollover_summary = summary[ + : MAX_SUMMARY_TOKENS * 4 + ] # ~4 chars/token cap + logger.info( + "Context rollover: handoff summary of ~%d tokens prepared in scratch chat %s", + count_tok(rollover_summary), + scratch_chat, + ) - session_id = await create_new_chat(tok["token"]) - await _db(save_session, sig, token_id, session_id, 0) - parent_message_id = 0 + session_id = await create_new_chat(tok["token"]) + await _db(save_session, sig, token_id, session_id, 0) + parent_message_id = 0 + except Exception as e: + # Classify here, RETRY OUTSIDE the lock: the retry + # re-enters handle_chat, which takes this same + # non-reentrant asyncio.Lock — recursing in here hung + # the request instead of rotating (see + # _setup_failure_outcome). + setup_retry = await _setup_failure_outcome( + e, + token_id, + _setup_retry_kwargs( + messages, + model, + thinking, + search, + stream, + tools, + is_anthropic, + req_model, + scope, + _attempt, + _auth_rotated, + _context_retried, + is_responses, + response_opts, + ), + ) else: token_id = sess["token_id"] session_id = sess["session_id"] parent_message_id = sess["parent_message_id"] + if setup_retry is not None: + if isinstance(setup_retry, dict): + return await handle_chat(**setup_retry) + return setup_retry + tok = await _db(get_token, token_id) if not tok: return JSONResponse({"error": "Token expired"}, status_code=503) @@ -1289,14 +1457,17 @@ async def handle_chat( ) reply_id.set(None) - gen = send_message( - session_id, - tok["token"], - prompt, - parent_message_id, - thinking, - search, - file_ids, + gen = _with_deadline( + send_message( + session_id, + tok["token"], + prompt, + parent_message_id, + thinking, + search, + file_ids, + ), + "upstream stream", ) gen = await _preflight_stream(gen) if stream: @@ -1674,6 +1845,11 @@ async def stream_response( await _db(mark_limited, token_id, code) await _db(delete_sessions_for_chat, token_id, session_id) logger.warning("stream_response upstream HTTP %s: %s", code, e) + elif _is_truncated_stream_error(e): + # Upstream hung up mid-stream. Log a warning, not a traceback: + # nothing here is fixable and the client gets a clean error frame. + await _db(delete_sessions_for_chat, token_id, session_id) + logger.warning("stream_response upstream stream truncated: %s", e) else: await _db(delete_sessions_for_chat, token_id, session_id) logger.exception("stream_response failed") @@ -1861,6 +2037,9 @@ async def stream_anthropic_response( await _db(mark_limited, token_id, code) await _db(delete_sessions_for_chat, token_id, session_id) logger.warning("stream_anthropic_response upstream HTTP %s: %s", code, e) + elif _is_truncated_stream_error(e): + await _db(delete_sessions_for_chat, token_id, session_id) + logger.warning("stream_anthropic_response upstream stream truncated: %s", e) else: await _db(delete_sessions_for_chat, token_id, session_id) logger.exception("stream_anthropic_response failed") @@ -2584,6 +2763,9 @@ def _delta_msg(text): await _db(mark_limited, token_id, code) await _db(delete_sessions_for_chat, token_id, session_id) logger.warning("stream_responses_response upstream HTTP %s: %s", code, e) + elif _is_truncated_stream_error(e): + await _db(delete_sessions_for_chat, token_id, session_id) + logger.warning("stream_responses_response upstream stream truncated: %s", e) else: await _db(delete_sessions_for_chat, token_id, session_id) logger.exception("stream_responses_response failed") @@ -3512,6 +3694,33 @@ def convert_anthropic_messages(messages): return openai_msgs +@app.post("/v1/messages/count_tokens") +@app.post("/messages/count_tokens") +async def anthropic_count_tokens(request: Request): + """Anthropic-compatible token count. + + The proxy had no such route, so Claude Code's ~16 parallel count probes + per turn all 404'd and the client fell back to its own estimates. The + count is the DeepSeek tokenizer over the same text the chat path sends: + flattened system prompt plus message text. Tools are not counted — that + only ever costs a few hundred tokens of estimate error. + """ + if not check_key(request): + return JSONResponse({"error": "Invalid API key"}, status_code=401) + body = await _json_body(request) + text = _messages_text(body.get("messages") or []) + system = body.get("system") + if isinstance(system, list): + system = " ".join( + c.get("text", "") + for c in system + if isinstance(c, dict) and c.get("type") == "text" + ) + if system: + text = str(system) + "\n" + text + return JSONResponse({"input_tokens": count_tok(text)}) + + @app.post("/v1/messages") @app.post("/messages") async def anthropic_messages(request: Request): diff --git a/functions.py b/functions.py index ba5b081..829b6cf 100644 --- a/functions.py +++ b/functions.py @@ -333,6 +333,40 @@ def __init__(self, status, body=""): super().__init__(f"HTTP {self.status}: {self.body}") +class UpstreamBizError(UpstreamError): + """HTTP 200 whose envelope carries a business-level error. + + DeepSeek answers several endpoints with 200 + `{"code": N, "data": null}` + (or `data.biz_data: null`) when the account is flagged, its web session + was killed, or the per-account "messages too frequent" wall is hit. The + old code indexed straight into biz_data and died with + `'NoneType' object is not subscriptable`, which is not an UpstreamError — + so the token never rotated and the dead account stayed pinned forever. + + Subclassing UpstreamError means `_upstream_http_code()` already reads + `.status`, and handle_chat's rotation path (401/403/429) picks it up with + no changes there. + """ + + +def _biz_data(payload, what="response"): + """Return `payload["data"]["biz_data"]` or raise UpstreamBizError. + + Auth-shaped business messages map to 401 so the token is quarantined / + rotated; everything else maps to 429 so it cools down and rotates. Both + are statuses handle_chat already understands. + """ + data = payload.get("data") if isinstance(payload, dict) else None + biz = data.get("biz_data") if isinstance(data, dict) else None + if biz: + return biz + payload = payload if isinstance(payload, dict) else {} + code = payload.get("code") + msg = payload.get("msg") or payload.get("message") or "empty biz_data" + status = 401 if re.search(r"login|auth|token|session|expire", str(msg), re.I) else 429 + raise UpstreamBizError(status, f"{what}: code={code} msg={msg}") + + def _read_cookie_file(): """Return valid (unexpired) cookies, or None.""" try: @@ -1995,7 +2029,7 @@ async def create_challange_pow(target_path, auth_token): ) async with response: data = await response.json() - return data["data"]["biz_data"]["challenge"] + return _biz_data(data, "chat/create_pow_challenge")["challenge"] def write_string_pow(text, alloc_func, memory, store): @@ -2039,7 +2073,7 @@ async def create_new_chat(auth_token): ) async with response: data = await response.json() - return data["data"]["biz_data"]["chat_session"]["id"] + return _biz_data(data, "chat_session/create")["chat_session"]["id"] _EMPTY_SSE_DEFAULT = ( @@ -2136,7 +2170,28 @@ def _scan_recent_sse_errors(parsed_events): return None -def _raise_empty_sse_response(recent_lines, parsed_events): +def _non_sse_biz_error(non_sse_tail): + """Return (code, msg) if the non-SSE body tail holds a JSON error envelope. + + Upstream answers some protocol failures on the SSE endpoint with a plain + `application/json` envelope (`{"code":26,"msg":"invalid message id",...}`) + instead of an event stream. The parser sees zero `data:` lines and the + failure surfaces as an anonymous empty stream; this recovers the reason so + the error message says what actually went wrong. + """ + for line in reversed(non_sse_tail or []): + try: + obj = json.loads(line) + except Exception: + continue + if isinstance(obj, dict) and (obj.get("code") or obj.get("msg")): + return obj.get("code"), obj.get("msg") or obj.get("message") + return None, None + + +def _raise_empty_sse_response( + recent_lines, parsed_events, resp_ct="", non_sse_tail=None +): err = _scan_recent_sse_errors(parsed_events) if err: code, msg = err @@ -2149,10 +2204,18 @@ def _raise_empty_sse_response(recent_lines, parsed_events): msg = _EMPTY_SSE_CONTEXT else: msg = _EMPTY_SSE_DEFAULT + # Log the response type and the raw non-SSE body so the next upstream + # protocol change is identifiable from one request, not an hour of guessing. + tail = non_sse_tail or [] logger.warning( - "DeepSeek SSE ended without assistant output; last events: %s", + "DeepSeek SSE ended without assistant output; ct=%s non_sse=%s; last events: %s", + resp_ct or "-", + (" | ".join(tail[-3:])[:400] or "-"), recent_lines[-_SSE_RECENT_LINES:], ) + biz_code, biz_msg = _non_sse_biz_error(tail) + if biz_code is not None or biz_msg: + raise UpstreamError(502, f"{msg}: code={biz_code} msg={biz_msg}") raise UpstreamError(502, msg) @@ -2228,8 +2291,6 @@ async def _send_message_unlocked( ): # Backup: WAF cookies not required with Android headers. Kept as fallback: # cookie = await get_cookies() - if parent_message_id == 0: - parent_message_id = None file_ids = file_ids_ or [] headers = get_headers( @@ -2237,7 +2298,6 @@ async def _send_message_unlocked( ) json_data = { "chat_session_id": chat_id, - "parent_message_id": parent_message_id, "model_type": DEFAULT_MODEL_TYPE, "prompt": message, "ref_file_ids": file_ids, @@ -2246,6 +2306,12 @@ async def _send_message_unlocked( "preempt": False, "action": None, } + # Since 2026-10-31 upstream rejects parent_message_id=0 with biz_code 26 + # "invalid message id", and edges disagree about whether null is tolerated. + # Omitting the key on the first turn is the only shape that works on all + # of them; later turns still chain off the previous response id. + if parent_message_id: + json_data["parent_message_id"] = parent_message_id think_open = False got_output = False @@ -2273,6 +2339,12 @@ async def _send_message_unlocked( recent_lines = [] parsed_events = [] + # Non-SSE body lines (no `data: ` prefix). Upstream answers some + # protocol failures with a plain application/json envelope on this + # endpoint; keeping the raw lines is what makes those diagnosable. + non_sse_tail = [] + # getattr: response doubles used by tests may not expose headers. + resp_ct = (getattr(resp, "headers", None) or {}).get("content-type") or "" line_buf = b"" def _check_incomplete(): @@ -2319,7 +2391,9 @@ async def _emit_from_event(data): yield "\n\n\n" _check_incomplete() if not got_output: - _raise_empty_sse_response(recent_lines, parsed_events) + _raise_empty_sse_response( + recent_lines, parsed_events, resp_ct, non_sse_tail + ) raise _SseStreamFinished() if "v" in data and isinstance(data["v"], dict) and "response" in data["v"]: @@ -2371,6 +2445,9 @@ async def _emit_from_event(data): async def _consume_sse_line(decoded_line: str): if not decoded_line.startswith("data: "): + non_sse_tail.append(_redact_sse_line(decoded_line)) + if len(non_sse_tail) > _SSE_RECENT_LINES: + del non_sse_tail[0] return _remember_sse_line(recent_lines, decoded_line) try: @@ -2417,7 +2494,9 @@ async def _consume_sse_line(decoded_line: str): yield "\n\n\n" _check_incomplete() if not got_output: - _raise_empty_sse_response(recent_lines, parsed_events) + _raise_empty_sse_response( + recent_lines, parsed_events, resp_ct, non_sse_tail + ) async def upload_file(file_bytes, file_name, file_content_type, auth_token): @@ -2452,9 +2531,9 @@ async def upload_file(file_bytes, file_name, file_content_type, auth_token): ) async with response: resp_json = await response.json() - file_id = resp_json["data"]["biz_data"]["id"] + js_data = _biz_data(resp_json, "file/upload_file") + file_id = js_data["id"] yield ("uploaded", file_id) - js_data = resp_json["data"]["biz_data"] status = js_data["status"] headers = get_headers(auth_token) deadline = time.time() + 300 @@ -2467,7 +2546,7 @@ async def upload_file(file_bytes, file_name, file_content_type, auth_token): # cookies=cookie, # Backup WAF fallback timeout=aiohttp.ClientTimeout(total=30), ) as resp: - js_data = (await resp.json())["data"]["biz_data"]["files"][0] + js_data = _biz_data(await resp.json(), "file/fetch_files")["files"][0] status = js_data["status"] if status == "SUCCESS": tp_data = datetime.fromtimestamp(js_data["updated_at"], UTC) @@ -2496,7 +2575,7 @@ async def get_file_content(auth_token, file_id): timeout=aiohttp.ClientTimeout(total=30), ) as resp: resp_json = await resp.json() - js_data = resp_json["data"]["biz_data"]["files"][0] + js_data = _biz_data(resp_json, "file/fetch_files")["files"][0] yield mimetypes.guess_type(js_data["file_name"])[0] deadline = time.time() + 60 while ( @@ -2511,7 +2590,7 @@ async def get_file_content(auth_token, file_id): # cookies=cookie, # Backup WAF fallback timeout=aiohttp.ClientTimeout(total=30), ) as resp: - js_data = (await resp.json())["data"]["biz_data"]["files"][0] + js_data = _biz_data(await resp.json(), "file/fetch_files")["files"][0] if not js_data.get("signed_path"): return file_path = "https://files.deepseeksvc.com/api" + js_data["signed_path"] + "&ty=r" diff --git a/tests/test_incomplete_status.py b/tests/test_incomplete_status.py index 4f3d8c6..1bec7ee 100644 --- a/tests/test_incomplete_status.py +++ b/tests/test_incomplete_status.py @@ -4,9 +4,9 @@ import json import pytest +from test_sse_upstream import _collect_send_message import functions -from test_sse_upstream import _collect_send_message def _sse(*events): diff --git a/tests/test_parent_message_id.py b/tests/test_parent_message_id.py new file mode 100644 index 0000000..50553cd --- /dev/null +++ b/tests/test_parent_message_id.py @@ -0,0 +1,149 @@ +"""DeepSeek 2026-10-31 protocol change: the first message must not carry +`parent_message_id`. + +Upstream started rejecting `parent_message_id: 0` with biz_code 26 +"invalid message id", delivered as `application/json` on an endpoint the +client parses as SSE -- so the failure surfaced as a silent empty stream. +Different edges also disagree on whether `null` is tolerated, so the only +shape that works everywhere is omitting the key on the first turn. +Follow-up turns keep the existing parent = previous response id chain. + +See PR #92 and the incident report on it. +""" + +import asyncio +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import functions # noqa: E402 + + +class _FakeStream: + def __init__(self, chunks): + self._chunks = list(chunks) + + def iter_any(self): + async def _gen(): + for chunk in self._chunks: + yield chunk + + return _gen() + + +class _FakeResponse: + def __init__(self, body: bytes, status=200, content_type="text/event-stream"): + self.status = status + self.headers = {"content-type": content_type} + self._body = body + self.content = _FakeStream([body] if body else []) + + async def text(self): + return self._body.decode("utf-8", errors="replace") + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + +def _run_send(body: bytes, parent, content_type="text/event-stream"): + """Drive send_message with a fake upstream; return (payloads, output|error).""" + + captured = [] + + async def fake_pow(target_path, auth_token): + return "pow" + + async def fake_post(path, headers=None, json=None, **kwargs): + captured.append(dict(json or {})) + return _FakeResponse(body, content_type=content_type) + + async def run(): + orig_pow = functions.solve_create_pow + orig_post = functions.post_with_failover + functions.solve_create_pow = fake_pow + functions.post_with_failover = fake_post + try: + parts = [] + async for chunk in functions.send_message("chat-1", "tok", "hello", parent): + parts.append(chunk) + return "".join(parts) + finally: + functions.solve_create_pow = orig_pow + functions.post_with_failover = orig_post + + try: + return captured, asyncio.run(run()) + except Exception as e: # noqa: BLE001 - the test inspects the error + return captured, e + + +_OK_BODY = ( + b'data: {"p":"response/fragments","o":"APPEND","v":[{"type":"RESPONSE","content":"hi"}]}\n' + b'data: {"p":"response/status","v":"FINISHED"}\n' +) + + +def test_first_message_omits_parent_message_id_key(): + """parent=0 is the first turn; the key must be absent, not 0 and not null.""" + payloads, out = _run_send(_OK_BODY, 0) + assert out == "hi" + assert "parent_message_id" not in payloads[0] + + +def test_followup_message_keeps_parent_message_id(): + """Later turns still chain off the previous response id.""" + payloads, out = _run_send(_OK_BODY, 12) + assert out == "hi" + assert payloads[0]["parent_message_id"] == 12 + + +def test_non_sse_json_error_is_reported_not_silent(): + """A JSON biz error on the SSE endpoint names its code and message.""" + body = b'{"code":26,"msg":"invalid message id","data":null}' + payloads, err = _run_send(body, 0, content_type="application/json") + + assert isinstance(err, functions.UpstreamError), err + text = str(err) + assert "no parseable SSE content" in text + assert "code=26" in text + assert "invalid message id" in text + # The body never reached the parser as content. + assert "data:" not in text + + +def test_non_sse_error_does_not_rotate_tokens(): + """A protocol error is not a token error: keep the 502, do not cool the pool.""" + body = b'{"code":26,"msg":"invalid message id","data":null}' + _, err = _run_send(body, 0, content_type="application/json") + assert isinstance(err, functions.UpstreamError) + assert err.status == 502 + + +TESTS = [ + test_first_message_omits_parent_message_id_key, + test_followup_message_keeps_parent_message_id, + test_non_sse_json_error_is_reported_not_silent, + test_non_sse_error_does_not_rotate_tokens, +] + + +def main(): + failed = 0 + for t in TESTS: + try: + t() + print(f"PASS {t.__name__}") + except Exception as e: # noqa: BLE001 + failed += 1 + print(f"FAIL {t.__name__}: {e}") + if failed: + sys.exit(1) + print(f"{len(TESTS)} tests passed") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/tests/test_setup_retry.py b/tests/test_setup_retry.py new file mode 100644 index 0000000..8bff49f --- /dev/null +++ b/tests/test_setup_retry.py @@ -0,0 +1,131 @@ +"""Regression tests for auth failures on the turn-PREPARATION hops. + +`create_new_chat` runs before the try block that owns token rotation, so an +upstream 401 there (HTTP 200 + `code=40003 INVALID_TOKEN`, which +`_biz_data` turns into `UpstreamBizError(401)`) used to escape handle_chat +entirely: the client got a bare 500 and the dead token was never marked, so +it kept being picked and every later turn failed the same way. + +The retry must NOT recurse while the create path still holds `_sig_locks[sig]` +— asyncio.Lock is not reentrant, so that variant hung the request instead of +rotating it. + +Run: python -m pytest tests/test_setup_retry.py +""" + +import asyncio +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import app as app_module # noqa: E402 +import functions # noqa: E402 + + +def _run_create_path_scenario(): + """First create_new_chat raises 401; the turn must rotate and succeed.""" + calls = {"create": 0, "limited": [], "picked": [], "send": 0} + sessions = {} + + async def fake_sig(messages, model, scope=""): + return "sig-setup-retry" + + def fake_find(s): + return sessions.get(s) + + def fake_pick(exclude=None): + calls["picked"].append(set(exclude or ())) + return 2 if exclude and 1 in exclude else 1 + + def fake_get_token(tid): + return {"id": tid, "token": f"tok-{tid}", "status": "ACTIVE", "alias": f"t{tid}"} + + async def fake_create_chat(token): + calls["create"] += 1 + if calls["create"] == 1: + raise functions.UpstreamBizError(401, "chat_session/create: code=40003 msg=INVALID_TOKEN") + return f"chat-{token}" + + def fake_send(chat_id, auth_token, message, parent, thinking=False, search=False, file_ids_=None): + calls["send"] += 1 + + async def ok_gen(): + yield "recovered after setup rotation" + + return ok_gen() + + async def fake_files(messages, token, last_user_only=False): + return [] + + async def fake_prompt(messages, tools, model, is_first, rollover_summary=None): + return "prompt" + + def fake_save(sig, token_id, session_id, parent): + sessions[sig] = { + "token_id": token_id, + "session_id": session_id, + "parent_message_id": parent, + } + + def fake_limited(tid, code=None): + calls["limited"].append((tid, code)) + + patches = [ + ("get_auth_token", lambda: "tok"), + ("generate_signature", fake_sig), + ("find_session", fake_find), + ("pick_token", fake_pick), + ("get_token", fake_get_token), + ("send_message", fake_send), + ("create_new_chat", fake_create_chat), + ("extract_and_upload_files", fake_files), + ("build_prompt", fake_prompt), + ("mark_limited", fake_limited), + ("mark_active", lambda tid: None), + ("delete_sessions_for_chat", lambda *a: None), + ("save_session", fake_save), + ("parse_tools", lambda t: ([], t)), + ("format_response", lambda text, model, messages, tools=None: text), + ("needs_rollover", lambda messages, tools: False), + ] + + async def run(): + app_module._chat_locks.clear() + app_module._sig_locks.clear() + saved = [(name, getattr(app_module, name)) for name, _ in patches] + for name, fn in patches: + setattr(app_module, name, fn) + try: + # A hard timeout turns the old deadlock into a failure, not a hang. + return await asyncio.wait_for( + app_module.handle_chat( + [{"role": "user", "content": "hi"}], "test-model" + ), + timeout=10, + ) + finally: + for name, fn in saved: + setattr(app_module, name, fn) + + return asyncio.run(run()), calls + + +def test_create_path_401_rotates_instead_of_500(): + result, calls = _run_create_path_scenario() + + assert result == "recovered after setup rotation", result + # The dead token was marked (401 -> quarantine door) ... + assert calls["limited"] == [(1, 401)], calls["limited"] + # ... and the retry went to a different account, not the same one again. + assert 1 in calls["picked"][-1], calls["picked"] + assert calls["send"] == 1 + + +def test_create_path_401_does_not_deadlock_on_the_signature_lock(): + """The retry re-enters handle_chat, which takes _sig_locks[sig] again.""" + # If the create-path retry recursed while still holding the sig lock, the + # scenario's wait_for would time out; reaching past it proves the release. + _run_create_path_scenario() + for sig, lock in app_module._sig_locks.items(): + assert not lock.locked(), f"signature lock {sig!r} left held" \ No newline at end of file diff --git a/tests/test_upstream_hardening.py b/tests/test_upstream_hardening.py new file mode 100644 index 0000000..e95c821 --- /dev/null +++ b/tests/test_upstream_hardening.py @@ -0,0 +1,261 @@ +"""Tests for the upstream-hardening patches. + +Covers: + * /v1/messages/count_tokens (was 404 — every Claude Code count probe) + * functions._biz_data / UpstreamBizError (HTTP 200 + null biz_data used to + die with TypeError inside the route; the token never rotated) + * app._with_deadline (no upper bound on one upstream turn) + * app._is_truncated_stream_error (upstream hang-up mid-stream) + +Run: python tests/test_upstream_hardening.py (pytest-compatible) +""" + +import asyncio +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import aiohttp # noqa: E402 + +import app as app_module # noqa: E402 +from functions import UpstreamBizError, UpstreamError, _biz_data # noqa: E402 + +# -------------------------------------------------------------------------- +# count_tokens route +# -------------------------------------------------------------------------- + + +def _client(): + from fastapi.testclient import TestClient + + return TestClient(app_module.app) + + +def test_count_tokens_requires_api_key(): + r = _client().post("/v1/messages/count_tokens", json={"messages": []}) + assert r.status_code == 401, r.status_code + + +def test_count_tokens_counts_messages_and_string_system(): + r = _client().post( + "/v1/messages/count_tokens", + headers={"x-api-key": app_module.API_KEY}, + json={ + "model": "claude-sonnet-4", + "system": "you are terse", + "messages": [{"role": "user", "content": "hello world"}], + }, + ) + assert r.status_code == 200, r.text + body = r.json() + assert set(body) == {"input_tokens"}, body + assert body["input_tokens"] > 0 + + +def test_count_tokens_flattens_list_system_and_content_blocks(): + r = _client().post( + "/v1/messages/count_tokens", + headers={"x-api-key": app_module.API_KEY}, + json={ + "system": [{"type": "text", "text": "abc def"}], + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "ghi"}, + {"type": "image", "source": {}}, + ], + } + ], + }, + ) + assert r.status_code == 200, r.text + assert r.json()["input_tokens"] > 0 + + +def test_count_tokens_matches_tokenizer_on_flattened_text(): + body = { + "system": "sys", + "messages": [{"role": "user", "content": "one two three"}], + } + r = _client().post( + "/v1/messages/count_tokens", + headers={"authorization": "Bearer " + app_module.API_KEY}, + json=body, + ) + expected = app_module.count_tok("sys\none two three") + assert r.json()["input_tokens"] == expected + + +# -------------------------------------------------------------------------- +# _biz_data +# -------------------------------------------------------------------------- + + +def test_biz_data_returns_biz_data_when_present(): + payload = {"code": 0, "data": {"biz_data": {"chat_session": {"id": "s1"}}}} + assert _biz_data(payload) == {"chat_session": {"id": "s1"}} + + +def test_biz_data_null_data_raises_and_maps_auth_to_401(): + payload = {"code": 40003, "msg": "user login expired", "data": None} + try: + _biz_data(payload, "chat_session/create") + except UpstreamBizError as e: + assert e.status == 401, e.status + assert isinstance(e, UpstreamError) # rotation path understands it + assert "chat_session/create" in str(e) + else: + raise AssertionError("expected UpstreamBizError") + + +def test_biz_data_null_biz_data_maps_non_auth_to_429(): + payload = {"code": 42900, "msg": "messages too frequent", "data": {"biz_data": None}} + try: + _biz_data(payload, "chat/create_pow_challenge") + except UpstreamBizError as e: + assert e.status == 429, e.status + else: + raise AssertionError("expected UpstreamBizError") + + +def test_biz_data_non_dict_payload_is_an_error_not_a_crash(): + for payload in (None, [], "oops", 0): + try: + _biz_data(payload) + except UpstreamBizError: + pass + else: + raise AssertionError(f"expected UpstreamBizError for {payload!r}") + + +def test_upstream_biz_error_status_is_seen_by_the_rotation_path(): + e = UpstreamBizError(401, "x") + assert app_module._upstream_http_code(e) == 401 + + +# -------------------------------------------------------------------------- +# _with_deadline +# -------------------------------------------------------------------------- + + +def _run(coro): + return asyncio.run(coro) + + +async def _fast_gen(n=3): + for i in range(n): + yield f"c{i}" + + +async def _slow_gen(): + for _ in range(100): + await asyncio.sleep(0.5) + yield "late" + + +def test_with_deadline_passes_chunks_through_unchanged(): + async def go(): + out = [] + async for chunk in app_module._with_deadline(_fast_gen(), "t", 5): + out.append(chunk) + return out + + assert _run(go()) == ["c0", "c1", "c2"] + + +def test_with_deadline_raises_upstream_504_on_hang(): + async def go(): + async for _ in app_module._with_deadline(_slow_gen(), "t", 0.2): + pass + + try: + _run(go()) + except UpstreamError as e: + assert e.status == 504, e.status + else: + raise AssertionError("expected UpstreamError(504)") + + +def test_with_deadline_does_not_deadlock_on_early_consumer_abort(): + async def go(): + gen = app_module._with_deadline(_slow_gen(), "t", 5) + first = await gen.__anext__() + assert first == "late" + await asyncio.wait_for(gen.aclose(), timeout=2) + + _run(go()) + + +def test_with_deadline_cancels_upstream_generator_on_abort(): + closed = {"v": False} + + async def tracked(): + try: + while True: + await asyncio.sleep(0.5) + yield "x" + finally: + closed["v"] = True + + async def go(): + gen = app_module._with_deadline(tracked(), "t", 5) + await gen.__anext__() + await asyncio.wait_for(gen.aclose(), timeout=2) + + _run(go()) + assert closed["v"] is True + + +# -------------------------------------------------------------------------- +# _is_truncated_stream_error +# -------------------------------------------------------------------------- + + +def test_truncated_stream_error_detects_aiohttp_payload_error(): + assert app_module._is_truncated_stream_error(aiohttp.ClientPayloadError("x")) + assert app_module._is_truncated_stream_error(aiohttp.ClientError("x")) + + +def test_truncated_stream_error_excludes_timeouts_and_other_errors(): + assert not app_module._is_truncated_stream_error(TimeoutError()) + assert not app_module._is_truncated_stream_error(RuntimeError("boom")) + assert not app_module._is_truncated_stream_error(UpstreamError(429, "x")) + + +TESTS = [ + test_count_tokens_requires_api_key, + test_count_tokens_counts_messages_and_string_system, + test_count_tokens_flattens_list_system_and_content_blocks, + test_count_tokens_matches_tokenizer_on_flattened_text, + test_biz_data_returns_biz_data_when_present, + test_biz_data_null_data_raises_and_maps_auth_to_401, + test_biz_data_null_biz_data_maps_non_auth_to_429, + test_biz_data_non_dict_payload_is_an_error_not_a_crash, + test_upstream_biz_error_status_is_seen_by_the_rotation_path, + test_with_deadline_passes_chunks_through_unchanged, + test_with_deadline_raises_upstream_504_on_hang, + test_with_deadline_does_not_deadlock_on_early_consumer_abort, + test_with_deadline_cancels_upstream_generator_on_abort, + test_truncated_stream_error_detects_aiohttp_payload_error, + test_truncated_stream_error_excludes_timeouts_and_other_errors, +] + + +def main(): + failed = 0 + for t in TESTS: + try: + t() + print(f"PASS {t.__name__}") + except AssertionError as e: + failed += 1 + print(f"FAIL {t.__name__}: {e}") + if failed: + sys.exit(1) + print(f"{len(TESTS)} tests passed") + + +if __name__ == "__main__": + main() \ No newline at end of file