hitl-continue-api-test.py

"""
Test Team HITL Continue Run API
===============================

This script tests the team continue run endpoint by:
1. Starting an AgentOS server with a team that has a tool requiring confirmation
2. Creating a run via the API (which will pause)
3. Continuing the run via the /continue API endpoint
"""

import json

from fastapi.testclient import TestClient

from agno.agent import Agent
from agno.db.postgres import PostgresDb
from agno.models.openai import OpenAIChat
from agno.os import AgentOS
from agno.team.team import Team
from agno.tools import tool


@tool(requires_confirmation=True)
def deploy_to_production(app_name: str, version: str) -> str:
    """Deploy an application to production."""
    return f"Successfully deployed {app_name} v{version} to production"


# Setup DB and Team
db = PostgresDb(db_url="postgresql+psycopg://ai:ai@localhost:5532/ai")

deploy_agent = Agent(
    name="Deploy Agent",
    model=OpenAIChat(id="gpt-4o"),
    tools=[deploy_to_production],
)

team = Team(
    id="devops-team",
    name="DevOps Team",
    members=[deploy_agent],
    model=OpenAIChat(id="gpt-4o"),
    db=db,

    store_member_responses=True,
)

# Create AgentOS and test client
agent_os = AgentOS(teams=[team])
app = agent_os.get_app()
client = TestClient(app)

if __name__ == "__main__":
    # Step 1: Create a run via API - this should pause
    print("--- Step 1: Creating team run via API (expects pause) ---")
    response = client.post(
        "/teams/devops-team/runs",
        data={"message": "Deploy the payments app version 2.1 to production", "stream": "false"},
    )
    print(f"Response status code: {response.status_code}")

    run_data = response.json()
    print(f"Run ID: {run_data.get('run_id')}")
    print(f"Session ID: {run_data.get('session_id')}")
    print(f"Status: {run_data.get('status')}")

    # Check if paused by status (is_paused is a computed property, not in JSON)
    is_paused = run_data.get("status") == "PAUSED"
    print(f"Is paused: {is_paused}")

    if not is_paused:
        print("WARNING: Run did not pause as expected!")
        print(f"Full response: {json.dumps(run_data, indent=2)}")
    else:
        print("Team paused as expected!")

        # Step 2: Get the requirements and confirm them
        print("\n--- Step 2: Processing requirements ---")
        requirements = run_data.get("requirements", [])
        print(f"Found {len(requirements)} requirements")

        # Confirm each requirement by setting confirmed=True on the tool_execution
        for req in requirements:
            if req.get("tool_execution"):
                tool_exec = req["tool_execution"]
                print(f"  Confirming tool: {tool_exec.get('tool_name')}({tool_exec.get('tool_args')})")
                # Set confirmed to True on the tool_execution
                tool_exec["confirmed"] = True
                tool_exec["confirmation_note"] = "Approved via API test"

        # Step 3: Continue the run via the new API endpoint
        print("\n--- Step 3: Continuing run via /continue API ---")
        run_id = run_data["run_id"]
        session_id = run_data["session_id"]

        continue_response = client.post(
            f"/teams/devops-team/runs/{run_id}/continue",
            data={
                "requirements": json.dumps(requirements),
                "session_id": session_id,
                "stream": "false",
            },
        )
        print(f"Continue response status code: {continue_response.status_code}")

        if continue_response.status_code == 200:
            continue_data = continue_response.json()
            print(f"Final status: {continue_data.get('status')}")
            print(f"Content: {str(continue_data.get('content', ''))[:200]}...")
            print("\n--- SUCCESS: Team continue run API works! ---")
        else:
            print(f"Continue failed: {continue_response.text}")