diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index aabb12ca..afb9475f 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -525,10 +525,17 @@ class BaseUpstreamProvider: session, max_cost_for_model, ) + remaining_balance_msats = fresh_key.balance # Merge cost into usage usage_chunk_data["usage"]["cost"] = cost_data.get( "total_usd", 0.0 ) + usage_chunk_data["usage"]["cost_sats"] = ( + cost_data.get("total_msats", 0) // 1000 + ) + usage_chunk_data["usage"]["remaining_balance_msats"] = ( + remaining_balance_msats + ) # Keep detailed cost in metadata usage_chunk_data["metadata"] = usage_chunk_data.get( "metadata", {} @@ -536,6 +543,12 @@ class BaseUpstreamProvider: usage_chunk_data["metadata"]["routstr"] = { "cost": cost_data } + usage_chunk_data["metadata"]["routstr"]["cost"][ + "sats_cost" + ] = cost_data.get("total_msats", 0) // 1000 + usage_chunk_data["metadata"]["routstr"]["cost"][ + "remaining_balance_msats" + ] = remaining_balance_msats yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode() usage_finalized = True except Exception as e: @@ -626,14 +639,31 @@ class BaseUpstreamProvider: 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", @@ -803,6 +833,7 @@ class BaseUpstreamProvider: session, max_cost_for_model, ) + remaining_balance_msats = fresh_key.balance # Merge cost into usage chunk if ( "response" in usage_chunk_data @@ -811,10 +842,22 @@ class BaseUpstreamProvider: usage_chunk_data["response"]["usage"]["cost"] = ( cost_data.get("total_usd", 0.0) ) + usage_chunk_data["response"]["usage"][ + "cost_sats" + ] = cost_data.get("total_msats", 0) // 1000 + usage_chunk_data["response"]["usage"][ + "remaining_balance_msats" + ] = remaining_balance_msats elif "usage" in usage_chunk_data: usage_chunk_data["usage"]["cost"] = cost_data.get( "total_usd", 0.0 ) + usage_chunk_data["usage"]["cost_sats"] = ( + cost_data.get("total_msats", 0) // 1000 + ) + usage_chunk_data["usage"][ + "remaining_balance_msats" + ] = remaining_balance_msats # Keep detailed cost in metadata usage_chunk_data["metadata"] = usage_chunk_data.get( @@ -823,6 +866,12 @@ class BaseUpstreamProvider: usage_chunk_data["metadata"]["routstr"] = { "cost": cost_data } + usage_chunk_data["metadata"]["routstr"]["cost"][ + "sats_cost" + ] = cost_data.get("total_msats", 0) // 1000 + usage_chunk_data["metadata"]["routstr"]["cost"][ + "remaining_balance_msats" + ] = remaining_balance_msats yield f"data: {json.dumps(usage_chunk_data)}\n\n".encode() usage_finalized = True except Exception: @@ -907,14 +956,31 @@ class BaseUpstreamProvider: 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 Responses API", diff --git a/routstr/upstream/gemini.py b/routstr/upstream/gemini.py index a008975a..8a416bfa 100644 --- a/routstr/upstream/gemini.py +++ b/routstr/upstream/gemini.py @@ -259,7 +259,15 @@ class GeminiUpstreamProvider(BaseUpstreamProvider): cost_data = await adjust_payment_for_tokens( key, openai_format_response, session, max_cost_for_model ) + await session.refresh(key) + remaining_balance_msats = key.balance openai_format_response["cost"] = cost_data + openai_format_response["cost"]["sats_cost"] = ( + cost_data.get("total_msats", 0) // 1000 + ) + openai_format_response["cost"]["remaining_balance_msats"] = ( + remaining_balance_msats + ) logger.info( "Gemini non-streaming payment completed",