jiang
be6317eb48
Add project-bound OAuth MCP queries and controlled analysis with scoped evidence, limits, timeouts and a nearby-model spatial query using EPSG:4326 inputs. Default nearby searches to 200 meters and five results. Integrate investigation messaging/presentation, health/readiness and typed history; remove duplicate web search and fix online simulation and compressed baseline writes. Include project migrations 023-025; keep local credentials and demo seeding excluded. Validation: 319 focused pytest checks; real PostGIS temporary-table regression; read-only Jinqian bridge radius comparison; cross-language MCP discovery/calls.
277 lines
13 KiB
Python
277 lines
13 KiB
Python
import json
|
|
import os
|
|
from uuid import uuid4
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from scripts.run_operational_worker import CredentialProvider, OperationalWorker, WorkerSettings
|
|
|
|
|
|
@pytest.fixture
|
|
def anyio_backend():
|
|
return "asyncio"
|
|
|
|
|
|
@pytest.mark.parametrize("environment,use_dev_url", [
|
|
("development", True), ("production", False), ("staging", False), (None, False),
|
|
])
|
|
def test_worker_dev_token_url_is_only_used_in_development(monkeypatch, environment, use_dev_url):
|
|
monkeypatch.setattr("scripts.run_operational_worker.load_dotenv", lambda *args, **kwargs: None)
|
|
for key in list(os.environ):
|
|
if key.startswith("OPERATIONAL_WORKER_") or key == "ENVIRONMENT":
|
|
monkeypatch.delenv(key)
|
|
if environment:
|
|
monkeypatch.setenv("ENVIRONMENT", environment)
|
|
monkeypatch.setenv("OPERATIONAL_WORKER_PROJECT_ID", str(uuid4()))
|
|
monkeypatch.setenv("OPERATIONAL_WORKER_CLIENT_ID", "test-service")
|
|
monkeypatch.setenv("OPERATIONAL_WORKER_CLIENT_SECRET", "TEST_ONLY_SECRET")
|
|
standard_url = "https://identity.test/token"
|
|
dev_url = "http://192.168.1.10:8080/token"
|
|
monkeypatch.setenv("OPERATIONAL_WORKER_TOKEN_URL", standard_url)
|
|
monkeypatch.setenv("OPERATIONAL_WORKER_DEV_TOKEN_URL", dev_url)
|
|
assert WorkerSettings.from_environment().token_url == (dev_url if use_dev_url else standard_url)
|
|
monkeypatch.delenv("OPERATIONAL_WORKER_DEV_TOKEN_URL")
|
|
assert WorkerSettings.from_environment().token_url == standard_url
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_all_worker_operations_share_investigation_creation_scope():
|
|
boundary = "2026-10-10T10:29:47.864313Z"
|
|
requests = []
|
|
|
|
def handler(request):
|
|
requests.append(request)
|
|
if request.url.path.endswith("/agent-messages/pending"):
|
|
return httpx.Response(200, json={"items": []})
|
|
if request.method == "GET":
|
|
assert request.url.params["task_type"] == "event_investigation"
|
|
assert request.url.params["task_created_from"] == boundary
|
|
return httpx.Response(200, json={"items": []})
|
|
assert request.url.path.endswith("leases/recover")
|
|
assert json.loads(request.content) == {
|
|
"limit": 100, "task_type": "event_investigation", "task_created_from": boundary,
|
|
}
|
|
return httpx.Response(200, json={"task_ids": []})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
|
|
bearer="TEST_ONLY", task_created_from=boundary)
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
assert await OperationalWorker(settings, client).tick() == 0
|
|
assert len(requests) == 5
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_worker_reconciles_messages_at_capacity_and_preserves_transport_retries():
|
|
pending = [{"message_id": str(uuid4())} for _ in range(3)]
|
|
failures, submitted = [], []
|
|
|
|
def handler(request):
|
|
path = request.url.path
|
|
if path.endswith("leases/recover"):
|
|
return httpx.Response(200, json={"items": []})
|
|
if path.endswith("agent-messages/pending"):
|
|
return httpx.Response(200, json={"items": pending})
|
|
if path.endswith("operational-task-runs"):
|
|
return httpx.Response(200, json={"items": [{"run_id": str(uuid4())}]
|
|
if request.url.params["status"] == "running" else []})
|
|
body = json.loads(request.content)
|
|
if path.endswith("task-messages/deliver"):
|
|
assert str(request.url) == "http://agent.test/api/v1/task-messages/deliver"
|
|
submitted.append(body["message_id"])
|
|
status = {pending[0]["message_id"]: 503, pending[1]["message_id"]: 404}.get(body["message_id"], 202)
|
|
return httpx.Response(status, json={})
|
|
assert path.endswith("/delivery") and request.url.host == "server.test"
|
|
failures.append((path, body))
|
|
return httpx.Response(200, json={})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
|
|
bearer="TEST_ONLY", max_concurrency=1)
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
assert await OperationalWorker(settings, client).tick() == 0
|
|
assert submitted == [row["message_id"] for row in pending]
|
|
assert len(failures) == 1 and pending[1]["message_id"] in failures[0][0]
|
|
assert failures[0][1]["status"] == "failed"
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_message_sweep_advances_past_idle_chats_across_bounded_ticks():
|
|
messages = [str(uuid4()) for _ in range(5)]
|
|
scans, delivered = [], []
|
|
|
|
def handler(request):
|
|
if request.method == "GET":
|
|
cursor = request.url.params.get("cursor")
|
|
page = int(cursor) if cursor else 0
|
|
scans.append(cursor)
|
|
return httpx.Response(200, json={
|
|
"items": [{"message_id": messages[page]}],
|
|
"next_cursor": str(page + 1) if page < 4 else None,
|
|
})
|
|
delivered.append(json.loads(request.content)["message_id"])
|
|
return httpx.Response(202, json={})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
worker = OperationalWorker(settings, client)
|
|
await worker._reconcile_messages({})
|
|
assert scans == [None, "1", "2", "3"] and worker._message_cursor == "4"
|
|
await worker._reconcile_messages({})
|
|
assert scans[-1] == "4" and worker._message_cursor is None
|
|
assert delivered == messages
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_persisted_queue_retries_with_backend_execution_identity():
|
|
task_id, run_id, execution_id = str(uuid4()), str(uuid4()), str(uuid4())
|
|
submitted = []
|
|
def handler(request):
|
|
assert request.headers["Authorization"] == "Bearer TEST_ONLY"
|
|
assert not request.url.path.endswith("schedules/dispatch-due")
|
|
if request.url.path.endswith("leases/recover"):
|
|
return httpx.Response(200, json={"items": []})
|
|
if request.method == "GET":
|
|
if request.url.params.get("status") != "queued":
|
|
return httpx.Response(200, json={"items": []})
|
|
return httpx.Response(200, json={"items": [{
|
|
"task_id": task_id, "run_id": run_id, "execution_request_id": execution_id,
|
|
"status": "queued", "revision": 1,
|
|
}]})
|
|
body = json.loads(request.content)
|
|
assert str(request.url) == "http://agent.test/api/v1/task-executions"
|
|
submitted.append(body)
|
|
return httpx.Response(202, json={"task_id": task_id, "run_id": run_id, "session_id": "execution-session"})
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
worker = OperationalWorker(settings, client)
|
|
assert await worker.tick() == 1
|
|
assert await worker.tick() == 1
|
|
assert submitted[0] == submitted[1]
|
|
assert submitted[0] == {
|
|
"task_id": task_id, "run_id": run_id, "execution_request_id": execution_id,
|
|
}
|
|
assert "credential" not in submitted[0]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
@pytest.mark.parametrize("run_status", ["claimed", "running", "failed"])
|
|
async def test_only_queued_runs_are_submitted(run_status):
|
|
def handler(request):
|
|
assert not request.url.path.endswith("schedules/dispatch-due")
|
|
if request.method == "GET":
|
|
requested = request.url.params.get("status")
|
|
if requested == run_status and run_status in {"claimed", "running"}:
|
|
return httpx.Response(200, json={"items": [{"run_id": str(uuid4()), "status": run_status}]})
|
|
return httpx.Response(200, json={"items": []})
|
|
assert request.url.host == "server.test"
|
|
return httpx.Response(200, json={"items": []})
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
assert await OperationalWorker(settings, client).tick() == 0
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_bounded_queue_continues_after_one_agent_handoff_fails():
|
|
runs = [{"task_id": str(uuid4()), "run_id": str(uuid4()),
|
|
"execution_request_id": str(uuid4()), "status": "queued",
|
|
"revision": 1} for _ in range(2)]
|
|
limits, submitted, failures = [], [], []
|
|
|
|
def handler(request):
|
|
if request.method == "GET":
|
|
if request.url.params.get("status") != "queued":
|
|
return httpx.Response(200, json={"items": []})
|
|
limits.append(int(request.url.params["limit"]))
|
|
return httpx.Response(200, json={"items": runs})
|
|
if request.url.host == "server.test":
|
|
if request.url.path.endswith("leases/recover"):
|
|
return httpx.Response(200, json={"items": []})
|
|
failures.append(json.loads(request.content))
|
|
return httpx.Response(200, json={})
|
|
body = json.loads(request.content)
|
|
submitted.append(body)
|
|
if body["task_id"] == runs[0]["task_id"]:
|
|
return httpx.Response(503, json={"detail": "temporary failure"})
|
|
return httpx.Response(202, json={"task_id": body["task_id"], "run_id": body["run_id"]})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
assert await OperationalWorker(settings, client).tick() == 1
|
|
assert limits == [settings.max_concurrency]
|
|
assert [body["task_id"] for body in submitted] == [run["task_id"] for run in runs]
|
|
assert [body["execution_request_id"] for body in submitted] == [run["execution_request_id"] for run in runs]
|
|
assert failures[0]["retryable"] is True
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_service_credentials_are_cached_and_refreshed_without_entering_task_payload(monkeypatch):
|
|
clock = [100.0]
|
|
monkeypatch.setattr("scripts.run_operational_worker.time.monotonic", lambda: clock[0])
|
|
requests = []
|
|
|
|
def handler(request):
|
|
requests.append(request)
|
|
assert request.url == "https://identity.test/token"
|
|
assert "authorization" not in request.headers
|
|
return httpx.Response(200, json={"access_token": f"TEST_ONLY_{len(requests)}", "expires_in": 60})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
|
|
token_url="https://identity.test/token", client_id="test-service",
|
|
client_secret="TEST_ONLY_SECRET")
|
|
provider = CredentialProvider(settings)
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
first = await provider.headers(client)
|
|
clock[0] = 140.0
|
|
assert await provider.headers(client) == first
|
|
assert len(requests) == 1
|
|
clock[0] = 146.0
|
|
refreshed = await provider.headers(client)
|
|
assert len(requests) == 2
|
|
assert first["Authorization"] == "Bearer TEST_ONLY_1"
|
|
assert refreshed == {"Authorization": "Bearer TEST_ONLY_2", "X-Project-Id": settings.project_id}
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_long_lived_oauth_token_is_refreshed_before_product_token_age_limit(monkeypatch):
|
|
clock, requests = [0.0], []
|
|
monkeypatch.setattr("scripts.run_operational_worker.time.monotonic", lambda: clock[0])
|
|
|
|
def handler(request):
|
|
requests.append(request)
|
|
return httpx.Response(200, json={"access_token": "TEST_ONLY", "expires_in": 3600})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()),
|
|
token_url="https://identity.test/token", client_id="test-service",
|
|
client_secret="TEST_ONLY_SECRET")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
provider = CredentialProvider(settings)
|
|
await provider.headers(client)
|
|
clock[0] = 301.0
|
|
await provider.headers(client)
|
|
assert len(requests) == 2
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_worker_refreshes_running_agent_credentials_after_restart():
|
|
run_id, project_id = str(uuid4()), str(uuid4())
|
|
renewed = []
|
|
|
|
def handler(request):
|
|
if request.url.path.endswith("leases/recover"):
|
|
return httpx.Response(200, json={"items": []})
|
|
if request.method == "GET":
|
|
items = [{"run_id": run_id, "agent_session_id": "existing-session"}] if request.url.params.get("status") == "running" else []
|
|
return httpx.Response(200, json={"items": items})
|
|
assert request.url.host == "agent.test"
|
|
assert str(request.url) == "http://agent.test/api/v1/session/existing-session/credentials"
|
|
assert request.headers["Authorization"] == "Bearer TEST_ONLY"
|
|
assert request.headers["X-Project-Id"] == project_id
|
|
assert json.loads(request.content) == {}
|
|
renewed.append(request)
|
|
return httpx.Response(202, json={"accepted": True})
|
|
|
|
settings = WorkerSettings("http://server.test", "http://agent.test", project_id, bearer="TEST_ONLY")
|
|
async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
|
|
assert await OperationalWorker(settings, client).tick() == 0
|
|
assert len(renewed) == 1
|