feat(mcp): queue long-running calculations
This commit is contained in:
+12
-1
@@ -145,5 +145,16 @@ TJWATER_MCP_ALLOWED_ORIGINS='[]'
|
|||||||
TJWATER_MCP_CONTEXT_TTL_SECONDS=600
|
TJWATER_MCP_CONTEXT_TTL_SECONDS=600
|
||||||
# 普通查询预算;计算工具由 Agent 通过 timeout_seconds 按次指定。
|
# 普通查询预算;计算工具由 Agent 通过 timeout_seconds 按次指定。
|
||||||
TJWATER_MCP_QUERY_TIMEOUT_SECONDS=30
|
TJWATER_MCP_QUERY_TIMEOUT_SECONDS=30
|
||||||
# 计算等待上限为 120–86400 秒,实际上限通过 tools/list 发布。
|
# 后台计算执行上限为 120–86400 秒,实际上限通过 tools/list 发布;不延长 HTTP 超时。
|
||||||
TJWATER_MCP_MAX_CALCULATION_TIMEOUT_SECONDS=3600
|
TJWATER_MCP_MAX_CALCULATION_TIMEOUT_SECONDS=3600
|
||||||
|
|
||||||
|
# 独立计算队列消费者;显式项目 UUID,以逗号分隔。不会在 API lifespan 内启动。
|
||||||
|
# 先在每个目标项目 BizDB 应用 026_calculation_jobs.sql,再启动计算服务。
|
||||||
|
CALCULATION_WORKER_PROJECT_IDS=""
|
||||||
|
# 所有配置项目合计并发,独立于 Uvicorn --workers 4;单实例锁防止重复消费者扩容。
|
||||||
|
CALCULATION_WORKER_MAX_CONCURRENCY=2
|
||||||
|
CALCULATION_WORKER_POLL_SECONDS=2
|
||||||
|
CALCULATION_WORKER_LEASE_SECONDS=30
|
||||||
|
CALCULATION_WORKER_HEARTBEAT_SECONDS=5
|
||||||
|
# 持久化本机恢复账本;容器模板已挂载独立 volume,不能当缓存删除。
|
||||||
|
CALCULATION_WORKER_STATE_DIR="temp/calculation-jobs"
|
||||||
|
|||||||
+1
-1
@@ -23,7 +23,7 @@ COPY app ./app
|
|||||||
COPY contracts ./contracts
|
COPY contracts ./contracts
|
||||||
COPY infra ./infra
|
COPY infra ./infra
|
||||||
COPY resources ./resources
|
COPY resources ./resources
|
||||||
COPY scripts/run_operational_worker.py scripts/sync_peripheral_work_orders.py ./scripts/
|
COPY scripts/run_operational_worker.py scripts/run_calculation_worker.py scripts/sync_peripheral_work_orders.py ./scripts/
|
||||||
COPY samples/peripheral ./samples/peripheral
|
COPY samples/peripheral ./samples/peripheral
|
||||||
RUN python -c "from pathlib import Path; from zipfile import ZipFile; model_dir = Path('app/algorithms/pipe_health_prediction/model'); zip_path = model_dir / 'my_survival_forest_model_quxi.zip'; joblib_name = 'my_survival_forest_model_quxi.joblib'; joblib_path = model_dir / joblib_name; assert zip_path.exists(), f'Model archive not found: {zip_path}'; archive = ZipFile(zip_path); archive.extract(joblib_name, model_dir); archive.close(); assert joblib_path.exists(), f'Model file not extracted: {joblib_path}'" && \
|
RUN python -c "from pathlib import Path; from zipfile import ZipFile; model_dir = Path('app/algorithms/pipe_health_prediction/model'); zip_path = model_dir / 'my_survival_forest_model_quxi.zip'; joblib_name = 'my_survival_forest_model_quxi.joblib'; joblib_path = model_dir / joblib_name; assert zip_path.exists(), f'Model archive not found: {zip_path}'; archive = ZipFile(zip_path); archive.extract(joblib_name, model_dir); archive.close(); assert joblib_path.exists(), f'Model file not extracted: {joblib_path}'" && \
|
||||||
rm -f app/algorithms/pipe_health_prediction/model/my_survival_forest_model_quxi.zip
|
rm -f app/algorithms/pipe_health_prediction/model/my_survival_forest_model_quxi.zip
|
||||||
|
|||||||
@@ -195,3 +195,7 @@ Gitea Actions 工作流位于 `.gitea/workflows/package.yml`,可手动触发
|
|||||||
Server 可选启用 MCP 2026-07-28,默认关闭。当前目录包含 67 项类型化工具,其中 42 项查询、12 项计算、13 项预览式写入。MCP 使用专用 OAuth resource Bearer、宿主绑定上下文和现有项目或任务权限,不提供自由命令执行。
|
Server 可选启用 MCP 2026-07-28,默认关闭。当前目录包含 67 项类型化工具,其中 42 项查询、12 项计算、13 项预览式写入。MCP 使用专用 OAuth resource Bearer、宿主绑定上下文和现有项目或任务权限,不提供自由命令执行。
|
||||||
|
|
||||||
写入工具默认只执行 `dry_run`。真实写入必须使用同一 `request_id`,带上预览返回的 `expected_digest`,并通过宿主原生许可和项目权限检查。超时或结果不明时先查询操作状态,不盲目重试。实施、启用和授权边界见 [MCP 实施与启用](docs/mcp-query-implementation.md) 和 [协议与旧 CLI 能力对照](docs/mcp-query-design.md)。
|
写入工具默认只执行 `dry_run`。真实写入必须使用同一 `request_id`,带上预览返回的 `expected_digest`,并通过宿主原生许可和项目权限检查。超时或结果不明时先查询操作状态,不盲目重试。实施、启用和授权边界见 [MCP 实施与启用](docs/mcp-query-implementation.md) 和 [协议与旧 CLI 能力对照](docs/mcp-query-design.md)。
|
||||||
|
|
||||||
|
## 后台计算队列
|
||||||
|
|
||||||
|
长计算通过项目数据库队列和独立计算 Worker 执行,API 保持4个进程,HTTP超时不延长。接入需要新增数据库迁移和独立服务,详见[计算队列部署与MCP契约](docs/calculation-queue.md)。
|
||||||
|
|||||||
@@ -3,6 +3,14 @@ import os
|
|||||||
import pandas as pd
|
import pandas as pd
|
||||||
|
|
||||||
|
|
||||||
|
def calculation_temp_directory(kind: str) -> str:
|
||||||
|
"""Keep native outputs inside the supervisor-owned job directory."""
|
||||||
|
root = os.environ.get("TJWATER_CALCULATION_JOB_DIR", "temp")
|
||||||
|
path = os.path.abspath(os.path.join(root, kind))
|
||||||
|
os.makedirs(path, exist_ok=True)
|
||||||
|
return path
|
||||||
|
|
||||||
|
|
||||||
def fill_time_gaps(
|
def fill_time_gaps(
|
||||||
data: pd.DataFrame,
|
data: pd.DataFrame,
|
||||||
time_col: str = "time",
|
time_col: str = "time",
|
||||||
|
|||||||
@@ -8,15 +8,14 @@ import sys
|
|||||||
import pandas as pd
|
import pandas as pd
|
||||||
import wntr
|
import wntr
|
||||||
|
|
||||||
from app.algorithms._utils import _cleanup_temp_files
|
from app.algorithms._utils import _cleanup_temp_files, calculation_temp_directory
|
||||||
|
|
||||||
_PIPE2LEAKNODE = None
|
_PIPE2LEAKNODE = None
|
||||||
_SIGNATURE_WORKER_DATA = {}
|
_SIGNATURE_WORKER_DATA = {}
|
||||||
|
|
||||||
|
|
||||||
def _make_temp_prefix(tag):
|
def _make_temp_prefix(tag):
|
||||||
temp_dir = os.path.abspath(os.path.join("temp", "burst_location"))
|
temp_dir = calculation_temp_directory("burst_location")
|
||||||
os.makedirs(temp_dir, exist_ok=True)
|
|
||||||
safe_tag = str(tag).replace(os.sep, "_").replace(" ", "_")
|
safe_tag = str(tag).replace(os.sep, "_").replace(" ", "_")
|
||||||
return os.path.join(temp_dir, f"{safe_tag}_{os.getpid()}")
|
return os.path.join(temp_dir, f"{safe_tag}_{os.getpid()}")
|
||||||
|
|
||||||
|
|||||||
@@ -15,7 +15,7 @@ from pymoo.optimize import minimize as pymoo_minimize
|
|||||||
from pymoo.termination.default import DefaultSingleObjectiveTermination
|
from pymoo.termination.default import DefaultSingleObjectiveTermination
|
||||||
|
|
||||||
|
|
||||||
from app.algorithms._utils import _cleanup_temp_files
|
from app.algorithms._utils import _cleanup_temp_files, calculation_temp_directory
|
||||||
|
|
||||||
_worker_data: dict[str, Any] = {}
|
_worker_data: dict[str, Any] = {}
|
||||||
DEFAULT_N_WORKERS = max(1, min(cpu_count() - 1, 4))
|
DEFAULT_N_WORKERS = max(1, min(cpu_count() - 1, 4))
|
||||||
@@ -92,7 +92,7 @@ def _worker_evaluate(raw_ratios: np.ndarray) -> float:
|
|||||||
demand_obj.base_value = original_val + per_node_leak
|
demand_obj.base_value = original_val + per_node_leak
|
||||||
modifications.append((demand_obj, original_val))
|
modifications.append((demand_obj, original_val))
|
||||||
|
|
||||||
temp_dir = os.path.abspath(os.path.join("temp", "leakage"))
|
temp_dir = calculation_temp_directory("leakage")
|
||||||
os.makedirs(temp_dir, exist_ok=True)
|
os.makedirs(temp_dir, exist_ok=True)
|
||||||
prefix = os.path.join(temp_dir, f"temp_{os.getpid()}")
|
prefix = os.path.join(temp_dir, f"temp_{os.getpid()}")
|
||||||
|
|
||||||
@@ -569,7 +569,7 @@ class LeakageProblem(Problem):
|
|||||||
modifications.append((demand_obj, original_val))
|
modifications.append((demand_obj, original_val))
|
||||||
|
|
||||||
# 结果保存在根目录的temp/leakage文件夹中
|
# 结果保存在根目录的temp/leakage文件夹中
|
||||||
temp_dir = os.path.abspath(os.path.join("temp", "leakage"))
|
temp_dir = calculation_temp_directory("leakage")
|
||||||
os.makedirs(temp_dir, exist_ok=True)
|
os.makedirs(temp_dir, exist_ok=True)
|
||||||
prefix = os.path.join(temp_dir, f"temp_{os.getpid()}")
|
prefix = os.path.join(temp_dir, f"temp_{os.getpid()}")
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,103 @@
|
|||||||
|
"""Durable project calculation queue. Every worker update is generation fenced."""
|
||||||
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from fastapi import HTTPException
|
||||||
|
from psycopg.types.json import Jsonb
|
||||||
|
|
||||||
|
|
||||||
|
class CalculationJobRepository:
|
||||||
|
def __init__(self, connection, project_id):
|
||||||
|
self.connection = connection
|
||||||
|
self.project_id = str(project_id)
|
||||||
|
|
||||||
|
async def enqueue(self, *, owner_user_id, request_key, tool, parameters, binding, timeout_seconds):
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"""INSERT INTO workbench.calculation_jobs
|
||||||
|
(job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds)
|
||||||
|
VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
|
||||||
|
ON CONFLICT(project_id,owner_user_id,request_key) DO NOTHING RETURNING *""",
|
||||||
|
(uuid4(), self.project_id, owner_user_id, request_key, tool,
|
||||||
|
Jsonb(parameters), Jsonb(binding), timeout_seconds),
|
||||||
|
)
|
||||||
|
row = await cursor.fetchone()
|
||||||
|
if row:
|
||||||
|
return dict(row)
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s AND request_key=%s",
|
||||||
|
(self.project_id, owner_user_id, request_key),
|
||||||
|
)
|
||||||
|
row = dict(await cursor.fetchone())
|
||||||
|
if row['tool'] != tool or row['parameters'] != parameters or row['binding'] != binding:
|
||||||
|
raise HTTPException(409, 'Calculation request_key was used with different inputs')
|
||||||
|
return row
|
||||||
|
|
||||||
|
async def get(self, job_id, owner_user_id=None):
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND job_id=%s
|
||||||
|
AND (%s::uuid IS NULL OR owner_user_id=%s)""",
|
||||||
|
(self.project_id, job_id, owner_user_id, owner_user_id),
|
||||||
|
)
|
||||||
|
row = await cursor.fetchone()
|
||||||
|
return dict(row) if row else None
|
||||||
|
|
||||||
|
async def list(self, *, owner_user_id, binding, limit, offset, status=None):
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s
|
||||||
|
AND binding @> %s AND (%s::text IS NULL OR status=%s) ORDER BY created_at DESC,job_id LIMIT %s OFFSET %s""",
|
||||||
|
(self.project_id, owner_user_id, Jsonb(binding), status, status, limit, offset),
|
||||||
|
)
|
||||||
|
return [dict(row) for row in await cursor.fetchall()]
|
||||||
|
|
||||||
|
async def cancel(self, job_id, owner_user_id):
|
||||||
|
await self.connection.execute(
|
||||||
|
"""UPDATE workbench.calculation_jobs SET status=CASE WHEN status='queued' THEN 'cancelled'
|
||||||
|
ELSE 'cancel_requested' END,updated_at=clock_timestamp()
|
||||||
|
WHERE project_id=%s AND job_id=%s AND owner_user_id=%s AND status IN ('queued','running')""",
|
||||||
|
(self.project_id, job_id, owner_user_id),
|
||||||
|
)
|
||||||
|
return await self.get(job_id, owner_user_id)
|
||||||
|
|
||||||
|
async def recover(self):
|
||||||
|
# A crashed write is never automatically replayed.
|
||||||
|
await self.connection.execute(
|
||||||
|
"""UPDATE workbench.calculation_jobs SET status='interrupted',
|
||||||
|
error='{"code":"WORKER_INTERRUPTED"}'::jsonb,updated_at=clock_timestamp()
|
||||||
|
WHERE project_id=%s AND status IN ('running','cancel_requested')
|
||||||
|
AND lease_expires_at < clock_timestamp()""", (self.project_id,),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def claim(self, owner, lease_seconds):
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"""UPDATE workbench.calculation_jobs j SET status='running',lease_owner=%s,
|
||||||
|
lease_generation=j.lease_generation+1,
|
||||||
|
lease_expires_at=clock_timestamp()+%s*interval '1 second',updated_at=clock_timestamp()
|
||||||
|
WHERE j.job_id=(SELECT job_id FROM workbench.calculation_jobs
|
||||||
|
WHERE project_id=%s AND status='queued' ORDER BY created_at,job_id
|
||||||
|
FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING j.*""",
|
||||||
|
(owner, lease_seconds, self.project_id),
|
||||||
|
)
|
||||||
|
row = await cursor.fetchone()
|
||||||
|
return dict(row) if row else None
|
||||||
|
|
||||||
|
async def heartbeat(self, row, owner, lease_seconds):
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"""UPDATE workbench.calculation_jobs SET lease_expires_at=clock_timestamp()+%s*interval '1 second',
|
||||||
|
updated_at=clock_timestamp() WHERE project_id=%s AND job_id=%s AND lease_owner=%s
|
||||||
|
AND lease_generation=%s AND lease_expires_at>clock_timestamp()
|
||||||
|
AND status IN ('running','cancel_requested') RETURNING status""",
|
||||||
|
(lease_seconds, self.project_id, row['job_id'], owner, row['lease_generation']),
|
||||||
|
)
|
||||||
|
result = await cursor.fetchone()
|
||||||
|
return result['status'] if result else None
|
||||||
|
|
||||||
|
async def finish(self, row, owner, status, *, result=None, error=None):
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"""UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,updated_at=clock_timestamp()
|
||||||
|
WHERE project_id=%s AND job_id=%s AND lease_owner=%s AND lease_generation=%s
|
||||||
|
AND lease_expires_at>clock_timestamp() AND status IN ('running','cancel_requested')
|
||||||
|
AND (status!='cancel_requested' OR %s IN ('cancelled','interrupted')
|
||||||
|
OR (%s='failed' AND %s='CLEANUP_FAILED')) RETURNING job_id""",
|
||||||
|
(status, Jsonb(result) if result is not None else None, Jsonb(error) if error else None,
|
||||||
|
self.project_id, row['job_id'], owner, row['lease_generation'], status, status, (error or {}).get('code')),
|
||||||
|
)
|
||||||
|
return await cursor.fetchone() is not None
|
||||||
@@ -75,6 +75,20 @@ class OperationalTaskRepository:
|
|||||||
status=409,
|
status=409,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def pending_calculation_jobs(
|
||||||
|
self, task_id: str, run_id: str, execution_request_id: str,
|
||||||
|
) -> list[str]:
|
||||||
|
"""Only the exact original run blocks final artifacts and completion."""
|
||||||
|
cursor = await self.connection.execute(
|
||||||
|
"""SELECT job_id FROM workbench.calculation_jobs
|
||||||
|
WHERE project_id=%s AND binding->>'task_id'=%s
|
||||||
|
AND binding->>'run_id'=%s AND binding->>'execution_request_id'=%s
|
||||||
|
AND status IN ('queued','running','cancel_requested')
|
||||||
|
ORDER BY created_at,job_id LIMIT 100""",
|
||||||
|
(self.project_id, task_id, run_id, execution_request_id),
|
||||||
|
)
|
||||||
|
return [str(row["job_id"]) for row in await cursor.fetchall()]
|
||||||
|
|
||||||
async def get_task(self, task_id: str) -> dict | None:
|
async def get_task(self, task_id: str) -> dict | None:
|
||||||
cursor = await self.connection.execute(
|
cursor = await self.connection.execute(
|
||||||
"""SELECT document FROM workbench.operational_tasks
|
"""SELECT document FROM workbench.operational_tasks
|
||||||
|
|||||||
+12
-7
@@ -347,20 +347,25 @@ class McpQueryApplication:
|
|||||||
)
|
)
|
||||||
result = QueryResult.model_validate(previous["result"])
|
result = QueryResult.model_validate(previous["result"])
|
||||||
else:
|
else:
|
||||||
result = (
|
from app.mcp.calculation_jobs import should_queue, submit_calculation
|
||||||
await self.executor(
|
if should_queue(tool, args):
|
||||||
scope, binding, tool, args, full_snapshot=True
|
result = await submit_calculation(scope, binding, tool, args, call_id)
|
||||||
|
else:
|
||||||
|
result = (
|
||||||
|
await self.executor(
|
||||||
|
scope, binding, tool, args, full_snapshot=True
|
||||||
|
)
|
||||||
|
if include_full
|
||||||
|
else await self.executor(scope, binding, tool, args)
|
||||||
)
|
)
|
||||||
if include_full
|
|
||||||
else await self.executor(scope, binding, tool, args)
|
|
||||||
)
|
|
||||||
# Never serialize an unbounded response into a model message or database receipt.
|
# Never serialize an unbounded response into a model message or database receipt.
|
||||||
if len(result.model_dump_json().encode()) > 8 * 1024 * 1024:
|
if len(result.model_dump_json().encode()) > 8 * 1024 * 1024:
|
||||||
return tool_error(
|
return tool_error(
|
||||||
"RESULT_TOO_LARGE",
|
"RESULT_TOO_LARGE",
|
||||||
"Narrow the query or request a smaller page",
|
"Narrow the query or request a smaller page",
|
||||||
)
|
)
|
||||||
if binding.task_id and name not in {
|
if binding.task_id and not should_queue(tool, args) and name not in {
|
||||||
|
"get_calculation_job", "get_calculation_result", "cancel_calculation_job", "list_calculation_jobs",
|
||||||
"get_task_context",
|
"get_task_context",
|
||||||
"get_task_evidence",
|
"get_task_evidence",
|
||||||
}:
|
}:
|
||||||
|
|||||||
@@ -0,0 +1,196 @@
|
|||||||
|
"""Durable MCP calculation submission and owner-scoped result retrieval.
|
||||||
|
|
||||||
|
Only safe server-owned binding fields are persisted. OAuth tokens and context
|
||||||
|
handles are never used to resume a calculation.
|
||||||
|
"""
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from uuid import uuid5, NAMESPACE_URL
|
||||||
|
from fastapi import HTTPException
|
||||||
|
from app.mcp.models import BoundContext, QueryResult, PageInfo
|
||||||
|
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
||||||
|
|
||||||
|
QUEUED_TOOLS = frozenset({
|
||||||
|
"run_burst_location", "run_leakage_identification", "run_burst_detection",
|
||||||
|
"run_simulation", "run_burst_simulation", "run_contaminant_simulation",
|
||||||
|
"analyze_valve_closure", "analyze_flushing", "analyze_pressure_regulation",
|
||||||
|
"optimize_sensor_placement", "save_simulation_scenario", "create_sensor_placement",
|
||||||
|
})
|
||||||
|
CONTROL_TOOLS = frozenset({"get_calculation_job", "get_calculation_result", "cancel_calculation_job", "list_calculation_jobs"})
|
||||||
|
BINDING_FIELDS = {"project_id", "subject", "client_id", "native_session_id", "purpose", "parent_session_id", "task_id", "run_id", "execution_request_id", "tool_cap", "permission_cap", "task_scope_cap"}
|
||||||
|
|
||||||
|
|
||||||
|
def durable_binding(binding: BoundContext) -> dict:
|
||||||
|
return binding.model_dump(mode="json", include=BINDING_FIELDS)
|
||||||
|
|
||||||
|
|
||||||
|
def restore_binding(document: dict) -> BoundContext:
|
||||||
|
# Expiration governs admission via a context handle; durable work has its own
|
||||||
|
# lifecycle and must revalidate identity, membership and run on every use.
|
||||||
|
return BoundContext.model_validate({**document, "context_handle": "durable-calculation", "expires_at": datetime.max.replace(tzinfo=timezone.utc)})
|
||||||
|
|
||||||
|
|
||||||
|
def should_queue(tool, args) -> bool:
|
||||||
|
return tool.name in QUEUED_TOOLS and (not tool.mutating or not args.dry_run)
|
||||||
|
|
||||||
|
|
||||||
|
def envelope(source: str, data, *, job_id="list", page=None) -> QueryResult:
|
||||||
|
return QueryResult(source=source, source_ref="calculation:" + str(job_id), queried_at=datetime.now(timezone.utc), semantics="task", data=data, page=page, truncated=bool(page and (page.offset or page.has_more)),
|
||||||
|
data_gaps=["Calculation status and queue receipts are not analytical evidence."])
|
||||||
|
|
||||||
|
|
||||||
|
def public_job(row: dict) -> dict:
|
||||||
|
fields = {"job_id", "status", "tool", "created_at", "started_at", "finished_at", "updated_at", "error"}
|
||||||
|
from fastapi.encoders import jsonable_encoder
|
||||||
|
return {**jsonable_encoder({key: row[key] for key in fields if key in row}), "poll_after_ms": 1500}
|
||||||
|
|
||||||
|
|
||||||
|
async def submit_calculation(scope, binding, tool, args, call_id: str) -> QueryResult:
|
||||||
|
from app.mcp.query_service import authorize_query
|
||||||
|
await authorize_query(scope, binding, tool, args)
|
||||||
|
if binding.purpose == "task_reader":
|
||||||
|
raise HTTPException(403, "Read-only task agents cannot start calculations")
|
||||||
|
if tool.mutating:
|
||||||
|
from app.mcp.mutation_service import (
|
||||||
|
preview_change, check_digest, mutation_arguments, retryable_capacity_rejection,
|
||||||
|
)
|
||||||
|
from app.infra.db.postgresql.mcp_queries import McpQueryRepository
|
||||||
|
mutation_repo = McpQueryRepository(scope.connection, scope.project.project_id)
|
||||||
|
mutation_key = "mutation:" + binding.subject + ":" + binding.native_session_id
|
||||||
|
approved_parameters = {"tool": tool.name, "arguments": mutation_arguments(args, preview=False)}
|
||||||
|
async with scope.connection.transaction():
|
||||||
|
await mutation_repo.lock_receipt(mutation_key, str(args.request_id))
|
||||||
|
previous = await mutation_repo.receipt(mutation_key, str(args.request_id))
|
||||||
|
if previous:
|
||||||
|
if previous["parameters"] != approved_parameters:
|
||||||
|
raise HTTPException(409, "request_id already has different parameters")
|
||||||
|
if previous["result"].get("status") == "completed":
|
||||||
|
return QueryResult.model_validate(previous["result"]["output"])
|
||||||
|
if not retryable_capacity_rejection(previous["result"]):
|
||||||
|
raise HTTPException(409, "Operation already started; inspect the existing calculation or mutation status")
|
||||||
|
check_digest(args, await preview_change(scope, binding, tool.name, args))
|
||||||
|
parameters = args.model_dump(mode="json", exclude={"context_handle"})
|
||||||
|
request_id = str(args.request_id) if tool.mutating else call_id
|
||||||
|
key = str(uuid5(NAMESPACE_URL, ":".join([binding.subject, binding.client_id, binding.native_session_id, str(binding.task_id), str(binding.run_id), request_id])))
|
||||||
|
repo = CalculationJobRepository(scope.connection, scope.project.project_id)
|
||||||
|
if binding.task_id:
|
||||||
|
from app.mcp.query_service import task_service
|
||||||
|
from app.mcp.contexts import verify_task
|
||||||
|
task_repo = task_service(scope).repo
|
||||||
|
# Share the operational transaction lock with finalization. Admission
|
||||||
|
# cannot race a pending-job check and then publish work into a closed run.
|
||||||
|
async with task_repo.transaction():
|
||||||
|
await task_repo.lock(f"task:{binding.task_id}")
|
||||||
|
await task_repo.lock(f"run:{binding.run_id}")
|
||||||
|
await verify_task(scope, binding)
|
||||||
|
row = await repo.enqueue(
|
||||||
|
owner_user_id=scope.project.user_id, request_key=key, tool=tool.name,
|
||||||
|
parameters=parameters, binding=durable_binding(binding), timeout_seconds=args.timeout_seconds,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
row = await repo.enqueue(
|
||||||
|
owner_user_id=scope.project.user_id, request_key=key, tool=tool.name,
|
||||||
|
parameters=parameters, binding=durable_binding(binding), timeout_seconds=args.timeout_seconds,
|
||||||
|
)
|
||||||
|
await scope.connection.commit()
|
||||||
|
return envelope(tool.name, public_job(row), job_id=row["job_id"])
|
||||||
|
|
||||||
|
|
||||||
|
def check_binding(binding, row):
|
||||||
|
original = row["binding"]
|
||||||
|
if any(str(original.get(key)) != str(getattr(binding, key)) for key in ("subject", "client_id", "native_session_id", "task_id", "run_id", "execution_request_id")):
|
||||||
|
raise HTTPException(403, "Calculation belongs to a different session or task run")
|
||||||
|
|
||||||
|
|
||||||
|
async def authorize_original(scope, row):
|
||||||
|
from app.mcp.catalog import TOOLS
|
||||||
|
from app.mcp.query_service import authorize_query
|
||||||
|
from app.mcp.contexts import verify_task
|
||||||
|
original = restore_binding(row["binding"])
|
||||||
|
tool = TOOLS[row["tool"]]
|
||||||
|
args = tool.model.model_validate({**row["parameters"], "context_handle": "durable-calculation-context-handle"})
|
||||||
|
if (str(scope.project.user_id) != str(row["owner_user_id"])
|
||||||
|
or str(scope.project.project_id) != str(original.project_id)):
|
||||||
|
raise HTTPException(403, "Calculation owner is no longer authorized")
|
||||||
|
if original.task_id:
|
||||||
|
await verify_task(scope, original)
|
||||||
|
await authorize_query(scope, original, tool, args)
|
||||||
|
return original, tool, args
|
||||||
|
|
||||||
|
|
||||||
|
async def read_calculation_job(scope, binding, tool, args) -> QueryResult:
|
||||||
|
repo = CalculationJobRepository(scope.connection, scope.project.project_id)
|
||||||
|
if tool.name == "list_calculation_jobs":
|
||||||
|
identity_scope = binding.model_dump(mode="json", include={
|
||||||
|
"subject", "client_id", "native_session_id", "task_id", "run_id", "execution_request_id",
|
||||||
|
})
|
||||||
|
rows = await repo.list(owner_user_id=scope.project.user_id, binding=identity_scope, limit=args.limit + 1, offset=args.offset, status=args.status)
|
||||||
|
return envelope(tool.name, [public_job(row) for row in rows[:args.limit]], page=PageInfo(limit=args.limit, offset=args.offset, has_more=len(rows) > args.limit))
|
||||||
|
row = await repo.get(args.job_id, owner_user_id=scope.project.user_id)
|
||||||
|
if row is None:
|
||||||
|
raise HTTPException(404, "Calculation was not found")
|
||||||
|
check_binding(binding, row)
|
||||||
|
await authorize_original(scope, row)
|
||||||
|
if tool.name == "cancel_calculation_job":
|
||||||
|
if binding.purpose == "task_reader":
|
||||||
|
raise HTTPException(403, "Read-only task agents cannot cancel calculations")
|
||||||
|
row = await repo.cancel(args.job_id, owner_user_id=scope.project.user_id)
|
||||||
|
await scope.connection.commit()
|
||||||
|
if tool.name != "get_calculation_result" or row["status"] != "succeeded":
|
||||||
|
return envelope(tool.name, public_job(row), job_id=row["job_id"])
|
||||||
|
result = QueryResult.model_validate(row["result"])
|
||||||
|
result = await attach_calculation_evidence(scope, row, result)
|
||||||
|
value = result.data
|
||||||
|
for part in args.data_path:
|
||||||
|
if isinstance(value, dict) and part in value:
|
||||||
|
value = value[part]
|
||||||
|
else:
|
||||||
|
raise HTTPException(422, "data_path must select an existing object field")
|
||||||
|
if isinstance(value, list):
|
||||||
|
result.data = value[args.offset:args.offset + args.limit]
|
||||||
|
result.page = PageInfo(limit=args.limit, offset=args.offset, total=len(value), has_more=args.offset + args.limit < len(value))
|
||||||
|
result.truncated = result.truncated or args.offset > 0 or result.page.has_more
|
||||||
|
elif args.offset:
|
||||||
|
raise HTTPException(422, "offset requires a collection data_path")
|
||||||
|
else:
|
||||||
|
result.data = value
|
||||||
|
if len(result.model_dump_json().encode()) > 11000:
|
||||||
|
raise HTTPException(413, "Select a narrower data_path or smaller result page")
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
async def execute_calculation_job(scope, row) -> QueryResult:
|
||||||
|
"""Worker entry point; success is persisted by the fenced supervisor."""
|
||||||
|
from app.mcp.query_service import execute_query
|
||||||
|
original, tool, args = await authorize_original(scope, row)
|
||||||
|
result = await execute_query(scope, original, tool, args)
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
async def attach_calculation_evidence(scope, row, result):
|
||||||
|
from app.mcp.query_service import task_service
|
||||||
|
from app.mcp.contexts import verify_task
|
||||||
|
original, tool, _ = await authorize_original(scope, row)
|
||||||
|
if original.task_id:
|
||||||
|
from app.domain.schemas.operational_task import TaskDataSnapshot
|
||||||
|
from app.infra.db.postgresql.mcp_queries import McpQueryRepository
|
||||||
|
repo = McpQueryRepository(scope.connection, scope.project.project_id)
|
||||||
|
key, call_id = "calculation:" + str(row["job_id"]), "result"
|
||||||
|
# Result lookup can leave an implicit SELECT transaction open. End it so
|
||||||
|
# evidence and its receipt commit together before returning a response.
|
||||||
|
await scope.connection.commit()
|
||||||
|
task_repo = task_service(scope).repo
|
||||||
|
# Match operational finalization's lock order: event lock, then task.
|
||||||
|
async with task_repo.transaction():
|
||||||
|
await repo.lock_receipt(key, call_id)
|
||||||
|
previous = await repo.receipt(key, call_id)
|
||||||
|
if previous:
|
||||||
|
return QueryResult.model_validate(previous["result"])
|
||||||
|
await task_repo.lock(f"task:{original.task_id}")
|
||||||
|
await verify_task(scope, original)
|
||||||
|
result.source_ref = key
|
||||||
|
evidence, _ = await task_service(scope).record_data_snapshot(original.task_id, TaskDataSnapshot(
|
||||||
|
source=tool.name, source_ref=key, queried_at=result.queried_at, origin=result.origin,
|
||||||
|
data=result.model_dump(mode="json"), truncated=result.truncated))
|
||||||
|
result.evidence_id = evidence.evidence_id
|
||||||
|
await repo.save_receipt(key, call_id, {"job_id": str(row["job_id"])}, result.model_dump(mode="json"))
|
||||||
|
return result
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
"""Bounded inputs for durable computation control."""
|
||||||
|
from typing import Literal
|
||||||
|
from uuid import UUID
|
||||||
|
from pydantic import Field
|
||||||
|
from app.mcp.models import StrictInput, PageInput, Id
|
||||||
|
|
||||||
|
CalculationStatus = Literal["queued", "running", "cancel_requested", "succeeded", "failed", "cancelled", "interrupted", "timed_out"]
|
||||||
|
|
||||||
|
|
||||||
|
class CalculationJobInput(StrictInput):
|
||||||
|
job_id: UUID
|
||||||
|
|
||||||
|
|
||||||
|
class CalculationResultInput(CalculationJobInput):
|
||||||
|
data_path: list[Id] = Field(default_factory=list, max_length=8)
|
||||||
|
limit: int = Field(default=20, ge=1, le=500)
|
||||||
|
offset: int = Field(default=0, ge=0, le=1000000)
|
||||||
|
|
||||||
|
|
||||||
|
class CalculationJobsInput(PageInput):
|
||||||
|
status: CalculationStatus | None = None
|
||||||
+13
-2
@@ -4,6 +4,7 @@ from dataclasses import dataclass
|
|||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
from mcp import types
|
from mcp import types
|
||||||
from app.mcp import models as m
|
from app.mcp import models as m
|
||||||
|
from app.mcp import calculation_models as jobs
|
||||||
from app.mcp import extended_models as e
|
from app.mcp import extended_models as e
|
||||||
from app.mcp.analysis_models import BurstDetectionInput, BurstLocationInput, ValveIsolationInput, LeakageIdentificationInput
|
from app.mcp.analysis_models import BurstDetectionInput, BurstLocationInput, ValveIsolationInput, LeakageIdentificationInput
|
||||||
from app.mcp.analysis_models import ScenarioInput, ContaminantSimulationInput, ValveClosureInput, FlushingInput, PressureRegulationInput, SensorPlacementInput
|
from app.mcp.analysis_models import ScenarioInput, ContaminantSimulationInput, ValveClosureInput, FlushingInput, PressureRegulationInput, SensorPlacementInput
|
||||||
@@ -36,11 +37,15 @@ class QueryTool:
|
|||||||
)
|
)
|
||||||
return types.Tool(
|
return types.Tool(
|
||||||
name=self.name,
|
name=self.name,
|
||||||
description=self.description + suffix,
|
description=self.description + suffix + (
|
||||||
|
"。执行返回持久job_id;排队回执不是分析证据,使用get_calculation_job查询状态、get_calculation_result分页读取成功结果,不重复提交。"
|
||||||
|
if self.name in {"run_burst_location", "run_leakage_identification", "run_burst_detection", "run_simulation", "run_burst_simulation", "run_contaminant_simulation", "analyze_valve_closure", "analyze_flushing", "analyze_pressure_regulation", "optimize_sensor_placement", "save_simulation_scenario", "create_sensor_placement"}
|
||||||
|
else ""
|
||||||
|
),
|
||||||
inputSchema=self.model.model_json_schema(),
|
inputSchema=self.model.model_json_schema(),
|
||||||
outputSchema=m.QueryResult.model_json_schema(),
|
outputSchema=m.QueryResult.model_json_schema(),
|
||||||
annotations=types.ToolAnnotations(
|
annotations=types.ToolAnnotations(
|
||||||
readOnlyHint=not (self.computation or self.mutating), destructiveHint=self.mutating or self.name == "run_simulation", openWorldHint=self.name in {"search_project_places", "reverse_geocode"}
|
readOnlyHint=not (self.computation or self.mutating or self.name == "cancel_calculation_job"), destructiveHint=self.mutating or self.name == "run_simulation", openWorldHint=self.name in {"search_project_places", "reverse_geocode"}
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -256,4 +261,10 @@ _specs += [
|
|||||||
("retry_operational_task", "预览或重新入队已终止任务;不自动启动新Agent会话", e.TaskRetryInput, "operational_task.write", "task", False, False, True),
|
("retry_operational_task", "预览或重新入队已终止任务;不自动启动新Agent会话", e.TaskRetryInput, "operational_task.write", "task", False, False, True),
|
||||||
("cancel_task_run", "预览或取消指定任务运行,校验运行版本", e.TaskCancelInput, "operational_task.write", "task", False, False, True),
|
("cancel_task_run", "预览或取消指定任务运行,校验运行版本", e.TaskCancelInput, "operational_task.write", "task", False, False, True),
|
||||||
]
|
]
|
||||||
|
_specs += [
|
||||||
|
("get_calculation_job", "读取当前身份与会话的持久计算状态;排队回执不是分析证据", jobs.CalculationJobInput, "webgis.view", "task"),
|
||||||
|
("get_calculation_result", "读取成功计算的原始结果;data_path选择嵌套集合,limit/offset按需分页", jobs.CalculationResultInput, "webgis.view", "task"),
|
||||||
|
("cancel_calculation_job", "请求取消本人当前会话的计算;运行中取消由独立计算进程处理,须继续核实最终状态", jobs.CalculationJobInput, "webgis.view", "task"),
|
||||||
|
("list_calculation_jobs", "分页读取本人当前会话及任务范围的计算状态", jobs.CalculationJobsInput, "webgis.view", "task"),
|
||||||
|
]
|
||||||
TOOLS = {row[0]: QueryTool(*row) for row in _specs}
|
TOOLS = {row[0]: QueryTool(*row) for row in _specs}
|
||||||
|
|||||||
+1
-1
@@ -70,7 +70,7 @@ async def bind_context(
|
|||||||
subject=principal.subject,
|
subject=principal.subject,
|
||||||
client_id=principal.client_id,
|
client_id=principal.client_id,
|
||||||
tool_cap=sorted(name for name, tool in TOOLS.items()
|
tool_cap=sorted(name for name, tool in TOOLS.items()
|
||||||
if (command.purpose != "task_reader" or not (tool.computation or tool.mutating))
|
if (command.purpose != "task_reader" or not (tool.computation or tool.mutating or name == "cancel_calculation_job"))
|
||||||
and (not tool.mutating or command.purpose == "ordinary_chat")),
|
and (not tool.mutating or command.purpose == "ordinary_chat")),
|
||||||
permission_cap=sorted(permissions_for_context(scope.project)),
|
permission_cap=sorted(permissions_for_context(scope.project)),
|
||||||
native_session_id=command.native_session_id,
|
native_session_id=command.native_session_id,
|
||||||
|
|||||||
@@ -131,6 +131,9 @@ async def execute_query(
|
|||||||
full_snapshot: bool = False,
|
full_snapshot: bool = False,
|
||||||
) -> QueryResult:
|
) -> QueryResult:
|
||||||
await authorize_query(scope, binding, tool, args)
|
await authorize_query(scope, binding, tool, args)
|
||||||
|
if tool.name in {"get_calculation_job", "get_calculation_result", "cancel_calculation_job", "list_calculation_jobs"}:
|
||||||
|
from app.mcp.calculation_jobs import read_calculation_job
|
||||||
|
return await read_calculation_job(scope, binding, tool, args)
|
||||||
if tool.mutating:
|
if tool.mutating:
|
||||||
from app.mcp.mutation_service import execute_mutation
|
from app.mcp.mutation_service import execute_mutation
|
||||||
return await execute_mutation(scope, binding, tool, args)
|
return await execute_mutation(scope, binding, tool, args)
|
||||||
|
|||||||
@@ -201,6 +201,8 @@ def temporary_project_name(project: str, purpose: str) -> str:
|
|||||||
def temporary_project_database(project: str, purpose: str):
|
def temporary_project_database(project: str, purpose: str):
|
||||||
"""Clone one project's runnable model into an isolated temporary database."""
|
"""Clone one project's runnable model into an isolated temporary database."""
|
||||||
temporary_name = temporary_project_name(project, purpose)
|
temporary_name = temporary_project_name(project, purpose)
|
||||||
|
from app.services.calculation_cleanup import record_calculation_database
|
||||||
|
record_calculation_database(temporary_name)
|
||||||
try:
|
try:
|
||||||
copy_project(
|
copy_project(
|
||||||
get_project_template_database_name(project),
|
get_project_template_database_name(project),
|
||||||
@@ -211,12 +213,15 @@ def temporary_project_database(project: str, purpose: str):
|
|||||||
finally:
|
finally:
|
||||||
if have_project(temporary_name):
|
if have_project(temporary_name):
|
||||||
delete_project(temporary_name)
|
delete_project(temporary_name)
|
||||||
|
record_calculation_database(temporary_name, removed=True)
|
||||||
|
|
||||||
|
|
||||||
@contextmanager
|
@contextmanager
|
||||||
def temporary_template_database(name_hint: str, purpose: str):
|
def temporary_template_database(name_hint: str, purpose: str):
|
||||||
"""Create an empty schema-only temporary database from the fixed template."""
|
"""Create an empty schema-only temporary database from the fixed template."""
|
||||||
temporary_name = temporary_project_name(name_hint, purpose)
|
temporary_name = temporary_project_name(name_hint, purpose)
|
||||||
|
from app.services.calculation_cleanup import record_calculation_database
|
||||||
|
record_calculation_database(temporary_name)
|
||||||
try:
|
try:
|
||||||
copy_project(
|
copy_project(
|
||||||
get_schema_template_database_name(),
|
get_schema_template_database_name(),
|
||||||
@@ -227,6 +232,7 @@ def temporary_template_database(name_hint: str, purpose: str):
|
|||||||
finally:
|
finally:
|
||||||
if have_project(temporary_name):
|
if have_project(temporary_name):
|
||||||
delete_project(temporary_name)
|
delete_project(temporary_name)
|
||||||
|
record_calculation_database(temporary_name, removed=True)
|
||||||
|
|
||||||
|
|
||||||
def have_project(name: str) -> bool:
|
def have_project(name: str) -> bool:
|
||||||
|
|||||||
@@ -0,0 +1,74 @@
|
|||||||
|
"""Record only job-owned native databases so abrupt cancellation is recoverable."""
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
import re
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
|
||||||
|
def _names(directory: Path) -> list[str]:
|
||||||
|
path = directory / "databases.json"
|
||||||
|
if not path.exists():
|
||||||
|
return []
|
||||||
|
if path.is_symlink() or path.stat().st_size > 65536:
|
||||||
|
raise ValueError("Invalid calculation database ledger")
|
||||||
|
names = json.loads(path.read_text())
|
||||||
|
if not isinstance(names, list) or len(names) > 100 or any(
|
||||||
|
not isinstance(name, str) or not re.fullmatch(r"tjw_tmp_[a-z0-9_]+_[a-f0-9]{32}", name)
|
||||||
|
or len(name) > 63 for name in names
|
||||||
|
):
|
||||||
|
raise ValueError("Invalid calculation database ledger")
|
||||||
|
return names
|
||||||
|
|
||||||
|
|
||||||
|
def _write(directory: Path, names: list[str]) -> None:
|
||||||
|
path = directory / "databases.json"
|
||||||
|
staging = directory / "databases.json.new"
|
||||||
|
with staging.open("w") as output:
|
||||||
|
json.dump(names, output)
|
||||||
|
output.flush()
|
||||||
|
os.fsync(output.fileno())
|
||||||
|
staging.replace(path)
|
||||||
|
|
||||||
|
|
||||||
|
def record_calculation_database(name: str, *, removed: bool = False) -> None:
|
||||||
|
"""Called before clone creation and after successful normal cleanup."""
|
||||||
|
root = os.environ.get("TJWATER_CALCULATION_JOB_DIR")
|
||||||
|
if not root:
|
||||||
|
return
|
||||||
|
directory = Path(root)
|
||||||
|
names = _names(directory)
|
||||||
|
if removed:
|
||||||
|
names = [value for value in names if value != name]
|
||||||
|
elif name not in names:
|
||||||
|
names.append(name)
|
||||||
|
_write(directory, names)
|
||||||
|
|
||||||
|
|
||||||
|
async def cleanup_calculation_databases(project_id: str, job_directory: str) -> None:
|
||||||
|
directory = Path(job_directory)
|
||||||
|
names = _names(directory)
|
||||||
|
if not names:
|
||||||
|
return
|
||||||
|
from app.infra.db.metadb.database import SessionLocal
|
||||||
|
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
|
||||||
|
from app.infra.db.project_routing import ActiveProjectRouting, activate_project_routing
|
||||||
|
from app.native.wndb.core.projects import have_project, delete_project
|
||||||
|
async with SessionLocal() as session:
|
||||||
|
metadata = MetadataRepository(session)
|
||||||
|
project = await metadata.get_project_by_id(UUID(str(project_id)))
|
||||||
|
route = await metadata.get_project_db_routing(UUID(str(project_id)), "biz_data")
|
||||||
|
if not project or not route or route.db_type != "postgresql":
|
||||||
|
raise ValueError("Calculation cleanup routing unavailable")
|
||||||
|
routing = ActiveProjectRouting(project.code, route.dsn, project_id=project.id)
|
||||||
|
await session.commit()
|
||||||
|
|
||||||
|
def cleanup():
|
||||||
|
with activate_project_routing(routing):
|
||||||
|
for name in names[:]:
|
||||||
|
if have_project(name):
|
||||||
|
delete_project(name)
|
||||||
|
names.remove(name)
|
||||||
|
_write(directory, names)
|
||||||
|
await asyncio.to_thread(cleanup)
|
||||||
@@ -0,0 +1,318 @@
|
|||||||
|
"""Independent, singleton calculation supervisor and isolated job subprocesses.
|
||||||
|
|
||||||
|
Never imported by API lifespan. Child output contains the original QueryResult;
|
||||||
|
algorithm stdout/stderr is discarded rather than interpreted as a protocol.
|
||||||
|
"""
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from contextlib import asynccontextmanager, contextmanager
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
import signal
|
||||||
|
import shutil
|
||||||
|
import sys
|
||||||
|
import tempfile
|
||||||
|
import time
|
||||||
|
from uuid import UUID, uuid4
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class WorkerSettings:
|
||||||
|
project_ids: tuple[str, ...]
|
||||||
|
max_concurrency: int = 2
|
||||||
|
poll_seconds: float = 2
|
||||||
|
lease_seconds: int = 30
|
||||||
|
heartbeat_seconds: float = 5
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_environment(cls):
|
||||||
|
projects = tuple(dict.fromkeys(str(UUID(p.strip())) for p in
|
||||||
|
os.environ['CALCULATION_WORKER_PROJECT_IDS'].split(',')))
|
||||||
|
result = cls(projects, int(os.getenv('CALCULATION_WORKER_MAX_CONCURRENCY', '2')),
|
||||||
|
float(os.getenv('CALCULATION_WORKER_POLL_SECONDS', '2')),
|
||||||
|
int(os.getenv('CALCULATION_WORKER_LEASE_SECONDS', '30')),
|
||||||
|
float(os.getenv('CALCULATION_WORKER_HEARTBEAT_SECONDS', '5')))
|
||||||
|
if not projects or not 1 <= result.max_concurrency <= 32:
|
||||||
|
raise ValueError('Configure explicit project IDs and concurrency between 1 and 32')
|
||||||
|
if not 0 < result.heartbeat_seconds < result.lease_seconds / 3 or result.poll_seconds <= 0:
|
||||||
|
raise ValueError('Heartbeat must be positive and below one third of the lease')
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def queue_repository(project_id):
|
||||||
|
"""Resolve current routing from metadata, never from a job payload."""
|
||||||
|
from psycopg import AsyncConnection
|
||||||
|
from psycopg.rows import dict_row
|
||||||
|
from app.infra.db.metadb.database import SessionLocal
|
||||||
|
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
|
||||||
|
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
||||||
|
async with SessionLocal() as session:
|
||||||
|
metadata = MetadataRepository(session)
|
||||||
|
project = await metadata.get_project_by_id(UUID(str(project_id)))
|
||||||
|
if not project or project.status != 'active':
|
||||||
|
raise ValueError('Calculation project is not active')
|
||||||
|
routing = await metadata.get_project_db_routing(UUID(str(project_id)), 'biz_data')
|
||||||
|
if not routing or routing.db_type != 'postgresql':
|
||||||
|
raise ValueError('Project business PostgreSQL routing is unavailable')
|
||||||
|
await session.commit()
|
||||||
|
async with await AsyncConnection.connect(routing.dsn, row_factory=dict_row) as conn:
|
||||||
|
yield CalculationJobRepository(conn, project_id)
|
||||||
|
|
||||||
|
|
||||||
|
@asynccontextmanager
|
||||||
|
async def execution_scope(row):
|
||||||
|
from fastapi import HTTPException
|
||||||
|
from app.mcp.scope import QueryMetadataRepository, QueryScope
|
||||||
|
from app.auth.project_dependencies import (resolve_project_context,
|
||||||
|
resolve_project_business_routing, get_project_pg_connection)
|
||||||
|
from app.infra.db.metadb.database import SessionLocal
|
||||||
|
from app.infra.db.project_routing import activate_project_routing
|
||||||
|
async with SessionLocal() as session:
|
||||||
|
metadata = QueryMetadataRepository(session)
|
||||||
|
user = await metadata.get_user_by_id(UUID(str(row['owner_user_id'])))
|
||||||
|
if not user or not user.is_active:
|
||||||
|
raise HTTPException(403, 'Calculation owner is inactive')
|
||||||
|
project = await resolve_project_context(str(row['project_id']), user, metadata)
|
||||||
|
routing = await resolve_project_business_routing(project, metadata)
|
||||||
|
await session.commit()
|
||||||
|
async with asynccontextmanager(get_project_pg_connection)(project, metadata) as conn:
|
||||||
|
await session.commit()
|
||||||
|
with activate_project_routing(routing):
|
||||||
|
yield QueryScope(project, user, conn, routing, metadata)
|
||||||
|
|
||||||
|
|
||||||
|
async def execute_child(input_path: str, output_path: str):
|
||||||
|
from app.mcp.calculation_jobs import execute_calculation_job
|
||||||
|
row = json.loads(Path(input_path).read_text())
|
||||||
|
try:
|
||||||
|
async with execution_scope(row) as scope:
|
||||||
|
result = await execute_calculation_job(scope, row)
|
||||||
|
payload = {'status': 'succeeded', 'result': result.model_dump(mode='json')}
|
||||||
|
except Exception as exc:
|
||||||
|
# Do not persist exception text: DSNs, SQL, credentials or model data can occur there.
|
||||||
|
category = 'CALCULATION_FAILED'
|
||||||
|
if getattr(exc, 'status_code', None) in (401, 403):
|
||||||
|
category = 'AUTHORIZATION_DENIED'
|
||||||
|
elif getattr(exc, 'status_code', None) == 409:
|
||||||
|
category = 'EXECUTION_CONFLICT'
|
||||||
|
elif getattr(exc, 'status_code', None) in (404, 410):
|
||||||
|
category = 'EXECUTION_CONTEXT_UNAVAILABLE'
|
||||||
|
payload = {'status': 'failed', 'error': {'code': category}}
|
||||||
|
encoded = json.dumps(payload, allow_nan=False)
|
||||||
|
if len(encoded.encode()) > 8 * 1024 * 1024:
|
||||||
|
encoded = json.dumps({'status': 'failed', 'error': {'code': 'RESULT_TOO_LARGE'}})
|
||||||
|
Path(output_path).write_text(encoded)
|
||||||
|
|
||||||
|
|
||||||
|
async def stop_process_group(process):
|
||||||
|
"""Always kill the group, even if its original leader already exited."""
|
||||||
|
try:
|
||||||
|
os.killpg(process.pid, signal.SIGTERM)
|
||||||
|
except ProcessLookupError:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
await asyncio.wait_for(process.wait(), 2)
|
||||||
|
except TimeoutError:
|
||||||
|
pass
|
||||||
|
# Native pools can ignore SIGTERM or survive the Python leader.
|
||||||
|
try:
|
||||||
|
os.killpg(process.pid, signal.SIGKILL)
|
||||||
|
except ProcessLookupError:
|
||||||
|
pass
|
||||||
|
await process.wait()
|
||||||
|
if process.stdin is not None:
|
||||||
|
process.stdin.close()
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def job_directory(row):
|
||||||
|
root = Path(os.getenv('CALCULATION_WORKER_STATE_DIR', 'temp/calculation-jobs'))
|
||||||
|
root.mkdir(parents=True, exist_ok=True)
|
||||||
|
directory = Path(tempfile.mkdtemp(prefix=f"{row.get('job_id', 'job')}-", dir=root))
|
||||||
|
directory.chmod(0o700)
|
||||||
|
(directory / 'scope.json').write_text(json.dumps({'project_id': str(row['project_id']), 'job_id': str(row.get('job_id', 'job')),
|
||||||
|
'lease_generation': row.get('lease_generation')}))
|
||||||
|
try:
|
||||||
|
yield str(directory.resolve())
|
||||||
|
finally:
|
||||||
|
# Failed external resource cleanup must remain operator-recoverable.
|
||||||
|
if not (directory / 'cleanup-failed').exists():
|
||||||
|
shutil.rmtree(directory)
|
||||||
|
|
||||||
|
|
||||||
|
async def cleanup_job_resources(row, directory):
|
||||||
|
from app.services.calculation_cleanup import cleanup_calculation_databases
|
||||||
|
try:
|
||||||
|
await cleanup_calculation_databases(row['project_id'], directory)
|
||||||
|
except BaseException:
|
||||||
|
(Path(directory) / 'cleanup-failed').touch()
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
async def recover_job_resources(project_ids):
|
||||||
|
"""Reconcile only this worker's explicit scopes after obtaining the singleton."""
|
||||||
|
root = Path(os.getenv('CALCULATION_WORKER_STATE_DIR', 'temp/calculation-jobs'))
|
||||||
|
if not root.exists():
|
||||||
|
return
|
||||||
|
for directory in root.iterdir():
|
||||||
|
if directory.is_symlink() or not directory.is_dir():
|
||||||
|
continue
|
||||||
|
scope_file = directory / 'scope.json'
|
||||||
|
if not scope_file.exists():
|
||||||
|
continue
|
||||||
|
row = json.loads(scope_file.read_text())
|
||||||
|
if str(row.get('project_id')) not in project_ids:
|
||||||
|
continue
|
||||||
|
job_id = str(UUID(row['job_id']))
|
||||||
|
if not directory.name.startswith(job_id + '-'):
|
||||||
|
raise RuntimeError('Calculation state directory identity is invalid')
|
||||||
|
async with queue_repository(row['project_id']) as repository:
|
||||||
|
stored = await repository.get(job_id)
|
||||||
|
if not stored or stored['lease_generation'] != row['lease_generation']:
|
||||||
|
raise RuntimeError('Calculation state ownership cannot be verified')
|
||||||
|
expiry = stored.get('lease_expires_at')
|
||||||
|
if expiry and expiry > datetime.now(timezone.utc) and stored['status'] in ('running', 'cancel_requested'):
|
||||||
|
raise RuntimeError('Previous calculation lease has not expired')
|
||||||
|
await cleanup_job_resources(row, str(directory.resolve()))
|
||||||
|
async with queue_repository(row['project_id']) as repository:
|
||||||
|
await repository.recover()
|
||||||
|
shutil.rmtree(directory)
|
||||||
|
|
||||||
|
|
||||||
|
class CalculationWorker:
|
||||||
|
def __init__(self, settings, *, repository=queue_repository):
|
||||||
|
self.settings = settings
|
||||||
|
self.repository = repository
|
||||||
|
self.owner = str(uuid4())
|
||||||
|
|
||||||
|
async def run_job(self, row):
|
||||||
|
from fastapi.encoders import jsonable_encoder
|
||||||
|
with job_directory(row) as directory:
|
||||||
|
input_path = Path(directory) / 'input.json'
|
||||||
|
output_path = Path(directory) / 'result.json'
|
||||||
|
input_path.write_text(json.dumps(jsonable_encoder(row)))
|
||||||
|
process = await asyncio.create_subprocess_exec(
|
||||||
|
sys.executable, '-m', 'app.services.calculation_worker',
|
||||||
|
'--guard', str(input_path), str(output_path), start_new_session=True,
|
||||||
|
stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.DEVNULL,
|
||||||
|
stderr=asyncio.subprocess.DEVNULL,
|
||||||
|
env={**os.environ, 'TMPDIR': directory, 'TJWATER_CALCULATION_JOB_DIR': directory},
|
||||||
|
)
|
||||||
|
started = time.monotonic()
|
||||||
|
cleaned = False
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
async with asyncio.timeout(self.settings.heartbeat_seconds):
|
||||||
|
async with self.repository(row['project_id']) as repository:
|
||||||
|
status = await repository.heartbeat(row, self.owner, self.settings.lease_seconds)
|
||||||
|
if status is None:
|
||||||
|
return # Finally stops all owned computation on a lost lease.
|
||||||
|
terminal = None
|
||||||
|
if status == 'cancel_requested':
|
||||||
|
terminal = {'status': 'cancelled'}
|
||||||
|
elif time.monotonic() - started >= row['timeout_seconds']:
|
||||||
|
terminal = {'status': 'timed_out', 'error': {'code': 'EXECUTION_TIMEOUT'}}
|
||||||
|
elif process.returncode is not None:
|
||||||
|
try:
|
||||||
|
if output_path.stat().st_size > 8 * 1024 * 1024:
|
||||||
|
raise ValueError('Result exceeds persistence limit')
|
||||||
|
terminal = json.loads(output_path.read_text())
|
||||||
|
except (OSError, ValueError):
|
||||||
|
terminal = {'status': 'failed', 'error': {'code': 'CHILD_PROCESS_FAILED'}}
|
||||||
|
if terminal:
|
||||||
|
await stop_process_group(process)
|
||||||
|
try:
|
||||||
|
await cleanup_job_resources(row, directory)
|
||||||
|
cleaned = True
|
||||||
|
except Exception:
|
||||||
|
terminal = {'status': 'failed', 'error': {'code': 'CLEANUP_FAILED'}}
|
||||||
|
cleaned = True
|
||||||
|
async with self.repository(row['project_id']) as repository:
|
||||||
|
finished = await repository.finish(row, self.owner, **terminal)
|
||||||
|
if not finished:
|
||||||
|
status = await repository.heartbeat(row, self.owner, self.settings.lease_seconds)
|
||||||
|
if status == 'cancel_requested':
|
||||||
|
await repository.finish(row, self.owner, 'cancelled')
|
||||||
|
return
|
||||||
|
await asyncio.sleep(self.settings.heartbeat_seconds)
|
||||||
|
finally:
|
||||||
|
await stop_process_group(process)
|
||||||
|
if not cleaned:
|
||||||
|
await cleanup_job_resources(row, directory)
|
||||||
|
|
||||||
|
async def run(self):
|
||||||
|
from sqlalchemy import text
|
||||||
|
from app.infra.db.metadb.database import engine
|
||||||
|
# Dedicated session-level metadata lock: independent API workers never
|
||||||
|
# participate; duplicate supervisors exit before claiming any jobs.
|
||||||
|
async with engine.connect() as lock:
|
||||||
|
acquired = await lock.scalar(text("SELECT pg_try_advisory_lock(1748531781, 2)"))
|
||||||
|
await lock.commit()
|
||||||
|
if not acquired:
|
||||||
|
raise RuntimeError('Another calculation supervisor is already active')
|
||||||
|
tasks = set()
|
||||||
|
project_cursor = 0
|
||||||
|
try:
|
||||||
|
# Allow any old supervisor to detect a broken lock connection and
|
||||||
|
# stop its children before this replacement claims capacity.
|
||||||
|
await asyncio.sleep(self.settings.lease_seconds)
|
||||||
|
await recover_job_resources(self.settings.project_ids)
|
||||||
|
while True:
|
||||||
|
async with asyncio.timeout(self.settings.heartbeat_seconds):
|
||||||
|
await lock.execute(text('SELECT 1'))
|
||||||
|
await lock.commit()
|
||||||
|
for task in tuple(tasks):
|
||||||
|
if task.done():
|
||||||
|
tasks.remove(task)
|
||||||
|
task.result() # Fail closed if heartbeat/persistence failed.
|
||||||
|
projects = self.settings.project_ids
|
||||||
|
ordered = projects[project_cursor:] + projects[:project_cursor]
|
||||||
|
project_cursor = (project_cursor + 1) % len(projects)
|
||||||
|
for project_id in ordered:
|
||||||
|
if len(tasks) >= self.settings.max_concurrency:
|
||||||
|
break
|
||||||
|
async with self.repository(project_id) as repository:
|
||||||
|
await repository.recover()
|
||||||
|
row = await repository.claim(self.owner, self.settings.lease_seconds)
|
||||||
|
if row:
|
||||||
|
tasks.add(asyncio.create_task(self.run_job(row)))
|
||||||
|
await asyncio.sleep(min(self.settings.poll_seconds, self.settings.heartbeat_seconds))
|
||||||
|
finally:
|
||||||
|
for task in tasks:
|
||||||
|
task.cancel()
|
||||||
|
await asyncio.gather(*tasks, return_exceptions=True)
|
||||||
|
await lock.execute(text('SELECT pg_advisory_unlock(1748531781, 2)'))
|
||||||
|
await lock.commit()
|
||||||
|
|
||||||
|
|
||||||
|
async def guard_child(input_path, output_path):
|
||||||
|
"""A parent-death pipe kills the whole group even when the supervisor is SIGKILLed."""
|
||||||
|
reader = asyncio.StreamReader()
|
||||||
|
protocol = asyncio.StreamReaderProtocol(reader)
|
||||||
|
transport, _ = await asyncio.get_running_loop().connect_read_pipe(lambda: protocol, sys.stdin)
|
||||||
|
process = await asyncio.create_subprocess_exec(
|
||||||
|
sys.executable, '-m', 'app.services.calculation_worker', input_path, output_path,
|
||||||
|
stdin=asyncio.subprocess.DEVNULL,
|
||||||
|
)
|
||||||
|
completed = asyncio.create_task(process.wait())
|
||||||
|
parent_gone = asyncio.create_task(reader.read(1))
|
||||||
|
try:
|
||||||
|
done, _ = await asyncio.wait((completed, parent_gone), return_when=asyncio.FIRST_COMPLETED)
|
||||||
|
if parent_gone in done:
|
||||||
|
os.killpg(os.getpgrp(), signal.SIGKILL)
|
||||||
|
await completed
|
||||||
|
finally:
|
||||||
|
parent_gone.cancel()
|
||||||
|
transport.close()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
if sys.argv[1] == '--guard':
|
||||||
|
asyncio.run(guard_child(sys.argv[2], sys.argv[3]))
|
||||||
|
else:
|
||||||
|
asyncio.run(execute_child(sys.argv[1], sys.argv[2]))
|
||||||
@@ -778,6 +778,10 @@ class OperationalTaskService:
|
|||||||
"Only a running run accepts Agent artifacts", status=409)
|
"Only a running run accepts Agent artifacts", status=409)
|
||||||
now = self._now()
|
now = self._now()
|
||||||
self._require_active_run_lease(active_run, now)
|
self._require_active_run_lease(active_run, now)
|
||||||
|
if command.report and (
|
||||||
|
command.report.stage == "final" or command.report.outcome == "insufficient_data"
|
||||||
|
):
|
||||||
|
await self._require_no_pending_calculations(task_id, active_run)
|
||||||
previous_report = await self.repo.get_latest_report(str(task_id))
|
previous_report = await self.repo.get_latest_report(str(task_id))
|
||||||
require(not previous_report or previous_report.get("execution_generation") != task["execution_generation"]
|
require(not previous_report or previous_report.get("execution_generation") != task["execution_generation"]
|
||||||
or previous_report["report"].get("outcome") != "insufficient_data",
|
or previous_report["report"].get("outcome") != "insufficient_data",
|
||||||
@@ -1866,6 +1870,15 @@ class OperationalTaskService:
|
|||||||
task_ids.append(task_id)
|
task_ids.append(task_id)
|
||||||
return DispatchDueSchedulesResult(task_ids=task_ids, dispatched_at=now)
|
return DispatchDueSchedulesResult(task_ids=task_ids, dispatched_at=now)
|
||||||
|
|
||||||
|
async def _require_no_pending_calculations(self, task_id: UUID, run: dict) -> None:
|
||||||
|
if hasattr(self.repo, "pending_calculation_jobs"):
|
||||||
|
pending = await self.repo.pending_calculation_jobs(
|
||||||
|
str(task_id), str(run["run_id"]), str(run["execution_request_id"]),
|
||||||
|
)
|
||||||
|
require(not pending, "PENDING_CALCULATIONS",
|
||||||
|
"Wait for submitted calculations before publishing a final report or completing the run",
|
||||||
|
refs=pending, status=409, recovery="retry_same_request")
|
||||||
|
|
||||||
async def _finish_run(
|
async def _finish_run(
|
||||||
self, task_id: UUID, run_id: UUID, command: AgentRunCommand,
|
self, task_id: UUID, run_id: UUID, command: AgentRunCommand,
|
||||||
*, outcome: str, failure_code: str | None = None,
|
*, outcome: str, failure_code: str | None = None,
|
||||||
@@ -1890,6 +1903,7 @@ class OperationalTaskService:
|
|||||||
task = await self._task(task_id)
|
task = await self._task(task_id)
|
||||||
final_outcome = outcome
|
final_outcome = outcome
|
||||||
if outcome == "succeeded":
|
if outcome == "succeeded":
|
||||||
|
await self._require_no_pending_calculations(task_id, run)
|
||||||
if hasattr(self.repo, "pending_run_messages"):
|
if hasattr(self.repo, "pending_run_messages"):
|
||||||
pending_messages = await self.repo.pending_run_messages(str(run_id))
|
pending_messages = await self.repo.pending_run_messages(str(run_id))
|
||||||
require(not pending_messages, "PENDING_AGENT_MESSAGES",
|
require(not pending_messages, "PENDING_AGENT_MESSAGES",
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ from collections.abc import Iterator
|
|||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from tempfile import NamedTemporaryFile
|
from tempfile import NamedTemporaryFile
|
||||||
|
import os
|
||||||
|
|
||||||
from app.native.wndb.inp.exporter import dump_inp
|
from app.native.wndb.inp.exporter import dump_inp
|
||||||
|
|
||||||
@@ -17,9 +18,11 @@ def temporary_project_inp(
|
|||||||
version: str = "2",
|
version: str = "2",
|
||||||
) -> Iterator[Path]:
|
) -> Iterator[Path]:
|
||||||
"""Export one project model to an isolated INP and remove it afterwards."""
|
"""Export one project model to an isolated INP and remove it afterwards."""
|
||||||
PROJECT_INP_DIRECTORY.mkdir(parents=True, exist_ok=True)
|
job_directory = os.environ.get("TJWATER_CALCULATION_JOB_DIR")
|
||||||
|
directory = Path(job_directory) / "inp" if job_directory else PROJECT_INP_DIRECTORY
|
||||||
|
directory.mkdir(parents=True, exist_ok=True)
|
||||||
with NamedTemporaryFile(
|
with NamedTemporaryFile(
|
||||||
dir=PROJECT_INP_DIRECTORY,
|
dir=directory,
|
||||||
prefix=f"{purpose}_",
|
prefix=f"{purpose}_",
|
||||||
suffix=".inp",
|
suffix=".inp",
|
||||||
delete=False,
|
delete=False,
|
||||||
|
|||||||
@@ -0,0 +1,86 @@
|
|||||||
|
# MCP 后台计算队列
|
||||||
|
|
||||||
|
API 保留 `uvicorn --workers 4`。长计算在项目 BizDB 持久化入队,由独立计算进程领取;所有 API 进程读取相同任务状态,不依赖粘性会话。HTTP 的 5 分钟空闲超时保留,Agent 的提交和状态查询采用 30 秒调用预算、35 秒宿主预算。`timeout_seconds` 对入队工具表示 Worker 的执行预算,不是 HTTP 等待时间,也不是模拟时间 `duration`。
|
||||||
|
|
||||||
|
```text
|
||||||
|
前端 / Agent → API × 4 → 项目 BizDB calculation_jobs
|
||||||
|
↑
|
||||||
|
独立计算 supervisor
|
||||||
|
↓
|
||||||
|
受管理的计算子进程组
|
||||||
|
```
|
||||||
|
|
||||||
|
## 工具契约
|
||||||
|
|
||||||
|
原工具名继续使用。以下工具提交后返回 `QueryResult`,其中 `semantics=task`,`data` 包含 `job_id`、`status`、`tool` 和 `poll_after_ms`;这只证明提交状态,不是分析结果或调查证据:
|
||||||
|
|
||||||
|
- `run_burst_location`、`run_leakage_identification`、`run_burst_detection`
|
||||||
|
- `run_simulation`、`run_burst_simulation`、`run_contaminant_simulation`
|
||||||
|
- `analyze_valve_closure`、`analyze_flushing`、`analyze_pressure_regulation`
|
||||||
|
- `optimize_sensor_placement`
|
||||||
|
- `save_simulation_scenario`、`create_sensor_placement` 的 `dry_run=false` 执行阶段
|
||||||
|
|
||||||
|
后两项保留预览、摘要校验和原有审批;提交与执行时均校验权限及预览摘要,已有写入回执继续约束重复请求。`analyze_valve_isolation`、`calculate_demand_allocation` 保持同步执行。普通查询和其他写入不进入此队列。
|
||||||
|
|
||||||
|
| MCP 工具 | 输入 | 返回 |
|
||||||
|
| --- | --- | --- |
|
||||||
|
| `get_calculation_job` | `job_id` | 状态、工具、时间和安全错误码 |
|
||||||
|
| `get_calculation_result` | `job_id`,可选 `data_path/limit/offset` | 成功后读取原分析结果;未成功返回任务状态 |
|
||||||
|
| `list_calculation_jobs` | 可选 `status/limit/offset` | 当前授权会话及调查轮次的任务列表 |
|
||||||
|
| `cancel_calculation_job` | `job_id` | 排队任务可直接取消;运行任务先进入 `cancel_requested` |
|
||||||
|
|
||||||
|
状态为 `queued/running/cancel_requested/succeeded/failed/cancelled/timed_out/interrupted`。任务取消不等于在线模型写入回滚;发生失败、超时或中断时,需要核查已有写入,不能自动重算整个窗口。
|
||||||
|
|
||||||
|
当前使用普通 MCP 业务工具承载队列契约,没有宣称已实现协议级 Tasks 扩展。Server 和 Agent 目录必须配套更新,旧 Agent 不应接收新的队列返回契约。
|
||||||
|
|
||||||
|
## 身份、幂等和结果
|
||||||
|
|
||||||
|
任务记录保存提交人的元数据用户 ID、项目、原工具参数和受限授权快照,不保存 Bearer、数据库 DSN 或短期 `context_handle`。执行和读取结果时重新查询当前用户、成员关系、项目路由、权限及调查运行租约。任务不会因为短期 MCP 句柄失效而丢失;也不会借持久化任务绕过撤销的权限。
|
||||||
|
|
||||||
|
普通计算使用宿主从原生工具调用记录解析出的稳定调用 ID;同一次传输重试复用任务,不将新的业务计算自动按参数合并。重写工具同时保留原审批 `request_id` 幂等语义。不同参数复用同一操作编号会被拒绝。
|
||||||
|
|
||||||
|
调查的排队回执不写成分析证据。成功结果读取时,按原计算来源、任务及运行验证,并幂等地保存证据。未完成的计算阻止提交最终报告、`insufficient_data` 或将运行标记成功;中间进度仍可保存。
|
||||||
|
|
||||||
|
Agent 调查执行在原生回合空闲后由宿主短轮询等待,继续维持调查心跳,计算进入终态后再触发模型读取结果。普通聊天展示任务编号,可在当前会话后续查询。异常流结束时若无法安全继续处理新结果,会明确失败而不是虚报完成;Server 中的计算结果仍保留。
|
||||||
|
|
||||||
|
## 并发、取消和恢复
|
||||||
|
|
||||||
|
一个 supervisor 默认全项目合计并发 2 项,独立于 API 的 4 个进程。它持有元数据库的会话级 advisory lock,第二个 supervisor 会拒绝启动;当前不是多台计算节点各自扩容的模式。各项目轮转领取,使用 PostgreSQL 短事务 `SKIP LOCKED`,租约代次约束心跳和结果写回。
|
||||||
|
|
||||||
|
每项任务有独立子进程组;取消、执行超时或失去租约时先停止进程组,再写终态。守护管道在 supervisor 被强杀后关闭,负责终止仍在运行的算法进程。API 重启不影响计算 supervisor;计算进程崩溃后,旧任务标为中断,不自动重放。
|
||||||
|
|
||||||
|
原生临时文件写入独立任务目录。临时模型数据库创建前登记到该目录的账本,正常结束或取消后只清理明确登记的临时数据库。清理失败返回 `CLEANUP_FAILED` 并保留恢复目录;启动时检查项目、任务和执行代次,确认旧租约过期后恢复清理。无法验证归属或清理失败时停止领取新工作。
|
||||||
|
|
||||||
|
## 部署
|
||||||
|
|
||||||
|
1. 先对每个目标项目的 BizDB 应用 `resources/sql/026_calculation_jobs.sql`,包括未来项目使用的相应业务模板。不要应用到 TimescaleDB。迁移为新增表及索引,不删除业务数据。
|
||||||
|
2. 配置 `CALCULATION_WORKER_PROJECT_IDS` 为明确的项目 UUID 列表,配置现有元数据库连接及数据库路由解密密钥。Worker 不需要新 OAuth 密钥;仅部署侧工作进程读取这些设置。
|
||||||
|
3. 同步发布 Server 和 Agent,在接入长计算前启动独立 Worker。不要从 API `lifespan` 启动,不要把调查 Worker 当作计算 Worker。
|
||||||
|
|
||||||
|
从 Server 仓库启动:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
conda run -n server python scripts/run_calculation_worker.py
|
||||||
|
```
|
||||||
|
|
||||||
|
systemd 模板:`infra/systemd/tjwater-calculation-worker.service`。Docker 使用同一 Server 镜像的独立服务:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
docker compose -f infra/docker/docker-compose.yml --profile calculations up -d calculation-worker
|
||||||
|
```
|
||||||
|
|
||||||
|
Compose 使用 `../../.env` 注入配置,需确保其中数据库地址在容器内可达。`calculation-state` volume 保存任务恢复账本,不能当作临时缓存删除。非容器启动可用 `CALCULATION_WORKER_STATE_DIR` 指定持久目录。
|
||||||
|
|
||||||
|
并发默认 2、轮询默认 2 秒、租约 30 秒、心跳 5 秒。算法内部并行度仍消耗额外 CPU/内存;调整并发必须依据实际资源测量。
|
||||||
|
|
||||||
|
回滚前暂停接收新的队列计算并等待或明确取消在途任务,核对有副作用任务的写入。保留任务表、结果和恢复目录;旧同步版本不支持继续处理这些作业,不能直接删除记录后重新提交。
|
||||||
|
|
||||||
|
## 验证
|
||||||
|
|
||||||
|
```bash
|
||||||
|
conda run -n server python -m pytest tests/unit/test_mcp_calculation_jobs.py tests/unit/test_calculation_worker.py tests/unit/test_calculation_cleanup.py tests/unit/test_mcp_query.py tests/unit/test_mcp_analysis.py tests/unit/test_operational_tasks.py -q
|
||||||
|
```
|
||||||
|
|
||||||
|
真实 PostgreSQL 并发测试为 `tests/integration/test_calculation_queue_storage.py`。它要求 `CALCULATION_QUEUE_TEST_DSN` 指向名称以 `tjwater_calc_test_` 开头的专用可丢弃数据库,不能指向业务库。覆盖四连接重复提交、参数冲突、并发领取、项目和用户隔离、取消竞争、失效代次拒绝和过期租约恢复。
|
||||||
|
|
||||||
|
这些验证不代替真实项目算法运行、浏览器断线重连或部署验收。生产接入还应观察超过 5 分钟的真实计算、恢复后结果读取和原生临时资源清理。
|
||||||
@@ -54,15 +54,15 @@ Agent 原生工具统一使用 `mcp_tjwater_` 前缀。Server目录共67项:42
|
|||||||
|
|
||||||
## 资源边界与发布
|
## 资源边界与发布
|
||||||
|
|
||||||
12项计算和13项预览式写入工具均提供 `timeout_seconds`,由Agent根据模型规模、模拟步数及已有耗时按次选择;它表示实际等待预算,与 `duration` 表示的模拟时间不同。未填写默认120秒。服务端上限由 `TJWATER_MCP_MAX_CALCULATION_TIMEOUT_SECONDS` 控制,默认3600秒,管理员可设120至86400秒;`tools/list` schema 发布实际有效上限,超限在计算前拒绝。
|
长计算采用[后台计算队列](calculation-queue.md)。10项分析计算与2项重计算写入的正式执行阶段快速返回任务编号,再使用 `get_calculation_job`、`get_calculation_result`、`list_calculation_jobs`、`cancel_calculation_job` 管理。原有预览与审批仍有效;排队回执不作为分析证据。
|
||||||
|
|
||||||
例如 `duration=3600, timeout_seconds=900` 表示模拟1小时、允许等待15分钟。Agent的上游HTTP/SDK预算为所选预算加5秒,整个宿主调用另留60秒处理认证、上下文绑定和清理;OpenCode连接上限从工具目录计算并写入宿主连接文件。本地MCP请求关闭Bun空闲断连,仍受上述调用截止时间控制。普通查询继续使用原30秒Server预算与35秒宿主预算。
|
`timeout_seconds` 对队列工具表示计算 Worker 的执行预算,与模拟时间 `duration` 和单次 HTTP 请求时限独立。各工具默认值与管理员上限继续由目录公布。Agent的提交/查询采用30秒调用预算和35秒宿主预算,不延长Bun的5分钟HTTP空闲超时。
|
||||||
|
|
||||||
计算在线程执行,不阻塞事件循环,单Server进程最多两项计算。超时不能杀死底层原生计算,槽位一直保留到计算真正退出;超时结果不会补写为调查成功结果,不自动反复重试。本次未增加后台计算作业队列。普通聊天大结果仍受现有12KB内联限制,调查结果可经证据分页读取。
|
`analyze_valve_isolation` 和 `calculate_demand_allocation` 保持同步。队列消费者独立于Uvicorn的4个API进程,默认全部配置项目合计并发2项,并通过元数据库单实例锁防止重复启动。取消、超时和租约失效需要停止计算子进程组;成功结果按原运行授权保存证据,未完成计算阻止调查提前结案。
|
||||||
|
|
||||||
在线重算逐个时刻写入基线,超时或中途失败可能已有部分时刻更新,不能把工具失败解释为整个窗口已回滚;不要自动重试整个窗口。
|
在线重算逐个时刻写入基线,超时或中途失败可能已有部分时刻更新,不能把工具失败解释为整个窗口已回滚;不要自动重试整个窗口。
|
||||||
|
|
||||||
无需新增数据库迁移或模型配置。须同步发布Server、Agent和前端证据枚举;先更新Server目录,再重启Agent,client模式还需重载原生连接。保持既有OAuth、MCP开关及数据库路由配置,不自动更改权限或重启调查。
|
队列接入需要在目标项目BizDB应用 `resources/sql/026_calculation_jobs.sql` 并启动独立计算Worker。Server与Agent需配套发布,client模式还需重载原生连接。既有OAuth与项目权限保持不变;完整部署与回滚步骤见后台计算队列文档。
|
||||||
|
|
||||||
本次验证限schema、权限、协议目录和mock业务函数,不运行真实模拟、优化或付费模型。真实项目的模板同步状态、运算耗时、SCADA覆盖和结果准确性仍需另行验收。
|
本次验证限schema、权限、协议目录和mock业务函数,不运行真实模拟、优化或付费模型。真实项目的模板同步状态、运算耗时、SCADA覆盖和结果准确性仍需另行验收。
|
||||||
|
|
||||||
|
|||||||
@@ -1,10 +1,10 @@
|
|||||||
# 产品 MCP 查询:本地实施与启用
|
# 产品 MCP 查询:本地实施与启用
|
||||||
|
|
||||||
实现使用 MCP 2026-07-28、Python mcp 2.3.0、TypeScript @modelcontextprotocol/client 与 @modelcontextprotocol/server 2.3.1。Server 提供 `/mcp` 和67项工具(42项查询、12项计算、13项预览式写入);Agent 将其投影合并到本地 MCP,由 OpenCode 原生发现为 `mcp_tjwater_*`,普通项目会话、调查主 Agent 和已验证的直接只读子 Agent 共享42项查询;12项计算仅向非reader会话开放,13项写入仅向项目成员的普通聊天开放。个人会话不能获取项目凭据。旧 CLI 只作为能力清单,没有运行时 CLI 桥接。
|
实现使用 MCP 2026-07-28、Python mcp 2.3.0、TypeScript @modelcontextprotocol/client 与 @modelcontextprotocol/server 2.3.1。Server 提供 `/mcp` 和71项工具(42项原查询、12项计算、13项预览式写入、4项计算任务管理);Agent 将其投影合并到本地 MCP,由 OpenCode 原生发现为 `mcp_tjwater_*`,普通项目会话、调查主 Agent 和已验证的直接只读子 Agent 共享42项查询;12项计算仅向非reader会话开放,13项写入仅向项目成员的普通聊天开放。个人会话不能获取项目凭据。旧 CLI 只作为能力清单,没有运行时 CLI 桥接。
|
||||||
|
|
||||||
`run_simulation` 是在线重算,复用 REST 在线计算服务而不复制数据库;按水力步长更新模型条件和实时基线,不创建正式方案。MCP 时长为秒,原 REST 时长仍为分钟。其 annotation 为 `readOnlyHint=false`、`destructiveHint=true`;只读 reader 不能调用,其他方案模拟继续使用隔离临时模型。详细边界见 [MCP 分析工具说明](mcp-analysis-tools.md)。
|
`run_simulation` 是在线重算,复用 REST 在线计算服务而不复制数据库;按水力步长更新模型条件和实时基线,不创建正式方案。MCP 时长为秒,原 REST 时长仍为分钟。其 annotation 为 `readOnlyHint=false`、`destructiveHint=true`;只读 reader 不能调用,其他方案模拟继续使用隔离临时模型。详细边界见 [MCP 分析工具说明](mcp-analysis-tools.md)。
|
||||||
|
|
||||||
计算工具支持Agent按次指定 `timeout_seconds`,默认120秒;计算上限默认1小时,由 `TJWATER_MCP_MAX_CALCULATION_TIMEOUT_SECONDS` 配置并在工具schema公布。Agent、OpenCode和本地HTTP调用预算随目录同步,普通查询仍保持原有短超时。等待超时不意味着底层线程已停止或在线基线已回滚,禁止自动重试整个写入窗口。
|
长计算使用持久化队列,原工具快速返回任务编号,新增状态、结果、列表和取消工具。`timeout_seconds` 对队列工具表示执行预算;HTTP保持短调用,不延长5分钟空闲超时。独立计算Worker、迁移与权限边界见[后台计算队列](calculation-queue.md)。
|
||||||
|
|
||||||
## 认证与数据边界
|
## 认证与数据边界
|
||||||
|
|
||||||
|
|||||||
@@ -21,6 +21,26 @@ services:
|
|||||||
- timescaledb
|
- timescaledb
|
||||||
- postgis
|
- postgis
|
||||||
|
|
||||||
|
# Enable after provisioning calculation storage and explicit project scopes.
|
||||||
|
calculation-worker:
|
||||||
|
profiles: [calculations]
|
||||||
|
image: ${TJWATER_SERVER_IMAGE:-tjwater-server:local}
|
||||||
|
init: true
|
||||||
|
restart: on-failure
|
||||||
|
command: ["python", "scripts/run_calculation_worker.py"]
|
||||||
|
env_file:
|
||||||
|
- ../../.env
|
||||||
|
environment:
|
||||||
|
- PYTHONPATH=/app
|
||||||
|
volumes:
|
||||||
|
- ../../app:/app/app
|
||||||
|
- ../../resources:/app/resources
|
||||||
|
- calculation-state:/app/temp/calculation-jobs
|
||||||
|
stop_grace_period: 30s
|
||||||
|
depends_on:
|
||||||
|
- timescaledb
|
||||||
|
- postgis
|
||||||
|
|
||||||
# ==========================================
|
# ==========================================
|
||||||
# Infrastructure Services
|
# Infrastructure Services
|
||||||
# ==========================================
|
# ==========================================
|
||||||
@@ -127,6 +147,9 @@ services:
|
|||||||
volumes:
|
volumes:
|
||||||
- ./mapservice/geoserver_data:/opt/geoserver_data
|
- ./mapservice/geoserver_data:/opt/geoserver_data
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
calculation-state:
|
||||||
|
|
||||||
networks:
|
networks:
|
||||||
default:
|
default:
|
||||||
driver: bridge
|
driver: bridge
|
||||||
|
|||||||
@@ -0,0 +1,17 @@
|
|||||||
|
[Unit]
|
||||||
|
Description=TJWater queued calculation worker
|
||||||
|
After=network-online.target
|
||||||
|
Wants=network-online.target
|
||||||
|
|
||||||
|
[Service]
|
||||||
|
Type=simple
|
||||||
|
WorkingDirectory=%h/NextTJWater/next-tjwater-server
|
||||||
|
ExecStart=%h/miniconda3/envs/server/bin/python scripts/run_calculation_worker.py
|
||||||
|
Restart=on-failure
|
||||||
|
RestartSec=5
|
||||||
|
KillMode=control-group
|
||||||
|
TimeoutStopSec=30
|
||||||
|
UMask=0077
|
||||||
|
|
||||||
|
[Install]
|
||||||
|
WantedBy=default.target
|
||||||
@@ -0,0 +1,23 @@
|
|||||||
|
-- Apply independently to each project's business database.
|
||||||
|
CREATE SCHEMA IF NOT EXISTS workbench;
|
||||||
|
CREATE TABLE IF NOT EXISTS workbench.calculation_jobs (
|
||||||
|
job_id uuid PRIMARY KEY,
|
||||||
|
project_id uuid NOT NULL,
|
||||||
|
owner_user_id uuid NOT NULL,
|
||||||
|
request_key text NOT NULL,
|
||||||
|
tool text NOT NULL,
|
||||||
|
parameters jsonb NOT NULL,
|
||||||
|
binding jsonb NOT NULL,
|
||||||
|
timeout_seconds integer NOT NULL CHECK (timeout_seconds BETWEEN 1 AND 86400),
|
||||||
|
status text NOT NULL DEFAULT 'queued' CHECK (status IN
|
||||||
|
('queued','running','cancel_requested','succeeded','failed','cancelled','timed_out','interrupted')),
|
||||||
|
result jsonb,
|
||||||
|
error jsonb,
|
||||||
|
lease_generation bigint NOT NULL DEFAULT 0,
|
||||||
|
lease_owner text,
|
||||||
|
lease_expires_at timestamptz,
|
||||||
|
created_at timestamptz NOT NULL DEFAULT clock_timestamp(),
|
||||||
|
updated_at timestamptz NOT NULL DEFAULT clock_timestamp(),
|
||||||
|
UNIQUE(project_id,owner_user_id,request_key)
|
||||||
|
);
|
||||||
|
CREATE INDEX IF NOT EXISTS calculation_jobs_queue ON workbench.calculation_jobs(project_id,status,created_at);
|
||||||
@@ -0,0 +1,30 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Run the independent calculation supervisor (never an API lifespan task)."""
|
||||||
|
import argparse
|
||||||
|
import asyncio
|
||||||
|
from pathlib import Path
|
||||||
|
import signal
|
||||||
|
import sys
|
||||||
|
|
||||||
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||||
|
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
from app.services.calculation_worker import CalculationWorker, WorkerSettings
|
||||||
|
|
||||||
|
|
||||||
|
async def main():
|
||||||
|
load_dotenv(Path(__file__).resolve().parents[1] / '.env', override=False)
|
||||||
|
task = asyncio.create_task(CalculationWorker(WorkerSettings.from_environment()).run())
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
for signum in (signal.SIGTERM, signal.SIGINT):
|
||||||
|
loop.add_signal_handler(signum, task.cancel)
|
||||||
|
try:
|
||||||
|
await task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
argparse.ArgumentParser(description=__doc__).parse_args()
|
||||||
|
asyncio.run(main())
|
||||||
@@ -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, client, _, _ = runtime
|
||||||
app.settings.query_timeout_seconds = 1
|
|
||||||
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||||
|
receipt = calculation_jobs.envelope("run_simulation", {"job_id": str(uuid4()), "status": "queued"})
|
||||||
async def execute(scope, bound, tool, args):
|
submit = AsyncMock(return_value=receipt)
|
||||||
await asyncio.sleep(1.05)
|
monkeypatch.setattr(calculation_jobs, "submit_calculation", submit)
|
||||||
return QueryResult(
|
app.executor = AsyncMock(side_effect=AssertionError("native work entered API process"))
|
||||||
source=tool.name, source_ref="test:calculation", semantics="analysis",
|
response = request(client, "tools/call", {"name": "run_simulation", "arguments": {
|
||||||
queried_at=datetime.now(timezone.utc), data={"completed": True},
|
"context_handle": binding().context_handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||||
)
|
|
||||||
|
|
||||||
app.executor = execute
|
|
||||||
handle = binding().context_handle
|
|
||||||
arguments = {
|
|
||||||
"context_handle": handle, "start_time": "2026-10-08T12:00:00+08:00",
|
|
||||||
"duration": 3600, "output_node_ids": ["J1"], "timeout_seconds": 2,
|
"duration": 3600, "output_node_ids": ["J1"], "timeout_seconds": 2,
|
||||||
}
|
}}).json()
|
||||||
response = request(client, "tools/call", {"name": "run_simulation", "arguments": arguments}).json()
|
assert response["result"]["structuredContent"]["data"]["status"] == "queued"
|
||||||
assert response["result"].get("isError") is not True
|
assert response["result"]["structuredContent"]["semantics"] == "task"
|
||||||
assert response["result"]["structuredContent"]["data"] == {"completed": True}
|
app.executor.assert_not_awaited()
|
||||||
arguments["timeout_seconds"] = 1
|
assert submit.await_args.args[3].timeout_seconds == 2
|
||||||
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"
|
|
||||||
|
|
||||||
|
|
||||||
def test_calculation_limit_is_advertised_and_enforced_before_execution(runtime):
|
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 == []
|
assert calls == []
|
||||||
|
|
||||||
|
|
||||||
def test_calculation_capacity_has_an_explicit_mcp_error(runtime, monkeypatch):
|
def test_queue_submission_errors_preserve_mcp_error_mapping(runtime, monkeypatch):
|
||||||
from threading import BoundedSemaphore
|
from app.mcp import calculation_jobs
|
||||||
from app.mcp import analysis_service
|
_, client, _, _ = runtime
|
||||||
|
|
||||||
app, client, _, _ = runtime
|
|
||||||
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||||
slots = BoundedSemaphore(2)
|
monkeypatch.setattr(calculation_jobs, "submit_calculation", AsyncMock(side_effect=HTTPException(409, "Calculation request_key was used with different inputs")))
|
||||||
monkeypatch.setattr(analysis_service, "_slots", slots)
|
response = request(client, "tools/call", {"name": "run_simulation", "arguments": {
|
||||||
async def execute(*unused):
|
"context_handle": binding().context_handle, "start_time": "2026-10-08T12:00:00+08:00",
|
||||||
return await analysis_service.bounded_calculation(lambda: pytest.fail("occupied calculation entered"))
|
"duration": 3600, "output_node_ids": ["J1"],
|
||||||
app.executor = execute
|
}}).json()
|
||||||
assert slots.acquire(blocking=False) and slots.acquire(blocking=False)
|
assert response["result"]["structuredContent"]["error"]["code"] == "QUERY_CONTEXT_STALE"
|
||||||
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"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_burst_location_default_inputs_pass_protocol_validation(runtime, monkeypatch):
|
def test_burst_location_default_inputs_pass_protocol_validation(runtime, monkeypatch):
|
||||||
app, client, calls, _ = runtime
|
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())
|
monkeypatch.setattr(application, "authorize_query", AsyncMock())
|
||||||
schema = next(tool["inputSchema"] for tool in request(client).json()["result"]["tools"]
|
schema = next(tool["inputSchema"] for tool in request(client).json()["result"]["tools"]
|
||||||
if tool["name"] == "run_burst_location")
|
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",
|
"scada_burst_end": "2026-10-08T13:00:00+08:00",
|
||||||
}}).json()
|
}}).json()
|
||||||
assert response["result"].get("isError") is not True
|
assert response["result"].get("isError") is not True
|
||||||
assert calls == [("session-1", "run_burst_location")]
|
assert calls == []
|
||||||
|
|
||||||
|
|
||||||
def settings():
|
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,
|
request_id=uuid4(), source_session_id="session-1", report=report,
|
||||||
))
|
))
|
||||||
assert error.value.code == "EVIDENCE_NOT_FOUND"
|
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