Source code for chariot.cli.awm.chat.sessions

"""Session and agent selection helpers for AWM chat."""

import click
import questionary

from chariot import awm
from chariot.cli.awm.common import (
    _BACK,
    _QUESTIONARY_STYLE,
    _flush_stdin,
)
from chariot.identity import whoami


def _current_user_id() -> str:
    """Fetch the authenticated user id from identity `/users/me`."""
    user_id = whoami().get("id")
    if not user_id:
        raise click.ClickException("Could not determine current user ID from identity API.")
    return str(user_id).strip()


def _session_schedule_indicator(session: awm.models.SessionDetails) -> str:
    if not session.schedule_id:
        return ""
    return "  [scheduled]"


def _session_timestamp(session: awm.models.SessionDetails) -> str:
    return session.created_at.strftime("%Y-%m-%d %H:%M")


def _session_row_title(
    session: awm.models.SessionDetails,
    *,
    time_w: int,
    id_w: int,
    count_w: int,
    include_agents: bool = False,
) -> str:
    row = (
        f"{_session_timestamp(session).ljust(time_w)}  "
        f"{session.session_id.ljust(id_w)}  "
        f"{str(session.user_activity_count).rjust(count_w)}"
    )
    if include_agents:
        agents = ", ".join(session.agents) if session.agents else "—"
        row += f"  {agents}"
    return row + _session_schedule_indicator(session)


def _session_column_widths(sessions: list[awm.models.SessionDetails]) -> tuple[int, int, int]:
    time_w = max(
        16, len("Session Time"), max((len(_session_timestamp(s)) for s in sessions), default=0)
    )
    id_w = max(28, len("Session ID"), max((len(s.session_id) for s in sessions), default=0))
    count_w = max(
        4, len("Msgs"), max((len(str(s.user_activity_count)) for s in sessions), default=0)
    )
    return time_w, id_w, count_w


def _list_user_sessions(
    workflow_id: str, agent_name: str | None = None
) -> list[awm.models.SessionDetails]:
    current_user_id = _current_user_id()
    return list(
        awm.activity.get_sessions(
            workflow_id,
            user_ids=[current_user_id],
            agents=[agent_name] if agent_name else None,
        )
    )


def _fetch_session_history(workflow_id: str, session_id: str) -> list[awm.models.Activity]:
    types = [
        "user_to_agent_message",
        "agent_to_user_message",
        "error",
        "thinking",
        "tool_call_requests",
        "tool_call_responses",
        "clear_context",
        "tool_call_approvals",
    ]
    # TODO(https://striveworks.atlassian.net/browse/AF2-419) handle longer sessions
    return list(
        awm.activity.get_activities(
            workflow_id,
            session_id=session_id,
            types=types,
            sort_order="asc",
            limit=100,
        )
    )




def extract_agent_names(config: dict) -> list[str]:
    """Extract valid chat agent names from workflow config.

    Supported locations:
    - config.orchestrator.name
    - config.orchestrator.service_agents[*].name
    - config.service_agents[*].name
    - config.serviceAgents[*].name
    - config["service agents"][*].name
    """
    if not isinstance(config, dict):
        return []

    names: list[str] = []

    def _add_name(value: object) -> None:
        if not value:
            return
        name = str(value).strip()
        if name and name not in names:
            names.append(name)

    def _add_agents_from_collection(collection: object) -> None:
        if isinstance(collection, list):
            for item in collection:
                if isinstance(item, dict):
                    _add_name(item.get("name"))
                else:
                    _add_name(item)
        elif isinstance(collection, dict):
            # Sometimes service agents are represented as keyed objects.
            for key, value in collection.items():
                if isinstance(value, dict) and value.get("name"):
                    _add_name(value.get("name"))
                else:
                    _add_name(key)

    orchestrator = config.get("orchestrator")
    if isinstance(orchestrator, dict):
        _add_name(orchestrator.get("name"))
        _add_agents_from_collection(orchestrator.get("service_agents"))
        _add_agents_from_collection(orchestrator.get("serviceAgents"))
        _add_agents_from_collection(orchestrator.get("service agents"))

    _add_agents_from_collection(config.get("service_agents"))
    _add_agents_from_collection(config.get("serviceAgents"))
    _add_agents_from_collection(config.get("service agents"))
    _add_agents_from_collection(config.get("agents"))
    return names



def _pick_agent(workflow: awm.models.Workflow) -> str | None:
    """Prompt the user to select an agent from the workflow config."""
    latest_workflow = awm.workflows.get_workflow(workflow.id)
    suggested = extract_agent_names(latest_workflow.config or {})

    if not suggested:
        click.echo(
            click.style(
                "  Could not determine valid agents from workflow config. "
                "Expected names under orchestrator/service_agents.",
                fg="red",
            )
        )
        return None

    if len(suggested) == 1:
        return suggested[0]

    _flush_stdin()
    agent_name = questionary.select(
        "Select agent:",
        choices=[questionary.Choice(title=name, value=name) for name in suggested]
        + [questionary.Choice(title="↩ Cancel", value=_BACK)],
        style=_QUESTIONARY_STYLE,
    ).ask()
    return None if (agent_name is None or agent_name == _BACK) else agent_name