From 5174c42e7c4f7bab5d40ae18ba76b8d46f42815f Mon Sep 17 00:00:00 2001 From: 9qeklajc Date: Sat, 28 Mar 2026 00:31:30 +0100 Subject: [PATCH] fix claude code and open code --- routstr/auth.py | 13 +- routstr/upstream/base.py | 1009 ++++++++++++++++++++------------------ 2 files changed, 523 insertions(+), 499 deletions(-) diff --git a/routstr/auth.py b/routstr/auth.py index f844f1e7..54dbe6a1 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -732,22 +732,20 @@ async def adjust_payment_for_tokens( update(ApiKey) .where(col(ApiKey.hashed_key) == billing_key.hashed_key) .values( - reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost, balance=col(ApiKey.balance) - cost.total_msats, total_spent=col(ApiKey.total_spent) + cost.total_msats, ) ) result = await session.exec(finalize_stmt) # type: ignore[call-overload] - # Also update total_spent and reserved_balance on the child key if it's different + # Also update total_spent and balance on the child key if it's different if billing_key.hashed_key != key.hashed_key: child_stmt = ( update(ApiKey) .where(col(ApiKey.hashed_key) == key.hashed_key) .values( total_spent=col(ApiKey.total_spent) + cost.total_msats, - reserved_balance=col(ApiKey.reserved_balance) - - deducted_max_cost, + balance=col(ApiKey.balance) - cost.total_msats, ) ) await session.exec(child_stmt) # type: ignore[call-overload] @@ -819,23 +817,20 @@ async def adjust_payment_for_tokens( update(ApiKey) .where(col(ApiKey.hashed_key) == billing_key.hashed_key) .values( - reserved_balance=col(ApiKey.reserved_balance) - - deducted_max_cost, balance=col(ApiKey.balance) - total_cost_msats, total_spent=col(ApiKey.total_spent) + total_cost_msats, ) ) await session.exec(finalize_stmt) # type: ignore[call-overload] - # Also update total_spent and reserved_balance on the child key if it's different + # Also update total_spent and balance on the child key if it's different if billing_key.hashed_key != key.hashed_key: child_stmt = ( update(ApiKey) .where(col(ApiKey.hashed_key) == key.hashed_key) .values( total_spent=col(ApiKey.total_spent) + total_cost_msats, - reserved_balance=col(ApiKey.reserved_balance) - - deducted_max_cost, + balance=col(ApiKey.balance) - total_cost_msats, ) ) await session.exec(child_stmt) # type: ignore[call-overload] diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index f11eb85e..5481f274 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -446,496 +446,6 @@ class BaseUpstreamProvider: error_type, message, mapped_status, request=request ) - async def handle_streaming_chat_completion( - self, response: httpx.Response, key: ApiKey, max_cost_for_model: int - ) -> StreamingResponse: - """Handle streaming chat completion responses with token usage tracking and cost adjustment. - - Args: - response: Streaming response from upstream - key: API key for the authenticated user - max_cost_for_model: Maximum cost deducted upfront for the model - - Returns: - StreamingResponse with cost data injected at the end - """ - logger.info( - "Processing streaming chat completion", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "key_balance": key.balance, - "response_status": response.status_code, - }, - ) - background_tasks = BackgroundTasks() - - async def stream_with_cost( - max_cost_for_model: int, - ) -> AsyncGenerator[bytes, None]: - stored_chunks: list[bytes] = [] - usage_finalized: bool = False - last_model_seen: str | None = None - - async def finalize_without_usage() -> bytes | None: - nonlocal usage_finalized - if usage_finalized: - return None - async with create_session() as new_session: - fresh_key = await new_session.get(key.__class__, key.hashed_key) - if not fresh_key: - logger.warning( - "Key not found when finalizing streaming payment", - extra={"key_hash": key.hashed_key[:8] + "..."}, - ) - usage_finalized = True - return None - try: - fallback: dict = { - "model": last_model_seen or "unknown", - "usage": None, - } - cost_data = await adjust_payment_for_tokens( - fresh_key, fallback, new_session, max_cost_for_model - ) - usage_finalized = True - logger.info( - "Finalized streaming payment without explicit usage", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "cost_data": cost_data, - "balance_after_adjustment": fresh_key.balance, - }, - ) - return f"data: {json.dumps({'cost': cost_data})}\n\n".encode() - except Exception as cost_error: - logger.error( - "Error finalizing payment without usage", - extra={ - "error": str(cost_error), - "error_type": type(cost_error).__name__, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - usage_finalized = True - return None - - try: - async for chunk in response.aiter_bytes(): - stored_chunks.append(chunk) - try: - for part in re.split(b"data: ", chunk): - if not part or part.strip() in (b"[DONE]", b""): - continue - try: - obj = json.loads(part) - if isinstance(obj, dict) and obj.get("model"): - last_model_seen = str(obj.get("model")) - except json.JSONDecodeError: - pass - except Exception: - pass - - yield chunk - - logger.debug( - "Streaming completed, analyzing usage data", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "chunks_count": len(stored_chunks), - }, - ) - - for i in range(len(stored_chunks) - 1, -1, -1): - chunk = stored_chunks[i] - if not chunk: - continue - try: - events = re.split(b"data: ", chunk) - for event_data in events: - if not event_data or event_data.strip() in (b"[DONE]", b""): - continue - try: - data = json.loads(event_data) - if isinstance(data, dict) and data.get("model"): - last_model_seen = str(data.get("model")) - if isinstance(data, dict) and isinstance( - data.get("usage"), dict - ): - async with create_session() as new_session: - fresh_key = await new_session.get( - key.__class__, key.hashed_key - ) - if fresh_key: - try: - cost_data = ( - await adjust_payment_for_tokens( - fresh_key, - data, - new_session, - max_cost_for_model, - ) - ) - usage_finalized = True - logger.info( - "Payment adjustment completed for streaming", - extra={ - "key_hash": key.hashed_key[:8] - + "...", - "cost_data": cost_data, - "model": last_model_seen, - "balance_after_adjustment": fresh_key.balance, - }, - ) - yield f"data: {json.dumps({'cost': cost_data})}\n\n".encode() - except Exception as cost_error: - logger.error( - "Error adjusting payment for streaming tokens", - extra={ - "error": str(cost_error), - "error_type": type( - cost_error - ).__name__, - "key_hash": key.hashed_key[:8] - + "...", - }, - ) - break - except json.JSONDecodeError: - continue - except Exception as e: - logger.error( - "Error processing streaming response chunk", - extra={ - "error": str(e), - "error_type": type(e).__name__, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - - if not usage_finalized: - maybe_cost_event = await finalize_without_usage() - if maybe_cost_event is not None: - yield maybe_cost_event - - except Exception as stream_error: - logger.warning( - "Streaming interrupted; finalizing without usage", - extra={ - "error": str(stream_error), - "error_type": type(stream_error).__name__, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - raise - finally: - if not usage_finalized: - await finalize_without_usage() - - # Remove inaccurate encoding headers from upstream response - response_headers = dict(response.headers) - response_headers.pop("content-encoding", None) - response_headers.pop("content-length", None) - - return StreamingResponse( - stream_with_cost(max_cost_for_model), - status_code=response.status_code, - headers=response_headers, - background=background_tasks, - ) - - async def handle_streaming_messages_completion( - self, response: httpx.Response, key: ApiKey, max_cost_for_model: int - ) -> StreamingResponse: - async def stream_with_cost( - max_cost_for_model: int, - ) -> AsyncGenerator[bytes, None]: - stored_chunks: list[bytes] = [] - usage_finalized: bool = False - last_model_seen: str | None = None - input_tokens: int = 0 - output_tokens: int = 0 - - async def finalize_without_usage() -> bytes | None: - nonlocal usage_finalized - if usage_finalized: - return None - async with create_session() as new_session: - fresh_key = await new_session.get(key.__class__, key.hashed_key) - if not fresh_key: - usage_finalized = True - return None - try: - fallback: dict = { - "model": last_model_seen or "unknown", - "usage": None, - } - cost_data = await adjust_payment_for_tokens( - fresh_key, fallback, new_session, max_cost_for_model - ) - usage_finalized = True - return f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() - except Exception: - usage_finalized = True - return None - - try: - async for chunk in response.aiter_bytes(): - stored_chunks.append(chunk) - try: - decoded_chunk = chunk.decode("utf-8", errors="ignore") - for line in decoded_chunk.split("\n"): - if line.startswith("data: "): - try: - data = json.loads(line[6:]) - if isinstance(data, dict): - msg = data.get("message", {}) - if msg and msg.get("model"): - last_model_seen = str(msg.get("model")) - - if usage := msg.get("usage"): - input_tokens += usage.get("input_tokens", 0) - output_tokens += usage.get( - "output_tokens", 0 - ) - - if usage := data.get("usage"): - input_tokens += usage.get("input_tokens", 0) - output_tokens += usage.get( - "output_tokens", 0 - ) - except json.JSONDecodeError: - pass - except Exception: - pass - - yield chunk - - usage_data = { - "input_tokens": input_tokens, - "output_tokens": output_tokens, - } - - if input_tokens > 0 or output_tokens > 0: - async with create_session() as new_session: - fresh_key = await new_session.get(key.__class__, key.hashed_key) - if fresh_key: - try: - combined_data = { - "model": last_model_seen or "unknown", - "usage": usage_data, - } - cost_data = await adjust_payment_for_tokens( - fresh_key, - combined_data, - new_session, - max_cost_for_model, - ) - usage_finalized = True - yield f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() - except Exception: - pass - - if not usage_finalized: - maybe_cost_event = await finalize_without_usage() - if maybe_cost_event is not None: - yield maybe_cost_event - - except Exception: - if not usage_finalized: - await finalize_without_usage() - raise - finally: - if not usage_finalized: - await finalize_without_usage() - - response_headers = dict(response.headers) - response_headers.pop("content-encoding", None) - response_headers.pop("content-length", None) - - return StreamingResponse( - stream_with_cost(max_cost_for_model), - status_code=response.status_code, - headers=response_headers, - ) - - async def handle_non_streaming_messages_completion( - self, - response: httpx.Response, - key: ApiKey, - session: AsyncSession, - deducted_max_cost: int, - path: str, - ) -> Response: - try: - content = await response.aread() - response_json = json.loads(content) - - if path.endswith("count_tokens") and "usage" not in response_json: - input_tokens = response_json.get("input_tokens", 0) - response_json["usage"] = {"input_tokens": input_tokens} - - cost_data = await adjust_payment_for_tokens( - key, response_json, session, deducted_max_cost - ) - response_json["cost"] = cost_data - - allowed_headers = { - "content-type", - "cache-control", - "date", - "vary", - "access-control-allow-origin", - "access-control-allow-methods", - "access-control-allow-headers", - "access-control-allow-credentials", - "access-control-expose-headers", - "access-control-max-age", - } - - response_headers = { - k: v - for k, v in response.headers.items() - if k.lower() in allowed_headers - } - - return Response( - content=json.dumps(response_json).encode(), - status_code=response.status_code, - headers=response_headers, - media_type="application/json", - ) - except Exception: - raise - - async def handle_non_streaming_chat_completion( - self, - response: httpx.Response, - key: ApiKey, - session: AsyncSession, - deducted_max_cost: int, - ) -> Response: - """Handle non-streaming chat completion responses with token usage tracking and cost adjustment. - - Args: - response: Response from upstream - key: API key for the authenticated user - session: Database session for updating balance - deducted_max_cost: Maximum cost deducted upfront - - Returns: - Response with cost data added to JSON body - """ - logger.info( - "Processing non-streaming chat completion", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "key_balance": key.balance, - "response_status": response.status_code, - }, - ) - - content: bytes | None = None - try: - content = await response.aread() - response_json = json.loads(content) - - logger.debug( - "Parsed response JSON", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "model": response_json.get("model", "unknown"), - "has_usage": "usage" in response_json, - }, - ) - - cost_data = await adjust_payment_for_tokens( - key, response_json, session, deducted_max_cost - ) - - await session.refresh(key) - remaining_balance_msats = key.balance - - # Merge cost into usage for OpenCode - if "usage" in response_json: - response_json["usage"]["cost"] = cost_data.get("total_usd", 0.0) - response_json["usage"]["cost_sats"] = ( - cost_data.get("total_msats", 0) // 1000 - ) - response_json["usage"]["remaining_balance_msats"] = ( - remaining_balance_msats - ) - - # Keep detailed cost - response_json["metadata"] = response_json.get("metadata", {}) - response_json["metadata"]["routstr"] = {"cost": cost_data} - response_json["metadata"]["routstr"]["cost"]["sats_cost"] = ( - cost_data.get("total_msats", 0) // 1000 - ) - response_json["metadata"]["routstr"]["cost"]["remaining_balance_msats"] = ( - remaining_balance_msats - ) - response_json["cost"] = cost_data - response_json["cost"]["sats_cost"] = cost_data.get("total_msats", 0) // 1000 - response_json["cost"]["remaining_balance_msats"] = remaining_balance_msats - - logger.info( - "Payment adjustment completed for non-streaming", - extra={ - "key_hash": key.hashed_key[:8] + "...", - "cost_data": cost_data, - "model": response_json.get("model", "unknown"), - "balance_after_adjustment": key.balance, - }, - ) - - allowed_headers = { - "content-type", - "cache-control", - "date", - "vary", - "access-control-allow-origin", - "access-control-allow-methods", - "access-control-allow-headers", - "access-control-allow-credentials", - "access-control-expose-headers", - "access-control-max-age", - } - - response_headers = { - k: v - for k, v in response.headers.items() - if k.lower() in allowed_headers - } - - return Response( - content=json.dumps(response_json).encode(), - status_code=response.status_code, - headers=response_headers, - media_type="application/json", - ) - except json.JSONDecodeError as e: - logger.error( - "Failed to parse JSON from upstream response", - extra={ - "error": str(e), - "key_hash": key.hashed_key[:8] + "...", - "content_preview": content[:200].decode(errors="ignore") - if content - else "empty", - }, - ) - raise - except Exception as e: - logger.error( - "Error processing non-streaming chat completion", - extra={ - "error": str(e), - "error_type": type(e).__name__, - "key_hash": key.hashed_key[:8] + "...", - }, - ) - raise - async def handle_streaming_responses_completion( self, response: httpx.Response, key: ApiKey, max_cost_for_model: int ) -> StreamingResponse: @@ -1253,6 +763,524 @@ class BaseUpstreamProvider: ) raise + async def handle_streaming_chat_completion( + self, response: httpx.Response, key: ApiKey, max_cost_for_model: int + ) -> StreamingResponse: + """Handle streaming chat completion responses with token usage tracking and cost adjustment. + + Args: + response: Streaming response from upstream + key: API key for the authenticated user + max_cost_for_model: Maximum cost deducted upfront for the model + + Returns: + StreamingResponse with cost data injected at the end + """ + logger.info( + "Processing streaming chat completion", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "key_balance": key.balance, + "response_status": response.status_code, + }, + ) + + async def stream_with_cost( + max_cost_for_model: int, + ) -> AsyncGenerator[bytes, None]: + usage_finalized: bool = False + last_model_seen: str | None = None + usage_chunk_data: dict | None = None + done_seen: bool = False + + async def finalize_db_only() -> None: + nonlocal usage_finalized + if usage_finalized: + return + async with create_session() as new_session: + fresh_key = await new_session.get(key.__class__, key.hashed_key) + if not fresh_key: + return + try: + await adjust_payment_for_tokens( + fresh_key, + {"model": last_model_seen or "unknown", "usage": None}, + new_session, + max_cost_for_model, + ) + usage_finalized = True + except Exception: + pass + + try: + async for chunk in response.aiter_bytes(): + # Split chunk into SSE events + parts = re.split(b"data: ", chunk) + for i, part in enumerate(parts): + if not part: + continue + + stripped_part = part.strip() + if not stripped_part: + continue + + if stripped_part == b"[DONE]": + done_seen = True + continue + + try: + obj = json.loads(part) + if isinstance(obj, dict): + if obj.get("model"): + last_model_seen = str(obj.get("model")) + + if isinstance(obj.get("usage"), dict): + # Hold this chunk back to merge cost later + usage_chunk_data = obj + continue + except json.JSONDecodeError: + pass + + prefix = ( + b"data: " if (i > 0 or chunk.startswith(b"data: ")) else b"" + ) + yield prefix + part + + # Stream finished, process usage if found + if usage_chunk_data: + async with create_session() as session: + fresh_key = await session.get(key.__class__, key.hashed_key) + if fresh_key: + try: + cost_data = await adjust_payment_for_tokens( + fresh_key, + usage_chunk_data, + session, + max_cost_for_model, + ) + # Merge cost into usage + usage_chunk_data["usage"]["cost"] = cost_data.get( + "total_usd", 0.0 + ) + # Keep detailed cost in metadata + usage_chunk_data["metadata"] = usage_chunk_data.get( + "metadata", {} + ) + usage_chunk_data["metadata"]["routstr"] = { + "cost": cost_data + } + yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode() + usage_finalized = True + + await session.refresh(key) + except Exception: + # Fallback: yield original usage chunk if adjustment fails + yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode() + + if not usage_finalized: + await finalize_db_only() + + if done_seen: + yield b"data: [DONE]\n\n" + + except Exception as stream_error: + logger.warning( + "Streaming interrupted; finalizing in background", + extra={ + "error": str(stream_error), + "key_hash": key.hashed_key[:8] + "...", + }, + ) + raise + finally: + if not usage_finalized: + await finalize_db_only() + + # Remove inaccurate encoding headers from upstream response + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return StreamingResponse( + stream_with_cost(max_cost_for_model), + status_code=response.status_code, + headers=response_headers, + ) + + async def handle_streaming_messages_completion( + self, response: httpx.Response, key: ApiKey, max_cost_for_model: int + ) -> StreamingResponse: + background_tasks = BackgroundTasks() + + async def stream_with_cost( + max_cost_for_model: int, + ) -> AsyncGenerator[bytes, None]: + stored_chunks: list[bytes] = [] + usage_finalized: bool = False + last_model_seen: str | None = None + input_tokens: int = 0 + output_tokens: int = 0 + + async def finalize_without_usage() -> bytes | None: + nonlocal usage_finalized + if usage_finalized: + return None + async with create_session() as new_session: + fresh_key = await new_session.get(key.__class__, key.hashed_key) + if not fresh_key: + usage_finalized = True + return None + try: + fallback: dict = { + "model": last_model_seen or "unknown", + "usage": None, + } + cost_data = await adjust_payment_for_tokens( + fresh_key, fallback, new_session, max_cost_for_model + ) + usage_finalized = True + await new_session.refresh(key) + return f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() + except Exception: + usage_finalized = True + return None + + try: + async for chunk in response.aiter_bytes(): + stored_chunks.append(chunk) + try: + decoded_chunk = chunk.decode("utf-8", errors="ignore") + for line in decoded_chunk.split("\n"): + if line.startswith("data: "): + try: + data = json.loads(line[6:]) + if isinstance(data, dict): + # Extract model + model = data.get("model") + if not model and isinstance( + data.get("message"), dict + ): + model = data["message"].get("model") + + if model: + # Strip trailing date if present (e.g., -20251222) + model_str = str(model) + if "-" in model_str: + parts = model_str.rsplit("-", 1) + if ( + len(parts) == 2 + and parts[1].isdigit() + ): + model_str = parts[0] + last_model_seen = model_str + + print("aaaaaaaaaa ", last_model_seen) + + # Extract usage + msg = data.get("message", {}) + if isinstance(msg, dict): + if usage := msg.get("usage"): + input_tokens += usage.get( + "input_tokens", 0 + ) + output_tokens += usage.get( + "output_tokens", 0 + ) + + if usage := data.get("usage"): + input_tokens += usage.get("input_tokens", 0) + output_tokens += usage.get( + "output_tokens", 0 + ) + + # if data.get("model"): + # last_model_seen = str(data.get("model")) + + msg = data.get("message", {}) + if isinstance(msg, dict): + # if msg.get("model"): + # last_model_seen = str(msg.get("model")) + if usage := msg.get("usage"): + input_tokens += usage.get( + "input_tokens", 0 + ) + output_tokens += usage.get( + "output_tokens", 0 + ) + + if usage := data.get("usage"): + input_tokens += usage.get("input_tokens", 0) + output_tokens += usage.get( + "output_tokens", 0 + ) + + # if data.get("model"): + # last_model_seen = str(data.get("model")) + + msg = data.get("message", {}) + if isinstance(msg, dict): + # if msg.get("model"): + # last_model_seen = str(msg.get("model")) + if usage := msg.get("usage"): + input_tokens += usage.get( + "input_tokens", 0 + ) + output_tokens += usage.get( + "output_tokens", 0 + ) + + if usage := data.get("usage"): + input_tokens += usage.get("input_tokens", 0) + output_tokens += usage.get( + "output_tokens", 0 + ) + except json.JSONDecodeError: + pass + except Exception: + pass + + yield chunk + + usage_data = { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + } + + if input_tokens > 0 or output_tokens > 0: + async with create_session() as new_session: + fresh_key = await new_session.get(key.__class__, key.hashed_key) + if fresh_key: + try: + combined_data = { + "model": last_model_seen or "unknown", + "usage": usage_data, + } + cost_data = await adjust_payment_for_tokens( + fresh_key, + combined_data, + new_session, + max_cost_for_model, + ) + usage_finalized = True + await new_session.refresh(key) + yield f"event: cost\ndata: {json.dumps({'cost': cost_data})}\n\n".encode() + except Exception: + pass + + if not usage_finalized: + maybe_cost_event = await finalize_without_usage() + if maybe_cost_event is not None: + yield maybe_cost_event + + except Exception: + if not usage_finalized: + await finalize_without_usage() + raise + finally: + if not usage_finalized: + await finalize_without_usage() + + response_headers = dict(response.headers) + response_headers.pop("content-encoding", None) + response_headers.pop("content-length", None) + + return StreamingResponse( + stream_with_cost(max_cost_for_model), + status_code=response.status_code, + headers=response_headers, + background=background_tasks, + ) + + async def handle_non_streaming_messages_completion( + self, + response: httpx.Response, + key: ApiKey, + session: AsyncSession, + deducted_max_cost: int, + path: str, + ) -> Response: + try: + content = await response.aread() + response_json = json.loads(content) + + if path.endswith("count_tokens") and "usage" not in response_json: + input_tokens = response_json.get("input_tokens", 0) + response_json["usage"] = {"input_tokens": input_tokens} + + cost_data = await adjust_payment_for_tokens( + key, response_json, session, deducted_max_cost + ) + response_json["cost"] = cost_data + + allowed_headers = { + "content-type", + "cache-control", + "date", + "vary", + "access-control-allow-origin", + "access-control-allow-methods", + "access-control-allow-headers", + "access-control-allow-credentials", + "access-control-expose-headers", + "access-control-max-age", + } + + response_headers = { + k: v + for k, v in response.headers.items() + if k.lower() in allowed_headers + } + + return Response( + content=json.dumps(response_json).encode(), + status_code=response.status_code, + headers=response_headers, + media_type="application/json", + ) + except Exception: + raise + + async def handle_non_streaming_chat_completion( + self, + response: httpx.Response, + key: ApiKey, + session: AsyncSession, + deducted_max_cost: int, + ) -> Response: + """Handle non-streaming chat completion responses with token usage tracking and cost adjustment. + + Args: + response: Response from upstream + key: API key for the authenticated user + session: Database session for updating balance + deducted_max_cost: Maximum cost deducted upfront + + Returns: + Response with cost data added to JSON body + """ + logger.info( + "Processing non-streaming chat completion", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "key_balance": key.balance, + "response_status": response.status_code, + }, + ) + + content: bytes | None = None + try: + content = await response.aread() + response_json = json.loads(content) + + logger.debug( + "Parsed response JSON", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "model": response_json.get("model", "unknown"), + "has_usage": "usage" in response_json, + }, + ) + + cost_data = await adjust_payment_for_tokens( + key, response_json, session, deducted_max_cost + ) + + # Filter cost data to exclude internal-only fields + clean_cost_data = { + k: v + for k, v in cost_data.items() + if k not in ["input_tokens", "output_tokens"] + } + + await session.refresh(key) + remaining_balance_msats = key.balance + + # Merge cost into usage for OpenCode + if "choices" in response_json: + if "usage" in response_json: + response_json["usage"]["cost"] = cost_data.get("total_usd", 0.0) + response_json["usage"]["cost_sats"] = ( + cost_data.get("total_msats", 0) // 1000 + ) + response_json["usage"]["remaining_balance_msats"] = ( + remaining_balance_msats + ) + + # Keep detailed cost + response_json["metadata"] = response_json.get("metadata", {}) + response_json["metadata"]["routstr"] = {"cost": clean_cost_data} + response_json["metadata"]["routstr"]["cost"]["sats_cost"] = ( + cost_data.get("total_msats", 0) // 1000 + ) + response_json["metadata"]["routstr"]["cost"][ + "remaining_balance_msats" + ] = remaining_balance_msats + response_json["cost"] = clean_cost_data + response_json["cost"]["sats_cost"] = ( + cost_data.get("total_msats", 0) // 1000 + ) + response_json["cost"]["remaining_balance_msats"] = ( + remaining_balance_msats + ) + + logger.info( + "Payment adjustment completed for non-streaming Responses API", + extra={ + "key_hash": key.hashed_key[:8] + "...", + "cost_data": cost_data, + "model": response_json.get("model", "unknown"), + "balance_after_adjustment": key.balance, + }, + ) + + allowed_headers = { + "content-type", + "cache-control", + "date", + "vary", + "access-control-allow-origin", + "access-control-allow-methods", + "access-control-allow-headers", + "access-control-allow-credentials", + "access-control-expose-headers", + "access-control-max-age", + } + + response_headers = { + k: v + for k, v in response.headers.items() + if k.lower() in allowed_headers + } + + return Response( + content=json.dumps(response_json).encode(), + status_code=response.status_code, + headers=response_headers, + media_type="application/json", + ) + except json.JSONDecodeError as e: + logger.error( + "Failed to parse JSON from upstream Responses API response", + extra={ + "error": str(e), + "key_hash": key.hashed_key[:8] + "...", + "content_preview": content[:200].decode(errors="ignore") + if content + else "empty", + }, + ) + raise + except Exception as e: + logger.error( + "Error processing non-streaming Responses API completion", + extra={ + "error": str(e), + "error_type": type(e).__name__, + "key_hash": key.hashed_key[:8] + "...", + }, + ) + raise + async def _finalize_generic_streaming_payment( self, key_hash: str, max_cost: int, path: str ) -> None: @@ -1281,6 +1309,7 @@ class BaseUpstreamProvider: "key_hash": key_hash[:8] + "...", }, ) + await session.refresh(key) except Exception as e: logger.error( "Error finalizing generic streaming payment in background",