mirror of
https://github.com/langgenius/dify.git
synced 2026-10-04 13:47:55 +08:00
feat(agent): record E2B sandbox runtime usage (#42554)
This commit is contained in:
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
|
||||
)
|
||||
Reference in new issue
Block a user