Federated Forge — Sub-Plan 01: Memoir Foundation¶

"""Schema models and gameplay-rule split logic for the Forge Memoir."""

from __future__ import annotations

import pytest

from kourai_common.federation.memoir_schema import SplitDecision


class TestSplitDecision:
    """SplitDecision is a frozen pydantic model with two booleans and an
    invariant that they cannot both be true."""

    def test_shared_eligible_only(self):
        d = SplitDecision(shared_eligible=True, private_only=False)
        assert d.shared_eligible is True
        assert d.private_only is False

    def test_private_only_only(self):
        d = SplitDecision(shared_eligible=False, private_only=True)
        assert d.shared_eligible is False
        assert d.private_only is True

    def test_neither_set_is_valid(self):
        # Used for entries that are still being constructed.
        d = SplitDecision(shared_eligible=False, private_only=False)
        assert d.shared_eligible is False
        assert d.private_only is False

    def test_both_true_is_rejected(self):
        with pytest.raises(ValueError, match="cannot both be true"):
            SplitDecision(shared_eligible=True, private_only=True)

"""Pydantic schemas for Forge Memoir entries.

Each Memoir entry has two faces — a narrative beat the visual novel can
replay, and a training tuple the federated-learning pipeline can consume.
The `SplitDecision` records which side of the council/bond split an entry
falls on. The `decide_split()` pure function applies the gameplay rules
from `docs/research/federated-forge/index.md`.
"""

from __future__ import annotations

from pydantic import BaseModel, ConfigDict, model_validator


class SplitDecision(BaseModel):
    """Whether a Memoir entry contributes to the federated council adapter,
    stays in the local bond adapter only, or has not been decided yet.

    Invariant: `shared_eligible` and `private_only` cannot both be true.
    Both false is valid for entries still under construction.
    """

    model_config = ConfigDict(frozen=True)

    shared_eligible: bool = False
    private_only: bool = False

    @model_validator(mode="after")
    def _at_most_one_true(self) -> SplitDecision:
        if self.shared_eligible and self.private_only:
            raise ValueError(
                "shared_eligible and private_only cannot both be true"
            )
        return self

from kourai_common.federation.memoir_schema import (
    EntrySource,
    decide_split,
)


class TestDecideSplit:
    """Pure function applying the gameplay rules from the design spec."""

    def test_specialist_proposed_output_is_shared(self):
        d = decide_split(EntrySource.SPECIALIST_PROPOSED, agent="kallos")
        assert d.shared_eligible is True
        assert d.private_only is False

    def test_player_revision_is_private(self):
        d = decide_split(EntrySource.PLAYER_REVISION, agent="kallos")
        assert d.shared_eligible is False
        assert d.private_only is True

    def test_federating_agent_interrupt_is_shared(self):
        d = decide_split(EntrySource.AGENT_INTERRUPT, agent="aidos")
        assert d.shared_eligible is True
        assert d.private_only is False

    def test_cupid_interrupt_is_private(self):
        d = decide_split(EntrySource.AGENT_INTERRUPT, agent="cupid")
        assert d.private_only is True
        assert d.shared_eligible is False

    def test_puck_interrupt_is_private(self):
        d = decide_split(EntrySource.AGENT_INTERRUPT, agent="puck")
        assert d.private_only is True
        assert d.shared_eligible is False

    def test_disagreement_resolution_is_shared(self):
        d = decide_split(EntrySource.DISAGREEMENT_RESOLUTION, agent="hephaestus")
        assert d.shared_eligible is True
        assert d.private_only is False

    def test_cupid_scene_is_private(self):
        d = decide_split(EntrySource.CUPID_SCENE, agent="cupid")
        assert d.private_only is True

    def test_puck_engagement_is_private(self):
        d = decide_split(EntrySource.PUCK_ENGAGEMENT, agent="puck")
        assert d.private_only is True

    def test_affinity_change_is_private(self):
        d = decide_split(EntrySource.AFFINITY_CHANGE, agent="cupid")
        assert d.private_only is True

    def test_player_profile_is_private(self):
        d = decide_split(EntrySource.PLAYER_PROFILE, agent=None)
        assert d.private_only is True

    def test_raw_transcript_is_private(self):
        d = decide_split(EntrySource.RAW_TRANSCRIPT, agent=None)
        assert d.private_only is True

    def test_task_description_is_private(self):
        d = decide_split(EntrySource.TASK_DESCRIPTION, agent=None)
        assert d.private_only is True

    def test_unknown_agent_for_interrupt_raises(self):
        with pytest.raises(ValueError, match="unknown agent"):
            decide_split(EntrySource.AGENT_INTERRUPT, agent="nobody")

from enum import StrEnum


class EntrySource(StrEnum):
    """Where a Memoir entry came from. The split decision is keyed off this
    enum and (for interrupts) the originating agent."""

    SPECIALIST_PROPOSED = "specialist_proposed"
    PLAYER_REVISION = "player_revision"
    AGENT_INTERRUPT = "agent_interrupt"
    DISAGREEMENT_RESOLUTION = "disagreement_resolution"
    CUPID_SCENE = "cupid_scene"
    PUCK_ENGAGEMENT = "puck_engagement"
    AFFINITY_CHANGE = "affinity_change"
    PLAYER_PROFILE = "player_profile"
    RAW_TRANSCRIPT = "raw_transcript"
    TASK_DESCRIPTION = "task_description"


# Agents that have a council adapter. Cupid and Puck are bond-only by
# construction (see spec, Goals/Non-goals section).
FEDERATING_AGENTS: frozenset[str] = frozenset(
    {
        "hephaestus",
        "metis",
        "techne",
        "dokimasia",
        "kallos",
        "mneme",
        "aidos",
        "aletheia",
    }
)
BOND_ONLY_AGENTS: frozenset[str] = frozenset({"cupid", "puck"})
ALL_AGENTS: frozenset[str] = FEDERATING_AGENTS | BOND_ONLY_AGENTS


def decide_split(source: EntrySource, *, agent: str | None) -> SplitDecision:
    """Apply the gameplay rules from the design spec.

    Patterns leave the forge; instances do not. Federating agents'
    proposed outputs and interrupts are shared-eligible; everything
    keyed to the player or the relationship layer is private-only.
    """
    if source in {
        EntrySource.PLAYER_REVISION,
        EntrySource.CUPID_SCENE,
        EntrySource.PUCK_ENGAGEMENT,
        EntrySource.AFFINITY_CHANGE,
        EntrySource.PLAYER_PROFILE,
        EntrySource.RAW_TRANSCRIPT,
        EntrySource.TASK_DESCRIPTION,
    }:
        return SplitDecision(shared_eligible=False, private_only=True)

    if source in {
        EntrySource.SPECIALIST_PROPOSED,
        EntrySource.DISAGREEMENT_RESOLUTION,
    }:
        return SplitDecision(shared_eligible=True, private_only=False)

    if source is EntrySource.AGENT_INTERRUPT:
        if agent is None or agent not in ALL_AGENTS:
            raise ValueError(f"unknown agent {agent!r} for AGENT_INTERRUPT")
        if agent in BOND_ONLY_AGENTS:
            return SplitDecision(shared_eligible=False, private_only=True)
        return SplitDecision(shared_eligible=True, private_only=False)

    raise ValueError(f"unhandled EntrySource {source!r}")

from kourai_common.federation.memoir_schema import (
    MemoirEntry,
    PlayerResponse,
    TrainingLabel,
)


class TestMemoirEntry:
    """The on-disk schema. Each entry carries narrative + training payloads
    and auto-populates its SplitDecision from EntrySource + agent."""

    def test_minimal_specialist_proposal_is_shared(self):
        entry = MemoirEntry(
            scene_id="session-1.turn-1",
            agent="kallos",
            source=EntrySource.SPECIALIST_PROPOSED,
            agent_proposed="some style edit",
        )
        assert entry.split.shared_eligible is True
        assert entry.split.private_only is False

    def test_cupid_scene_is_private(self):
        entry = MemoirEntry(
            scene_id="session-1.scene-cupid-7",
            agent="cupid",
            source=EntrySource.CUPID_SCENE,
            narrative_beat="cupid_late_night_check_in",
        )
        assert entry.split.private_only is True

    def test_pipeline_turn_with_player_response(self):
        entry = MemoirEntry(
            scene_id="session-1.turn-2",
            agent="kallos",
            source=EntrySource.SPECIALIST_PROPOSED,
            agent_proposed="lint fix",
            player_response=PlayerResponse(
                kind="modified", delta="player edits", felt="right"
            ),
            training_label=TrainingLabel(
                preference_pair=[
                    {"text": "lint fix", "score": 0},
                    {"text": "lint fix with player edits", "score": 1},
                ],
                weight=1.0,
            ),
        )
        assert entry.player_response.kind == "modified"
        assert entry.training_label.weight == 1.0

    def test_unknown_agent_rejected(self):
        with pytest.raises(ValueError, match="unknown agent"):
            MemoirEntry(
                scene_id="session-1.turn-3",
                agent="nobody",
                source=EntrySource.SPECIALIST_PROPOSED,
                agent_proposed="x",
            )

    def test_round_trip_json(self):
        original = MemoirEntry(
            scene_id="session-1.turn-1",
            agent="kallos",
            source=EntrySource.SPECIALIST_PROPOSED,
            agent_proposed="hello",
        )
        as_json = original.model_dump_json()
        restored = MemoirEntry.model_validate_json(as_json)
        assert restored == original

    def test_optional_context_block_round_trips(self):
        # Spec shows a `context` block with task_type, transcript_hash,
        # preceding_agents. We accept any dict for forward compatibility;
        # plan-02 will lock in the host-emitted shape.
        original = MemoirEntry(
            scene_id="session-1.turn-1",
            agent="kallos",
            source=EntrySource.SPECIALIST_PROPOSED,
            agent_proposed="hello",
            context={
                "task_type": "style_review",
                "transcript_hash": "sha256:abc",
                "preceding_agents": ["techne"],
            },
        )
        restored = MemoirEntry.model_validate_json(original.model_dump_json())
        assert restored.context["task_type"] == "style_review"
        assert restored.context["preceding_agents"] == ["techne"]

from typing import Any, Literal


class PlayerResponse(BaseModel):
    """How the player responded to an agent's proposal.

    `kind` is the high-level reaction; `delta` carries free-form payload
    (a diff, a comment, etc); `felt` is the optional Likert-style affect
    tag the host may choose to gather.
    """

    model_config = ConfigDict(frozen=True)

    kind: Literal["accepted", "modified", "rejected", "deferred"]
    delta: str | None = None
    felt: Literal["right", "off", "unsure"] | None = None


class TrainingLabel(BaseModel):
    """The FL-pipeline view of an entry. `preference_pair` is two scored
    candidates the trainer can consume as a DPO pair; `weight` lets the
    host emphasize or downweight specific entries (e.g. interrupted turns
    weigh less than fully-completed ones).
    """

    model_config = ConfigDict(frozen=True)

    preference_pair: list[dict[str, Any]] | None = None
    weight: float = 1.0


class MemoirEntry(BaseModel):
    """One Memoir entry on disk and on the wire.

    Fields are split into:

    - **identity / context** — `scene_id`, `agent`, `source`,
      `narrative_beat`
    - **narrative payload** — `agent_proposed`, `player_response`,
      `affinity_delta`
    - **training payload** — `training_label`
    - **derived** — `split`, decided automatically from `source` + `agent`
      via `decide_split()` if not provided

    The dual-face contract: every entry can be replayed by the VN via
    `narrative_beat`, AND consumed by the FL pipeline via `training_label`.
    """

    model_config = ConfigDict(frozen=True)

    scene_id: str
    agent: str
    source: EntrySource

    context: dict[str, Any] | None = None

    narrative_beat: str | None = None
    agent_proposed: str | None = None
    player_response: PlayerResponse | None = None
    affinity_delta: float = 0.0
    training_label: TrainingLabel | None = None

    split: SplitDecision = SplitDecision()

    @model_validator(mode="after")
    def _populate_split(self) -> MemoirEntry:
        if self.split.shared_eligible or self.split.private_only:
            return self  # caller provided one explicitly; trust them
        decided = decide_split(self.source, agent=self.agent)
        # Pydantic frozen models need object.__setattr__ to mutate after init.
        object.__setattr__(self, "split", decided)
        return self

"""Memoir reader/writer scoped to one ForgeSession workdir."""

from __future__ import annotations

import pytest

from kourai_common.federation.memoir import Memoir, MemoirError
from kourai_common.federation.memoir_schema import (
    EntrySource,
    MemoirEntry,
)


class TestMemoirAppend:
    """Append-only writes produce one JSON-encoded line per entry."""

    def test_append_creates_file(self, tmp_path):
        memoir = Memoir(tmp_path)
        entry = MemoirEntry(
            scene_id="s1.t1",
            agent="kallos",
            source=EntrySource.SPECIALIST_PROPOSED,
            agent_proposed="x",
        )
        memoir.append(entry)
        assert (tmp_path / "memoir.jsonl").exists()

    def test_append_writes_one_line_per_entry(self, tmp_path):
        memoir = Memoir(tmp_path)
        for i in range(3):
            memoir.append(
                MemoirEntry(
                    scene_id=f"s1.t{i}",
                    agent="kallos",
                    source=EntrySource.SPECIALIST_PROPOSED,
                    agent_proposed=f"x{i}",
                )
            )
        contents = (tmp_path / "memoir.jsonl").read_text()
        lines = contents.strip().split("\n")
        assert len(lines) == 3

    def test_append_does_not_truncate(self, tmp_path):
        m1 = Memoir(tmp_path)
        m1.append(
            MemoirEntry(
                scene_id="s1.t1",
                agent="kallos",
                source=EntrySource.SPECIALIST_PROPOSED,
                agent_proposed="first",
            )
        )

        m2 = Memoir(tmp_path)
        m2.append(
            MemoirEntry(
                scene_id="s1.t2",
                agent="kallos",
                source=EntrySource.SPECIALIST_PROPOSED,
                agent_proposed="second",
            )
        )

        contents = (tmp_path / "memoir.jsonl").read_text()
        assert "first" in contents
        assert "second" in contents

    def test_append_to_missing_directory_raises(self, tmp_path):
        bogus = tmp_path / "does_not_exist"
        memoir = Memoir(bogus)
        with pytest.raises(MemoirError, match="not a directory"):
            memoir.append(
                MemoirEntry(
                    scene_id="s1.t1",
                    agent="kallos",
                    source=EntrySource.SPECIALIST_PROPOSED,
                    agent_proposed="x",
                )
            )

"""Forge Memoir reader/writer.

The Memoir is the canonical record of one ForgeSession — a JSONL file
where each line is a `MemoirEntry`. Every entry has two faces: a
narrative beat the visual novel can replay, and a training tuple the
federated-learning pipeline can consume.

Append-only by contract. One host owns one Memoir at a time; no
cross-process coordination is provided in this iteration.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Iterator

    from kourai_common.federation.memoir_schema import MemoirEntry


MEMOIR_FILENAME = "memoir.jsonl"


class MemoirError(Exception):
    """Raised on Memoir read/write failures."""


class Memoir:
    """Reader/writer for one ForgeSession's memoir.jsonl file."""

    def __init__(self, workdir: Path) -> None:
        self.workdir = Path(workdir)
        self.path = self.workdir / MEMOIR_FILENAME

    def append(self, entry: MemoirEntry) -> None:
        if not self.workdir.is_dir():
            raise MemoirError(f"workdir {self.workdir} is not a directory")
        with self.path.open("a", encoding="utf-8") as f:
            f.write(entry.model_dump_json())
            f.write("\n")

class TestMemoirRead:
    """`entries()` yields one MemoirEntry per JSONL line, preserving order."""

    def test_empty_memoir_yields_nothing(self, tmp_path):
        memoir = Memoir(tmp_path)
        assert list(memoir.entries()) == []

    def test_round_trip_preserves_order(self, tmp_path):
        memoir = Memoir(tmp_path)
        originals = [
            MemoirEntry(
                scene_id=f"s1.t{i}",
                agent="kallos",
                source=EntrySource.SPECIALIST_PROPOSED,
                agent_proposed=f"x{i}",
            )
            for i in range(5)
        ]
        for entry in originals:
            memoir.append(entry)

        restored = list(memoir.entries())
        assert restored == originals

    def test_round_trip_preserves_split(self, tmp_path):
        memoir = Memoir(tmp_path)
        cupid_entry = MemoirEntry(
            scene_id="s1.cupid",
            agent="cupid",
            source=EntrySource.CUPID_SCENE,
            narrative_beat="cupid_check_in",
        )
        kallos_entry = MemoirEntry(
            scene_id="s1.kallos",
            agent="kallos",
            source=EntrySource.SPECIALIST_PROPOSED,
            agent_proposed="lint fix",
        )
        memoir.append(cupid_entry)
        memoir.append(kallos_entry)

        restored = list(memoir.entries())
        assert restored[0].split.private_only is True
        assert restored[1].split.shared_eligible is True

    def test_malformed_line_raises(self, tmp_path):
        memoir = Memoir(tmp_path)
        memoir.path.write_text("not valid json\n")
        with pytest.raises(MemoirError, match="malformed"):
            list(memoir.entries())

import json

    def entries(self) -> Iterator[MemoirEntry]:
        """Yield each entry from the on-disk file in append order.

        Returns an empty iterator if the file does not exist yet.
        """
        from kourai_common.federation.memoir_schema import MemoirEntry

        if not self.path.exists():
            return
        with self.path.open("r", encoding="utf-8") as f:
            for line_number, line in enumerate(f, start=1):
                stripped = line.strip()
                if not stripped:
                    continue
                try:
                    yield MemoirEntry.model_validate_json(stripped)
                except (json.JSONDecodeError, ValueError) as e:
                    raise MemoirError(
                        f"malformed entry on line {line_number}: {e}"
                    ) from e