fix: page ready agents during health-check startup

agentex/src/temporal/run_healthcheck_workflow.py

@@ -18,6 +18,7 @@
 from src.utils.logging import make_logger
 
 logger = make_logger(__name__)
+READY_AGENT_PAGE_SIZE = 200
 
 
 async def main() -> None:
@@ -45,37 +46,50 @@ async def main() -> None:
         logger.error("Temporal is not configured, skipping workflow creation")
         return
 
-    # Initialize repository and list agents
+    # Initialize repository and list ready agents
     engine = database_async_read_write_engine()
     session_maker = database_async_read_write_session_maker(engine)
     read_only_session_maker = database_async_read_only_session_maker(engine)
     agent_repo = AgentRepository(session_maker, read_only_session_maker)
-    agents = await agent_repo.list()
 
     adapter = TemporalAdapter(temporal_client=global_dependencies.temporal_client)
     logger.info(f"Adding Health Check workflows to task queue: {task_queue}")
-    # Try to add health check workflows to task queue for each agent
-    for agent in agents:
-        if agent.status != AgentStatus.READY:
-            logger.info(
-                f"Agent {agent.id} is not ready, skipping health check workflow"
-            )
-            continue
-        try:
-            await adapter.start_workflow(
-                workflow_id=f"healthcheck_workflow_{agent.id}",
-                workflow=HealthCheckWorkflow,
-                args=[{"agent_id": agent.id, "acp_url": agent.acp_url}],
-                task_queue=task_queue,
-            )
-        except TemporalWorkflowAlreadyExistsError:
-            # Expected if workflow is already running for existing agent registration
-            logger.info(f"Health check workflow already exists for agent {agent.id}")
-        except Exception as e:
-            # Unexpected error, don't raise here to continue with the next agent
-            logger.error(
-                f"Failed to start health check workflow for agent {agent.id}: {e}"
-            )
+
+    page_number = 1
+    while True:
+        agents = await agent_repo.list(
+            filters={"status": AgentStatus.READY},
+            limit=READY_AGENT_PAGE_SIZE,
+            page_number=page_number,
+            order_by="id",
+            order_direction="asc",
+        )
+        if not agents:
+            break
+
+        # Try to add health check workflows to task queue for each ready agent
+        for agent in agents:
+            try:
+                await adapter.start_workflow(
+                    workflow_id=f"healthcheck_workflow_{agent.id}",
+                    workflow=HealthCheckWorkflow,
+                    args=[{"agent_id": agent.id, "acp_url": agent.acp_url}],
+                    task_queue=task_queue,
+                )
+            except TemporalWorkflowAlreadyExistsError:
+                # Expected if workflow is already running for existing agent registration
+                logger.info(
+                    f"Health check workflow already exists for agent {agent.id}"
+                )
+            except Exception as e:
+                # Unexpected error, don't raise here to continue with the next agent
+                logger.error(
+                    f"Failed to start health check workflow for agent {agent.id}: {e}"
+                )
+
+        if len(agents) < READY_AGENT_PAGE_SIZE:
+            break
+        page_number += 1
 
 
 if __name__ == "__main__":

agentex/tests/unit/temporal/test_run_healthcheck_workflow.py

@@ -0,0 +1,151 @@
+from types import SimpleNamespace
+
+import pytest
+from src.temporal import run_healthcheck_workflow
+
+
+async def _run_main_with_pages(monkeypatch, pages):
+    class FakeGlobalDependencies:
+        temporal_client = object()
+
+        async def load(self):
+            return None
+
+    class FakeAgentRepository:
+        def __init__(self, *args):
+            self.calls = []
+
+        async def list(
+            self,
+            filters,
+            limit,
+            page_number,
+            order_by,
+            order_direction,
+        ):
+            self.calls.append(
+                {
+                    "filters": filters,
+                    "limit": limit,
+                    "page_number": page_number,
+                    "order_by": order_by,
+                    "order_direction": order_direction,
+                }
+            )
+            return pages.get(page_number, [])
+
+    class FakeTemporalAdapter:
+        def __init__(self, temporal_client):
+            self.temporal_client = temporal_client
+            self.started_workflows = []
+
+        async def start_workflow(self, **kwargs):
+            self.started_workflows.append(kwargs)
+
+    fake_repo = FakeAgentRepository()
+    fake_adapter = FakeTemporalAdapter(object())
+    fake_env = SimpleNamespace(
+        ENABLE_HEALTH_CHECK_WORKFLOW=True,
+        AGENTEX_SERVER_TASK_QUEUE="agentex-server",
+    )
+
+    monkeypatch.setattr(
+        run_healthcheck_workflow,
+        "GlobalDependencies",
+        FakeGlobalDependencies,
+    )
+    monkeypatch.setattr(
+        run_healthcheck_workflow.EnvironmentVariables,
+        "refresh",
+        lambda: fake_env,
+    )
+    monkeypatch.setattr(
+        run_healthcheck_workflow.TemporalClientFactory,
+        "is_temporal_configured",
+        lambda env: True,
+    )
+    monkeypatch.setattr(
+        run_healthcheck_workflow,
+        "database_async_read_write_engine",
+        lambda: object(),
+    )
+    monkeypatch.setattr(
+        run_healthcheck_workflow,
+        "database_async_read_write_session_maker",
+        lambda engine: object(),
+    )
+    monkeypatch.setattr(
+        run_healthcheck_workflow,
+        "database_async_read_only_session_maker",
+        lambda engine: object(),
+    )
+    monkeypatch.setattr(
+        run_healthcheck_workflow,
+        "AgentRepository",
+        lambda *args: fake_repo,
+    )
+    monkeypatch.setattr(
+        run_healthcheck_workflow,
+        "TemporalAdapter",
+        lambda temporal_client: fake_adapter,
+    )
+
+    await run_healthcheck_workflow.main()
+
+    return fake_repo, fake_adapter
+
+
+def _agents(count: int, prefix: str = "agent"):
+    return [
+        SimpleNamespace(id=f"{prefix}-{i}", acp_url=f"http://{prefix}-{i}")
+        for i in range(count)
+    ]
+
+
+def _expected_call(page_number: int):
+    return {
+        "filters": {"status": run_healthcheck_workflow.AgentStatus.READY},
+        "limit": run_healthcheck_workflow.READY_AGENT_PAGE_SIZE,
+        "page_number": page_number,
+        "order_by": "id",
+        "order_direction": "asc",
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.unit
+async def test_main_pages_ready_agents_for_healthcheck(monkeypatch):
+    final_agent = SimpleNamespace(id="agent-final", acp_url="http://agent-final")
+    fake_repo, fake_adapter = await _run_main_with_pages(
+        monkeypatch,
+        {
+            1: _agents(run_healthcheck_workflow.READY_AGENT_PAGE_SIZE),
+            2: [final_agent],
+        },
+    )
+
+    assert fake_repo.calls == [_expected_call(1), _expected_call(2)]
+    assert len(fake_adapter.started_workflows) == (
+        run_healthcheck_workflow.READY_AGENT_PAGE_SIZE + 1
+    )
+    assert (
+        fake_adapter.started_workflows[-1]["workflow_id"]
+        == "healthcheck_workflow_agent-final"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.unit
+async def test_main_checks_empty_page_for_exact_page_size(monkeypatch):
+    fake_repo, fake_adapter = await _run_main_with_pages(
+        monkeypatch,
+        {
+            1: _agents(run_healthcheck_workflow.READY_AGENT_PAGE_SIZE),
+            2: [],
+        },
+    )
+
+    assert fake_repo.calls == [_expected_call(1), _expected_call(2)]
+    assert len(fake_adapter.started_workflows) == (
+        run_healthcheck_workflow.READY_AGENT_PAGE_SIZE
+    )