feat(mcp): queue long-running calculations
This commit is contained in:
@@ -0,0 +1,70 @@
|
||||
"""Opt-in real PostgreSQL queue races, using an explicitly disposable database."""
|
||||
import asyncio
|
||||
import os
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from psycopg import AsyncConnection
|
||||
from psycopg.conninfo import conninfo_to_dict
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
||||
|
||||
|
||||
def test_concurrent_submission_claims_fencing_and_project_isolation():
|
||||
dsn = os.getenv("CALCULATION_QUEUE_TEST_DSN")
|
||||
if not dsn:
|
||||
pytest.skip("Set CALCULATION_QUEUE_TEST_DSN to a disposable PostgreSQL database")
|
||||
assert conninfo_to_dict(dsn)["dbname"].startswith("tjwater_calc_test_")
|
||||
|
||||
async def run():
|
||||
project, other, owner = uuid4(), uuid4(), uuid4()
|
||||
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
||||
await conn.execute((Path(__file__).parents[2] / "resources/sql/026_calculation_jobs.sql").read_text())
|
||||
|
||||
async def submit(key, parameters=None, selected_project=project):
|
||||
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
||||
return await CalculationJobRepository(conn, selected_project).enqueue(
|
||||
owner_user_id=owner, request_key=key, tool="run_burst_location",
|
||||
parameters=parameters or {"x": 1}, binding={"native_session_id": "session"}, timeout_seconds=600)
|
||||
|
||||
# Four independent API connections, same logical submission.
|
||||
submissions = await asyncio.gather(*(submit("same") for _ in range(4)))
|
||||
assert len({str(row["job_id"]) for row in submissions}) == 1
|
||||
with pytest.raises(HTTPException) as conflict:
|
||||
await submit("same", {"x": 2})
|
||||
assert conflict.value.status_code == 409
|
||||
second = await submit("second")
|
||||
foreign = await submit("foreign", selected_project=other)
|
||||
|
||||
async def claim(worker):
|
||||
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
||||
return await CalculationJobRepository(conn, project).claim(worker, 60)
|
||||
|
||||
claimed = await asyncio.gather(*(claim(f"worker-{index}") for index in range(4)))
|
||||
rows = [row for row in claimed if row]
|
||||
assert len(rows) == 2
|
||||
assert len({str(row["job_id"]) for row in rows}) == 2
|
||||
assert str(foreign["job_id"]) not in {str(row["job_id"]) for row in rows}
|
||||
|
||||
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
||||
repo = CalculationJobRepository(conn, project)
|
||||
assert await repo.get(foreign["job_id"], owner) is None
|
||||
assert await repo.get(second["job_id"], uuid4()) is None
|
||||
row = rows[0]
|
||||
stale = {**row, "lease_generation": row["lease_generation"] - 1}
|
||||
assert await repo.heartbeat(stale, row["lease_owner"], 60) is None
|
||||
assert not await repo.finish(stale, row["lease_owner"], "succeeded", result={})
|
||||
cancelled = await repo.cancel(row["job_id"], owner)
|
||||
assert cancelled["status"] == "cancel_requested"
|
||||
assert not await repo.finish(row, row["lease_owner"], "succeeded", result={})
|
||||
assert await repo.finish(row, row["lease_owner"], "cancelled")
|
||||
expired = rows[1]
|
||||
await conn.execute("UPDATE workbench.calculation_jobs SET lease_expires_at=now()-interval '1 second' WHERE job_id=%s", (expired["job_id"],))
|
||||
await repo.recover()
|
||||
assert (await repo.get(expired["job_id"], owner))["status"] == "interrupted"
|
||||
assert not await repo.finish(expired, expired["lease_owner"], "succeeded", result={})
|
||||
assert await repo.claim("replacement", 60) is None
|
||||
asyncio.run(run())
|
||||
@@ -0,0 +1,25 @@
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.calculation_cleanup import _names, record_calculation_database
|
||||
from app.algorithms._utils import calculation_temp_directory
|
||||
|
||||
|
||||
def test_job_owned_database_ledger_survives_until_successful_cleanup(tmp_path, monkeypatch):
|
||||
monkeypatch.setenv("TJWATER_CALCULATION_JOB_DIR", str(tmp_path))
|
||||
name = "tjw_tmp_run_project_" + "a" * 32
|
||||
record_calculation_database(name)
|
||||
record_calculation_database(name)
|
||||
assert _names(tmp_path) == [name]
|
||||
assert json.loads((tmp_path / "databases.json").read_text()) == [name]
|
||||
record_calculation_database(name, removed=True)
|
||||
assert _names(tmp_path) == []
|
||||
assert calculation_temp_directory("leakage") == str(tmp_path / "leakage")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("names", [["production"], ["tjw_tmp_not_random"], {"database": "production"}])
|
||||
def test_cleanup_rejects_non_job_database_names(tmp_path, names):
|
||||
(tmp_path / "databases.json").write_text(json.dumps(names))
|
||||
with pytest.raises(ValueError):
|
||||
_names(tmp_path)
|
||||
@@ -0,0 +1,241 @@
|
||||
"""Worker boundaries use real managed subprocesses, without external databases."""
|
||||
import asyncio
|
||||
from functools import wraps
|
||||
from types import SimpleNamespace
|
||||
from contextlib import asynccontextmanager
|
||||
import os
|
||||
from pathlib import Path
|
||||
import signal
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from app.services.calculation_worker import CalculationWorker, WorkerSettings, stop_process_group
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_external_cleanup(monkeypatch):
|
||||
async def cleanup(*args):
|
||||
pass
|
||||
monkeypatch.setattr('app.services.calculation_worker.cleanup_job_resources', cleanup)
|
||||
|
||||
|
||||
def run_async(function):
|
||||
@wraps(function)
|
||||
def wrapper(*args, **kwargs):
|
||||
return asyncio.run(function(*args, **kwargs))
|
||||
return wrapper
|
||||
|
||||
|
||||
def test_worker_settings_require_explicit_valid_scope(monkeypatch):
|
||||
monkeypatch.setenv('CALCULATION_WORKER_PROJECT_IDS', '6b39b128-2ba8-454f-9bbe-e9553902ab53')
|
||||
settings = WorkerSettings.from_environment()
|
||||
assert settings.max_concurrency == 2
|
||||
monkeypatch.setenv('CALCULATION_WORKER_HEARTBEAT_SECONDS', '30')
|
||||
with pytest.raises(ValueError):
|
||||
WorkerSettings.from_environment()
|
||||
|
||||
|
||||
@run_async
|
||||
async def test_process_group_cleanup_kills_native_style_descendant(tmp_path):
|
||||
pid_file = tmp_path / 'pid'
|
||||
code = ('import subprocess,time,pathlib,sys; '
|
||||
'p=subprocess.Popen([sys.executable,"-c","import time; time.sleep(60)"]); '
|
||||
'pathlib.Path(sys.argv[1]).write_text(str(p.pid)); time.sleep(60)')
|
||||
process = await asyncio.create_subprocess_exec(sys.executable, '-c', code,
|
||||
str(pid_file), start_new_session=True)
|
||||
try:
|
||||
for _ in range(100):
|
||||
if pid_file.exists():
|
||||
break
|
||||
await asyncio.sleep(.02)
|
||||
child_pid = int(pid_file.read_text())
|
||||
await stop_process_group(process)
|
||||
assert process.returncode is not None
|
||||
# Linux may retain a dead grandchild as a zombie until PID 1 reaps it.
|
||||
stat_path = Path(f'/proc/{child_pid}/stat')
|
||||
for _ in range(100):
|
||||
if not stat_path.exists() or stat_path.read_text().split()[2] == 'Z':
|
||||
break
|
||||
await asyncio.sleep(.01)
|
||||
assert not stat_path.exists() or stat_path.read_text().split()[2] == 'Z'
|
||||
finally:
|
||||
await stop_process_group(process)
|
||||
|
||||
|
||||
@run_async
|
||||
@pytest.mark.parametrize('lease_status, expected', [('cancel_requested', 'cancelled'), (None, None)])
|
||||
async def test_cancel_and_lost_lease_stop_process_before_persisting(monkeypatch, lease_status, expected):
|
||||
processes = []
|
||||
original = asyncio.create_subprocess_exec
|
||||
async def spawn(*args, **kwargs):
|
||||
process = await original(sys.executable, '-c', 'import time; time.sleep(60)',
|
||||
start_new_session=True)
|
||||
processes.append(process)
|
||||
return process
|
||||
monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
|
||||
finishes = []
|
||||
class Repository:
|
||||
async def heartbeat(self, *args):
|
||||
return lease_status
|
||||
async def finish(self, row, owner, status, **kwargs):
|
||||
assert processes[0].returncode is not None
|
||||
finishes.append(status)
|
||||
return True
|
||||
@asynccontextmanager
|
||||
async def repository(project):
|
||||
yield Repository()
|
||||
worker = CalculationWorker(WorkerSettings(('project',), heartbeat_seconds=.05), repository=repository)
|
||||
await worker.run_job({'project_id': 'project', 'job_id': 'job', 'timeout_seconds': 60})
|
||||
assert finishes == ([expected] if expected else [])
|
||||
assert processes[0].returncode is not None
|
||||
|
||||
|
||||
@run_async
|
||||
async def test_timeout_stops_process_before_recording(monkeypatch):
|
||||
original = asyncio.create_subprocess_exec
|
||||
processes = []
|
||||
async def spawn(*args, **kwargs):
|
||||
p = await original(sys.executable, '-c', 'import time; time.sleep(60)', start_new_session=True)
|
||||
processes.append(p)
|
||||
return p
|
||||
monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
|
||||
statuses = []
|
||||
class Repository:
|
||||
async def heartbeat(self, *args):
|
||||
return 'running'
|
||||
async def finish(self, row, owner, status, **kwargs):
|
||||
assert processes[0].returncode is not None
|
||||
statuses.append(status)
|
||||
return True
|
||||
@asynccontextmanager
|
||||
async def repository(project):
|
||||
yield Repository()
|
||||
worker = CalculationWorker(WorkerSettings(('project',), heartbeat_seconds=.02), repository=repository)
|
||||
await worker.run_job({'project_id': 'project', 'timeout_seconds': 0})
|
||||
assert statuses == ['timed_out']
|
||||
|
||||
|
||||
@run_async
|
||||
async def test_duplicate_supervisor_never_claims(monkeypatch):
|
||||
from app.infra.db.metadb import database
|
||||
class Lock:
|
||||
async def scalar(self, statement):
|
||||
assert 'pg_try_advisory_lock' in str(statement)
|
||||
return False
|
||||
async def commit(self):
|
||||
pass
|
||||
@asynccontextmanager
|
||||
async def connection():
|
||||
yield Lock()
|
||||
monkeypatch.setattr(database, 'engine', SimpleNamespace(connect=connection))
|
||||
@asynccontextmanager
|
||||
async def repository(project):
|
||||
pytest.fail('A duplicate supervisor must never open the queue')
|
||||
yield
|
||||
with pytest.raises(RuntimeError, match='already active'):
|
||||
await CalculationWorker(WorkerSettings(('project',)), repository=repository).run()
|
||||
|
||||
|
||||
@run_async
|
||||
async def test_repository_terminal_updates_are_fenced():
|
||||
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
||||
statements = []
|
||||
class Cursor:
|
||||
async def fetchone(self):
|
||||
return None
|
||||
class Connection:
|
||||
async def execute(self, sql, parameters):
|
||||
statements.append((sql, parameters))
|
||||
return Cursor()
|
||||
repository = CalculationJobRepository(Connection(), 'project')
|
||||
row = {'job_id': 'job', 'lease_generation': 7}
|
||||
assert await repository.finish(row, 'owner', 'succeeded', result={'original': True}) is False
|
||||
sql, parameters = statements[-1]
|
||||
assert 'lease_generation=%s' in sql and 'lease_owner=%s' in sql
|
||||
assert 'lease_expires_at>clock_timestamp()' in sql
|
||||
assert "status!='cancel_requested'" in sql
|
||||
assert 7 in parameters
|
||||
|
||||
|
||||
@run_async
|
||||
async def test_guard_parent_death_pipe_kills_entire_group(tmp_path):
|
||||
pid_file = tmp_path / 'guard-child'
|
||||
child = 'import os,time,pathlib,sys; pathlib.Path(sys.argv[1]).write_text(str(os.getpid())); time.sleep(60)'
|
||||
launcher = '''import asyncio,sys
|
||||
from app.services.calculation_worker import guard_child
|
||||
original = asyncio.create_subprocess_exec
|
||||
async def spawn(*args, **kwargs):
|
||||
return await original(sys.executable, '-c', sys.argv[2], sys.argv[1])
|
||||
asyncio.create_subprocess_exec = spawn
|
||||
asyncio.run(guard_child('unused','unused'))
|
||||
'''
|
||||
guard = await asyncio.create_subprocess_exec(sys.executable, '-c', launcher,
|
||||
str(pid_file), child, stdin=asyncio.subprocess.PIPE, start_new_session=True)
|
||||
try:
|
||||
for _ in range(100):
|
||||
if pid_file.exists():
|
||||
break
|
||||
await asyncio.sleep(.02)
|
||||
child_pid = int(pid_file.read_text())
|
||||
guard.stdin.close()
|
||||
await asyncio.wait_for(guard.wait(), 3)
|
||||
assert guard.returncode == -signal.SIGKILL
|
||||
stat = Path(f'/proc/{child_pid}/stat')
|
||||
for _ in range(100):
|
||||
if not stat.exists() or stat.read_text().split()[2] == 'Z':
|
||||
break
|
||||
await asyncio.sleep(.01)
|
||||
assert not stat.exists() or stat.read_text().split()[2] == 'Z'
|
||||
finally:
|
||||
await stop_process_group(guard)
|
||||
|
||||
|
||||
@run_async
|
||||
async def test_cleanup_failure_is_not_reported_cancelled(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv('CALCULATION_WORKER_STATE_DIR', str(tmp_path))
|
||||
async def cleanup(row, directory):
|
||||
(Path(directory) / 'cleanup-failed').touch()
|
||||
raise RuntimeError('external resource unavailable')
|
||||
monkeypatch.setattr('app.services.calculation_worker.cleanup_job_resources', cleanup)
|
||||
original = asyncio.create_subprocess_exec
|
||||
async def spawn(*args, **kwargs):
|
||||
return await original(sys.executable, '-c', 'import time; time.sleep(60)', start_new_session=True)
|
||||
monkeypatch.setattr(asyncio, 'create_subprocess_exec', spawn)
|
||||
results = []
|
||||
class Repository:
|
||||
async def heartbeat(self, *args):
|
||||
return 'cancel_requested'
|
||||
async def finish(self, row, owner, status, **kwargs):
|
||||
results.append((status, kwargs['error']['code']))
|
||||
return True
|
||||
@asynccontextmanager
|
||||
async def repository(project):
|
||||
yield Repository()
|
||||
worker = CalculationWorker(WorkerSettings(('project',)), repository=repository)
|
||||
await worker.run_job({'project_id': 'project', 'timeout_seconds': 60})
|
||||
assert results == [('failed', 'CLEANUP_FAILED')]
|
||||
assert len(list(tmp_path.glob('*/cleanup-failed'))) == 1
|
||||
|
||||
|
||||
@run_async
|
||||
async def test_startup_recovery_checks_persisted_generation(monkeypatch, tmp_path):
|
||||
import json
|
||||
from app.services.calculation_worker import recover_job_resources
|
||||
project = '6b39b128-2ba8-454f-9bbe-e9553902ab53'
|
||||
job = '2d3a6231-a8ee-4acb-bb13-58328d29b88f'
|
||||
directory = tmp_path / (job + '-owned')
|
||||
directory.mkdir()
|
||||
(directory / 'scope.json').write_text(json.dumps({
|
||||
'project_id': project, 'job_id': job, 'lease_generation': 1}))
|
||||
monkeypatch.setenv('CALCULATION_WORKER_STATE_DIR', str(tmp_path))
|
||||
class Repository:
|
||||
async def get(self, job_id):
|
||||
return {'lease_generation': 2}
|
||||
@asynccontextmanager
|
||||
async def repository(project):
|
||||
yield Repository()
|
||||
monkeypatch.setattr('app.services.calculation_worker.queue_repository', repository)
|
||||
with pytest.raises(RuntimeError, match='ownership'):
|
||||
await recover_job_resources((project,))
|
||||
assert directory.exists()
|
||||
@@ -0,0 +1,139 @@
|
||||
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()
|
||||
@@ -37,31 +37,22 @@ PERMISSIONS = {
|
||||
}
|
||||
|
||||
|
||||
def test_calculation_budget_outlives_query_default_and_still_times_out(runtime, monkeypatch):
|
||||
def test_calculation_submission_returns_receipt_without_running_native_code(runtime, monkeypatch):
|
||||
from app.mcp import calculation_jobs
|
||||
app, client, _, _ = runtime
|
||||
app.settings.query_timeout_seconds = 1
|
||||
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||
|
||||
async def execute(scope, bound, tool, args):
|
||||
await asyncio.sleep(1.05)
|
||||
return QueryResult(
|
||||
source=tool.name, source_ref="test:calculation", semantics="analysis",
|
||||
queried_at=datetime.now(timezone.utc), data={"completed": True},
|
||||
)
|
||||
|
||||
app.executor = execute
|
||||
handle = binding().context_handle
|
||||
arguments = {
|
||||
"context_handle": handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||
receipt = calculation_jobs.envelope("run_simulation", {"job_id": str(uuid4()), "status": "queued"})
|
||||
submit = AsyncMock(return_value=receipt)
|
||||
monkeypatch.setattr(calculation_jobs, "submit_calculation", submit)
|
||||
app.executor = AsyncMock(side_effect=AssertionError("native work entered API process"))
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": {
|
||||
"context_handle": binding().context_handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||
"duration": 3600, "output_node_ids": ["J1"], "timeout_seconds": 2,
|
||||
}
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": arguments}).json()
|
||||
assert response["result"].get("isError") is not True
|
||||
assert response["result"]["structuredContent"]["data"] == {"completed": True}
|
||||
arguments["timeout_seconds"] = 1
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": arguments}).json()
|
||||
assert response["result"]["isError"] is True
|
||||
assert response["result"]["structuredContent"]["error"]["code"] == "QUERY_TIMEOUT"
|
||||
}}).json()
|
||||
assert response["result"]["structuredContent"]["data"]["status"] == "queued"
|
||||
assert response["result"]["structuredContent"]["semantics"] == "task"
|
||||
app.executor.assert_not_awaited()
|
||||
assert submit.await_args.args[3].timeout_seconds == 2
|
||||
|
||||
|
||||
def test_calculation_limit_is_advertised_and_enforced_before_execution(runtime):
|
||||
@@ -80,34 +71,23 @@ def test_calculation_limit_is_advertised_and_enforced_before_execution(runtime):
|
||||
assert calls == []
|
||||
|
||||
|
||||
def test_calculation_capacity_has_an_explicit_mcp_error(runtime, monkeypatch):
|
||||
from threading import BoundedSemaphore
|
||||
from app.mcp import analysis_service
|
||||
|
||||
app, client, _, _ = runtime
|
||||
def test_queue_submission_errors_preserve_mcp_error_mapping(runtime, monkeypatch):
|
||||
from app.mcp import calculation_jobs
|
||||
_, client, _, _ = runtime
|
||||
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||
slots = BoundedSemaphore(2)
|
||||
monkeypatch.setattr(analysis_service, "_slots", slots)
|
||||
async def execute(*unused):
|
||||
return await analysis_service.bounded_calculation(lambda: pytest.fail("occupied calculation entered"))
|
||||
app.executor = execute
|
||||
assert slots.acquire(blocking=False) and slots.acquire(blocking=False)
|
||||
try:
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": {
|
||||
"context_handle": binding().context_handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||
"duration": 3600, "output_node_ids": ["J1"],
|
||||
}}).json()
|
||||
finally:
|
||||
slots.release()
|
||||
slots.release()
|
||||
error = response["result"]["structuredContent"]["error"]
|
||||
assert response["result"]["isError"] is True
|
||||
assert error["code"] == "CALCULATION_CAPACITY_EXCEEDED"
|
||||
assert "not started" in error["message"]
|
||||
monkeypatch.setattr(calculation_jobs, "submit_calculation", AsyncMock(side_effect=HTTPException(409, "Calculation request_key was used with different inputs")))
|
||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": {
|
||||
"context_handle": binding().context_handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||
"duration": 3600, "output_node_ids": ["J1"],
|
||||
}}).json()
|
||||
assert response["result"]["structuredContent"]["error"]["code"] == "QUERY_CONTEXT_STALE"
|
||||
|
||||
|
||||
def test_burst_location_default_inputs_pass_protocol_validation(runtime, monkeypatch):
|
||||
app, client, calls, _ = runtime
|
||||
from app.mcp import calculation_jobs
|
||||
executor = AsyncMock(return_value=calculation_jobs.envelope("run_burst_location", {"job_id": str(uuid4()), "status": "queued"}))
|
||||
monkeypatch.setattr(calculation_jobs, "submit_calculation", executor)
|
||||
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||
schema = next(tool["inputSchema"] for tool in request(client).json()["result"]["tools"]
|
||||
if tool["name"] == "run_burst_location")
|
||||
@@ -118,7 +98,7 @@ def test_burst_location_default_inputs_pass_protocol_validation(runtime, monkeyp
|
||||
"scada_burst_end": "2026-10-08T13:00:00+08:00",
|
||||
}}).json()
|
||||
assert response["result"].get("isError") is not True
|
||||
assert calls == [("session-1", "run_burst_location")]
|
||||
assert calls == []
|
||||
|
||||
|
||||
def settings():
|
||||
|
||||
@@ -1842,3 +1842,95 @@ async def test_interactive_report_publication_is_idempotent_and_versioned_with_e
|
||||
request_id=uuid4(), source_session_id="session-1", report=report,
|
||||
))
|
||||
assert error.value.code == "EVIDENCE_NOT_FOUND"
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("report_factory", [_current_report, _insufficient_report])
|
||||
async def test_final_report_rejected_while_original_run_has_pending_calculations(report_factory):
|
||||
now = datetime(2026, 9, 23, 8, tzinfo=timezone.utc)
|
||||
context = _context()
|
||||
task = _task_document(context, now, lease=now + timedelta(minutes=1))
|
||||
repository = RunTaskRepository(task, _active_run(task, now))
|
||||
pending_job = str(uuid4())
|
||||
repository.pending_calculation_jobs = AsyncMock(return_value=[pending_job])
|
||||
service = OperationalTaskService(repository, context, set(), clock=lambda: now)
|
||||
with pytest.raises(DomainError) as error:
|
||||
await service.publish_artifacts(task['task_id'], AgentArtifactsCommand(
|
||||
request_id=uuid4(), source_session_id='session-1', run_id=repository.run['run_id'], report=report_factory(),
|
||||
))
|
||||
assert error.value.code == 'PENDING_CALCULATIONS'
|
||||
assert error.value.status == 409
|
||||
assert error.value.recovery == 'retry_same_request'
|
||||
assert error.value.refs == [pending_job]
|
||||
repository.pending_calculation_jobs.assert_awaited_once_with(
|
||||
str(task['task_id']), str(repository.run['run_id']), str(repository.run['execution_request_id']),
|
||||
)
|
||||
assert not repository.reports
|
||||
assert repository.run['status'] in {'claimed', 'running'}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_completion_rejected_until_calculations_reach_terminal_state():
|
||||
now = datetime(2026, 9, 23, 8, tzinfo=timezone.utc)
|
||||
context = _context()
|
||||
task = _task_document(context, now, lease=now + timedelta(minutes=1))
|
||||
repository = RunTaskRepository(task, _active_run(task, now))
|
||||
repository.pending_calculation_jobs = AsyncMock(return_value=[str(uuid4())])
|
||||
service = OperationalTaskService(repository, context, set(), clock=lambda: now)
|
||||
command = AgentRunCommand(request_id=uuid4(), expected_revision=repository.run['revision'], source_session_id='session-1')
|
||||
with pytest.raises(DomainError) as error:
|
||||
await service.complete_run(task['task_id'], repository.run['run_id'], command)
|
||||
assert error.value.code == 'PENDING_CALCULATIONS'
|
||||
assert repository.run['status'] in {'claimed', 'running'}
|
||||
repository.pending_calculation_jobs.return_value = []
|
||||
# Once terminal, the normal final-artifact requirement runs again.
|
||||
completed = await service.complete_run(task['task_id'], repository.run['run_id'], command)
|
||||
assert repository.run['failure_code'] == 'ARTIFACT_INCOMPLETE'
|
||||
assert completed.lifecycle_status == 'failed'
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_pending_calculation_query_matches_exact_original_execution_binding():
|
||||
job_id = uuid4()
|
||||
connection = SimpleNamespace(execute=AsyncMock())
|
||||
connection.execute.return_value.fetchall = AsyncMock(return_value=[{'job_id': job_id}])
|
||||
project_id, task_id, run_id, request_id = (uuid4() for _ in range(4))
|
||||
repo = OperationalTaskRepository(connection, project_id)
|
||||
assert await repo.pending_calculation_jobs(str(task_id), str(run_id), str(request_id)) == [str(job_id)]
|
||||
sql, parameters = connection.execute.await_args.args
|
||||
assert parameters == (str(project_id), str(task_id), str(run_id), str(request_id))
|
||||
assert "status IN ('queued','running','cancel_requested')" in sql
|
||||
assert "binding->>'execution_request_id'=%s" in sql
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interim_report_can_explain_progress_while_calculation_is_pending():
|
||||
now = datetime(2026, 9, 23, 8, tzinfo=timezone.utc)
|
||||
context = _context()
|
||||
task = _task_document(context, now, lease=now + timedelta(minutes=1))
|
||||
repository = RunTaskRepository(task, _active_run(task, now))
|
||||
repository.pending_calculation_jobs = AsyncMock(return_value=[str(uuid4())])
|
||||
service = OperationalTaskService(repository, context, set(), clock=lambda: now)
|
||||
result = await service.publish_artifacts(task['task_id'], AgentArtifactsCommand(
|
||||
request_id=uuid4(), source_session_id='session-1', run_id=repository.run['run_id'],
|
||||
report={**_current_report(), 'stage': 'interim'},
|
||||
))
|
||||
assert result.report_revision == 1
|
||||
repository.pending_calculation_jobs.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_missing_calculation_migration_fails_finalization_closed(monkeypatch):
|
||||
from psycopg.errors import UndefinedTable
|
||||
from app.infra.db.postgresql import operational_tasks as storage
|
||||
@asynccontextmanager
|
||||
async def transaction():
|
||||
yield
|
||||
connection = SimpleNamespace(transaction=transaction, execute=AsyncMock(side_effect=UndefinedTable('calculation_jobs missing')))
|
||||
monkeypatch.setattr(storage, 'lock_event_transaction', AsyncMock())
|
||||
repo = OperationalTaskRepository(connection, uuid4())
|
||||
with pytest.raises(DomainError) as error:
|
||||
async with repo.transaction():
|
||||
await repo.pending_calculation_jobs(str(uuid4()), str(uuid4()), str(uuid4()))
|
||||
assert error.value.code == 'DEPENDENCY_UNAVAILABLE'
|
||||
assert error.value.status == 503
|
||||
|
||||
Reference in New Issue
Block a user