diff --git a/plugin/src/claude_smart/publish.py b/plugin/src/claude_smart/publish.py index fe98cfa..4885cd7 100644 --- a/plugin/src/claude_smart/publish.py +++ b/plugin/src/claude_smart/publish.py @@ -8,7 +8,8 @@ from __future__ import annotations -from typing import Literal +import uuid +from typing import Any, Literal from claude_smart import state from claude_smart.reflexio_adapter import Adapter @@ -16,6 +17,50 @@ PublishStatus = Literal["nothing", "ok", "failed"] +def _publish_request_id(session_id: str, start: int, end: int) -> str: + """Return a stable request ID for one buffered publish range.""" + + name = f"claude-smart:{session_id}:{start}:{end}" + return str(uuid.uuid5(uuid.NAMESPACE_URL, name)) + + +def _prepare_publish_batch( + session_id: str, +) -> tuple[int, int, list[dict[str, Any]]] | None: + """Freeze and return one canonical unpublished record range.""" + + while True: + records = state.read_all(session_id) + published, available = state.unpublished_slice(records) + if not available: + return None + + publish_end = state.pending_publish_end(records, published) + if publish_end is None: + publish_end = state.safe_publish_end(records, published) + _, proposed = state.unpublished_slice( + records[:publish_end], published_override=published + ) + if not proposed: + return None + state.append( + session_id, + {"publish_attempt": {"start": published, "end": publish_end}}, + ) + + canonical = state.read_all(session_id) + if state.published_record_offset(canonical) != published: + continue + canonical_end = state.pending_publish_end(canonical, published) + if canonical_end is None: + continue + _, interactions = state.unpublished_slice( + canonical[:canonical_end], published_override=published + ) + if interactions: + return published, canonical_end, interactions + + def publish_unpublished( *, session_id: str, @@ -56,20 +101,24 @@ def publish_unpublished( was unreachable. On ``"failed"`` the watermark is not advanced, so the next hook retries the same batch. """ - records = state.read_all(session_id) - _, interactions = state.unpublished_slice(records) - if not interactions: + batch = _prepare_publish_batch(session_id) + if batch is None: return ("nothing", 0) + published_record_offset, publish_end, interactions = batch + client = adapter if adapter is not None else Adapter() ok = client.publish( session_id=session_id, project_id=project_id, + request_id=_publish_request_id( + session_id, published_record_offset, publish_end + ), interactions=interactions, force_extraction=force_extraction, override_learning_stall=override_learning_stall, skip_aggregation=skip_aggregation, ) if ok: - state.append(session_id, {"published_up_to": len(records)}) + state.append(session_id, {"published_up_to": publish_end}) return ("ok", len(interactions)) return ("failed", len(interactions)) diff --git a/plugin/src/claude_smart/reflexio_adapter.py b/plugin/src/claude_smart/reflexio_adapter.py index 8077d84..ea0ca53 100644 --- a/plugin/src/claude_smart/reflexio_adapter.py +++ b/plugin/src/claude_smart/reflexio_adapter.py @@ -109,6 +109,7 @@ def publish( *, session_id: str, project_id: str, + request_id: str | None = None, interactions: Sequence[dict[str, Any]], force_extraction: bool = False, override_learning_stall: bool = False, @@ -121,9 +122,57 @@ def publish( if client is None: return False try: + interaction_list = list(interactions) + if _needs_raw_retrieved_learning_publish(interaction_list): + raw_request = getattr(client, "_make_request", None) + if request_id is not None and callable(raw_request): + payload = { + "request_id": request_id, + "user_id": project_id, + "interaction_data_list": interaction_list, + "agent_version": runtime.agent_version(), + "session_id": session_id, + "skip_aggregation": skip_aggregation, + "force_extraction": force_extraction, + "evaluation_only": False, + "override_learning_stall": override_learning_stall, + } + try: + raw_request( + "POST", + "/api/publish_interaction", + json=payload, + params=None, + ) + return True + except Exception as exc: # noqa: BLE001 + _LOGGER.warning( + "Could not confirm retrieved-learning links; retrying " + "the base interaction with the same request ID: %s", + exc, + ) + fallback_payload = { + **payload, + "interaction_data_list": _without_retrieved_learnings( + interaction_list + ), + } + raw_request( + "POST", + "/api/publish_interaction", + json=fallback_payload, + params=None, + ) + return True + else: + _LOGGER.warning( + "Stable raw publishing is unavailable; publishing " + "without optional retrieved-learning links" + ) + interaction_list = _without_retrieved_learnings(interaction_list) kwargs = { "user_id": project_id, - "interactions": list(interactions), + "interactions": interaction_list, "agent_version": runtime.agent_version(), "session_id": session_id, "wait_for_response": False, @@ -476,6 +525,32 @@ def _supports_keyword(callable_obj: Any, keyword: str) -> bool: return keyword in signature.parameters +def _needs_raw_retrieved_learning_publish( + interactions: Sequence[dict[str, Any]], +) -> bool: + """Return whether links require a caller-stable raw request. + + The public client does not accept a caller-supplied request ID, and older + clients also strip this field. The authenticated helper preserves both. + """ + return any(item.get("retrieved_learnings") for item in interactions) + + +def _without_retrieved_learnings( + interactions: Sequence[dict[str, Any]], +) -> list[dict[str, Any]]: + """Copy interactions without the optional compatibility-only field.""" + + return [ + { + key: value + for key, value in interaction.items() + if key != "retrieved_learnings" + } + for interaction in interactions + ] + + def _filter_rejected_agent_playbooks(items: list[Any]) -> list[Any]: """Drop rejected shared skills defensively, even if an older backend ignores filters.""" return [ diff --git a/plugin/src/claude_smart/state.py b/plugin/src/claude_smart/state.py index 775b946..0ff98b6 100644 --- a/plugin/src/claude_smart/state.py +++ b/plugin/src/claude_smart/state.py @@ -9,6 +9,10 @@ to the next assistant turn at ``Stop`` time - ``{"published_up_to": N}`` — high-water mark so Stop / SessionEnd don't re-publish rows already sent to reflexio +- ``{"retrieved_learning_refs": [...]}`` — identity pairs shown to the agent, + attached by event order to the next Assistant turn +- ``{"publish_attempt": {"start": N, "end": M}}`` — frozen retry boundary for + one in-flight publish The buffer exists for offline resilience: when reflexio is unreachable, Stop appends without publishing and the next successful hook drains. @@ -38,6 +42,9 @@ _VALID_CITATION_KINDS = frozenset( {"playbook", "profile", "user_playbook", "agent_playbook"} ) +_VALID_RETRIEVED_PLAYBOOK_KINDS = frozenset({"user_playbook", "agent_playbook"}) +_VALID_RETRIEVED_KINDS = _VALID_RETRIEVED_PLAYBOOK_KINDS | {"profile"} +_RETRIEVED_LEARNINGS_PUBLISH_CAP = 1000 def _truncate_tool_data_field(value: Any) -> Any: @@ -84,17 +91,37 @@ def injected_path(session_id: str) -> Path: return state_dir() / f"{session_id}.injected.jsonl" +def _normalize_registry_ref(entry: dict[str, Any]) -> dict[str, str] | None: + """Return the identity pair carried by one citation-registry entry.""" + + learning_id = entry.get("real_id") + kind = entry.get("kind") + if not isinstance(learning_id, str) or not learning_id: + return None + if kind == "profile": + wire_kind = "profile" + elif ( + kind == "playbook" + and entry.get("source_kind") in _VALID_RETRIEVED_PLAYBOOK_KINDS + ): + wire_kind = entry["source_kind"] + else: + return None + return {"kind": wire_kind, "learning_id": learning_id} + + def append_injected(session_id: str, entries: Iterable[dict[str, Any]]) -> None: - """Append citation-registry entries to the per-session injected-items file. + """Record injected items for citations and the next Assistant publish. - Each entry maps a short ``id`` (4-hex-char) back to the skill or - preference it came from so the Stop hook can resolve citation ids into - human-readable titles for the dashboard. - Silently no-ops when ``entries`` is empty. + The citation sidecar keeps the display metadata used by Stop. The ordered + session event stores identity pairs only, so publishing never has to join + two files or infer event order from timestamps. """ + records = list(entries) if not records: return + path = injected_path(session_id) path.parent.mkdir(parents=True, exist_ok=True) with path.open("a", encoding="utf-8") as fh: @@ -103,8 +130,22 @@ def append_injected(session_id: str, entries: Iterable[dict[str, Any]]) -> None: fcntl.flock(fh.fileno(), fcntl.LOCK_EX) except OSError as exc: _LOGGER.debug("flock failed on %s: %s", path, exc) - for rec in records: - fh.write(json.dumps(rec, ensure_ascii=False) + "\n") + for record in records: + fh.write(json.dumps(record, ensure_ascii=False) + "\n") + + refs: list[dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + for record in records: + ref = _normalize_registry_ref(record) + if ref is None: + continue + key = (ref["kind"], ref["learning_id"]) + if key in seen: + continue + seen.add(key) + refs.append(ref) + if refs: + append(session_id, {"retrieved_learning_refs": refs}) def read_injected(session_id: str) -> dict[str, dict[str, Any]]: @@ -114,6 +155,7 @@ def read_injected(session_id: str) -> dict[str, dict[str, Any]]: (identical content produces the same hash-derived id, so the extra record only refreshes metadata). """ + path = injected_path(session_id) if not path.exists(): return {} @@ -229,44 +271,127 @@ def _to_wire_citations(cited_items: Any) -> list[dict[str, str]]: return out +def published_record_offset(records: list[dict[str, Any]]) -> int: + """Return the highest valid publish watermark in the session buffer.""" + + limit = len(records) + return max( + ( + value + for record in records + if isinstance((value := record.get("published_up_to")), int) + and 0 <= value <= limit + ), + default=0, + ) + + +def pending_publish_end( + records: list[dict[str, Any]], published: int +) -> int | None: + """Return the frozen end of the current publish attempt, if any.""" + + for marker_index, record in enumerate(records[published:], start=published): + attempt = record.get("publish_attempt") + if not isinstance(attempt, dict): + continue + start = attempt.get("start") + end = attempt.get("end") + if ( + start == published + and isinstance(end, int) + and published < end <= marker_index + ): + _, interactions = unpublished_slice( + records[:end], published_override=published + ) + if interactions: + return end + return None + + +def _wire_retrieved_ref(value: Any) -> dict[str, str] | None: + """Validate one identity pair read from the ordered session buffer.""" + + if not isinstance(value, dict): + return None + kind = value.get("kind") + learning_id = value.get("learning_id") + if kind not in _VALID_RETRIEVED_KINDS: + return None + if not isinstance(learning_id, str) or not learning_id: + return None + return {"kind": kind, "learning_id": learning_id} + + +def safe_publish_end(records: list[dict[str, Any]], published: int) -> int: + """Stop before identity refs that do not yet have an Assistant target.""" + + unmatched_start: int | None = None + for index, record in enumerate(records[published:], start=published): + refs = record.get("retrieved_learning_refs") + if ( + unmatched_start is None + and isinstance(refs, list) + and any(_wire_retrieved_ref(value) is not None for value in refs) + ): + unmatched_start = index + if record.get("role") == "Assistant": + unmatched_start = None + return len(records) if unmatched_start is None else unmatched_start + + def unpublished_slice( records: Iterable[dict[str, Any]], + *, + published_override: int | None = None, ) -> tuple[int, list[dict[str, Any]]]: - """Split records into (last-published index, unpublished turn records). + """Return the current watermark and wire-ready unpublished interactions. - Walks the records in order, tracking the most recent ``published_up_to`` - marker and collecting turn records (anything with a ``role``) that come - after it. Tool records are folded into the closest following Assistant - turn's ``tools_used``. - - Returns: - tuple[int, list[dict]]: ``(published_up_to, interactions)``. The - integer is the watermark after which all turns are unpublished; - the list is formatted for ``InteractionData`` construction. + Metadata records are interpreted by their numeric boundary rather than + their physical position in the append-only file. This preserves turns that + arrive while an earlier network request is in flight. """ - published = 0 + + buffered = list(records) + if published_override is None: + published = published_record_offset(buffered) + elif 0 <= published_override <= len(buffered): + published = published_override + else: + raise ValueError("published_override is outside the supplied record range") pending_tools: list[dict[str, Any]] = [] + pending_refs: list[dict[str, str]] = [] turns: list[dict[str, Any]] = [] - for idx, rec in enumerate(records): - if "published_up_to" in rec: - published = rec["published_up_to"] - pending_tools = [] - turns = [] + attached_total = skipped = truncated = 0 + + for index, record in enumerate(buffered): + if index < published: continue - if idx < published: + + refs = record.get("retrieved_learning_refs") + if "role" not in record and isinstance(refs, list): + for value in refs: + ref = _wire_retrieved_ref(value) + if ref is None: + skipped += 1 + else: + pending_refs.append(ref) continue - role = rec.get("role") + + role = record.get("role") if role == "Assistant_tool": - tool_input = rec.get("tool_input") or {} - tool_output = rec.get("tool_output") or "" + tool_input = record.get("tool_input") or {} + tool_output = record.get("tool_output") or "" tool_entry: dict[str, Any] = { - "tool_name": rec.get("tool_name", ""), - "status": rec.get("status", "success"), + "tool_name": record.get("tool_name", ""), + "status": record.get("status", "success"), } tool_data: dict[str, Any] = {} if tool_input: tool_data["input"] = { - k: _truncate_tool_data_field(v) for k, v in tool_input.items() + key: _truncate_tool_data_field(value) + for key, value in tool_input.items() } if tool_output: tool_data["output"] = _truncate_tool_data_field(tool_output) @@ -274,20 +399,45 @@ def unpublished_slice( tool_entry["tool_data"] = tool_data pending_tools.append(tool_entry) continue - if role in {"User", "Assistant"}: - # ``cited_items`` is local-only metadata (dashboard "used" badge); - # map it onto the wire's ``citations`` field — reflexio uses those - # to drive skill/preference reflection in the publish flow. - turn = { - k: v for k, v in rec.items() if k not in {"role", "ts", "cited_items"} - } - turn["role"] = role - if role == "Assistant": - citations = _to_wire_citations(rec.get("cited_items")) - if citations: - turn["citations"] = citations - if pending_tools: - turn["tools_used"] = pending_tools - pending_tools = [] - turns.append(turn) + + if role not in {"User", "Assistant"}: + continue + + turn = { + key: value + for key, value in record.items() + if key not in {"role", "ts", "cited_items"} + } + turn["role"] = role + if role == "Assistant": + citations = _to_wire_citations(record.get("cited_items")) + if citations: + turn["citations"] = citations + if pending_tools: + turn["tools_used"] = pending_tools + pending_tools = [] + + retrieved: list[dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + for ref in pending_refs: + key = (ref["kind"], ref["learning_id"]) + if key in seen: + continue + seen.add(key) + if attached_total >= _RETRIEVED_LEARNINGS_PUBLISH_CAP: + truncated += 1 + continue + retrieved.append(ref) + attached_total += 1 + pending_refs = [] + if retrieved: + turn["retrieved_learnings"] = retrieved + turns.append(turn) + + if skipped: + _LOGGER.debug("Skipped %d malformed retrieved-learning refs", skipped) + if truncated: + _LOGGER.warning( + "Dropped %d retrieved-learning refs at the publish request cap", truncated + ) return published, turns diff --git a/tests/host_learning_harness.py b/tests/host_learning_harness.py index 154837f..1510f89 100644 --- a/tests/host_learning_harness.py +++ b/tests/host_learning_harness.py @@ -77,6 +77,7 @@ def publish( *, session_id: str, project_id: str, + request_id: str, interactions: list[dict[str, Any]], force_extraction: bool = False, override_learning_stall: bool = False, @@ -88,6 +89,7 @@ def publish( { "session_id": session_id, "project_id": project_id, + "request_id": request_id, "interactions": interactions, "force_extraction": force_extraction, "override_learning_stall": override_learning_stall, @@ -245,6 +247,7 @@ def run_host_learning_happy_path( publish_call = fake_reflexio.publish_calls[0] assert publish_call["session_id"] == session_id assert publish_call["project_id"] == project.name + assert publish_call["request_id"] assert publish_call["force_extraction"] is False assert publish_call["skip_aggregation"] is False assert publish_call["host"] == host @@ -269,6 +272,14 @@ def run_host_learning_happy_path( "title": f"{host} project skill", } ] + assert { + (item["kind"], item["learning_id"]) + for item in interactions[1]["retrieved_learnings"] + } == { + ("user_playbook", f"user-{host}"), + ("agent_playbook", f"agent-{host}"), + ("profile", f"profile-{host}"), + } _, unpublished = state.unpublished_slice(state.read_all(session_id)) assert unpublished == [] diff --git a/tests/test_adapter.py b/tests/test_adapter.py index 03e90ac..44c35d9 100644 --- a/tests/test_adapter.py +++ b/tests/test_adapter.py @@ -26,12 +26,18 @@ def __init__( self._agent_playbook_resp = agent_playbook_resp self._profile_resp = profile_resp self.published_kwargs: dict[str, Any] = {} + self.raw_request: dict[str, Any] = {} def publish_interaction(self, **kwargs): self.published_kwargs = kwargs if not self._publish_ok: raise RuntimeError("reflexio unreachable") + def _make_request(self, method, path, **kwargs): + self.raw_request = {"method": method, "path": path, **kwargs} + if not self._publish_ok: + raise RuntimeError("reflexio unreachable") + def search_user_playbooks(self, **_kw): return self._user_playbook_resp @@ -275,6 +281,107 @@ def publish_interaction( assert "override_learning_stall" not in client.published_kwargs +def test_link_publish_uses_stable_raw_request() -> None: + client = _FakeClient() + adapter = _adapter_with(client) + + ok = adapter.publish( + session_id="s1", + project_id="p1", + request_id="request-1", + interactions=[ + { + "role": "Assistant", + "content": "done", + "retrieved_learnings": [ + {"kind": "profile", "learning_id": "profile-1"} + ], + } + ], + ) + + assert ok is True + assert client.published_kwargs == {} + assert client.raw_request["method"] == "POST" + assert client.raw_request["path"] == "/api/publish_interaction" + assert client.raw_request["json"]["request_id"] == "request-1" + assert client.raw_request["json"]["interaction_data_list"][0][ + "retrieved_learnings" + ] == [{"kind": "profile", "learning_id": "profile-1"}] + + +def test_missing_raw_request_helper_falls_back_without_optional_links( + monkeypatch, +) -> None: + client = _FakeClient() + monkeypatch.setattr(client, "_make_request", None) + adapter = _adapter_with(client) + + ok = adapter.publish( + session_id="s1", + project_id="p1", + interactions=[ + { + "role": "Assistant", + "content": "done", + "retrieved_learnings": [ + {"kind": "profile", "learning_id": "profile-1"} + ], + } + ], + ) + + assert ok is True + assert client.published_kwargs["interactions"] == [ + {"role": "Assistant", "content": "done"} + ] + + +def test_ambiguous_raw_failure_retries_base_with_same_request_id() -> None: + class ResponseLossClient(_FakeClient): + def __init__(self): + super().__init__() + self.raw_requests = [] + + def _make_request(self, method, path, **kwargs): + self.raw_requests.append(kwargs["json"]) + if len(self.raw_requests) == 1: + raise TimeoutError("response was lost after send") + + def publish_interaction(self, **kwargs): + raise AssertionError("fallback must preserve the stable request ID") + + client = ResponseLossClient() + adapter = _adapter_with(client) + + ok = adapter.publish( + session_id="s1", + project_id="p1", + request_id="stable-request", + interactions=[ + { + "role": "Assistant", + "content": "done", + "retrieved_learnings": [ + {"kind": "profile", "learning_id": "profile-1"} + ], + } + ], + ) + + assert ok is True + assert len(client.raw_requests) == 2 + assert {request["request_id"] for request in client.raw_requests} == { + "stable-request" + } + assert client.raw_requests[0]["interaction_data_list"][0][ + "retrieved_learnings" + ] == [{"kind": "profile", "learning_id": "profile-1"}] + assert client.raw_requests[1]["interaction_data_list"] == [ + {"role": "Assistant", "content": "done"} + ] + + def test_publish_returns_false_when_client_raises() -> None: a = _adapter_with(_FakeClient(publish_ok=False)) ok = a.publish( diff --git a/tests/test_events.py b/tests/test_events.py index fb0414f..9dc2b3c 100644 --- a/tests/test_events.py +++ b/tests/test_events.py @@ -990,8 +990,16 @@ def test_user_prompt_injects_context_when_hits_present( calls: list[dict[str, Any]] = [] _stub_user_prompt_adapter( monkeypatch, - playbooks=[{"content": "run uv sync after edits", "trigger": "pyproject.toml"}], - profiles=[{"content": "prefers anyio over asyncio"}], + playbooks=[ + { + "user_playbook_id": "playbook-1", + "content": "run uv sync after edits", + "trigger": "pyproject.toml", + } + ], + profiles=[ + {"profile_id": "profile-1", "content": "prefers anyio over asyncio"} + ], calls=calls, ) buf = io.StringIO() @@ -1011,8 +1019,10 @@ def test_user_prompt_injects_context_when_hits_present( assert calls[0]["session_id"] == "s1" # The prompt is still buffered for downstream extraction. records = state.read_all("s1") - assert records[-1]["role"] == "User" - assert records[-1]["content"] == "edit pyproject.toml" + turns = [record for record in records if "role" in record] + assert turns[-1]["role"] == "User" + assert turns[-1]["content"] == "edit pyproject.toml" + assert records[-1]["retrieved_learning_refs"] # The citation registry is populated so Stop can resolve text citation ids. registry = state.read_injected("s1") assert registry, "expected injected registry to hold at least one entry" diff --git a/tests/test_publish.py b/tests/test_publish.py new file mode 100644 index 0000000..56abaa9 --- /dev/null +++ b/tests/test_publish.py @@ -0,0 +1,347 @@ +"""Publish payload tests for locally injected learning links.""" + +from __future__ import annotations + +import threading +from collections.abc import Callable +from typing import Any, cast + +from claude_smart import publish, state +from claude_smart.reflexio_adapter import Adapter + + +class _RecordingAdapter: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def publish(self, **kwargs: Any) -> bool: + self.calls.append(kwargs) + return True + + +class _SequencedAdapter(_RecordingAdapter): + def __init__(self, results: list[bool]) -> None: + super().__init__() + self.results = iter(results) + + def publish(self, **kwargs: Any) -> bool: + self.calls.append(kwargs) + return next(self.results) + + +class _CallbackAdapter(_RecordingAdapter): + def __init__(self, callback: Callable[[], None]) -> None: + super().__init__() + self.callback = callback + + def publish(self, **kwargs: Any) -> bool: + self.calls.append(kwargs) + self.callback() + return True + + +def _append_assistant(session_id: str, ts: int, content: str = "done") -> None: + state.append( + session_id, + {"role": "Assistant", "ts": ts, "content": content, "user_id": "user"}, + ) + + +def _append_user(session_id: str, ts: int, content: str = "question") -> None: + state.append( + session_id, + {"role": "User", "ts": ts, "content": content, "user_id": "user"}, + ) + + +def _inject_profile(session_id: str, learning_id: str, ts: int) -> None: + state.append_injected( + session_id, + [ + { + "id": f"rank-{learning_id}", + "kind": "profile", + "real_id": learning_id, + "title": "Preference", + "content": "Use the preference.", + "ts": ts, + } + ], + ) + + +def _publish(session_id: str, adapter: _RecordingAdapter) -> tuple[str, int]: + return publish.publish_unpublished( + session_id=session_id, + project_id="project", + force_extraction=False, + skip_aggregation=False, + adapter=cast(Adapter, adapter), + ) + + +def test_publish_attaches_identity_pairs_in_injection_order(session_dir) -> None: + state.append_injected( + "s1", + [ + {"id": "p", "kind": "profile", "real_id": "profile-1", "ts": 10}, + { + "id": "u", + "kind": "playbook", + "source_kind": "user_playbook", + "real_id": "11", + "ts": 11, + }, + { + "id": "a", + "kind": "playbook", + "source_kind": "agent_playbook", + "real_id": "22", + "ts": 12, + }, + {"id": "invalid", "kind": "profile", "ts": 13}, + {"id": "duplicate", "kind": "profile", "real_id": "profile-1"}, + ], + ) + _append_assistant("s1", 20) + adapter = _RecordingAdapter() + + assert _publish("s1", adapter) == ("ok", 1) + assert adapter.calls[0]["interactions"][0]["retrieved_learnings"] == [ + {"kind": "profile", "learning_id": "profile-1"}, + {"kind": "user_playbook", "learning_id": "11"}, + {"kind": "agent_playbook", "learning_id": "22"}, + ] + + +def test_retrieved_learnings_attach_once_across_publishes(session_dir) -> None: + _inject_profile("s1", "p1", 5) + _append_assistant("s1", 10, "first") + adapter = _RecordingAdapter() + assert _publish("s1", adapter) == ("ok", 1) + + _inject_profile("s1", "p2", 15) + _append_assistant("s1", 20, "second") + assert _publish("s1", adapter) == ("ok", 1) + + assert adapter.calls[0]["interactions"][0]["retrieved_learnings"] == [ + {"kind": "profile", "learning_id": "p1"} + ] + assert adapter.calls[1]["interactions"][0]["retrieved_learnings"] == [ + {"kind": "profile", "learning_id": "p2"} + ] + assert adapter.calls[0]["request_id"] != adapter.calls[1]["request_id"] + + +def test_failed_publish_retries_frozen_batch_before_new_turns(session_dir) -> None: + _inject_profile("s1", "p1", 5) + _append_assistant("s1", 10, "first") + adapter = _SequencedAdapter([False, True, True]) + + assert _publish("s1", adapter) == ("failed", 1) + + _inject_profile("s1", "p2", 15) + _append_assistant("s1", 20, "second") + + assert _publish("s1", adapter) == ("ok", 1) + assert _publish("s1", adapter) == ("ok", 1) + assert _publish("s1", adapter) == ("nothing", 0) + + assert [item["content"] for item in adapter.calls[0]["interactions"]] == [ + "first" + ] + assert adapter.calls[1]["interactions"] == adapter.calls[0]["interactions"] + assert adapter.calls[1]["request_id"] == adapter.calls[0]["request_id"] + assert [item["content"] for item in adapter.calls[2]["interactions"]] == [ + "second" + ] + assert adapter.calls[2]["request_id"] != adapter.calls[1]["request_id"] + + +def test_success_does_not_hide_turn_appended_during_publish(session_dir) -> None: + _inject_profile("s1", "p1", 5) + _append_assistant("s1", 10, "first") + + def append_next_turn() -> None: + _inject_profile("s1", "p2", 15) + _append_assistant("s1", 20, "second") + + first_adapter = _CallbackAdapter(append_next_turn) + assert _publish("s1", first_adapter) == ("ok", 1) + + second_adapter = _RecordingAdapter() + assert _publish("s1", second_adapter) == ("ok", 1) + assert second_adapter.calls[0]["interactions"] == [ + { + "role": "Assistant", + "content": "second", + "user_id": "user", + "retrieved_learnings": [ + {"kind": "profile", "learning_id": "p2"} + ], + } + ] + + +def test_prefix_uses_canonical_watermark_when_success_marker_is_later( + session_dir, +) -> None: + _append_assistant("s1", 1, "first") + + def append_unfinished_turn() -> None: + _append_user("s1", 2, "second question") + _inject_profile("s1", "p2", 3) + + first_adapter = _CallbackAdapter(append_unfinished_turn) + assert _publish("s1", first_adapter) == ("ok", 1) + + second_adapter = _RecordingAdapter() + assert _publish("s1", second_adapter) == ("ok", 1) + assert [item["content"] for item in second_adapter.calls[0]["interactions"]] == [ + "second question" + ] + + _append_assistant("s1", 4, "second") + assert _publish("s1", second_adapter) == ("ok", 1) + assert second_adapter.calls[1]["interactions"][0]["retrieved_learnings"] == [ + {"kind": "profile", "learning_id": "p2"} + ] + + +def test_early_publish_keeps_refs_for_the_next_assistant(session_dir) -> None: + _append_user("s1", 1) + _inject_profile("s1", "p1", 2) + adapter = _RecordingAdapter() + + assert _publish("s1", adapter) == ("ok", 1) + assert "retrieved_learnings" not in adapter.calls[0]["interactions"][0] + + _append_assistant("s1", 3) + assert _publish("s1", adapter) == ("ok", 1) + assert adapter.calls[1]["interactions"][0]["retrieved_learnings"] == [ + {"kind": "profile", "learning_id": "p1"} + ] + + +def test_overlapping_publishers_adopt_one_frozen_batch( + session_dir, monkeypatch +) -> None: + _inject_profile("s1", "p1", 1) + _append_assistant("s1", 2, "first") + + first_marker_ready = threading.Event() + allow_first_marker = threading.Event() + second_request_ready = threading.Event() + allow_second_request = threading.Event() + original_append = state.append + + def append_with_pause(session_id: str, record: dict[str, Any]) -> None: + if "publish_attempt" in record and threading.current_thread().name == "first": + first_marker_ready.set() + assert allow_first_marker.wait(timeout=5) + original_append(session_id, record) + + monkeypatch.setattr(state, "append", append_with_pause) + + first_adapter = _RecordingAdapter() + first_result: list[tuple[str, int]] = [] + first_thread = threading.Thread( + name="first", target=lambda: first_result.append(_publish("s1", first_adapter)) + ) + first_thread.start() + assert first_marker_ready.wait(timeout=5) + + _inject_profile("s1", "p2", 3) + _append_assistant("s1", 4, "second") + + class BlockingAdapter(_RecordingAdapter): + def publish(self, **kwargs: Any) -> bool: + self.calls.append(kwargs) + second_request_ready.set() + assert allow_second_request.wait(timeout=5) + return True + + second_adapter = BlockingAdapter() + second_result: list[tuple[str, int]] = [] + second_thread = threading.Thread( + name="second", + target=lambda: second_result.append(_publish("s1", second_adapter)), + ) + second_thread.start() + assert second_request_ready.wait(timeout=5) + + allow_first_marker.set() + first_thread.join(timeout=5) + assert not first_thread.is_alive() + allow_second_request.set() + second_thread.join(timeout=5) + assert not second_thread.is_alive() + + assert first_result == [("ok", 2)] + assert second_result == [("ok", 2)] + assert first_adapter.calls[0]["request_id"] == second_adapter.calls[0][ + "request_id" + ] + assert first_adapter.calls[0]["interactions"] == second_adapter.calls[0][ + "interactions" + ] + + +def test_invalid_empty_attempt_marker_is_skipped(session_dir) -> None: + state.append( + "s1", + { + "retrieved_learning_refs": [ + {"kind": "profile", "learning_id": "p1"} + ] + }, + ) + state.append("s1", {"publish_attempt": {"start": 0, "end": 1}}) + _append_assistant("s1", 2) + adapter = _RecordingAdapter() + + assert _publish("s1", adapter) == ("ok", 1) + assert adapter.calls[0]["interactions"][0]["retrieved_learnings"] == [ + {"kind": "profile", "learning_id": "p1"} + ] + + +def test_same_second_injections_follow_event_order(session_dir) -> None: + _inject_profile("s1", "p1", 10) + _append_assistant("s1", 10, "first") + adapter = _RecordingAdapter() + assert _publish("s1", adapter) == ("ok", 1) + + _inject_profile("s1", "p2", 10) + _append_assistant("s1", 10, "second") + assert _publish("s1", adapter) == ("ok", 1) + + assert adapter.calls[1]["interactions"][0]["retrieved_learnings"] == [ + {"kind": "profile", "learning_id": "p2"} + ] + + +def test_publish_without_injected_learnings_keeps_typed_payload(session_dir) -> None: + _append_assistant("s1", 10) + adapter = _RecordingAdapter() + + assert _publish("s1", adapter) == ("ok", 1) + assert "retrieved_learnings" not in adapter.calls[0]["interactions"][0] + + +def test_publish_keeps_first_links_at_request_cap(session_dir) -> None: + state.append_injected( + "s1", + [ + {"id": str(index), "kind": "profile", "real_id": str(index), "ts": 1} + for index in range(1002) + ], + ) + _append_assistant("s1", 10) + adapter = _RecordingAdapter() + + assert _publish("s1", adapter) == ("ok", 1) + retrieved = adapter.calls[0]["interactions"][0]["retrieved_learnings"] + assert len(retrieved) == 1000 + assert retrieved[0]["learning_id"] == "0" + assert retrieved[-1]["learning_id"] == "999" diff --git a/tests/test_state.py b/tests/test_state.py index 4436794..062af73 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -139,13 +139,34 @@ def test_append_injected_roundtrip(session_dir) -> None: state.append_injected( "s1", [ - {"id": "s1-ab12", "kind": "playbook", "title": "t1", "content": "c1"}, - {"id": "p1-cd34", "kind": "profile", "title": "t2", "content": "c2"}, + { + "id": "s1-ab12", + "kind": "playbook", + "source_kind": "user_playbook", + "real_id": "11", + "title": "t1", + "content": "c1", + }, + { + "id": "p1-cd34", + "kind": "profile", + "real_id": "p1", + "title": "t2", + "content": "c2", + }, ], ) registry = state.read_injected("s1") assert registry["s1-ab12"]["title"] == "t1" assert registry["p1-cd34"]["kind"] == "profile" + assert state.read_all("s1") == [ + { + "retrieved_learning_refs": [ + {"kind": "user_playbook", "learning_id": "11"}, + {"kind": "profile", "learning_id": "p1"}, + ] + } + ] def test_append_injected_empty_iter_is_noop(session_dir) -> None: