diff --git a/api/routes.py b/api/routes.py index 279b3ca0..6033ced7 100644 --- a/api/routes.py +++ b/api/routes.py @@ -12320,6 +12320,14 @@ def handle_post(handler, parsed) -> bool: branch_title = f"{source_title} (fork)" # Create new session inheriting workspace/model/profile + from api.session_ops import truncate_context_for_display_keep + + fork_keep = keep_count if keep_count is not None else len(source_messages) + forked_context = truncate_context_for_display_keep( + getattr(source, "context_messages", None), + source_messages, + fork_keep, + ) branch = Session( workspace=source.workspace, model=source.model, @@ -12332,8 +12340,8 @@ def handle_post(handler, parsed) -> bool: enabled_toolsets=getattr(source, "enabled_toolsets", None), context_length=getattr(source, "context_length", None), threshold_tokens=getattr(source, "threshold_tokens", None), - # context_messages — deep copy so the branch has independent context - context_messages=copy.deepcopy(getattr(source, "context_messages", None) or []), + # context_messages — truncated to fork prefix (not full parent copy) + context_messages=copy.deepcopy(forked_context), # Gateway routing — inherit from source gateway_routing=copy.deepcopy(getattr(source, "gateway_routing", None)), # Context engine — inherit state so branch's context engine starts correctly diff --git a/api/session_ops.py b/api/session_ops.py index 9f33075f..696826c6 100644 --- a/api/session_ops.py +++ b/api/session_ops.py @@ -67,6 +67,29 @@ def _truncation_watermark_for(messages): return 0.0 +def truncate_context_for_display_keep( + context_messages: list | None, + full_messages: list | None, + keep: int, +) -> list: + """Align model context with display prefix ``full_messages[:keep]``.""" + if keep <= 0: + return [] + ctx = context_messages if isinstance(context_messages, list) else [] + msgs = full_messages if isinstance(full_messages, list) else [] + if not ctx: + return [] + if len(ctx) == len(msgs): + return ctx[:keep] + if len(msgs) == 0: + return [] + # Context tail aligns with full display transcript; preserve leading-only rows. + prefix_len = max(0, len(ctx) - len(msgs)) + prefix = ctx[:prefix_len] + suffix = ctx[prefix_len:] + return prefix + suffix[:keep] + + def retry_last(session_id: str) -> dict[str, Any]: """Truncate the session to before the last user message, return its text. diff --git a/tests/test_issue_branch_context_at_fork.py b/tests/test_issue_branch_context_at_fork.py new file mode 100644 index 00000000..0cbfbeb6 --- /dev/null +++ b/tests/test_issue_branch_context_at_fork.py @@ -0,0 +1,34 @@ +"""Branch must not inherit full parent context_messages (#5096 A).""" + +import copy + +from api.session_ops import truncate_context_for_display_keep + + +def test_truncate_context_for_display_keep_drops_parent_tail(): + msgs = [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "after-fork"}, + ] + # Same-length context as display — fork at keep=2 matches messages[:2] + ctx = copy.deepcopy(msgs) + out = truncate_context_for_display_keep(ctx, msgs, 2) + assert len(out) == 2 + assert out[-1]["content"] == "a1" + assert "after-fork" not in [m["content"] for m in out] + + +def test_truncate_context_preserves_leading_compaction_row(): + msgs = [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + ] + ctx = [ + {"role": "user", "content": "compaction-ref-only"}, + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + ] + out = truncate_context_for_display_keep(ctx, msgs, 2) + assert len(out) == 3 + assert out[0]["content"] == "compaction-ref-only" \ No newline at end of file