fix(backend): Charge user credits before its block execution (#9427)

### Changes 🏗️

Instead of letting the user to execution the block then break it
post-execution.
We can charge the user first and execute it afterward.
The trade-offs:
* We can't charge a block that is charged based on the execution time.
* We will also charge failed block executions.

### Checklist 📋

#### For code changes:
- [ ] I have clearly listed my changes in the PR description
- [ ] I have made a test plan
- [ ] I have tested my changes according to the test plan:
  <!-- Put your test plan here: -->
  - [ ] ...

<details>
  <summary>Example test plan</summary>
  
  - [ ] Create from scratch and execute an agent with at least 3 blocks
- [ ] Import an agent from file upload, and confirm it executes
correctly
  - [ ] Upload agent to marketplace
- [ ] Import an agent from marketplace and confirm it executes correctly
  - [ ] Edit an agent from monitor, and confirm it executes correctly
</details>

#### For configuration changes:
- [ ] `.env.example` is updated or already compatible with my changes
- [ ] `docker-compose.yml` is updated or already compatible with my
changes
- [ ] I have included a list of my configuration changes in the PR
description (under **Changes**)

<details>
  <summary>Examples of configuration changes</summary>

  - Changing ports
  - Adding new services that need to communicate with each other
  - Secrets or environment variable changes
  - New or infrastructure changes such as databases
</details>
This commit is contained in:
Zamil Majdy
2025-02-05 16:39:41 +01:00
committed by GitHub
parent 22536de71f
commit 1d30e401fe
2 changed files with 8 additions and 11 deletions

View File

@@ -232,7 +232,7 @@ class UserCreditBase(ABC):
if amount < 0 and user_balance < abs(amount):
raise ValueError(
f"Insufficient balance for user {user_id}, balance: {user_balance}, amount: {amount}"
f"Insufficient balance of ${user_balance/100} to run the block that costs ${abs(amount)/100}"
)
# Create the transaction

View File

@@ -163,6 +163,7 @@ def execute_node(
# AgentExecutorBlock specially separate the node input_data & its input_default.
if isinstance(node_block, AgentExecutorBlock):
input_data = {**node.input_default, "data": input_data}
data.data = input_data
# Execute the node
input_data_str = json.dumps(input_data)
@@ -192,6 +193,11 @@ def execute_node(
output_size = 0
try:
# Charge the user for the execution before running the block.
# TODO: We assume the block is executed within 0 seconds.
# This is fine because for now, there is no block that is charged by time.
db_client.spend_credits(data, input_size + output_size, 0)
for output_name, output_data in node_block.execute(
input_data, **extra_exec_kwargs
):
@@ -210,16 +216,7 @@ def execute_node(
):
yield execution
# Update execution status and spend credits
res = update_execution(ExecutionStatus.COMPLETED)
s = input_size + output_size
t = (
(res.end_time - res.start_time).total_seconds()
if res.end_time and res.start_time
else 0
)
data.data = input_data
db_client.spend_credits(data, s, t)
update_execution(ExecutionStatus.COMPLETED)
except Exception as e:
error_msg = str(e)