agent-platform/services/team-service/app/application/services.py

from dataclasses import dataclass
from datetime import datetime, timedelta

from core_events import EventPublishContract, EventServiceClient, EventServiceClientError
from core_domain import AgentRunContract, TeamMemberContract
from core_shared import JSONValue, try_build_redis_client
from core_shared.task_queue import TaskQueuePublisher

from app.bootstrap.settings import TeamServiceSettings
from app.db.models import TeamDefinition, TeamRun, TeamVersion
from app.domain.repositories import (
    TeamDefinitionRepository,
    TeamRunRepository,
    TeamVersionRepository)
from app.infrastructure.agent_client import AgentServiceClient, AgentServiceClientError
from app.schemas.team import (
    TeamConfigCreateRequestDto,
    TeamConfigUpdateRequestDto,
    TeamCreateRequest,
    TeamCreateRequestDto,
    TeamDeleteRequestDto,
    TeamRunCreateRequest,
    TeamRunCreateRequestDto,
    TeamRunExecuteRequest,
    TeamRunStatusUpdateRequestDto,
    TeamRunStatusUpdateRequest,
    TeamStatusUpdateRequest,
    TeamUpdateRequestDto,
    TeamVersionCreateRequest)


@dataclass(frozen=True)
class TeamMemberRunResult:
    member: TeamMemberContract
    run: AgentRunContract


class TeamApplicationService:
    def __init__(
        self,
        *,
        team_repository: TeamDefinitionRepository,
        team_version_repository: TeamVersionRepository,
        team_run_repository: TeamRunRepository,
        agent_client: AgentServiceClient | None = None,
        event_client: EventServiceClient | None = None,
        task_queue_publisher: TaskQueuePublisher | None = None) -> None:
        self.team_repository = team_repository
        self.team_version_repository = team_version_repository
        self.team_run_repository = team_run_repository
        self.agent_client = agent_client
        self.event_client = event_client
        self.task_queue_publisher = task_queue_publisher

    def create_team(self, payload: TeamCreateRequest) -> TeamDefinition:
        return self.team_repository.create(
            code=payload.code,
            name=payload.name,
            description=payload.description,
            team_type=payload.team_type,
            owner_user_id=payload.owner_user_id,
            metadata_json=payload.metadata_json)

    def create_team_from_contract(self, payload: TeamCreateRequestDto) -> TeamDefinition:
        return self.create_team(
            TeamCreateRequest(
                code=self._build_team_code(payload.name),
                name=payload.name,
                description=payload.description,
                team_type=payload.teamType,
                owner_user_id=payload.ownerUserId,
                metadata_json=payload.metadata))

    def list_teams(self) -> list[TeamDefinition]:
        return self.team_repository.list_all()

    def get_team(self, *, team_id: str) -> TeamDefinition | None:
        return self.team_repository.get_by_id(team_id=team_id)

    def update_team_from_contract(self, payload: TeamUpdateRequestDto) -> TeamDefinition | None:
        entity = self.team_repository.get_by_id(team_id=payload.teamId)
        if entity is None:
            return None
        if payload.name is not None:
            entity.name = payload.name
            entity.code = self._build_team_code(payload.name)
        if payload.description is not None:
            entity.description = payload.description
        if payload.teamType is not None:
            entity.team_type = payload.teamType
        if payload.status is not None:
            entity.status = payload.status
        if payload.ownerUserId is not None:
            entity.owner_user_id = payload.ownerUserId
        if payload.metadata is not None:
            entity.metadata_json = payload.metadata
        return self.team_repository.save(entity)

    def delete_team_from_contract(self, payload: TeamDeleteRequestDto) -> bool:
        entity = self.team_repository.get_by_id(team_id=payload.teamId)
        if entity is None:
            return False
        self.team_repository.delete(entity)
        return True

    def update_team_status(
        self,
        *,
        team_id: str,
        payload: TeamStatusUpdateRequest) -> TeamDefinition | None:
        return self.team_repository.update_status(
            team_id=team_id,
            status=payload.status)

    def create_team_version(self, payload: TeamVersionCreateRequest) -> TeamVersion:
        team = self.team_repository.get_by_id(team_id=payload.team_id)
        if team is None:
            raise ValueError(f"team not found: {payload.team_id}")
        if not payload.member_refs:
            raise ValueError("team version requires at least one member")

        return self.team_version_repository.create(
            team_id=payload.team_id,
            status=payload.status,
            coordination_mode=payload.coordination_mode,
            objective=payload.objective,
            member_refs_json=[item.model_dump(mode="json") for item in payload.member_refs],
            policy_json=payload.policy_json)

    def create_team_config_from_contract(self, payload: TeamConfigCreateRequestDto) -> TeamVersion:
        return self.create_team_version(
            TeamVersionCreateRequest(
                team_id=payload.teamId,
                status=payload.status,
                coordination_mode=payload.coordinationMode,
                objective=payload.objective,
                member_refs=self._normalize_member_refs(payload.memberRefs),
                policy_json=payload.policy))

    def list_team_versions(self, *, team_id: str) -> list[TeamVersion]:
        return self.team_version_repository.list_by_team(team_id=team_id)

    def list_team_configs(self, *, team_id: str | None = None) -> list[TeamVersion]:
        if team_id is not None:
            return self.team_version_repository.list_by_team(team_id=team_id)
        return self.team_version_repository.list_all()

    def get_team_config(self, *, config_id: str) -> TeamVersion | None:
        return self.team_version_repository.get_by_id(team_version_id=config_id)

    def update_team_config_from_contract(self, payload: TeamConfigUpdateRequestDto) -> TeamVersion | None:
        entity = self.team_version_repository.get_by_id(team_version_id=payload.configId)
        if entity is None:
            return None
        if payload.status is not None:
            entity.status = payload.status
            entity.published_time = datetime.utcnow() if payload.status == "published" else entity.published_time
        if payload.coordinationMode is not None:
            entity.coordination_mode = payload.coordinationMode
        if payload.objective is not None:
            entity.objective = payload.objective
        if payload.memberRefs is not None:
            normalized = self._normalize_member_refs(payload.memberRefs)
            if not normalized:
                raise ValueError("team config requires at least one member")
            entity.member_refs_json = [item.model_dump(mode="json") for item in normalized]
        if payload.policy is not None:
            entity.policy_json = payload.policy
        return self.team_version_repository.save(entity)

    def delete_team_config(self, *, config_id: str) -> bool:
        entity = self.team_version_repository.get_by_id(team_version_id=config_id)
        if entity is None:
            return False
        self.team_version_repository.delete(entity)
        return True

    def create_team_run(self, payload: TeamRunCreateRequest) -> TeamRun:
        team_version = self._resolve_team_version(
            team_id=payload.team_id,
            team_version_id=payload.team_version_id)
        if team_version is None:
            raise ValueError("published team version not found")

        team_run = self.team_run_repository.create(
            team_id=payload.team_id,
            team_version_id=team_version.id,
            session_id=payload.session_id,
            input_text=payload.input_text,
            input_json=payload.input_json)
        self._publish_event(
            event_type="team.run.created",
            team_run=team_run,
            payload_json={"team_run_id": team_run.id, "status": team_run.status})
        if self.task_queue_publisher is not None:
            self.task_queue_publisher.publish_team_run(
                team_run_id=team_run.id)
        return team_run

    def create_team_run_from_contract(self, payload: TeamRunCreateRequestDto) -> TeamRun:
        return self.create_team_run(
            TeamRunCreateRequest(
                team_id=payload.teamId,
                team_version_id=payload.teamConfigId,
                session_id=payload.sessionId,
                input_text=payload.inputText,
                input_json=payload.inputJson))

    def list_team_runs(
        self,
        *,
        team_id: str | None = None,
        session_id: str | None = None) -> list[TeamRun]:
        return self.team_run_repository.list_by_scope(
            team_id=team_id,
            session_id=session_id)

    def get_team_run(self, *, team_run_id: str) -> TeamRun | None:
        return self.team_run_repository.get_by_id(team_run_id=team_run_id)

    def delete_team_run(self, *, team_run_id: str) -> bool:
        entity = self.team_run_repository.get_by_id(team_run_id=team_run_id)
        if entity is None:
            return False
        self.team_run_repository.delete(entity)
        return True

    def update_team_run_status(
        self,
        *,
        team_run_id: str,
        payload: TeamRunStatusUpdateRequest) -> TeamRun | None:
        entity = self.team_run_repository.get_by_id(
            team_run_id=team_run_id)
        if entity is None:
            return None
        return self.team_run_repository.update_status(
            team_run_id=team_run_id,
            status=payload.status,
            worker_key=payload.worker_key,
            output_text=payload.output_text,
            output_json=payload.output_json,
            error_code=payload.error_code,
            error_message=payload.error_message)

    def update_team_run_status_from_contract(
        self,
        payload: TeamRunStatusUpdateRequestDto) -> TeamRun | None:
        return self.update_team_run_status(
            team_run_id=payload.teamRunId,
            payload=TeamRunStatusUpdateRequest(
                status=payload.status,
                worker_key=payload.workerKey,
                output_text=payload.outputText,
                output_json=payload.outputJson,
                error_code=payload.errorCode,
                error_message=payload.errorMessage))

    def execute_team_run(
        self,
        *,
        team_run_id: str,
        payload: TeamRunExecuteRequest) -> TeamRun | None:
        team_run = self.team_run_repository.get_by_id(
            team_run_id=team_run_id)
        if team_run is None:
            return None

        team_version = self.team_version_repository.get_by_id(
            team_version_id=team_run.team_version_id)
        if team_version is None:
            failed_run = self.team_run_repository.update_status(
                team_run_id=team_run.id,
                status="failed",
                worker_key=payload.worker_key,
                error_code="team_version_missing",
                error_message=f"team version not found: {team_run.team_version_id}")
            if failed_run is not None:
                self._publish_event(
                    event_type="team.run.failed",
                    team_run=failed_run,
                    payload_json={
                        "team_run_id": failed_run.id,
                        "status": failed_run.status,
                        "error_code": "team_version_missing",
                    })
            return failed_run

        running_run = self.team_run_repository.update_status(
            team_run_id=team_run.id,
            status="running",
            worker_key=payload.worker_key)
        if running_run is None:
            return None

        members = self._read_team_members(team_version)
        if not members:
            return self.team_run_repository.update_status(
                team_run_id=team_run.id,
                status="failed",
                worker_key=payload.worker_key,
                error_code="team_members_missing",
                error_message="team version has no valid members")

        try:
            member_results = self._execute_members(
                team_run=team_run,
                team_version=team_version,
                members=members,
                worker_key=payload.worker_key,
                dry_run=payload.dry_run)
        except AgentServiceClientError as exc:
            return self.team_run_repository.update_status(
                team_run_id=team_run.id,
                status="failed",
                worker_key=payload.worker_key,
                error_code="agent_service_error",
                error_message=str(exc))

        failed_results = [item for item in member_results if item.run.status != "completed"]
        output_text = self._build_team_output_text(
            team_version=team_version,
            member_results=member_results)
        output_json: dict[str, JSONValue] = {
            "dry_run": payload.dry_run,
            "coordination_mode": team_version.coordination_mode,
            "team_version_id": team_version.id,
            "member_run_count": len(member_results),
            "member_results": [
                self._member_result_to_json(item) for item in member_results
            ],
        }
        if failed_results:
            failed_run = self.team_run_repository.update_status(
                team_run_id=team_run.id,
                status="failed",
                worker_key=payload.worker_key,
                output_text=output_text,
                output_json=output_json,
                error_code="member_run_failed",
                error_message=f"{len(failed_results)} member run(s) failed")
            if failed_run is not None:
                self._publish_event(
                    event_type="team.run.failed",
                    team_run=failed_run,
                    payload_json={
                        "team_run_id": failed_run.id,
                        "status": failed_run.status,
                        "failed_member_count": len(failed_results),
                    })
            return failed_run

        completed_run = self.team_run_repository.update_status(
            team_run_id=team_run.id,
            status="completed",
            worker_key=payload.worker_key,
            output_text=output_text,
            output_json=output_json)
        if completed_run is not None:
            self._publish_event(
                event_type="team.run.completed",
                team_run=completed_run,
                payload_json={
                    "team_run_id": completed_run.id,
                    "status": completed_run.status,
                    "member_run_count": len(member_results),
                })
        return completed_run

    def execute_next_claimed_team_run(
        self,
        *,
        worker_key: str,
        lease_seconds: int,
        stale_running_seconds: int,
        dry_run: bool) -> tuple[TeamRun, int] | None:
        released_lease_count = self.team_run_repository.release_expired_leases(
            now_time=datetime.utcnow(),
            stale_running_seconds=stale_running_seconds)
        claimed_team_run = self.team_run_repository.claim_next_queued(
            worker_key=worker_key,
            lease_expire_time=datetime.utcnow() + timedelta(seconds=lease_seconds))
        if claimed_team_run is None:
            return None

        result = self.execute_team_run(
            team_run_id=claimed_team_run.id,
            payload=TeamRunExecuteRequest(
                worker_key=worker_key,
                dry_run=dry_run))
        if result is None:
            return None
        return result, released_lease_count

    def _resolve_team_version(
        self,
        *,
        team_id: str,
        team_version_id: str | None) -> TeamVersion | None:
        if team_version_id is not None:
            return self.team_version_repository.get_by_id(
                team_version_id=team_version_id)
        return self.team_version_repository.get_latest_published(
            team_id=team_id)

    def _execute_members(
        self,
        *,
        team_run: TeamRun,
        team_version: TeamVersion,
        members: list[TeamMemberContract],
        worker_key: str | None,
        dry_run: bool) -> list[TeamMemberRunResult]:
        if self.agent_client is None:
            raise AgentServiceClientError("agent service client is not configured")

        member_results: list[TeamMemberRunResult] = []
        prior_outputs: list[dict[str, JSONValue]] = []
        for member in self._order_members(members):
            member_input_json = self._build_member_input_json(
                team_run=team_run,
                team_version=team_version,
                member=member,
                prior_outputs=prior_outputs)
            created_run = self.agent_client.create_agent_run(
                agent_id=member.agent_id,
                agent_version_id=member.agent_version_id,
                session_id=team_run.session_id,
                input_text=self._build_member_input_text(
                    team_run=team_run,
                    team_version=team_version,
                    member=member),
                input_json=member_input_json)
            executed_run = self.agent_client.execute_agent_run(
                agent_run_id=created_run.id,
                worker_key=worker_key,
                dry_run=dry_run)
            member_results.append(TeamMemberRunResult(member=member, run=executed_run))
            prior_outputs.append(
                {
                    "member_key": member.member_key,
                    "member_role": member.role,
                    "member_name": member.name,
                    "agent_id": member.agent_id,
                    "agent_run_id": executed_run.id,
                    "status": executed_run.status,
                    "output_text": executed_run.output_text,
                    "output_json": executed_run.output_json or {},
                }
            )
        return member_results

    def _read_team_members(self, team_version: TeamVersion) -> list[TeamMemberContract]:
        members: list[TeamMemberContract] = []
        for item in team_version.member_refs_json:
            try:
                members.append(TeamMemberContract.model_validate(item))
            except ValueError:
                continue
        return members

    def _order_members(self, members: list[TeamMemberContract]) -> list[TeamMemberContract]:
        role_priority = {
            "planner": 0,
            "supervisor": 1,
            "specialist": 2,
            "executor": 3,
            "reviewer": 4,
        }
        return sorted(members, key=lambda item: role_priority.get(item.role, 10))

    def _build_member_input_text(
        self,
        *,
        team_run: TeamRun,
        team_version: TeamVersion,
        member: TeamMemberContract) -> str:
        lines = [
            f"Team objective: {team_version.objective or 'No objective provided.'}",
            f"Member role: {member.role}",
        ]
        if member.responsibility:
            lines.append(f"Responsibility: {member.responsibility}")
        if team_run.input_text:
            lines.append(f"User task: {team_run.input_text}")
        return "\n".join(lines)

    def _build_member_input_json(
        self,
        *,
        team_run: TeamRun,
        team_version: TeamVersion,
        member: TeamMemberContract,
        prior_outputs: list[dict[str, JSONValue]]) -> dict[str, JSONValue]:
        input_json: dict[str, JSONValue] = dict(team_run.input_json or {})
        input_json.update(
            {
                "team_id": team_run.team_id,
                "team_run_id": team_run.id,
                "team_version_id": team_version.id,
                "team_objective": team_version.objective,
                "member_key": member.member_key,
                "member_role": member.role,
                "member_responsibility": member.responsibility,
                "prior_member_outputs": prior_outputs,
            }
        )
        configured_input = member.config_json.get("input_json")
        if isinstance(configured_input, dict):
            input_json.update(
                {str(item_key): item_value for item_key, item_value in configured_input.items()}
            )
        return input_json

    def _build_team_output_text(
        self,
        *,
        team_version: TeamVersion,
        member_results: list[TeamMemberRunResult]) -> str:
        lines = [
            f"Team objective: {team_version.objective or 'No objective provided.'}",
            f"Coordination mode: {team_version.coordination_mode}",
            "Team conversation:",
        ]
        for index, item in enumerate(member_results, start=1):
            output_text = item.run.output_text or item.run.error_message or ""
            speaker = item.member.name or item.member.member_key
            lines.append(
                f"{index}. {speaker} ({item.member.role}) "
                f"status={item.run.status}: {output_text}")
        return "\n".join(lines)

    def _member_result_to_json(self, result: TeamMemberRunResult) -> dict[str, JSONValue]:
        return {
            "member_key": result.member.member_key,
            "member_role": result.member.role,
            "member_name": result.member.name,
            "member_responsibility": result.member.responsibility,
            "agent_run_id": result.run.id,
            "agent_id": result.run.agent_id,
            "agent_version_id": result.run.agent_version_id,
            "status": result.run.status,
            "output_text": result.run.output_text,
            "output_json": result.run.output_json or {},
            "error_code": result.run.error_code,
            "error_message": result.run.error_message,
        }

    def _publish_event(
        self,
        *,
        event_type: str,
        team_run: TeamRun,
        payload_json: dict[str, JSONValue]) -> None:
        if self.event_client is None:
            return
        try:
            self.event_client.publish_event(
                EventPublishContract(
                    event_type=event_type,
                    source_service="team-service",
                    aggregate_type="team_run",
                    aggregate_id=team_run.id,
                    correlation_id=team_run.session_id,
                    payload_json={
                        **payload_json,
                        "team_id": team_run.team_id,
                        "team_version_id": team_run.team_version_id,
                    })
            )
        except EventServiceClientError:
            return

    def _build_team_code(self, name: str) -> str:
        base = "".join(
            char.lower() if char.isalnum() else "_"
            for char in name
        ).strip("_") or "team"
        return base[:64]

    def _normalize_member_refs(self, member_refs: list[dict[str, JSONValue]]) -> list[TeamMemberContract]:
        members: list[TeamMemberContract] = []
        for index, item in enumerate(member_refs, start=1):
            role = item.get("role")
            normalized_role = "executor" if role == "worker" else role
            member = {
                **item,
                "member_key": item.get("member_key") or item.get("memberKey") or f"member_{index}",
                "agent_id": item.get("agent_id") or item.get("agentId"),
                "agent_version_id": item.get("agent_version_id") or item.get("agentVersionId"),
                "role": normalized_role or "specialist",
                "config_json": item.get("config_json") or item.get("configJson") or {},
            }
            members.append(TeamMemberContract.model_validate(member))
        return members


def build_team_application_service(
    *,
    team_repository: TeamDefinitionRepository,
    team_version_repository: TeamVersionRepository,
    team_run_repository: TeamRunRepository,
    settings: TeamServiceSettings) -> TeamApplicationService:
    redis_client = try_build_redis_client(settings.redis_url)
    return TeamApplicationService(
        team_repository=team_repository,
        team_version_repository=team_version_repository,
        team_run_repository=team_run_repository,
        agent_client=AgentServiceClient(
            base_url=settings.agent_service_url,
            timeout_seconds=settings.agent_service_timeout_seconds),
        event_client=EventServiceClient(
            base_url=settings.event_service_url,
            timeout_seconds=settings.event_service_timeout_seconds),
        task_queue_publisher=(
            TaskQueuePublisher(client=redis_client) if redis_client is not None else None
        ))