integrations/hcm21_mcp_server.py

"""MCP tool server for HCM:21 — lets an MCP client (e.g. Nous Hermes-Agent)
drive the OpenEnv HR environment natively, with zero bespoke glue.

Exposes the environment as five tools:
    hcm21_reset(seed, size, scenario)  -> start an episode
    hcm21_step(action_type, ...)       -> take one action (phase-gated)
    hcm21_available_actions()          -> valid action_types right now
    hcm21_state()                      -> full state snapshot
    hcm21_score()                      -> final 0-1 score (once the episode ends)

The env runs in-process (no HTTP), so a Hermes skill can scan -> plan -> produce
-> control across 6 quarters by calling these tools in a loop, then read the
final score to learn from the trajectory.

Run (stdio transport, the default Hermes/MCP expects):
    pip install -e ".[server]" fastmcp
    python -m integrations.hcm21_mcp_server

Register with Hermes by adding an MCP server entry pointing at that command
(see RUNBOOK.md, step 5).
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from fastmcp import FastMCP

from hr_env.models import HRAction
from hr_env.server.environment import HRProductivityEnvironment
from hr_env.server.scoring import compute_final_score

mcp = FastMCP("hcm21")

# Single in-process episode per server connection.
_env = HRProductivityEnvironment()
_started = False
_final_score: Optional[float] = None


def _obs_to_dict(obs: Any) -> Dict[str, Any]:
    """Compact, agent-friendly view of an observation."""
    return {
        "message": getattr(obs, "message", ""),
        "current_quarter": getattr(obs, "current_quarter", None),
        "current_phase": getattr(obs, "current_phase", None),
        "step_in_quarter": getattr(obs, "step_in_quarter", None),
        "available_actions": getattr(obs, "available_actions", []),
        "warnings": getattr(obs, "warnings", []),
        "reward": getattr(obs, "reward", None),
        "done": getattr(obs, "done", False),
        "data": getattr(obs, "data", None),
    }


@mcp.tool()
def hcm21_reset(seed: int = 42, size: int = 300, scenario: Optional[str] = None) -> Dict[str, Any]:
    """Start a new 6-quarter episode and return the baseline observation.

    scenario: one of high_eng_turnover | budget_cuts | rapid_growth |
    balanced_optimization, or omit for the default company.
    """
    global _started, _final_score
    kwargs: Dict[str, Any] = {"size": size}
    if scenario:
        kwargs["scenario"] = scenario
    obs = _env.reset(seed=seed, **kwargs)
    _started = True
    _final_score = None
    return _obs_to_dict(obs)


@mcp.tool()
def hcm21_step(
    action_type: str,
    department: Optional[str] = None,
    employee_ids: Optional[List[str]] = None,
    metric_name: Optional[str] = None,
    amount: Optional[float] = None,
    count: Optional[int] = None,
    parameters: Optional[Dict[str, Any]] = None,
    rationale: Optional[str] = None,
) -> Dict[str, Any]:
    """Take one action. Actions are phase-gated — call hcm21_available_actions
    first if unsure. When the episode ends, `done` is true and the final score is
    available via hcm21_score().
    """
    global _final_score
    if not _started:
        return {"error": "Call hcm21_reset first."}
    action = HRAction(
        action_type=action_type,
        department=department,
        employee_ids=employee_ids,
        metric_name=metric_name,
        amount=amount,
        count=count,
        parameters=parameters,
        rationale=rationale,
    )
    obs = _env.step(action)
    result = _obs_to_dict(obs)
    if result.get("done"):
        _final_score = compute_final_score(_env.state.metric_history)
        result["final_score"] = _final_score
    return result


@mcp.tool()
def hcm21_available_actions() -> Dict[str, Any]:
    """Return the action_types currently valid in the active phase."""
    if not _started:
        return {"error": "Call hcm21_reset first."}
    st = _env.state
    return {
        "current_phase": st.current_phase,
        "current_quarter": st.current_quarter,
        "available_actions": _env._phase_manager.get_available_actions(),
    }


@mcp.tool()
def hcm21_state() -> Dict[str, Any]:
    """Return the full environment state (company snapshot, metric history, events)."""
    if not _started:
        return {"error": "Call hcm21_reset first."}
    st = _env.state
    return {
        "current_quarter": st.current_quarter,
        "current_phase": st.current_phase,
        "total_steps": st.total_steps,
        "company_snapshot": st.company_snapshot,
        "metric_history": st.metric_history,
        "budget_remaining": st.budget_remaining,
        "recent_events": st.recent_events,
    }


@mcp.tool()
def hcm21_score() -> Dict[str, Any]:
    """Return the final 0-1 episode score (only available once the episode is done)."""
    if _final_score is None:
        return {"error": "Episode not finished. Final score is computed after Q6."}
    return {"final_score": _final_score}


if __name__ == "__main__":
    mcp.run()