diff --git a/src/pf1e_simulator/combat.py b/src/pf1e_simulator/combat.py index 585641f..8b3a15f 100644 --- a/src/pf1e_simulator/combat.py +++ b/src/pf1e_simulator/combat.py @@ -358,7 +358,7 @@ class Action: "attack", "move", "wait", "swift", "free", "immediate", "full_round", "full_attack", "charge", "withdraw", "5foot_step", "cast_spell", "maneuver", "channel_energy", "hex", "deed", "spell_combat", - "inspire_courage", + "inspire_courage", "retrieve_weapon", ] target_id: str | None = None spell_name: str | None = None @@ -370,6 +370,7 @@ class Action: channel_mode: Literal["heal", "harm"] | None = None # channel energy polarity for this use direction: tuple[int, int] | None = None # steering for directional maneuvers trick_condition: TrickCondition | None = None # dirty trick condition choice + weapon_id: str | None = None # weapon id for retrieve_weapon / repair_weapon @dataclass(frozen=True) @@ -1475,6 +1476,8 @@ class CombatEngine: state, target, action.maneuver, direction=action.direction, trick_condition=action.trick_condition, ) + elif kind == "retrieve_weapon": + self._retrieve_weapon(state, action.weapon_id) elif kind in ("channel_energy", "hex", "deed", "spell_combat", "inspire_courage"): self._execute_class_ability(state, target, action) elif kind == "charge": @@ -1621,6 +1624,30 @@ class CombatEngine: } handlers[maneuver]() + def _retrieve_weapon(self, state: CombatantState, weapon_id: str | None) -> None: + """Restore a stolen weapon to its owner (remove it from ``weapon_stolen``). + + No-op (with a log line) when ``weapon_id`` is None or the id was never + stolen. Once restored, the weapon is usable again via ``usable_attacks()``. + """ + if weapon_id is None: + self._log( + f"round {self._current_round} {state.combatant.id}: " + f"retrieve: nothing to retrieve" + ) + return + if weapon_id not in state.weapon_stolen: + self._log( + f"round {self._current_round} {state.combatant.id}: " + f"retrieve {weapon_id}: nothing to retrieve" + ) + return + state.weapon_stolen.remove(weapon_id) + self._log( + f"round {self._current_round} {state.combatant.id}: " + f"retrieve {weapon_id} (restored)" + ) + def _maneuver_cmb(self, state: CombatantState) -> int: return ( state.combatant.cmb diff --git a/tests/test_retrieve_weapon.py b/tests/test_retrieve_weapon.py new file mode 100644 index 0000000..f7ff08a --- /dev/null +++ b/tests/test_retrieve_weapon.py @@ -0,0 +1,153 @@ +"""Tests for the ``retrieve_weapon`` action (recovering a stolen weapon). + +L7 of the ``llm-policy-contract`` plan: ``Action`` gains a +``retrieve_weapon`` kind and a ``weapon_id`` field so a policy can restore a +weapon that a steal maneuver took from its owner. The weapon is removed from +``state.weapon_stolen`` and becomes usable again via ``usable_attacks()``. + +No default policy emits ``retrieve_weapon``; it is only accepted by +``_execute`` (usage via custom policy / LLM). +""" + +from __future__ import annotations + +from pf1e_simulator.combat import Action, CombatantState, CombatEngine +from pf1e_simulator.grid import Grid +from pf1e_simulator.map import MapSpec, TerrainType +from pf1e_simulator.models import ( + AbilityScores, + ACProfile, + AttackSpec, + Combatant, + DamageComponent, + Saves, +) +from pf1e_simulator.rng import ScriptedRng, SeededRng + + +def make_combatant(cid: str, *, attack_bonus: int = 2) -> Combatant: + attack = AttackSpec( + id=f"{cid}-w", + name="sword", + kind="melee", + attack_bonus=attack_bonus, + damage=[DamageComponent(formula="1d4", types=["slashing"])], + ) + return Combatant( + id=cid, + name=cid, + level=1, + size="Medium", + abilities=AbilityScores( + str_score=10, dex_score=10, con_score=10, + int_score=10, wis_score=10, cha_score=10, + ), + hp_max=6, + ac=ACProfile(total=13, touch=13, flat_footed=13), + bab=attack_bonus, + initiative_mod=0, + speed_land_ft=30, + saves=Saves(fort=0, ref=0, will=0), + attacks=[attack], + ) + + +def make_state(combatant: Combatant, side: str, pos: tuple[int, int]) -> CombatantState: + return CombatantState(combatant=combatant, side=side, pos=pos, hp=combatant.hp_max) + + +def make_grid() -> Grid: + legend = {".": TerrainType(type="floor", move_cost=1)} + spec = MapSpec(name="test", terrain=tuple(["." * 8] * 8), legend=legend) + return Grid.from_spec(spec) + + +def make_engine(queue: list[int], states: list[CombatantState]) -> CombatEngine: + return CombatEngine(ScriptedRng(queue), make_grid(), states) + + +def _weapon_ids(state: CombatantState) -> list[str]: + return [w.id for w in state.usable_attacks()] + + +# --- Test 1: retrieve restores a stolen weapon ------------------------------- + +def test_retrieve_restores_stolen_weapon() -> None: + a_spec = make_combatant("a", attack_bonus=2) + b_spec = make_combatant("b") + a = make_state(a_spec, "players", (1, 1)) + b = make_state(b_spec, "monsters", (1, 2)) + engine = make_engine([12], [a, b]) + engine._execute(a, Action(kind="maneuver", target_id="b", maneuver="steal")) + assert "b-w" in b.weapon_stolen + assert "b-w" not in _weapon_ids(b) + engine._execute(b, Action(kind="retrieve_weapon", weapon_id="b-w")) + assert not b.weapon_stolen + assert "b-w" in _weapon_ids(b) + + +# --- Test 2: weapon_id=None -> nothing to retrieve, state unchanged ---------- + +def test_retrieve_none_logs_nothing_to_retrieve() -> None: + a_spec = make_combatant("a", attack_bonus=2) + b_spec = make_combatant("b") + a = make_state(a_spec, "players", (1, 1)) + b = make_state(b_spec, "monsters", (1, 2)) + b.weapon_stolen.add("b-w") + engine = make_engine([], [a, b]) + engine._execute(b, Action(kind="retrieve_weapon")) + assert "nothing to retrieve" in engine._transcript[-1] + assert b.weapon_stolen == {"b-w"} + + +# --- Test 3: unknown id (never stolen) -> log, state unchanged --------------- + +def test_retrieve_unknown_id_logs_nothing_to_retrieve() -> None: + a_spec = make_combatant("a", attack_bonus=2) + b_spec = make_combatant("b") + a = make_state(a_spec, "players", (1, 1)) + b = make_state(b_spec, "monsters", (1, 2)) + b.weapon_stolen.add("b-w") + engine = make_engine([], [a, b]) + engine._execute(b, Action(kind="retrieve_weapon", weapon_id="ghost")) + assert "nothing to retrieve" in engine._transcript[-1] + assert b.weapon_stolen == {"b-w"} + + +# --- Test 4: real steal then retrieve restores the weapon -------------------- + +def test_steal_then_retrieve_restores_weapon() -> None: + a_spec = make_combatant("a", attack_bonus=2) + b_spec = make_combatant("b") + a = make_state(a_spec, "players", (1, 1)) + b = make_state(b_spec, "monsters", (1, 2)) + engine = make_engine([12], [a, b]) + engine._execute(a, Action(kind="maneuver", target_id="b", maneuver="steal")) + assert "b-w" in b.weapon_stolen + assert engine.weapon_for(b, a) is None + engine._execute(b, Action(kind="retrieve_weapon", weapon_id="b-w")) + assert not b.weapon_stolen + assert engine.weapon_for(b, a) is not None + assert "retrieve b-w (restored)" in engine._transcript[-1] + + +# --- Test 5: determinism ----------------------------------------------------- + +def _run_with_retrieve() -> tuple[str, ...]: + a_spec = make_combatant("a", attack_bonus=5) + b_spec = make_combatant("b") + a = make_state(a_spec, "players", (1, 1)) + b = make_state(b_spec, "monsters", (1, 2)) + b.weapon_stolen.add("b-w") + + def policy(engine: CombatEngine, state: CombatantState) -> tuple[Action, ...]: + if state.combatant.id == "b": + return (Action(kind="retrieve_weapon", weapon_id="b-w"),) + return (Action(kind="wait"),) + + engine = CombatEngine(SeededRng(42), make_grid(), [a, b], policy=policy) + return engine.run().transcript + + +def test_determinism_same_seed_same_transcript() -> None: + assert _run_with_retrieve() == _run_with_retrieve()