feat(agent): record E2B sandbox runtime usage (#42554)

This commit is contained in:
zyssyz123 authored and GitHub committed 2026-09-21 07:11:12 +00:00
1 parent 1fa3245e83
commit af4f19840c
29 files changed
+3501 -5

No files matched your search

@@ -0,0 +1,219 @@
"""Authenticated HTTP contract backed by the real accounting service and DB."""
from collections.abc import Iterator
from typing import TypedDict
from uuid import uuid4
import pytest
import sqlalchemy as sa
from flask import Flask
from flask.testing import FlaskClient
from pydantic import JsonValue
from sqlalchemy.orm import Session, sessionmaker
from controllers.inner_api import bp as inner_api_bp
from controllers.inner_api import inner_api_ns
from models.agent_sandbox_usage import AgentSandboxUsageEvent
from tests.unit_tests.config_override import config_overrides_context
PROJECT = "431de237-596f-4d59-8a85-20a9846bf243"
class _EventPayload(TypedDict):
id: str
source: str
type: str
timestamp: str
payload: dict[str, JsonValue]
class _BatchPayload(TypedDict):
project_id: str
events: list[_EventPayload]
@pytest.fixture
def client() -> Iterator[FlaskClient]:
app = Flask(__name__)
app.config["TESTING"] = True
app.register_blueprint(inner_api_bp)
with config_overrides_context(
PLUGIN_DAEMON_KEY="test-daemon",
INNER_API_KEY_FOR_PLUGIN="test-inner",
AGENT_SANDBOX_METERING_ENABLED=True,
AGENT_SANDBOX_METERING_PROJECT_ID=PROJECT,
AGENT_SANDBOX_METERING_START_AT="2026-09-20T00:00:00Z",
):
yield app.test_client()
def payload() -> _BatchPayload:
event_id = str(uuid4())
return {
"project_id": PROJECT,
"events": [
{
"id": event_id,
"source": "provider",
"type": "sandbox.lifecycle.paused",
"timestamp": "2026-09-20T00:01:00Z",
"payload": {
"id": event_id,
"version": "v2",
"type": "sandbox.lifecycle.paused",
"timestamp": "2026-09-20T00:01:00Z",
"sandboxId": "sandbox-test",
"sandboxExecutionId": "execution-test",
"sandboxTeamId": PROJECT,
"eventData": {
"execution": {
"started_at": "2026-09-20T00:00:02Z",
"execution_time": 12345,
"memory_mb": 1024,
"vcpu_count": 2,
}
},
},
}
],
}
def test_state_requires_inner_auth(client: FlaskClient) -> None:
response = client.get(f"/inner/api/agent/sandbox-usage/state?project_id={PROJECT}")
assert response.status_code == 404
def test_state_exposes_persisted_scope_without_secret_or_environment(client: FlaskClient) -> None:
response = client.get(
f"/inner/api/agent/sandbox-usage/state?project_id={PROJECT}", headers={"X-Inner-Api-Key": "test-inner"}
)
assert response.status_code == 200
assert response.json == {
"enabled": True,
"project_id": PROJECT,
"started_at": "2026-09-20T00:00:00Z",
"checkpoint_at": None,
"full_scan_at": None,
"diagnostics": {
"unresolved_events": 0,
"conflict_events": 0,
"open_executions": 0,
"unattributed_executions": 0,
},
}
@pytest.mark.parametrize("query", ["", "?project_id=", f"?project_id={PROJECT}&unexpected=value"])
def test_state_rejects_invalid_query_before_activation(
client: FlaskClient, sqlite_session_factory: sessionmaker[Session], query: str
) -> None:
response = client.get(f"/inner/api/agent/sandbox-usage/state{query}", headers={"X-Inner-Api-Key": "test-inner"})
assert response.status_code == 400
body = response.get_json()
assert isinstance(body, dict)
assert body["code"] == "sandbox_usage_invalid_request"
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxUsageEvent)) == 0
def test_state_rejects_project_outside_allowlist_without_activation(
client: FlaskClient, sqlite_session_factory: sessionmaker[Session]
) -> None:
response = client.get(
"/inner/api/agent/sandbox-usage/state?project_id=other-project", headers={"X-Inner-Api-Key": "test-inner"}
)
assert response.status_code == 403
body = response.get_json()
assert isinstance(body, dict)
assert body["code"] == "sandbox_metering_project_not_allowed"
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxUsageEvent)) == 0
def test_state_reports_start_change_conflict_and_preserves_activation(client: FlaskClient) -> None:
url = f"/inner/api/agent/sandbox-usage/state?project_id={PROJECT}"
headers = {"X-Inner-Api-Key": "test-inner"}
original = client.get(url, headers=headers)
assert original.status_code == 200
with config_overrides_context(AGENT_SANDBOX_METERING_START_AT="2026-09-21T00:00:00Z"):
changed = client.get(url, headers=headers)
assert changed.status_code == 409
body = changed.get_json()
assert isinstance(body, dict)
assert body["code"] == "sandbox_metering_start_is_immutable"
restored = client.get(url, headers=headers)
assert restored.status_code == 200
assert restored.json == original.json
def test_state_reports_invalid_start_as_configuration_error(client: FlaskClient) -> None:
with config_overrides_context(AGENT_SANDBOX_METERING_START_AT="2026-09-20T00:00:00"):
response = client.get(
f"/inner/api/agent/sandbox-usage/state?project_id={PROJECT}", headers={"X-Inner-Api-Key": "test-inner"}
)
assert response.status_code == 503
body = response.get_json()
assert isinstance(body, dict)
assert body["code"] == "sandbox_metering_start_not_configured"
def test_events_ack_only_after_persisted_and_duplicate_is_idempotent(
client: FlaskClient, sqlite_session_factory: sessionmaker[Session]
) -> None:
body = payload()
for expected in (
{"accepted": 1, "duplicates": 0, "conflicts": 0, "ignored": 0},
{"accepted": 0, "duplicates": 1, "conflicts": 0, "ignored": 0},
):
response = client.post(
"/inner/api/agent/sandbox-usage/events", json=body, headers={"X-Inner-Api-Key": "test-inner"}
)
assert response.status_code == 200
assert response.json == expected
with sqlite_session_factory() as session:
assert (
session.scalar(
sa.select(sa.func.count())
.select_from(AgentSandboxUsageEvent)
.where(AgentSandboxUsageEvent.source_event_id == body["events"][0]["id"])
)
== 1
)
def test_ingestion_rejects_other_project_and_oversized_batch(client: FlaskClient) -> None:
body = payload()
body["project_id"] = "other-project"
response = client.post(
"/inner/api/agent/sandbox-usage/events", json=body, headers={"X-Inner-Api-Key": "test-inner"}
)
assert response.status_code == 403
body = payload()
body["events"] *= 101
response = client.post(
"/inner/api/agent/sandbox-usage/events", json=body, headers={"X-Inner-Api-Key": "test-inner"}
)
assert response.status_code == 400
def test_ingestion_payload_limit(client: FlaskClient) -> None:
body = payload()
body["events"][0]["payload"] = {"oversized": "x" * (1024 * 1024)}
response = client.post(
"/inner/api/agent/sandbox-usage/events", json=body, headers={"X-Inner-Api-Key": "test-inner"}
)
assert response.status_code == 413
def test_schema_and_http_exclude_business_operation_envelope_fields(client: FlaskClient) -> None:
schema = inner_api_ns.models["SandboxUsageEvent"].__schema__
assert set(schema["properties"]) == {"id", "source", "type", "timestamp", "sandbox_id", "execution_id", "payload"}
assert schema["additionalProperties"] is False
body = payload()
response = client.post(
"/inner/api/agent/sandbox-usage/events",
json={**body, "events": [{**body["events"][0], "allocation_id": str(uuid4())}]},
headers={"X-Inner-Api-Key": "test-inner"},
)
assert response.status_code == 400
@@ -162,6 +162,7 @@ class TestCelerySSLConfiguration:
# Mock all the scheduler configs
mock_config.CELERY_BEAT_SCHEDULER_TIME = 1
mock_config.AGENT_SANDBOX_METERING_ENABLED = False
mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = False
mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5
mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False
@@ -212,6 +213,7 @@ class TestCelerySSLConfiguration:
mock_config.CELERY_TASK_ANNOTATIONS = {}
mock_config.CELERY_BEAT_SCHEDULER_TIME = 1
mock_config.AGENT_SANDBOX_METERING_ENABLED = False
mock_config.ENABLE_CONVERSATION_CLEANUP_TASK = True
mock_config.CONVERSATION_CLEANUP_TASK_INTERVAL = 5
mock_config.ENABLE_CLEAN_EMBEDDING_CACHE_TASK = False
@@ -0,0 +1,179 @@
"""The scheduled collector reuses existing Celery workers and bounded HTTP."""
from collections.abc import Iterator
from datetime import timedelta
from unittest.mock import MagicMock
import httpx
import pytest
from celery.contrib.testing.app import setup_default_app
from configs.extra.agent_backend_config import AgentBackendConfig
from dify_app import DifyApp
from extensions import ext_celery
from schedule.collect_agent_sandbox_usage import collect_agent_sandbox_usage
from tests.unit_tests.config_override import config_overrides_context
PROJECT = "431de237-596f-4d59-8a85-20a9846bf243"
@pytest.fixture(autouse=True)
def restore_celery_app_state() -> Iterator[None]:
# init_app changes Celery's process-wide current/default app; close() alone
# leaves later shared tasks pointing at the now-unregistered test app.
with setup_default_app(collect_agent_sandbox_usage.app):
yield
@pytest.fixture(autouse=True)
def configured_collection() -> Iterator[None]:
with config_overrides_context(
AGENT_SANDBOX_METERING_ENABLED=True,
AGENT_SANDBOX_METERING_PROJECT_ID=PROJECT,
AGENT_BACKEND_BASE_URL="http://agent_backend:5001/",
AGENT_BACKEND_API_TOKEN="backend-token",
):
yield
def response(payload: object, status: int = 200) -> httpx.Response:
return httpx.Response(
status, json=payload, request=httpx.Request("POST", "http://agent_backend:5001/internal/e2b/usage/collect")
)
def test_task_uses_existing_queue_deadlines_and_single_authenticated_call(monkeypatch: pytest.MonkeyPatch) -> None:
post = MagicMock(return_value=response({"completed": True}))
monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post)
assert collect_agent_sandbox_usage() is True
post.assert_called_once_with(
"http://agent_backend:5001/internal/e2b/usage/collect",
headers={"Authorization": "Bearer backend-token"},
json={"project_id": PROJECT},
max_retries=0,
timeout=260,
)
dispatch = MagicMock()
monkeypatch.setattr(collect_agent_sandbox_usage.app, "send_task", dispatch)
collect_agent_sandbox_usage.apply_async()
assert dispatch.call_args.kwargs["queue"] == "ops_trace"
assert collect_agent_sandbox_usage.soft_time_limit == 270
assert collect_agent_sandbox_usage.time_limit == 300
def test_disabled_job_makes_no_http_request(monkeypatch: pytest.MonkeyPatch) -> None:
post = MagicMock()
monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post)
with config_overrides_context(AGENT_SANDBOX_METERING_ENABLED=False):
assert collect_agent_sandbox_usage() is False
post.assert_not_called()
@pytest.mark.parametrize(
"missing", ["AGENT_BACKEND_BASE_URL", "AGENT_BACKEND_API_TOKEN", "AGENT_SANDBOX_METERING_PROJECT_ID"]
)
def test_incomplete_configuration_only_fails_collection_job(missing: str, monkeypatch: pytest.MonkeyPatch) -> None:
post = MagicMock()
monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post)
with config_overrides_context(**{missing: ""}):
with pytest.raises(ValueError, match="requires backend URL, token, and project"):
collect_agent_sandbox_usage()
post.assert_not_called()
@pytest.mark.parametrize("payload", [{"completed": False}, {"completed": "true"}, {"unexpected": True}])
def test_incomplete_or_invalid_scan_is_a_job_failure_without_retry(
payload: dict[str, object], monkeypatch: pytest.MonkeyPatch
) -> None:
post = MagicMock(return_value=response(payload))
monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post)
with pytest.raises((ValueError, RuntimeError)):
collect_agent_sandbox_usage()
assert post.call_count == 1
@pytest.mark.parametrize("status", [403, 503, 504])
def test_backend_failure_does_not_trigger_http_retry(status: int, monkeypatch: pytest.MonkeyPatch) -> None:
post = MagicMock(return_value=response({"detail": "unavailable"}, status))
monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post)
with pytest.raises(httpx.HTTPStatusError):
collect_agent_sandbox_usage()
assert post.call_count == 1
def test_transport_failure_is_logged_without_credentials_or_raw_message(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
post = MagicMock(side_effect=httpx.ReadTimeout("do-not-log-secret"))
monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post)
with pytest.raises(httpx.ReadTimeout):
collect_agent_sandbox_usage()
assert post.call_count == 1
assert "Background sandbox usage collection failed" in caplog.text
assert "do-not-log-secret" not in caplog.text
assert "backend-token" not in caplog.text
@pytest.mark.parametrize("enabled", [False, True])
@pytest.mark.parametrize("interval", [90, "90"])
def test_beat_registers_only_enabled_collection_with_expiring_schedule(
enabled: bool, interval: int | str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(ext_celery, "setup_workflow_warm_shutdown_handler", lambda: None)
with config_overrides_context(
AGENT_SANDBOX_METERING_ENABLED=enabled,
AGENT_SANDBOX_METERING_INTERVAL_SECONDS=interval,
DISABLE_TELEMETRY=True,
):
celery = ext_celery.init_app(DifyApp(__name__))
try:
assert "schedule.collect_agent_sandbox_usage" in celery.conf.imports
schedule = celery.conf.beat_schedule.get("collect_agent_sandbox_usage")
if enabled:
assert schedule == {
"task": "schedule.collect_agent_sandbox_usage.collect_agent_sandbox_usage",
"schedule": timedelta(seconds=90),
"options": {"expires": 90},
}
else:
assert schedule is None
finally:
celery.close()
@pytest.mark.parametrize("raw_interval", ["0", "-1", "not-a-number", "1.5", "9" * 30])
def test_invalid_interval_env_only_disables_optional_schedule(
raw_interval: str, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
monkeypatch.setenv("AGENT_SANDBOX_METERING_INTERVAL_SECONDS", raw_interval)
config = AgentBackendConfig()
assert raw_interval == config.AGENT_SANDBOX_METERING_INTERVAL_SECONDS
monkeypatch.setattr(ext_celery, "setup_workflow_warm_shutdown_handler", lambda: None)
with config_overrides_context(
AGENT_SANDBOX_METERING_INTERVAL_SECONDS=config.AGENT_SANDBOX_METERING_INTERVAL_SECONDS,
ENABLE_CONVERSATION_CLEANUP_TASK=True,
CONVERSATION_CLEANUP_TASK_INTERVAL=2,
DISABLE_TELEMETRY=True,
):
app = DifyApp(__name__)
celery = ext_celery.init_app(app)
try:
assert app.extensions["celery"] is celery
assert "collect_agent_sandbox_usage" not in celery.conf.beat_schedule
assert "schedule.collect_agent_sandbox_usage" in celery.conf.imports
assert celery.conf.beat_schedule["conversation_cleanup_sweeper"] == {
"task": "tasks.delete_conversation_task.sweep_deleted_conversations",
"schedule": timedelta(minutes=2),
}
finally:
celery.close()
assert "Skipping sandbox usage schedule" in caplog.text
assert all(raw_interval not in record.getMessage() for record in caplog.records)
def test_manual_collection_does_not_depend_on_beat_interval(monkeypatch: pytest.MonkeyPatch) -> None:
post = MagicMock(return_value=response({"completed": True}))
monkeypatch.setattr("schedule.collect_agent_sandbox_usage.ssrf_proxy.post", post)
with config_overrides_context(AGENT_SANDBOX_METERING_INTERVAL_SECONDS="invalid"):
assert collect_agent_sandbox_usage() is True
post.assert_called_once()
@@ -0,0 +1,720 @@
"""Accounting invariants tested against real SQLAlchemy persistence."""
from collections.abc import Iterator
from copy import deepcopy
from datetime import datetime
from uuid import uuid4
import pytest
import sqlalchemy as sa
from pydantic import JsonValue
from sqlalchemy import event as sqlalchemy_event
from sqlalchemy.dialects import postgresql
from sqlalchemy.engine import Connection, Engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.sql import ClauseElement
from models.agent import (
Agent,
AgentConfigVersionKind,
AgentScope,
AgentSource,
AgentWorkspace,
AgentWorkspaceBinding,
AgentWorkspaceOwnerType,
)
from models.agent_sandbox_usage import AgentSandboxExecution, AgentSandboxUsageEvent
from models.model import App
from services.agent.runtime_usage_service import SandboxUsageError, SandboxUsageEvent, SandboxUsageService
from tests.unit_tests.config_override import config_overrides_context
PROJECT = "431de237-596f-4d59-8a85-20a9846bf243"
START = "2026-09-20T00:00:00Z"
@pytest.fixture(autouse=True)
def metering_config(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
monkeypatch.setattr("services.agent.runtime_usage_service.naive_utc_now", lambda: datetime(2026, 9, 21))
with config_overrides_context(
AGENT_SANDBOX_METERING_ENABLED=True,
AGENT_SANDBOX_METERING_PROJECT_ID=PROJECT,
AGENT_SANDBOX_METERING_START_AT=START,
):
yield
def provider_event(
*,
event_id: str = "event-1",
execution_id: str = "execution-1",
sandbox_id: str = "sandbox-1",
duration: int | None = 12345,
started_at: str | None = "2026-09-20T00:00:02Z",
kind: str = "paused",
metadata: dict[str, JsonValue] | None = None,
) -> SandboxUsageEvent:
return SandboxUsageEvent.model_validate(
{
"id": event_id,
"source": "provider",
"type": f"sandbox.lifecycle.{kind}",
"sandbox_id": sandbox_id,
"execution_id": execution_id,
"payload": {
"id": event_id,
"version": "v2",
"type": f"sandbox.lifecycle.{kind}",
"timestamp": "2026-09-20T00:00:14.551745815Z",
"sandboxId": sandbox_id,
"sandboxExecutionId": execution_id,
"sandboxTeamId": PROJECT,
"sandboxTemplateId": "template-1",
"sandboxBuildId": "build-1",
"eventData": {
"execution": {
"started_at": started_at,
"execution_time": duration,
"memory_mb": 1024,
"vcpu_count": 2,
},
"sandbox_metadata": metadata or {},
},
},
}
)
def ingest(*events: SandboxUsageEvent) -> dict[str, int]:
return SandboxUsageService.ingest(project_id=PROJECT, events=list(events))
def execution(factory: sessionmaker[Session]) -> AgentSandboxExecution:
with factory() as session:
return session.scalars(sa.select(AgentSandboxExecution)).one()
def test_duplicate_terminal_and_api_webhook_aliases_count_once(sqlite_session_factory: sessionmaker[Session]) -> None:
event = provider_event()
assert ingest(event)["accepted"] == 1
alias = deepcopy(event.payload)
for camel, snake in (
("sandboxId", "sandbox_id"),
("sandboxExecutionId", "sandbox_execution_id"),
("sandboxTeamId", "sandbox_team_id"),
("sandboxTemplateId", "sandbox_template_id"),
("sandboxBuildId", "sandbox_build_id"),
("eventData", "event_data"),
):
alias[snake] = alias.pop(camel)
assert ingest(event.model_copy(update={"payload": alias}))["duplicates"] == 1
assert ingest(provider_event(event_id="killed-event", kind="killed"))["accepted"] == 1
row = execution(sqlite_session_factory)
assert row.metered_duration_ms == 12345
assert row.quality == "metered"
assert row.started_at == datetime(2026, 9, 20, 0, 0, 2)
assert row.terminal_event_at == datetime(2026, 9, 20, 0, 0, 14, 551745)
def test_same_sandbox_resume_is_a_new_execution(sqlite_session_factory: sessionmaker[Session]) -> None:
ingest(provider_event(), provider_event(event_id="second", execution_id="execution-2", duration=6789))
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.sum(AgentSandboxExecution.metered_duration_ms))) == 19134
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxExecution)) == 2
def test_state_reports_pending_and_conflict_counts_without_owner_ids() -> None:
created = provider_event(event_id="created", kind="created", duration=None, started_at=None)
created.payload["eventData"] = dict[str, JsonValue]()
ingest(created)
assert SandboxUsageService.get_state(project_id=PROJECT)["diagnostics"] == {
"unresolved_events": 0,
"conflict_events": 0,
"open_executions": 1,
"unattributed_executions": 1,
}
ingest(provider_event())
ingest(provider_event(duration=1))
diagnostics = SandboxUsageService.get_state(project_id=PROJECT)["diagnostics"]
assert diagnostics == {
"unresolved_events": 0,
"conflict_events": 1,
"open_executions": 0,
"unattributed_executions": 1,
}
def test_unknown_duration_stays_null_and_enrichment_is_replayable(
sqlite_session_factory: sessionmaker[Session],
) -> None:
missing = provider_event(duration=None)
ingest(missing)
assert ingest(missing)["duplicates"] == 1
assert execution(sqlite_session_factory).metered_duration_ms is None
assert execution(sqlite_session_factory).quality == "pending"
assert ingest(provider_event())["accepted"] == 1
assert execution(sqlite_session_factory).metered_duration_ms == 12345
assert ingest(missing)["duplicates"] == 1
with sqlite_session_factory() as session:
event = session.scalars(
sa.select(AgentSandboxUsageEvent).where(AgentSandboxUsageEvent.source_event_id == "event-1")
).one()
assert event.payload["eventData"]["execution"]["execution_time"] is None
assert event.payload_versions
assert event.resolution["canonical"]["duration_ms"] == 12345
def test_conflicting_event_keeps_both_payloads_and_quarantines_usage(
sqlite_session_factory: sessionmaker[Session],
) -> None:
ingest(provider_event())
assert ingest(provider_event(duration=999))["conflicts"] == 1
assert execution(sqlite_session_factory).quality == "conflict"
assert execution(sqlite_session_factory).metered_duration_ms == 12345
with sqlite_session_factory() as session:
event = session.scalars(
sa.select(AgentSandboxUsageEvent).where(AgentSandboxUsageEvent.source_event_id == "event-1")
).one()
assert event.payload_versions[0]["payload"]["eventData"]["execution"]["execution_time"] == 999
ingest(provider_event(event_id="new-terminal"))
assert execution(sqlite_session_factory).quality == "conflict"
def test_conflicting_executions_never_add_or_overwrite_duration(sqlite_session_factory: sessionmaker[Session]) -> None:
ingest(provider_event())
assert ingest(provider_event(event_id="other-terminal", duration=1))["conflicts"] == 1
assert execution(sqlite_session_factory).metered_duration_ms == 12345
def test_late_start_does_not_reopen_closed_execution(sqlite_session_factory: sessionmaker[Session]) -> None:
ingest(provider_event())
ingest(provider_event(event_id="start", kind="resumed", duration=None))
row = execution(sqlite_session_factory)
assert (row.state, row.quality, row.metered_duration_ms) == ("closed", "metered", 12345)
def test_activation_survives_restarts_and_config_changes_are_rejected() -> None:
state = SandboxUsageService.get_state(project_id=PROJECT)
assert state["started_at"] == START
assert SandboxUsageService.get_state(project_id=PROJECT) == state
with config_overrides_context(AGENT_SANDBOX_METERING_START_AT="2026-09-21T00:00:00Z"):
with pytest.raises(SandboxUsageError, match="immutable"):
SandboxUsageService.get_state(project_id=PROJECT)
def test_no_old_data_even_when_terminal_arrives_after_activation(sqlite_session_factory: sessionmaker[Session]) -> None:
assert ingest(provider_event(started_at="2026-09-19T23:59:59Z"))["ignored"] == 1
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxExecution)) == 0
assert (
session.scalar(
sa.select(sa.func.count())
.select_from(AgentSandboxUsageEvent)
.where(AgentSandboxUsageEvent.source == "provider")
)
== 0
)
def test_t0_filter_cannot_hide_contradictory_existing_metered_event(
sqlite_session_factory: sessionmaker[Session],
) -> None:
ingest(provider_event())
assert ingest(provider_event(started_at="2026-09-19T23:59:59Z"))["conflicts"] == 1
assert execution(sqlite_session_factory).quality == "conflict"
def test_unknown_start_proven_before_t0_removes_pending_execution(
sqlite_session_factory: sessionmaker[Session],
) -> None:
unknown = provider_event(started_at=None, duration=None)
ingest(unknown)
assert ingest(provider_event(started_at="2026-09-19T23:59:59Z"))["ignored"] == 1
ingest(provider_event(event_id="late-start", kind="created", started_at=None, duration=None))
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxExecution)) == 0
def test_unusable_provider_data_is_visible_but_not_metered(sqlite_session_factory: sessionmaker[Session]) -> None:
event = provider_event(started_at=None)
assert ingest(event)["accepted"] == 1
with sqlite_session_factory() as session:
assert session.scalars(sa.select(AgentSandboxExecution)).one().quality == "pending"
row = session.scalars(
sa.select(AgentSandboxUsageEvent).where(AgentSandboxUsageEvent.source == "provider")
).one()
assert row.projection_status == "unresolved"
def test_created_without_execution_data_then_terminal_and_late_start(
sqlite_session_factory: sessionmaker[Session],
) -> None:
created = provider_event(event_id="created", kind="created", duration=None, started_at=None)
created.payload["eventData"] = dict[str, JsonValue]()
assert ingest(created)["accepted"] == 1
row = execution(sqlite_session_factory)
assert row.started_at is None
assert row.state == "open"
assert row.metered_duration_ms is None
ingest(provider_event())
assert execution(sqlite_session_factory).quality == "metered"
assert ingest(created)["duplicates"] == 1
late = created.model_copy(deep=True, update={"id": "late"})
late.payload["id"] = "late"
ingest(late)
row = execution(sqlite_session_factory)
assert (row.state, row.quality, row.metered_duration_ms) == ("closed", "metered", 12345)
def test_missing_provider_project_is_rejected() -> None:
event = provider_event()
event.payload.pop("sandboxTeamId")
with pytest.raises(SandboxUsageError, match="provider_project_mismatch"):
ingest(event)
def test_retention_gap_is_saved_without_advancing_checkpoint(sqlite_session_factory: sessionmaker[Session]) -> None:
with config_overrides_context(AGENT_SANDBOX_METERING_START_AT="2026-09-01T00:00:00Z"):
event = SandboxUsageEvent(
id="gap",
source="application",
type="collector_retention_gap",
timestamp=START,
payload={
"uncovered_start": "2026-09-01T00:00:00Z",
"uncovered_end": "2026-09-13T00:00:00Z",
"assumed_retention_seconds": 604800,
},
)
assert ingest(event)["accepted"] == 1
assert SandboxUsageService.get_state(project_id=PROJECT)["checkpoint_at"] is None
with sqlite_session_factory() as session:
row = session.scalars(
sa.select(AgentSandboxUsageEvent).where(AgentSandboxUsageEvent.source_event_id == "gap")
).one()
assert row.projection_status == "unresolved"
assert row.projection_error_code == "provider_retention_gap"
def test_checkpoint_only_advances_after_completed_scan() -> None:
event = SandboxUsageEvent(
id="checkpoint-1",
source="application",
type="collector_checkpoint",
timestamp=START,
payload={
"scan_started_at": START,
"window_start": START,
"window_end": START,
"completed": True,
"mode": "full",
"pages": 1,
"events": 0,
},
)
ingest(event)
state = SandboxUsageService.get_state(project_id=PROJECT)
assert state["checkpoint_at"] == START
assert state["full_scan_at"] == START
with pytest.raises(SandboxUsageError, match="invalid_collector_checkpoint"):
ingest(event.model_copy(update={"id": "partial", "payload": {**event.payload, "completed": False}}))
def test_forbidden_project_and_reserved_application_events_fail_without_partial_batch(
sqlite_session_factory: sessionmaker[Session],
) -> None:
with pytest.raises(SandboxUsageError, match="not_allowed"):
SandboxUsageService.ingest(project_id="other-project", events=[provider_event()])
reserved = SandboxUsageEvent(id="forged", source="application", type="allocation_registered", payload={})
with pytest.raises(SandboxUsageError, match="unsupported_application_event"):
ingest(provider_event(), reserved)
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxExecution)) == 0
def test_provider_team_must_match_configured_project() -> None:
event = provider_event()
event.payload["sandboxTeamId"] = "other-project"
with pytest.raises(SandboxUsageError, match="provider_project_mismatch"):
ingest(event)
def test_disabled_ingestion_does_not_ack(
sqlite_session_factory: sessionmaker[Session],
) -> None:
with config_overrides_context(AGENT_SANDBOX_METERING_ENABLED=False):
assert SandboxUsageService.get_state(project_id=PROJECT)["enabled"] is False
with pytest.raises(SandboxUsageError, match="disabled"):
ingest(provider_event())
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxUsageEvent)) == 0
@pytest.mark.parametrize(
("timestamp", "error"),
[
(123, "invalid_timestamp"),
("not-a-time", "invalid_timestamp"),
("2026-09-20T00:00:14", "timestamp_requires_timezone"),
],
)
def test_invalid_provider_time_rolls_back_entire_batch(
timestamp: JsonValue, error: str, sqlite_session_factory: sessionmaker[Session]
) -> None:
bad = provider_event(event_id="bad-time", execution_id="bad-execution")
bad.payload["timestamp"] = timestamp
with pytest.raises(SandboxUsageError, match=error):
ingest(provider_event(), bad)
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxExecution)) == 0
@pytest.mark.parametrize("quantity", [-1, True, 1.5, 2**63])
def test_invalid_duration_is_rejected_before_any_accounting(quantity: JsonValue) -> None:
event = provider_event()
data = event.payload["eventData"]
assert isinstance(data, dict)
details = data["execution"]
assert isinstance(details, dict)
details["execution_time"] = quantity
with pytest.raises(SandboxUsageError, match="invalid_provider_quantity"):
ingest(event)
@pytest.mark.parametrize("field", ["id", "sandboxId", "sandboxExecutionId"])
def test_raw_identity_cannot_disagree_with_delivery_envelope(field: str) -> None:
event = provider_event()
event.payload[field] = "different-id"
with pytest.raises(SandboxUsageError, match="provider_envelope_mismatch"):
ingest(event)
def test_invalid_provider_identifier_without_envelope_hint_is_rejected() -> None:
event = provider_event().model_copy(update={"sandbox_id": None})
event.payload["sandboxId"] = ""
with pytest.raises(SandboxUsageError, match="invalid_provider_identifier"):
ingest(event)
@pytest.mark.parametrize(
("configured_start", "error"),
[
("invalid", "sandbox_metering_start_not_configured"),
("2026-09-20T00:00:00.001Z", "sandbox_metering_start_requires_whole_second"),
],
)
def test_invalid_activation_configuration_cannot_create_ledger_scope(configured_start: str, error: str) -> None:
with config_overrides_context(AGENT_SANDBOX_METERING_START_AT=configured_start):
with pytest.raises(SandboxUsageError, match=error) as result:
SandboxUsageService.get_state(project_id=PROJECT)
assert result.value.status_code == 503
@pytest.mark.parametrize(
("key", "value"),
[
("dify.usage_allocation_id", 42),
("dify.usage_allocation_id", "invalid-uuid"),
("dify.binding_id", 42),
("dify.binding_id", "invalid-uuid"),
],
)
def test_invalid_owner_metadata_does_not_discard_metered_usage(
key: str, value: JsonValue, sqlite_session_factory: sessionmaker[Session]
) -> None:
assert ingest(provider_event(metadata={key: value}))["accepted"] == 1
row = execution(sqlite_session_factory)
assert row.quality == "metered"
assert row.metered_duration_ms == 12345
assert row.tenant_id is None
assert row.attribution_status == "unresolved"
@pytest.mark.parametrize("change", ["sandbox", "started_at", "resources"])
def test_execution_identity_and_resource_conflicts_preserve_original_usage(
change: str, sqlite_session_factory: sessionmaker[Session]
) -> None:
ingest(provider_event())
different = provider_event(event_id="conflicting-terminal")
if change == "sandbox":
different = provider_event(event_id="conflicting-terminal", sandbox_id="other-sandbox")
elif change == "started_at":
different = provider_event(event_id="conflicting-terminal", started_at="2026-09-20T00:00:03Z")
else:
data = different.payload["eventData"]
assert isinstance(data, dict)
details = data["execution"]
assert isinstance(details, dict)
details["vcpu_count"] = 4
assert ingest(different)["conflicts"] == 1
row = execution(sqlite_session_factory)
assert row.quality == "conflict"
assert row.sandbox_id == "sandbox-1"
assert row.vcpu_count == 2
assert row.metered_duration_ms == 12345
def test_unsupported_provider_contract_is_visible_without_fabricated_execution(
sqlite_session_factory: sessionmaker[Session],
) -> None:
event = provider_event()
event.payload["version"] = "unsupported"
assert ingest(event)["accepted"] == 1
with sqlite_session_factory() as session:
assert session.scalar(sa.select(sa.func.count()).select_from(AgentSandboxExecution)) == 0
row = session.scalars(
sa.select(AgentSandboxUsageEvent).where(AgentSandboxUsageEvent.source_event_id == event.id)
).one()
assert row.projection_status == "unresolved"
assert row.projection_error_code == "incomplete_provider_execution"
def business_owner(
session: Session,
*,
binding_id: str | None = None,
tenant_id: str | None = None,
sandbox_id: str = "sandbox-1",
) -> AgentWorkspaceBinding:
tenant_id = tenant_id or str(uuid4())
app_id, agent_id, workspace_id = (str(uuid4()) for _ in range(3))
session.add_all(
[
App(
id=app_id,
tenant_id=tenant_id,
name="Metering test",
description="",
mode="agent",
enable_site=False,
enable_api=False,
max_active_requests=None,
),
Agent(
id=agent_id, tenant_id=tenant_id, name=str(uuid4()), scope=AgentScope.ROSTER, source=AgentSource.ROSTER
),
AgentWorkspace(
id=workspace_id,
tenant_id=tenant_id,
app_id=app_id,
owner_type=AgentWorkspaceOwnerType.CONVERSATION,
owner_id=str(uuid4()),
owner_scope_key="root",
backend_workspace_ref=sandbox_id,
),
]
)
binding = AgentWorkspaceBinding(
id=binding_id or str(uuid4()),
tenant_id=tenant_id,
app_id=app_id,
agent_id=agent_id,
workspace_id=workspace_id,
agent_config_version_id=str(uuid4()),
agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT,
backend_binding_ref=sandbox_id,
)
session.add(binding)
session.commit()
return binding
def owner_metadata(binding: AgentWorkspaceBinding) -> dict[str, JsonValue]:
return {
"dify.binding_id": binding.id,
"dify.tenant_id": binding.tenant_id,
"dify.agent_id": binding.agent_id,
"dify.workspace_id": binding.workspace_id,
}
def test_provider_usage_without_business_rows_remains_metered(sqlite_session_factory: sessionmaker[Session]) -> None:
ingest(provider_event(metadata={"dify.binding_id": str(uuid4()), "dify.tenant_id": str(uuid4())}))
row = execution(sqlite_session_factory)
assert row.quality == "metered"
assert row.metered_duration_ms == 12345
assert row.attribution_status == "unresolved"
assert row.tenant_id is None
assert row.allocation_id is None
def test_background_attribution_uses_verified_business_chain(
sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
binding = business_owner(sqlite_session)
ingest(provider_event(metadata=owner_metadata(binding)))
row = execution(sqlite_session_factory)
assert (row.tenant_id, row.app_id, row.agent_id, row.binding_id, row.workspace_id) == (
binding.tenant_id,
binding.app_id,
binding.agent_id,
binding.id,
binding.workspace_id,
)
assert row.attribution_status == "resolved"
assert row.allocation_id is None
with sqlite_session_factory() as session:
assert (
session.scalar(
sa.select(sa.func.count())
.select_from(AgentSandboxUsageEvent)
.where(AgentSandboxUsageEvent.event_type == "allocation_registered")
)
== 0
)
def test_duplicate_provider_replay_can_resolve_business_rows_committed_later(
sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
binding_id, tenant_id = str(uuid4()), str(uuid4())
event = provider_event(metadata={"dify.binding_id": binding_id, "dify.tenant_id": tenant_id})
ingest(event)
assert execution(sqlite_session_factory).attribution_status == "unresolved"
binding = business_owner(sqlite_session, binding_id=binding_id, tenant_id=tenant_id)
assert ingest(event)["duplicates"] == 1
row = execution(sqlite_session_factory)
assert row.binding_id == binding.id
assert row.attribution_status == "resolved"
assert row.metered_duration_ms == 12345
@pytest.mark.parametrize("missing", ["binding", "workspace", "app", "agent"])
def test_deleted_business_rows_do_not_erase_retained_owner_or_usage(
missing: str, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
binding = business_owner(sqlite_session)
event = provider_event(metadata=owner_metadata(binding))
ingest(event)
if missing == "binding":
sqlite_session.delete(binding)
elif missing == "workspace":
sqlite_session.execute(sa.delete(AgentWorkspace).where(AgentWorkspace.id == binding.workspace_id))
elif missing == "app":
sqlite_session.execute(sa.delete(App).where(App.id == binding.app_id))
else:
sqlite_session.execute(sa.delete(Agent).where(Agent.id == binding.agent_id))
sqlite_session.commit()
assert ingest(event)["duplicates"] == 1
row = execution(sqlite_session_factory)
assert (row.binding_id, row.tenant_id) == (binding.id, binding.tenant_id)
assert row.attribution_status == "resolved"
assert row.metered_duration_ms == 12345
@pytest.mark.parametrize(
"broken_link",
[
"tenant_metadata",
"agent_metadata",
"workspace_metadata",
"binding_sandbox",
"workspace_sandbox",
"app_tenant",
"agent_tenant",
],
)
def test_unverified_business_chain_never_guesses_an_owner(
broken_link: str, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
binding = business_owner(sqlite_session)
metadata = owner_metadata(binding)
if broken_link == "tenant_metadata":
metadata["dify.tenant_id"] = str(uuid4())
elif broken_link in {"agent_metadata", "workspace_metadata"}:
metadata[f"dify.{broken_link.removesuffix('_metadata')}_id"] = str(uuid4())
elif broken_link == "binding_sandbox":
binding.backend_binding_ref = "different-sandbox"
elif broken_link == "workspace_sandbox":
sqlite_session.execute(
sa.update(AgentWorkspace)
.where(AgentWorkspace.id == binding.workspace_id)
.values(backend_workspace_ref="different-sandbox")
)
elif broken_link == "app_tenant":
sqlite_session.execute(sa.update(App).where(App.id == binding.app_id).values(tenant_id=str(uuid4())))
else:
sqlite_session.execute(sa.update(Agent).where(Agent.id == binding.agent_id).values(tenant_id=str(uuid4())))
sqlite_session.commit()
ingest(provider_event(metadata=metadata))
row = execution(sqlite_session_factory)
assert row.quality == "metered"
assert row.tenant_id is None
assert row.attribution_status == "unresolved"
@pytest.mark.parametrize("bad_value", [None, 42, "invalid-uuid"])
def test_bad_attribution_labels_do_not_invalidate_measured_runtime(
bad_value: JsonValue, sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
binding = business_owner(sqlite_session)
ingest(provider_event(metadata=owner_metadata(binding)))
changed = owner_metadata(binding)
changed["dify.tenant_id"] = bad_value
assert ingest(provider_event(metadata=changed))["accepted"] == 1
row = execution(sqlite_session_factory)
assert row.quality == "metered"
assert row.attribution_status == "resolved"
assert row.tenant_id == binding.tenant_id
assert row.metered_duration_ms == 12345
def test_owner_conflict_does_not_reassign_history_or_discard_project_usage(
sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
first, second = business_owner(sqlite_session), business_owner(sqlite_session)
ingest(provider_event(metadata=owner_metadata(first)))
ingest(provider_event(event_id="other-owner", metadata=owner_metadata(second)))
row = execution(sqlite_session_factory)
assert row.attribution_status == "conflict"
assert row.binding_id == first.id
assert row.tenant_id == first.tenant_id
assert row.quality == "metered"
assert row.metered_duration_ms == 12345
def test_same_sandbox_cannot_change_owner_across_executions(
sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]
) -> None:
first, second = business_owner(sqlite_session), business_owner(sqlite_session)
ingest(provider_event(metadata=owner_metadata(first)))
resumed = provider_event(event_id="new-execution", execution_id="execution-2", metadata=owner_metadata(second))
ingest(resumed)
# A subsequent replay with the original owner cannot erase the conflicting
# association. Both physical executions remain counted at project level.
ingest(provider_event(event_id="later-event", execution_id="execution-2", metadata=owner_metadata(first)))
with sqlite_session_factory() as session:
rows = list(
session.scalars(sa.select(AgentSandboxExecution).order_by(AgentSandboxExecution.provider_execution_id))
)
assert [(row.attribution_status, row.binding_id) for row in rows] == [
("resolved", first.id),
("conflict", None),
]
assert [row.quality for row in rows] == ["metered", "metered"]
assert sum(row.metered_duration_ms or 0 for row in rows) == 24690
@pytest.mark.parametrize("event_type", ["operation_requested", "operation_observed", "allocation_registered"])
def test_application_operations_are_no_longer_accepted(event_type: str) -> None:
event = SandboxUsageEvent(id="legacy-operation", source="application", type=event_type, payload={})
with pytest.raises(SandboxUsageError, match="unsupported_application_event"):
ingest(event)
def test_steady_state_lookup_only_selects_without_locking_activation(sqlite_session: Session) -> None:
engine = sqlite_session.get_bind()
assert isinstance(engine, Engine)
initial = SandboxUsageService.get_state(project_id=PROJECT)
statements: list[str] = []
def observe(_connection: Connection, statement: ClauseElement, *_args: object) -> None:
statements.append(str(statement.compile(dialect=postgresql.dialect())))
sqlalchemy_event.listen(engine, "before_execute", observe)
try:
assert SandboxUsageService.get_state(project_id=PROJECT) == initial
finally:
sqlalchemy_event.remove(engine, "before_execute", observe)
assert statements
assert all(statement.lstrip().startswith("SELECT") for statement in statements)
assert all("FOR UPDATE" not in statement for statement in statements)
@@ -637,3 +637,67 @@ def test_workspace_collection_final_delete_failure_propagates(
assert exc_info.value is error
client.destroy_execution_binding_sync.assert_called_once()
def test_binding_creation_ignores_unavailable_metering_configuration(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
apply_config_overrides(
monkeypatch,
AGENT_SANDBOX_METERING_ENABLED=True,
AGENT_SANDBOX_METERING_PROJECT_ID="",
AGENT_SANDBOX_METERING_START_AT="invalid",
)
client = _backend_client()
monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client))
binding = AgentWorkspaceService.create_binding(
session=sqlite_session,
scope=_scope(),
agent_id="agent-1",
base_home_snapshot_id=None,
agent_config_version_id="config-1",
agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT,
)
sqlite_session.commit()
assert binding.backend_binding_ref == "binding-ref"
client.create_execution_binding_sync.assert_called_once()
def test_binding_lookups_ignore_bad_metering_configuration_and_remain_tenant_scoped(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
apply_config_overrides(
monkeypatch,
AGENT_SANDBOX_METERING_ENABLED=True,
AGENT_SANDBOX_METERING_PROJECT_ID="",
AGENT_SANDBOX_METERING_START_AT="invalid",
)
workspace, binding = _workspace(), _binding()
sqlite_session.add_all([workspace, binding])
sqlite_session.commit()
assert (
AgentWorkspaceService.get_active_binding(
session=sqlite_session,
tenant_id=binding.tenant_id,
binding_id=binding.id,
expected_owner_scope=_scope(),
)
is binding
)
assert (
AgentWorkspaceService.resolve_active_binding_for_scope(
session=sqlite_session,
scope=_scope(),
agent_id=binding.agent_id,
)
is binding
)
assert (
AgentWorkspaceService.get_active_binding(
session=sqlite_session,
tenant_id="another-tenant",
binding_id=binding.id,
expected_owner_scope=_scope(),
)
is None
)