agent-system/orchestrator/main.py

import json
import threading
import time

from redis_client import get_redis_client
from result_listener import listen_for_results
from task_builder import build_task


def send_task(redis_client, task):
    redis_client.lpush("tasks", json.dumps(task))
    print(f"sent {task['action']} to {task['target']} ({task['task_id']})")


def dispatch(redis_client, command, ha_task_ids):
    if command == "ha":
        task = build_task("piha", "docker_ps", {})
        ha_task_ids.add(task["task_id"])
        send_task(redis_client, task)
    elif command == "all":
        send_task(redis_client, build_task("piha", "docker_ps", {}))
        send_task(redis_client, build_task("vps", "docker_ps", {}))
    else:
        send_task(redis_client, build_task("piha", "docker_ps", {}))


def main():
    redis_client = get_redis_client()
    ha_task_ids = set()

    listener = threading.Thread(
        target=listen_for_results,
        args=(redis_client, ha_task_ids),
        daemon=True,
    )
    listener.start()

    print("[orchestrator] ready")
    while True:
        try:
            command = input("> ").strip()
        except EOFError:
            print("[orchestrator] stdin closed, waiting...")
            time.sleep(2)
            continue
        except KeyboardInterrupt:
            print()
            continue

        if not command:
            continue

        dispatch(redis_client, command, ha_task_ids)


if __name__ == "__main__":
    main()