diff --git a/CHANGELOG.md b/CHANGELOG.md index fa0812663..aea6746b0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,8 @@ Only write entries that are worth mentioning to users. ## Unreleased +- Session: `/undo` and fork now map selected wire turns to real context user turns before truncating history, so local slash-command turns no longer make resumed sessions land several turns too early or show history the agent cannot see + ## 1.47.0 (2026-06-05) - Shell: Guide users to the new standalone Kimi Code — adds a `/upgrade` command that installs it (migrating your config & sessions automatically), a welcome-screen nudge, and a once-per-day tip shown on exit diff --git a/src/kimi_cli/session_fork.py b/src/kimi_cli/session_fork.py index 33d6e49ec..cb6dbfe81 100644 --- a/src/kimi_cli/session_fork.py +++ b/src/kimi_cli/session_fork.py @@ -207,6 +207,57 @@ def truncate_context_at_turn(context_path: Path, turn_index: int) -> list[str]: return lines +def _context_turn_index_for_wire_turn( + wire_path: Path, + context_path: Path, + wire_turn_index: int, +) -> int: + """Map a wire TurnBegin index to the corresponding real context user turn.""" + context_user_texts = _context_user_texts(context_path) + if not context_user_texts: + return -1 + + context_turn_index = -1 + next_context_turn = 0 + for turn in enumerate_turns(wire_path): + if turn.index > wire_turn_index: + break + if next_context_turn >= len(context_user_texts): + continue + + # A wire turn can be missing from context (local slash commands), and + # context can contain user turns without matching TurnBegin text (Stop + # hook reasons or expanded skill prompts). Match the next wire turn to + # the next same-text context user, preserving any context-only turns in + # between. + for matched_context_turn in range(next_context_turn, len(context_user_texts)): + if turn.user_text == context_user_texts[matched_context_turn]: + context_turn_index = matched_context_turn + next_context_turn = matched_context_turn + 1 + break + + return context_turn_index + + +def _context_user_texts(context_path: Path) -> list[str]: + if not context_path.exists(): + return [] + + texts: list[str] = [] + with open(context_path, encoding="utf-8") as f: + for line in f: + stripped = line.strip() + if not stripped: + continue + try: + record: dict[str, Any] = json.loads(stripped) + except json.JSONDecodeError: + continue + if record.get("role") == "user" and not _is_checkpoint_user_message(record): + texts.append(_extract_user_text(record.get("content", ""))) + return texts + + # --------------------------------------------------------------------------- # Full fork operation # --------------------------------------------------------------------------- @@ -241,7 +292,12 @@ async def fork_session( if turn_index is not None: truncated_wire_lines = truncate_wire_at_turn(wire_path, turn_index) - truncated_context_lines = truncate_context_at_turn(context_path, turn_index) + context_turn_index = _context_turn_index_for_wire_turn( + wire_path, + context_path, + turn_index, + ) + truncated_context_lines = truncate_context_at_turn(context_path, context_turn_index) else: # Copy all content truncated_wire_lines = _read_all_lines(wire_path) diff --git a/tests/core/test_session_fork.py b/tests/core/test_session_fork.py index b608a7db9..5005d59a0 100644 --- a/tests/core/test_session_fork.py +++ b/tests/core/test_session_fork.py @@ -309,6 +309,79 @@ async def test_fork_at_turn(self, isolated_share_dir: Path, work_dir: KaosPath): # 2 turns * (user + assistant) = 4 assert len(ctx_lines) == 4 + async def test_fork_maps_wire_turns_to_real_context_users( + self, isolated_share_dir: Path, work_dir: KaosPath + ): + from kimi_cli.session import Session + + source = await Session.create(work_dir) + _write_wire_file(source.dir, ["real 0", "/usage", "/sessions", "real 1", "real 2"]) + _write_context_file(source.dir, ["real 0", "real 1", "real 2"]) + + new_id = await fork_session( + source_session_dir=source.dir, + work_dir=work_dir, + turn_index=3, + title_prefix="Undo", + source_title="My Session", + ) + + new_session = await Session.find(work_dir, new_id) + assert new_session is not None + + context_lines = ( + (new_session.dir / "context.jsonl").read_text(encoding="utf-8").strip().split("\n") + ) + user_texts = [ + _extract_user_text(json.loads(line)["content"]) + for line in context_lines + if json.loads(line)["role"] == "user" + ] + assert user_texts == ["real 0", "real 1"] + + async def test_fork_preserves_context_only_user_turns_when_mapping_later_wire_turn( + self, isolated_share_dir: Path, work_dir: KaosPath + ): + from kimi_cli.session import Session + + source = await Session.create(work_dir) + _write_wire_file(source.dir, ["real 0", "/skill:demo", "real 1"]) + _write_context_file( + source.dir, + [ + "real 0", + "Stop hook reason", + "Expanded skill prompt", + "real 1", + ], + ) + + new_id = await fork_session( + source_session_dir=source.dir, + work_dir=work_dir, + turn_index=2, + title_prefix="Undo", + source_title="My Session", + ) + + new_session = await Session.find(work_dir, new_id) + assert new_session is not None + + context_lines = ( + (new_session.dir / "context.jsonl").read_text(encoding="utf-8").strip().split("\n") + ) + user_texts = [ + _extract_user_text(json.loads(line)["content"]) + for line in context_lines + if json.loads(line)["role"] == "user" + ] + assert user_texts == [ + "real 0", + "Stop hook reason", + "Expanded skill prompt", + "real 1", + ] + async def test_fork_all_turns(self, isolated_share_dir: Path, work_dir: KaosPath): from kimi_cli.session import Session