Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

2026-09-30 09:23:17 +00:00
1 parent 7082ab955a
commit 848f0d8f8a
@@ -0,0 +1,433 @@
+from __future__ import annotations
+
+import uuid
+from datetime import UTC, datetime
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+from pydantic import SecretStr
+from sqlalchemy import create_engine, select
+from sqlalchemy.orm import Session
+from sqlalchemy.pool import StaticPool
+
+from modelforge_api.api.routes.agent import (
+    get_agent_service,
+    get_decommission_service,
+)
+from modelforge_api.api.routes.agent import (
+    router as agent_router,
+)
+from modelforge_api.domain.agent_protocol import (
+    AgentMetadata,
+    HeartbeatRequest,
+    NodeManagementUpdate,
+)
+from modelforge_api.domain.node_decommission import NodeDecommissionExecute
+from modelforge_api.domain.node_lifecycle import (
+    NodeLifecycleChange,
+    NodeLifecyclePreviewRequest,
+)
+from modelforge_api.main import agent_protocol_error, node_decommission_error
+from modelforge_api.persistence.models import (
+    Accelerator,
+    AuditEvent,
+    Base,
+    ComputeNode,
+    GatewayRequest,
+    GpuLease,
+    NodeCredential,
+)
+from modelforge_api.services.node_agent import (
+    AgentAuthorizationError,
+    AgentConflictError,
+    NodeAgentService,
+    secret_hash,
+)
+from modelforge_api.services.node_decommission import (
+    NodeDecommissionError,
+    NodeDecommissionService,
+)
+from modelforge_api.settings import Settings, get_settings
+
+
+@pytest.fixture
+def session() -> Session:
+    engine = create_engine(
+        "sqlite+pysqlite:///:memory:",
+        connect_args={"check_same_thread": False},
+        poolclass=StaticPool,
+    )
+    Base.metadata.create_all(engine)
+    with Session(engine) as value:
+        yield value
+
+
+def _node(session: Session, *, enabled: bool = True, liveness: str = "offline") -> ComputeNode:
+    node = ComputeNode(
+        key=f"node-{uuid.uuid4()}",
+        hostname="node-lifecycle-test",
+        display_name="Lifecycle test node",
+        enabled=enabled,
+        status="active",
+        liveness_state=liveness,
+        inventory={"source": "test"},
+        observation_source="remote_agent",
+        agent_protocol_version=1,
+        agent_capabilities=["hardware.inventory"],
+        agent_version="1.3.0-rc.1",
+        last_heartbeat_at=datetime.now(UTC),
+    )
+    session.add(node)
+    session.commit()
+    return node
+
+
+def _request(preview, *, target_enabled: bool | None = None) -> NodeLifecycleChange:
+    return NodeLifecycleChange(
+        target_enabled=preview.target_enabled if target_enabled is None else target_enabled,
+        expected_generation=preview.node_generation,
+        dependency_digest=preview.dependency_digest,
+        operator="test-operator",
+        reason="Planned node maintenance with workloads drained.",
+        confirmation=preview.persisted_identity,
+    )
+
+
+def test_disable_requires_preview_confirmation_and_audits_transition(session: Session) -> None:
+    node = _node(session, liveness="online")
+    service = NodeDecommissionService(session)
+
+    preview = service.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
+    )
+    assert preview.can_apply is True
+    assert preview.current_enabled is True
+    assert preview.target_enabled is False
+    assert preview.blockers == []
+
+    result = service.apply_enabled_change(node.id, _request(preview))
+
+    session.refresh(node)
+    assert result.enabled is False
+    assert result.liveness == "disabled"
+    assert node.enabled is False
+    assert node.liveness_state == "disabled"
+    event = session.scalar(
+        select(AuditEvent).where(
+            AuditEvent.action == "NODE_DISABLED",
+            AuditEvent.resource_id == str(node.id),
+        )
+    )
+    assert event is not None
+    assert event.details["reason"] == "Planned node maintenance with workloads drained."
+
+
+def test_disable_fails_closed_for_active_gateway_request_and_gpu_lease(session: Session) -> None:
+    node = _node(session)
+    accelerator = Accelerator(
+        compute_node_id=node.id,
+        device_index=0,
+        device_uuid="GPU-node-lifecycle",
+        name="Test GPU",
+    )
+    session.add(accelerator)
+    session.flush()
+    gateway = GatewayRequest(
+        request_id=uuid.uuid4(),
+        capability_key="assistant.general",
+        capability_version=1,
+        compute_node_id=node.id,
+        status="running",
+        priority="production",
+        input_sha256="a" * 64,
+        input_count=1,
+    )
+    lease = GpuLease(
+        accelerator_id=accelerator.id,
+        deployment_id=uuid.uuid4(),
+        priority="production",
+        reserved_vram_mb=1024,
+        state="active",
+    )
+    session.add_all([gateway, lease])
+    session.commit()
+    service = NodeDecommissionService(session)
+
+    preview = service.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
+    )
+    assert preview.can_apply is False
+    assert {item.code for item in preview.blockers} >= {
+        "active_gateway_request",
+        "active_gpu_lease",
+    }
+    with pytest.raises(NodeDecommissionError) as raised:
+        service.apply_enabled_change(node.id, _request(preview))
+    assert raised.value.code == "node_lifecycle_blocked"
+    session.refresh(node)
+    assert node.enabled is True
+
+
+def test_disable_rejects_preview_when_work_arrives_after_impact_review(session: Session) -> None:
+    node = _node(session)
+    service = NodeDecommissionService(session)
+    preview = service.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
+    )
+    request = _request(preview)
+    session.add(
+        GatewayRequest(
+            request_id=uuid.uuid4(),
+            capability_key="assistant.general",
+            capability_version=1,
+            compute_node_id=node.id,
+            status="queued",
+            priority="production",
+            input_sha256="b" * 64,
+            input_count=1,
+        )
+    )
+    session.commit()
+
+    with pytest.raises(NodeDecommissionError) as raised:
+        service.apply_enabled_change(node.id, request)
+    assert raised.value.code == "node_lifecycle_preview_stale"
+    session.refresh(node)
+    assert node.enabled is True
+    assert node.liveness_state == "offline"
+
+
+def test_node_generation_change_invalidates_lifecycle_preview(session: Session) -> None:
+    node = _node(session)
+    service = NodeDecommissionService(session)
+    preview = service.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
+    )
+    request = _request(preview)
+    node.display_name = "Renamed after preview"
+    session.commit()
+
+    with pytest.raises(NodeDecommissionError) as raised:
+        service.apply_enabled_change(node.id, request)
+    assert raised.value.code == "node_generation_conflict"
+    session.refresh(node)
+    assert node.enabled is True
+
+
+def test_enable_requires_active_remote_credential_and_preserves_unknown_readiness(
+    session: Session,
+) -> None:
+    node = _node(session, enabled=False, liveness="disabled")
+    session.add(
+        NodeCredential(
+            compute_node_id=node.id,
+            secret_hash="c" * 64,
+            scope="node.publish",
+        )
+    )
+    session.commit()
+    service = NodeDecommissionService(session)
+
+    preview = service.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=True)
+    )
+    assert preview.can_apply is True
+    result = service.apply_enabled_change(node.id, _request(preview))
+
+    session.refresh(node)
+    assert result.enabled is True
+    assert result.liveness == "offline"
+    assert result.agent_evidence_state == "awaiting_fresh_authenticated_heartbeat"
+    assert node.enabled is True
+    assert node.liveness_state == "offline"
+    assert node.agent_protocol_version is None
+    assert node.agent_capabilities == []
+    assert node.agent_version is None
+    assert node.last_heartbeat_at is None
+    assert "awaiting fresh authenticated heartbeat" in node.last_connection_error
+
+
+def test_enable_without_active_credential_and_enable_of_tombstone_are_blocked(
+    session: Session,
+) -> None:
+    node = _node(session, enabled=False, liveness="disabled")
+    service = NodeDecommissionService(session)
+    preview = service.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=True)
+    )
+    assert preview.can_apply is False
+    assert "active_node_credential_missing" in {item.code for item in preview.blockers}
+    with pytest.raises(NodeDecommissionError) as raised:
+        service.apply_enabled_change(node.id, _request(preview))
+    assert raised.value.code == "node_lifecycle_blocked"
+    with pytest.raises(AgentConflictError, match="impact preview and typed confirmation"):
+        NodeAgentService(session, Settings(_env_file=None)).update_node(
+            node.id, NodeManagementUpdate(enabled=True)
+        )
+
+
+def test_direct_node_patch_rejects_enabled_but_keeps_metadata_updates(session: Session) -> None:
+    node = _node(session)
+    service = NodeAgentService(session, Settings(_env_file=None))
+
+    updated = service.update_node(node.id, NodeManagementUpdate(display_name="Renamed safely"))
+    assert updated.display_name == "Renamed safely"
+    with pytest.raises(AgentConflictError, match="impact preview and typed confirmation"):
+        service.update_node(node.id, NodeManagementUpdate(enabled=False))
+    session.refresh(node)
+    assert node.enabled is True
+
+
+def test_agent_heartbeat_and_post_body_auth_refuse_node_disabled_after_prebody_auth(
+    session: Session,
+) -> None:
+    node = _node(session)
+    credential_id = uuid.uuid4()
+    raw_credential = f"mfnode_{credential_id}_race-test-secret"
+    session.add(
+        NodeCredential(
+            id=credential_id,
+            compute_node_id=node.id,
+            secret_hash=secret_hash(raw_credential),
+            scope="node.publish",
+        )
+    )
+    session.commit()
+
+    # Keep the pre-body identity cached like a request whose lifecycle changes while its body is
+    # being read. A later admin transaction disables the same row before request dispatch.
+    agent_session = Session(session.get_bind(), expire_on_commit=False)
+    agent_service = NodeAgentService(agent_session, Settings(_env_file=None))
+    (_credential, stale_node), evidence = agent_service.authenticate_with_evidence(
+        f"Bearer {raw_credential}"
+    )
+    agent_session.commit()
+    assert stale_node.enabled is True
+
+    lifecycle_session = Session(session.get_bind())
+    lifecycle = NodeDecommissionService(lifecycle_session)
+    preview = lifecycle.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=False)
+    )
+    lifecycle.apply_enabled_change(node.id, _request(preview))
+    lifecycle_session.close()
+
+    with pytest.raises(AgentAuthorizationError, match="compute node is disabled"):
+        agent_service.reuse_authentication(evidence, f"Bearer {raw_credential}")
+
+    heartbeat = HeartbeatRequest(
+        identity_key=node.key,
+        metadata=AgentMetadata(
+            agent_version="1.3.0-rc.1",
+            protocol_version=1,
+            supported_capabilities=["hardware.inventory"],
+            started_at=datetime.now(UTC),
+        ),
+        observed_at=datetime.now(UTC),
+    )
+    with pytest.raises(AgentAuthorizationError, match="compute node is disabled"):
+        agent_service.heartbeat(stale_node, heartbeat)
+
+    session.refresh(node)
+    assert node.enabled is False
+    assert node.liveness_state == "disabled"
+    agent_session.close()
+
+
+def test_decommissioned_identity_cannot_be_reenabled(session: Session) -> None:
+    node = _node(session, enabled=False, liveness="disabled")
+    service = NodeDecommissionService(session)
+    retirement = service.preview(node.id)
+    service.execute(
+        node.id,
+        NodeDecommissionExecute(
+            expected_generation=retirement.node_generation,
+            preview_digest=retirement.dependency_digest,
+            idempotency_key=f"decom-{uuid.uuid4()}",
+            operator="test-operator",
+            reason="This test node has left service permanently.",
+            confirmation=retirement.persisted_identity,
+        ),
+    )
+
+    preview = service.preview_enabled_change(
+        node.id, NodeLifecyclePreviewRequest(target_enabled=True)
+    )
+    assert preview.can_apply is False
+    assert "node_decommissioned" in {item.code for item in preview.blockers}
+    with pytest.raises(NodeDecommissionError) as raised:
+        service.apply_enabled_change(node.id, _request(preview))
+    assert raised.value.code == "node_lifecycle_blocked"
+
+
+def test_admin_api_requires_preview_for_enabled_patch_but_preserves_metadata_patch(
+    session: Session,
+) -> None:
+    node = _node(session)
+    settings = Settings(_env_file=None, operator_api_key=SecretStr("admin-key"))
+    agent = NodeAgentService(session, settings)
+    lifecycle = NodeDecommissionService(session)
+    test_app = FastAPI()
+    test_app.include_router(agent_router)
+    test_app.add_exception_handler(AgentConflictError, agent_protocol_error)
+    test_app.add_exception_handler(NodeDecommissionError, node_decommission_error)
+    test_app.dependency_overrides[get_agent_service] = lambda: agent
+    test_app.dependency_overrides[get_decommission_service] = lambda: lifecycle
+    test_app.dependency_overrides[get_settings] = lambda: settings
+
+    try:
+        with TestClient(test_app) as client:
+            path = f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle/preview"
+            assert client.post(path, json={"target_enabled": False}).status_code == 401
+            headers = {"X-ModelForge-Admin-Token": "admin-key"}
+            preview_response = client.post(path, headers=headers, json={"target_enabled": False})
+            assert preview_response.status_code == 200
+            preview = preview_response.json()
+            assert preview["can_apply"] is True
+
+            bypass = client.patch(
+                f"/api/v1/admin/hardware/nodes/{node.id}",
+                headers=headers,
+                json={"enabled": False},
+            )
+            assert bypass.status_code == 409
+            assert bypass.json()["error"]["code"] == "agent_conflict"
+            metadata = client.patch(
+                f"/api/v1/admin/hardware/nodes/{node.id}",
+                headers=headers,
+                json={"display_name": "Updated metadata"},
+            )
+            assert metadata.status_code == 200
+            assert node.display_name == "Updated metadata"
+
+            stale_change = {
+                "target_enabled": False,
+                "expected_generation": preview["node_generation"],
+                "dependency_digest": preview["dependency_digest"],
+                "operator": "test-operator",
+                "reason": "Planned node maintenance after draining workloads.",
+                "confirmation": preview["persisted_identity"],
+            }
+            stale = client.post(
+                f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle",
+                headers=headers,
+                json=stale_change,
+            )
+            assert stale.status_code == 409
+            assert stale.json()["error"]["code"] == "node_generation_conflict"
+
+            refreshed_preview = client.post(
+                path, headers=headers, json={"target_enabled": False}
+            ).json()
+            change = {**stale_change, "expected_generation": refreshed_preview["node_generation"], "dependency_digest": refreshed_preview["dependency_digest"]}
+            applied = client.post(
+                f"/api/v1/admin/hardware/nodes/{node.id}/lifecycle",
+                headers=headers,
+                json=change,
+            )
+            assert applied.status_code == 200
+            assert applied.json()["enabled"] is False
+            assert applied.json()["liveness"] == "disabled"
+    finally:
+        test_app.dependency_overrides.clear()