feat(mcp): queue long-running calculations

This commit is contained in:
2026-10-29 18:25:15 +08:00
parent 2865811f1b
commit 63af998045
32 changed files with 1588 additions and 72 deletions
+26 -46
View File
@@ -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():