feat: extend operational investigation workflows

tests/unit/test_operational_worker.py

@@ -90,8 +90,7 @@ async def test_worker_reconciles_messages_at_capacity_and_preserves_transport_re
     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"
+    assert failures == []
 
 
 @pytest.mark.anyio
@@ -126,6 +125,8 @@ 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):
+        if request.url.path == "/health":
+            return httpx.Response(200, json={"ok": True, "agent_instance_id": "agent-a"})
         assert request.headers["Authorization"] == "Bearer TEST_ONLY"
         assert not request.url.path.endswith("schedules/dispatch-due")
         if request.url.path.endswith("leases/recover"):
@@ -140,7 +141,10 @@ async def test_persisted_queue_retries_with_backend_execution_identity():
         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"})
+        return httpx.Response(202, json={
+            "task_id": task_id, "run_id": run_id, "session_id": "execution-session",
+            "agent_instance_id": "agent-a",
+        })
     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)
@@ -149,6 +153,7 @@ async def test_persisted_queue_retries_with_backend_execution_identity():
     assert submitted[0] == submitted[1]
     assert submitted[0] == {
         "task_id": task_id, "run_id": run_id, "execution_request_id": execution_id,
+        "agent_instance_id": "agent-a",
     }
     assert "credential" not in submitted[0]
 
@@ -178,6 +183,8 @@ async def test_bounded_queue_continues_after_one_agent_handoff_fails():
     limits, submitted, failures = [], [], []
 
     def handler(request):
+        if request.url.path == "/health":
+            return httpx.Response(200, json={"ok": True, "agent_instance_id": "agent-a"})
         if request.method == "GET":
             if request.url.params.get("status") != "queued":
                 return httpx.Response(200, json={"items": []})
@@ -192,7 +199,10 @@ async def test_bounded_queue_continues_after_one_agent_handoff_fails():
         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"]})
+        return httpx.Response(202, json={
+            "task_id": body["task_id"], "run_id": body["run_id"],
+            "agent_instance_id": "agent-a",
+        })
 
     settings = WorkerSettings("http://server.test", "http://agent.test", str(uuid4()), bearer="TEST_ONLY")
     async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
@@ -257,10 +267,16 @@ async def test_worker_refreshes_running_agent_credentials_after_restart():
     renewed = []
 
     def handler(request):
+        if request.url.path == "/health":
+            return httpx.Response(200, json={"ok": True, "agent_instance_id": "agent-a"})
         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 []
+            items = [{
+                "run_id": run_id,
+                "agent_session_id": "existing-session",
+                "agent_instance_id": "agent-a",
+            }] 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"
@@ -274,3 +290,34 @@ async def test_worker_refreshes_running_agent_credentials_after_restart():
     async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
         assert await OperationalWorker(settings, client).tick() == 0
     assert len(renewed) == 1
+
+
+@pytest.mark.anyio
+async def test_worker_only_restores_sessions_owned_by_target_agent_instance():
+    own_run_id, foreign_run_id, project_id = str(uuid4()), str(uuid4()), str(uuid4())
+    renewed = []
+
+    def handler(request):
+        if request.url.path == "/health":
+            return httpx.Response(200, json={"ok": True, "agent_instance_id": "agent-a"})
+        if request.url.path.endswith("leases/recover"):
+            return httpx.Response(200, json={"items": []})
+        if request.method == "GET":
+            items = [{
+                "run_id": own_run_id,
+                "agent_session_id": "own-session",
+                "agent_instance_id": "agent-a",
+            }, {
+                "run_id": foreign_run_id,
+                "agent_session_id": "foreign-session",
+                "agent_instance_id": "agent-b",
+            }] if request.url.params.get("status") == "running" else []
+            return httpx.Response(200, json={"items": items})
+        assert str(request.url) == "http://agent.test/api/v1/session/own-session/credentials"
+        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