jiang
2e4f6cb93b
Add isolated burst and pollution impact evidence, parallel work-order workflows, persisted notifications and graph updates. Checkpoint current code. Known follow-ups remain for large pollution result previews and graceful geometry failure.
248 lines
14 KiB
Python
248 lines
14 KiB
Python
import asyncio
|
|
from datetime import datetime, timezone
|
|
from uuid import uuid4
|
|
from types import SimpleNamespace
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
from app.mcp.models import BoundContext
|
|
from app.mcp.calculation_jobs import durable_binding, restore_binding, check_binding, should_queue, QUEUED_TOOLS
|
|
|
|
|
|
def binding():
|
|
return BoundContext(context_handle='secret-handle', project_id=uuid4(), subject='user', client_id='host', native_session_id='session', purpose='ordinary_chat', expires_at=datetime.now(timezone.utc))
|
|
|
|
|
|
def test_durable_binding_never_stores_expiring_secret_handle():
|
|
original = binding()
|
|
safe = durable_binding(original)
|
|
assert 'context_handle' not in safe
|
|
assert 'expires_at' not in safe
|
|
assert restore_binding(safe).subject == original.subject
|
|
|
|
|
|
def test_original_session_boundary_is_enforced():
|
|
original = binding()
|
|
row = {'binding': durable_binding(original)}
|
|
check_binding(original, row)
|
|
original.native_session_id = 'other'
|
|
with pytest.raises(HTTPException) as error:
|
|
check_binding(original, row)
|
|
assert error.value.status_code == 403
|
|
|
|
|
|
def test_only_executed_expensive_mutations_are_queued():
|
|
tool = SimpleNamespace(name='save_simulation_scenario', mutating=True)
|
|
assert not should_queue(tool, SimpleNamespace(dry_run=True))
|
|
assert should_queue(tool, SimpleNamespace(dry_run=False))
|
|
assert 'analyze_valve_isolation' not in QUEUED_TOOLS
|
|
assert 'calculate_demand_allocation' not in QUEUED_TOOLS
|
|
assert len(QUEUED_TOOLS) == 12
|
|
|
|
|
|
def test_result_pages_keep_original_semantics_and_provenance(monkeypatch):
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs
|
|
from app.mcp.calculation_models import CalculationResultInput
|
|
from app.mcp.models import QueryResult
|
|
original = binding()
|
|
result = QueryResult(source='run_burst_location', source_ref='calculation:one', queried_at=datetime.now(timezone.utc), semantics='analysis', data={'result': {'items': list(range(10))}}, evidence_id='evidence-one')
|
|
row = {'job_id': uuid4(), 'binding': durable_binding(original), 'status': 'succeeded', 'result': result.model_dump(mode='json')}
|
|
repo = SimpleNamespace(get=AsyncMock(return_value=row))
|
|
monkeypatch.setattr(jobs, 'CalculationJobRepository', lambda *a: repo)
|
|
monkeypatch.setattr(jobs, 'authorize_original', AsyncMock())
|
|
monkeypatch.setattr(jobs, 'attach_calculation_evidence', AsyncMock(side_effect=lambda scope, row, result: result))
|
|
scope = SimpleNamespace(connection=None, project=SimpleNamespace(project_id=original.project_id, user_id=uuid4()))
|
|
args = CalculationResultInput(context_handle='x' * 32, job_id=row['job_id'], data_path=['result', 'items'], limit=3, offset=2)
|
|
page = asyncio.run(jobs.read_calculation_job(scope, original, SimpleNamespace(name='get_calculation_result'), args))
|
|
assert page.data == [2, 3, 4]
|
|
assert page.semantics == 'analysis'
|
|
assert page.source == 'run_burst_location'
|
|
assert page.evidence_id == 'evidence-one'
|
|
assert page.truncated and page.page.has_more and page.page.total == 10
|
|
|
|
|
|
def test_pending_result_is_task_receipt_without_evidence(monkeypatch):
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs
|
|
from app.mcp.calculation_models import CalculationResultInput
|
|
original = binding()
|
|
row = {'job_id': uuid4(), 'binding': durable_binding(original), 'status': 'running', 'tool': 'run_simulation'}
|
|
monkeypatch.setattr(jobs, 'CalculationJobRepository', lambda *a: SimpleNamespace(get=AsyncMock(return_value=row)))
|
|
monkeypatch.setattr(jobs, 'authorize_original', AsyncMock())
|
|
attach = AsyncMock()
|
|
monkeypatch.setattr(jobs, 'attach_calculation_evidence', attach)
|
|
scope = SimpleNamespace(connection=None, project=SimpleNamespace(project_id=original.project_id, user_id=uuid4()))
|
|
result = asyncio.run(jobs.read_calculation_job(scope, original, SimpleNamespace(name='get_calculation_result'), CalculationResultInput(context_handle='x' * 32, job_id=row['job_id'])))
|
|
assert result.semantics == 'task' and result.evidence_id is None
|
|
assert result.data['status'] == 'running'
|
|
attach.assert_not_awaited()
|
|
|
|
|
|
def test_mutation_submission_cannot_bypass_existing_approval_receipt(monkeypatch):
|
|
from contextlib import asynccontextmanager
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs, query_service
|
|
from app.mcp.catalog import TOOLS
|
|
from app.mcp.extended_models import SensorCreateInput
|
|
from app.mcp.mutation_service import mutation_arguments
|
|
from app.infra.db.postgresql import mcp_queries
|
|
args = SensorCreateInput(context_handle='x' * 32, request_id=uuid4(), dry_run=False, expected_digest='a' * 64, name='approved', method='kmeans', sensor_count=2)
|
|
@asynccontextmanager
|
|
async def transaction():
|
|
yield
|
|
previous = {'parameters': {'tool': 'create_sensor_placement', 'arguments': mutation_arguments(args, preview=False)}, 'result': {'status': 'started'}}
|
|
receipt = SimpleNamespace(lock_receipt=AsyncMock(), receipt=AsyncMock(return_value=previous))
|
|
monkeypatch.setattr(mcp_queries, 'McpQueryRepository', lambda *a: receipt)
|
|
monkeypatch.setattr(query_service, 'authorize_query', AsyncMock())
|
|
queue = AsyncMock()
|
|
monkeypatch.setattr(jobs, 'CalculationJobRepository', queue)
|
|
scope = SimpleNamespace(connection=SimpleNamespace(transaction=transaction), project=SimpleNamespace(project_id=uuid4(), user_id=uuid4()))
|
|
with pytest.raises(HTTPException) as error:
|
|
asyncio.run(jobs.submit_calculation(scope, binding(), TOOLS['create_sensor_placement'], args, 'host-call'))
|
|
assert error.value.status_code == 409
|
|
assert 'already started' in str(error.value.detail)
|
|
queue.assert_not_called()
|
|
changed = args.model_copy(update={'name': 'changed'})
|
|
with pytest.raises(HTTPException) as conflict:
|
|
asyncio.run(jobs.submit_calculation(scope, binding(), TOOLS['create_sensor_placement'], changed, 'host-call'))
|
|
assert conflict.value.detail == 'request_id already has different parameters'
|
|
queue.assert_not_called()
|
|
|
|
|
|
def test_evidence_receipt_commits_in_root_operational_transaction(monkeypatch):
|
|
from contextlib import asynccontextmanager
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs, query_service, contexts
|
|
from app.infra.db.postgresql import mcp_queries
|
|
from app.mcp.models import QueryResult
|
|
original = binding().model_copy(update={'task_id': uuid4(), 'run_id': uuid4()})
|
|
result = QueryResult(source='run_simulation', source_ref='original', queried_at=datetime.now(timezone.utc), semantics='analysis', data={})
|
|
events = []
|
|
async def commit():
|
|
events.append('commit-reads')
|
|
@asynccontextmanager
|
|
async def transaction():
|
|
events.append('begin-root-operational')
|
|
yield
|
|
events.append('commit-evidence')
|
|
task_repo = SimpleNamespace(transaction=transaction, lock=AsyncMock())
|
|
receipt_repo = SimpleNamespace(lock_receipt=AsyncMock(), receipt=AsyncMock(return_value=None), save_receipt=AsyncMock())
|
|
service = SimpleNamespace(repo=task_repo, record_data_snapshot=AsyncMock(return_value=(SimpleNamespace(evidence_id='saved'), None)))
|
|
monkeypatch.setattr(jobs, 'authorize_original', AsyncMock(return_value=(original, SimpleNamespace(name='run_simulation'), None)))
|
|
monkeypatch.setattr(contexts, 'verify_task', AsyncMock())
|
|
monkeypatch.setattr(query_service, 'task_service', lambda scope: service)
|
|
monkeypatch.setattr(mcp_queries, 'McpQueryRepository', lambda *a: receipt_repo)
|
|
scope = SimpleNamespace(connection=SimpleNamespace(commit=commit), project=SimpleNamespace(project_id=original.project_id))
|
|
attached = asyncio.run(jobs.attach_calculation_evidence(scope, {'job_id': uuid4()}, result))
|
|
assert attached.evidence_id == 'saved'
|
|
assert events == ['commit-reads', 'begin-root-operational', 'commit-evidence']
|
|
receipt_repo.save_receipt.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.parametrize('job_budget', [3600, 7200])
|
|
def test_http_timeout_does_not_limit_queued_execution(monkeypatch, job_budget):
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs, query_service
|
|
from app.mcp.catalog import TOOLS
|
|
tool = TOOLS['run_simulation']
|
|
args = tool.model.model_validate({
|
|
'context_handle': 'x' * 32, 'timeout_seconds': 1,
|
|
'start_time': '2026-10-08T12:00:00+08:00', 'duration': 3600,
|
|
'output_node_ids': ['J1'],
|
|
})
|
|
repo = SimpleNamespace(enqueue=AsyncMock(return_value={
|
|
'job_id': uuid4(), 'status': 'queued', 'tool': tool.name,
|
|
}))
|
|
monkeypatch.setattr(jobs, 'CalculationJobRepository', lambda *a: repo)
|
|
monkeypatch.setattr(query_service, 'authorize_query', AsyncMock())
|
|
scope = SimpleNamespace(connection=SimpleNamespace(commit=AsyncMock()),
|
|
project=SimpleNamespace(project_id=uuid4(), user_id=uuid4()))
|
|
result = asyncio.run(jobs.submit_calculation(
|
|
scope, binding(), tool, args, 'host-call', job_timeout_seconds=job_budget,
|
|
))
|
|
assert result.data['status'] == 'queued'
|
|
assert repo.enqueue.await_args.kwargs['timeout_seconds'] == job_budget
|
|
assert 'timeout_seconds' not in repo.enqueue.await_args.kwargs['parameters']
|
|
|
|
|
|
def test_expired_result_is_safe_explicit_error(monkeypatch):
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs
|
|
from app.mcp.calculation_models import CalculationResultInput
|
|
original = binding()
|
|
row = {'job_id': uuid4(), 'binding': durable_binding(original), 'status': 'succeeded',
|
|
'result': None, 'result_expires_at': datetime(2020, 1, 1, tzinfo=timezone.utc)}
|
|
repo = SimpleNamespace(get=AsyncMock(return_value=row))
|
|
monkeypatch.setattr(jobs, 'CalculationJobRepository', lambda *a: repo)
|
|
monkeypatch.setattr(jobs, 'authorize_original', AsyncMock())
|
|
attach = AsyncMock()
|
|
monkeypatch.setattr(jobs, 'attach_calculation_evidence', attach)
|
|
scope = SimpleNamespace(connection=None, project=SimpleNamespace(project_id=original.project_id, user_id=uuid4()))
|
|
with pytest.raises(HTTPException) as error:
|
|
asyncio.run(jobs.read_calculation_job(scope, original, SimpleNamespace(name='get_calculation_result'),
|
|
CalculationResultInput(context_handle='x' * 32, job_id=row['job_id'])))
|
|
assert error.value.status_code == 410
|
|
assert jobs.public_job(row)['result_available'] is False
|
|
assert repo.get.await_args.kwargs['include_result'] is True
|
|
attach.assert_not_awaited()
|
|
|
|
|
|
def test_legacy_execution_timeout_does_not_block_owner_authorization(monkeypatch):
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs, query_service
|
|
from app.mcp.catalog import TOOLS
|
|
original = binding()
|
|
owner = uuid4()
|
|
row = {'binding': durable_binding(original), 'tool': 'run_simulation',
|
|
'owner_user_id': owner, 'parameters': {'timeout_seconds': 3600,
|
|
'start_time': '2026-10-29T00:00:00Z', 'output_node_ids': ['1']}}
|
|
authorize = AsyncMock()
|
|
monkeypatch.setattr(query_service, 'authorize_query', authorize)
|
|
scope = SimpleNamespace(project=SimpleNamespace(project_id=original.project_id, user_id=owner))
|
|
_, tool, args = asyncio.run(jobs.authorize_original(scope, row))
|
|
assert tool is TOOLS['run_simulation']
|
|
assert args.timeout_seconds <= 240
|
|
assert row['parameters']['timeout_seconds'] == 3600
|
|
authorize.assert_awaited_once()
|
|
|
|
|
|
def test_large_burst_result_keeps_saved_map_command_and_full_evidence(monkeypatch):
|
|
from copy import deepcopy
|
|
from unittest.mock import AsyncMock
|
|
from app.mcp import calculation_jobs as jobs
|
|
from app.mcp.calculation_models import CalculationResultInput
|
|
from app.mcp.models import QueryResult
|
|
original = binding()
|
|
original.task_id = uuid4()
|
|
full = {"parameters": {"burst_pipe_ids": ["P1"]}, "result": {
|
|
"impact": {"phase": "after_burst_before_intervention", "status": "ready",
|
|
"compared_node_count": 1000, "comparisons": [{"node_id": f"J{i}", "pressure_drop_m": 4} for i in range(1000)],
|
|
"grades": [{"group_id": "drop_3_5", "node_ids": [f"J{i}" for i in range(1000)]}],
|
|
"unavailable_node_ids": []},
|
|
"coverage": {"status": "ready", "groups": {"drop_3_5": {"node_ids": [f"J{i}" for i in range(1000)],
|
|
"mapped_node_count": 1000, "missing_node_ids": []}}},
|
|
"map_payload": {"type": "FeatureCollection", "features": [{"large": "recorded-geometry"}]},
|
|
}}
|
|
result = QueryResult(source='run_burst_simulation', source_ref='mcp:burst', queried_at=datetime.now(timezone.utc),
|
|
semantics='analysis', data=full, evidence_id='saved-burst')
|
|
row = {'job_id': uuid4(), 'binding': durable_binding(original), 'status': 'succeeded', 'result': result.model_dump(mode='json')}
|
|
persisted = deepcopy(row)
|
|
monkeypatch.setattr(jobs, 'CalculationJobRepository', lambda *a: SimpleNamespace(get=AsyncMock(return_value=row)))
|
|
monkeypatch.setattr(jobs, 'authorize_original', AsyncMock())
|
|
monkeypatch.setattr(jobs, 'attach_calculation_evidence', AsyncMock(side_effect=lambda scope, row, result: result))
|
|
scope = SimpleNamespace(connection=None, project=SimpleNamespace(project_id=original.project_id, user_id=uuid4()))
|
|
args = CalculationResultInput(context_handle='x' * 32, job_id=row['job_id'])
|
|
response = asyncio.run(jobs.read_calculation_job(scope, original, SimpleNamespace(name='get_calculation_result'), args))
|
|
assert response.truncated and len(response.model_dump_json().encode()) < 11000
|
|
assert response.data['result']['impact']['compared_node_count'] == 1000
|
|
assert response.data['result']['impact']['grades'][0]['node_count'] == 1000
|
|
assert response.data['presentation']['command']['evidence_id'] == 'saved-burst'
|
|
assert response.data['presentation']['command']['action'] == 'set_overlay'
|
|
assert response.data['resource_uri'].endswith('/evidence/saved-burst')
|
|
assert 'comparisons' not in response.data['result']['impact']
|
|
assert row == persisted
|
|
# Explicit path pagination still reads every original node without summary replacement.
|
|
page = asyncio.run(jobs.read_calculation_job(scope, original, SimpleNamespace(name='get_calculation_result'),
|
|
args.model_copy(update={'data_path': ['result', 'impact', 'comparisons'], 'offset': 990, 'limit': 10})))
|
|
assert page.data[0]['node_id'] == 'J990' and page.page.total == 1000
|