Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion autogpt_platform/backend/backend/api/external/v1/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,15 @@
from backend.data import user as user_db
from backend.data.auth.base import APIAuthorizationInfo
from backend.data.block import BlockInput, CompletedBlockOutput
from backend.data.execution import ExecutionContext
from backend.executor.utils import (
add_graph_execution,
charge_for_direct_block_execution,
)
from backend.integrations.webhooks.graph_lifecycle_hooks import on_graph_activate
from backend.util.exceptions import InsufficientBalanceError
from backend.util.settings import Settings
from backend.util.timezone_utils import get_user_timezone_or_utc

from .integrations import integrations_router
from .tools import tools_router
Expand Down Expand Up @@ -103,6 +105,10 @@ async def execute_graph_block(
if obj.disabled:
raise HTTPException(status_code=403, detail=f"Block #{block_id} is disabled.")

user = await user_db.get_user_by_id(auth.user_id)
if not user:
raise HTTPException(status_code=404, detail="User not found.")

try:
await charge_for_direct_block_execution(
user_id=auth.user_id, block=obj, input_data=data, source="external"
Expand All @@ -112,8 +118,16 @@ async def execute_graph_block(
status_code=status.HTTP_402_PAYMENT_REQUIRED, detail=str(e)
) from e

# Direct block execution has no graph; build a minimal ExecutionContext
# carrying the caller's identity + timezone so blocks that depend on
# those (e.g. time blocks) get correct data.
execution_context = ExecutionContext(
user_id=auth.user_id,
user_timezone=get_user_timezone_or_utc(user.timezone),
)

output = defaultdict(list)
async for name, data in obj.execute(data):
async for name, data in obj.execute(data, execution_context=execution_context):
output[name].append(data)
return output

Expand Down
53 changes: 51 additions & 2 deletions autogpt_platform/backend/backend/api/external/v1/routes_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,24 +40,42 @@ def _stub_block(
block_id: str = "00000000-0000-0000-0000-000000000001",
name: str = "TestBlock",
disabled: bool = False,
capture: dict | None = None,
):
"""Build a minimal block stub for get_block(...) replacement.

Async-iterable execute() yields one (name, value) pair so the route's
`async for` loop can iterate without touching real block logic.
`async for` loop can iterate without touching real block logic. If
``capture`` is supplied, the stub stores the kwargs it was called with
so tests can assert on them (e.g. ``execution_context``).
"""
block = MagicMock()
block.id = block_id
block.name = name
block.disabled = disabled

async def _execute(_data):
async def _execute(_data, **kwargs):
if capture is not None:
capture.update(kwargs)
yield "result", "ok"

block.execute = _execute
return block


@pytest.fixture(autouse=True)
def stub_user_lookup(monkeypatch: pytest.MonkeyPatch, test_user_id: str):
"""The direct-block-execute route fetches the user to build an
ExecutionContext; stub it so block-route tests don't need a real DB."""
user = MagicMock()
user.id = test_user_id
user.timezone = "UTC"
monkeypatch.setattr(
"backend.api.external.v1.routes.user_db.get_user_by_id",
AsyncMock(return_value=user),
)


def test_zero_balance_returns_402_on_paid_block(monkeypatch: pytest.MonkeyPatch):
"""Zero-credit user calling a paid block must be rejected before execution."""
block = _stub_block(name="PaidBlock")
Expand Down Expand Up @@ -134,6 +152,37 @@ def test_free_block_runs_without_charging(monkeypatch: pytest.MonkeyPatch):
spend_mock.assert_not_awaited()


def test_execute_block_forwards_execution_context(
monkeypatch: pytest.MonkeyPatch, test_user_id: str
):
"""Regression for #12648: blocks that read ``execution_context`` (e.g.
time blocks) crashed because this route didn't forward one. The route
must build an ExecutionContext carrying the caller's user_id +
timezone and pass it through to ``Block.execute``."""
captured: dict = {}
block = _stub_block(name="FreeBlock", capture=captured)
monkeypatch.setattr("backend.blocks.get_block", lambda _: block)
monkeypatch.setattr(
"backend.executor.utils.block_usage_cost", lambda *_a, **_k: (0, {})
)

user = MagicMock()
user.id = test_user_id
user.timezone = "Europe/Amsterdam"
monkeypatch.setattr(
"backend.api.external.v1.routes.user_db.get_user_by_id",
AsyncMock(return_value=user),
)

response = client.post(f"/blocks/{block.id}/execute", json={})

assert response.status_code == 200, f"got {response.status_code}: {response.text}"
assert "execution_context" in captured
ctx = captured["execution_context"]
assert ctx.user_id == test_user_id
assert ctx.user_timezone == "Europe/Amsterdam"


def test_disabled_block_still_403(monkeypatch: pytest.MonkeyPatch):
"""Pre-existing behavior: disabled blocks return 403, not bypassed by new gates."""
block = _stub_block(name="DisabledBlock", disabled=True)
Expand Down
11 changes: 10 additions & 1 deletion autogpt_platform/backend/backend/api/features/v1.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@
sync_subscription_schedule_from_stripe,
sync_tier_from_checkout_session,
)
from backend.data.execution import ExecutionContext
from backend.data.execution_cost_summary import (
UserExecutionCostSummary,
get_user_cost_summary,
Expand Down Expand Up @@ -503,13 +504,21 @@ async def execute_graph_block(
except InsufficientBalanceError as e:
raise HTTPException(status_code=HTTP_402_PAYMENT_REQUIRED, detail=str(e)) from e

# Direct block execution has no graph; build a minimal ExecutionContext
# carrying the caller's identity + timezone so blocks that depend on
# those (e.g. time blocks) get correct data.
execution_context = ExecutionContext(
user_id=user_id,
user_timezone=get_user_timezone_or_utc(user.timezone),
)

start_time = time.time()
try:
output = defaultdict(list)
async for name, data in obj.execute(
data,
user_id=user_id,
# Note: graph_exec_id and graph_id are not available for direct block execution
execution_context=execution_context,
):
output[name].append(data)

Expand Down
46 changes: 46 additions & 0 deletions autogpt_platform/backend/backend/api/features/v1_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,52 @@ async def mock_execute(*args, **kwargs):
)


def test_execute_graph_block_forwards_execution_context(
mocker: pytest_mock.MockFixture,
test_user_id: str,
) -> None:
"""Regression for #12648: blocks that read execution_context (e.g. time
blocks) crashed because the direct-block-execute route didn't forward
one. The route must construct an ExecutionContext carrying the caller's
user_id + timezone and pass it through to ``Block.execute``."""
captured_kwargs: dict = {}

mock_block = Mock()
mock_block.disabled = False
mock_block.name = "TestBlock"

async def mock_execute(*args, **kwargs):
captured_kwargs.update(kwargs)
yield "output", {"data": "ok"}

mock_block.execute = mock_execute

mocker.patch(
"backend.api.features.v1.get_block",
return_value=mock_block,
)

mock_user = Mock()
mock_user.timezone = "America/New_York"
mocker.patch(
"backend.api.features.v1.get_user_by_id",
return_value=mock_user,
)

mocker.patch(
"backend.api.features.v1.execution_utils.block_usage_cost",
return_value=(0, {}),
)

response = client.post("/blocks/test-block/execute", json={"x": "y"})

assert response.status_code == 200
assert "execution_context" in captured_kwargs
ctx = captured_kwargs["execution_context"]
assert ctx.user_id == test_user_id
assert ctx.user_timezone == "America/New_York"


def test_execute_graph_block_charges_when_cost_positive(
mocker: pytest_mock.MockFixture,
) -> None:
Expand Down
41 changes: 24 additions & 17 deletions autogpt_platform/backend/backend/blocks/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -683,9 +683,17 @@ def get_info(self) -> BlockInfo:
uiType=self.block_type.value,
)

async def execute(self, input_data: BlockInput, **kwargs) -> BlockOutput:
async def execute(
self,
input_data: BlockInput,
*,
execution_context: "ExecutionContext",
**kwargs,
) -> BlockOutput:
try:
async for output_name, output_data in self._execute(input_data, **kwargs):
async for output_name, output_data in self._execute(
input_data, execution_context=execution_context, **kwargs
):
yield output_name, output_data
except Exception as ex:
if isinstance(ex, BlockError):
Expand Down Expand Up @@ -767,21 +775,19 @@ async def is_block_exec_need_review(
)
return False, reviewed_data

async def _execute(self, input_data: BlockInput, **kwargs) -> BlockOutput:
# Check for review requirement only if running within a graph execution context
# Direct block execution (e.g., from chat) skips the review process
has_graph_context = all(
key in kwargs
for key in (
"node_exec_id",
"graph_exec_id",
"graph_id",
"execution_context",
)
)
if has_graph_context:
async def _execute(
self,
input_data: BlockInput,
*,
execution_context: "ExecutionContext",
**kwargs,
) -> BlockOutput:
# Review is only meaningful inside a graph execution. Direct block
# execution (e.g. from the /blocks/{id}/execute API) has no graph
# context and skips the review path.
if execution_context.graph_exec_id is not None:
should_pause, input_data = await self.is_block_exec_need_review(
input_data, **kwargs
input_data, execution_context=execution_context, **kwargs
)
if should_pause:
return
Expand All @@ -791,7 +797,7 @@ async def _execute(self, input_data: BlockInput, **kwargs) -> BlockOutput:
# that would fail JSON schema required checks. We still validate the
# non-credential fields so blocks that execute for real during dry-run
# (e.g. AgentExecutorBlock) get proper input validation.
is_dry_run = getattr(kwargs.get("execution_context"), "dry_run", False)
is_dry_run = execution_context.dry_run
if is_dry_run:
# Credential fields may be absent (LLM-built agents often skip
# wiring them) or nullified earlier in the pipeline. Validate
Expand Down Expand Up @@ -874,6 +880,7 @@ async def _execute(self, input_data: BlockInput, **kwargs) -> BlockOutput:
# Use the validated input data
async for output_name, output_data in self.run(
self.input_schema(**{k: v for k, v in input_data.items() if v is not None}),
execution_context=execution_context,
**kwargs,
):
if output_name == "error":
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,11 @@ async def test_merge_pr_error_path():
"credentials": TEST_CREDENTIALS_INPUT,
}
with pytest.raises(BlockExecutionError, match="PR not mergeable"):
async for _ in block.execute(input_data, credentials=TEST_CREDENTIALS):
async for _ in block.execute(
input_data,
credentials=TEST_CREDENTIALS,
execution_context=ExecutionContext(),
):
pass


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
import pytest

from backend.blocks.google.sheets import GoogleSheetsReadBlock
from backend.data.execution import ExecutionContext
from backend.util.exceptions import BlockExecutionError


Expand All @@ -30,7 +31,7 @@ async def test_sheets_read_missing_credentials_yields_clean_error():
}

with pytest.raises(BlockExecutionError, match="Missing credentials"):
async for _ in block.execute(input_data):
async for _ in block.execute(input_data, execution_context=ExecutionContext()):
pass


Expand All @@ -44,7 +45,7 @@ async def test_sheets_read_no_spreadsheet_still_hits_credentials_guard():
input_data = {"range": "Sheet1!A1:B2"} # no spreadsheet, no credentials

with pytest.raises(BlockExecutionError, match="Missing credentials"):
async for _ in block.execute(input_data):
async for _ in block.execute(input_data, execution_context=ExecutionContext()):
pass


Expand Down Expand Up @@ -79,7 +80,7 @@ async def test_sheets_read_upstream_chained_value_skips_guard(mocker):
}

with pytest.raises(Exception) as exc_info:
async for _ in block.execute(input_data):
async for _ in block.execute(input_data, execution_context=ExecutionContext()):
pass

# The guard should skip (chained data present) and let us reach run(),
Expand Down Expand Up @@ -121,7 +122,7 @@ async def test_sheets_read_upstream_chained_with_explicit_none_cred_id_skips_gua
}

with pytest.raises(Exception) as exc_info:
async for _ in block.execute(input_data):
async for _ in block.execute(input_data, execution_context=ExecutionContext()):
pass

# The guard must not raise "Missing credentials" for this shape.
Expand Down
Loading
Loading