merge: output-shaper replays the pinned tail when shaping stops (headroomlabs-ai#2524)

Author: nangsontay

.github/workflows/docker.yml

@@ -204,10 +204,10 @@ jobs:
       - name: Upload digest marker
         uses: actions/upload-artifact@v7
         with:
-          # Variant + arch in the artifact name so the manifest job can
-          # download with `pattern: digests-<variant>-*` to gather all
-          # arches for one variant. `root` substitutes the empty-string
-          # variant since GHA artifact names can't end in a hyphen.
+          # Variant + arch uniquely identify the marker. The manifest job
+          # downloads both architecture artifacts by exact name; a glob such
+          # as `digests-code-*` would also match code-nonroot/code-slim.
+          # `root` substitutes the empty-string variant.
           name: digests-${{ matrix.variant.name || 'root' }}-${{ matrix.arch.name }}
           path: ${{ runner.temp }}/digests/*
           if-no-files-found: error
@@ -273,12 +273,17 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Download per-arch digests for this variant
+      - name: Download amd64 digest for this variant
         uses: actions/download-artifact@v8
         with:
-          pattern: digests-${{ matrix.variant.name || 'root' }}-*
+          name: digests-${{ matrix.variant.name || 'root' }}-amd64
+          path: ${{ runner.temp }}/digests
+
+      - name: Download arm64 digest for this variant
+        uses: actions/download-artifact@v8
+        with:
+          name: digests-${{ matrix.variant.name || 'root' }}-arm64
           path: ${{ runner.temp }}/digests
-          merge-multiple: true
 
       # Same tag rules as the pre-fan-out workflow — preserve every
       # tag flavor (semver, ref, sha-prefixed, version-suffixed,

.release-please-config.json

@@ -7,8 +7,8 @@
   "bump-patch-for-minor-pre-major": false,
   "draft": false,
   "prerelease": false,
-  "separate-pull-requests": false,
-  "pull-request-title-pattern": "chore: release ${version}",
+  "separate-pull-requests": true,
+  "pull-request-title-pattern": "chore: release${component} ${version}",
   "packages": {
     ".": {
       "package-name": "headroom-ai",

headroom/proxy/handlers/anthropic.py

@@ -2988,18 +2988,32 @@ def _count_tool_tokens(value: object) -> int:
                 )
                 from headroom.proxy.output_shaper import (
                     OutputShaperSettings,
+                    apply_verbosity_steering,
                     classify_turn,
                     resolve_verbosity_level,
                     shape_request,
                 )
 
+                # The level an established frozen prefix was built at outranks
+                # the CURRENT configuration. Dropping the tail from a prefix
+                # built with it busts the whole provider cache exactly like
+                # appending one to a prefix built without it — so the pin is
+                # read before the enablement and arm gates and replayed below
+                # when those no longer hold (shaper switched off live, rollout
+                # channel moved, holdout drift into control). Only an explicit
+                # request bypass, handled by the enclosing guard, may sacrifice
+                # the cache.
+                _pinned = getattr(prefix_tracker, "output_shaping_level", None)
+                _pinned_established = frozen_message_count > 0 and _pinned is not None
+
                 _shaper_settings = OutputShaperSettings.from_env(
                     enabled=(
                         self.config.rollout.is_enabled("proxy_output_shaper")
                         if getattr(self.config, "rollout", None) is not None
                         else None
                     )
                 )
+                _arm = "control"
                 if _shaper_settings.enabled:
                     # Conversation-stable holdout assignment: a whole
                     # conversation is treatment or control. This keeps the A/B
@@ -3027,48 +3041,56 @@ def _count_tool_tokens(value: object) -> int:
                     # outcome funnel can feed the savings ledger from any path.
                     transforms_applied.append(stratum_label(_arm, _stratum))
 
-                    if _arm == "treatment":
-                        # The system prompt is the head of the provider's cache
-                        # prefix, so the shaping tail has to stay byte-identical
-                        # for every turn of a conversation whose frozen prefix
-                        # is established. All three drifts invalidate that whole
-                        # prefix: appending the tail to a prefix built without
-                        # it, dropping the tail from one built with it, and
-                        # swapping it for another level's text. The level the
-                        # conversation was established at is therefore pinned on
-                        # the tracker and replayed until the prefix thaws.
-                        _pinned = getattr(prefix_tracker, "output_shaping_level", None)
-                        _level, _src = resolve_verbosity_level(_shaper_settings)
-                        if frozen_message_count > 0 and _pinned is None:
+                if _shaper_settings.enabled and _arm == "treatment":
+                    # The system prompt is the head of the provider's cache
+                    # prefix, so the shaping tail has to stay byte-identical
+                    # for every turn of a conversation whose frozen prefix
+                    # is established. All three drifts invalidate that whole
+                    # prefix: appending the tail to a prefix built without
+                    # it, dropping the tail from one built with it, and
+                    # swapping it for another level's text. The level the
+                    # conversation was established at is therefore pinned on
+                    # the tracker and replayed until the prefix thaws.
+                    _level, _src = resolve_verbosity_level(_shaper_settings)
+                    if frozen_message_count > 0 and _pinned is None:
+                        logger.info(
+                            f"[{request_id}] OutputShaper: skipped — frozen "
+                            f"prefix ({frozen_message_count} messages) predates "
+                            "the shaping tail; preserving provider cache"
+                        )
+                    else:
+                        if _pinned_established and _pinned != _level:
                             logger.info(
-                                f"[{request_id}] OutputShaper: skipped — frozen "
-                                f"prefix ({frozen_message_count} messages) predates "
-                                "the shaping tail; preserving provider cache"
+                                f"[{request_id}] OutputShaper: L{_level}/{_src} "
+                                f"pinned to L{_pinned} — the frozen prefix "
+                                f"({frozen_message_count} messages) carries "
+                                "that tail"
                             )
-                        else:
-                            if (
-                                frozen_message_count > 0
-                                and _pinned is not None
-                                and _pinned != _level
-                            ):
-                                logger.info(
-                                    f"[{request_id}] OutputShaper: L{_level}/{_src} "
-                                    f"pinned to L{_pinned} — the frozen prefix "
-                                    f"({frozen_message_count} messages) carries "
-                                    "that tail"
-                                )
-                                _level, _src = _pinned, "pinned"
-                            shape_result = shape_request(
-                                body, _shaper_settings, level_override=_level
+                            _level, _src = _pinned, "pinned"
+                        shape_result = shape_request(body, _shaper_settings, level_override=_level)
+                        prefix_tracker.output_shaping_level = _level
+                        if shape_result.changed:
+                            body_mutation_tracker.mark_mutated("output_shaper")
+                            transforms_applied.extend(shape_result.labels or [])
+                            logger.info(
+                                f"[{request_id}] OutputShaper(L{_level}/{_src}): "
+                                f"{shape_result.labels}"
                             )
-                            prefix_tracker.output_shaping_level = _level
-                            if shape_result.changed:
-                                body_mutation_tracker.mark_mutated("output_shaper")
-                                transforms_applied.extend(shape_result.labels or [])
-                                logger.info(
-                                    f"[{request_id}] OutputShaper(L{_level}/{_src}): "
-                                    f"{shape_result.labels}"
-                                )
+                elif _pinned_established:
+                    # Replay only. The conversation is no longer being shaped
+                    # (disabled, or control), but its frozen prefix already
+                    # carries the L{_pinned} tail, so those exact bytes have to
+                    # keep going out. Deliberately narrower than the treatment
+                    # path: no experiment label (attribution belongs to the live
+                    # assignment, not to a cache replay) and no effort routing
+                    # (that lever never enters the cached prefix).
+                    if _pinned > 0 and apply_verbosity_steering(body, _pinned):
+                        body_mutation_tracker.mark_mutated("output_shaper")
+                    logger.info(
+                        f"[{request_id}] OutputShaper: replaying pinned L{_pinned} — "
+                        f"the frozen prefix ({frozen_message_count} messages) carries "
+                        "that tail; shaping is off or in control"
+                    )
 
             # Unit 2: mark end of pre-upstream phase. Everything after this
             # point is upstream I/O or post-response bookkeeping.

tests/test_proxy_anthropic_cache_stability.py

@@ -671,6 +671,109 @@ def _post() -> None:
             assert fake_tracker.output_shaping_level == 2
 
 
+@pytest.mark.parametrize("drift", ["shaper_disabled", "arm_control"])
+def test_output_shaper_replays_pinned_tail_when_shaping_stops(monkeypatch, drift: str) -> None:
+    """Shaping stopping mid-conversation must not drop an established tail.
+
+    The pin is consulted before the enablement and arm gates: a conversation
+    that froze with an L2 tail keeps sending those exact system bytes even once
+    the shaper is disabled or its arm drifts into control. Dropping the tail
+    invalidates the whole provider prefix, the same bust as adding one.
+    Experiment attribution stays with the live assignment — a replay never
+    records a treatment label.
+    """
+    from headroom.proxy import runtime_env
+    from headroom.proxy.output_savings_policy import _STRATUM_LABEL
+    from headroom.proxy.output_verbosity_policy import STEERING_SENTINEL
+    from headroom.rollout import resolve_rollout
+
+    runtime_env.clear_overrides()
+    monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1")
+    # Rollout-gated as in the tests above: beta is required for the shaper to
+    # run at all.
+    monkeypatch.setenv("HEADROOM_ROLLOUT_CHANNEL", "beta")
+    monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "2")
+
+    captured = {}
+    with _make_proxy_client() as client:
+        proxy = client.app.state.proxy
+        proxy.config.optimize = False
+        proxy.config.image_optimize = False
+        proxy.config.ccr_proactive_expansion = False
+
+        fake_tracker = _FakePrefixTracker(frozen_count=0)
+        proxy.session_tracker_store.compute_session_id = (
+            lambda request, model, messages, **_kwargs: "stable-session"
+        )
+        proxy.session_tracker_store.get_or_create = lambda session_id, provider: fake_tracker
+
+        async def _fake_retry(method, url, headers, body, stream=False, **kwargs):  # noqa: ANN001
+            captured["body"] = body
+            return httpx.Response(
+                200,
+                json={
+                    "id": "msg_shape_3",
+                    "type": "message",
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": "ok"}],
+                    "usage": {
+                        "input_tokens": 20,
+                        "output_tokens": 3,
+                        "cache_read_input_tokens": 0,
+                        "cache_creation_input_tokens": 0,
+                    },
+                },
+            )
+
+        proxy._retry_request = _fake_retry
+
+        def _post() -> str:
+            response = client.post(
+                "/v1/messages",
+                headers={"x-api-key": "test-key", "anthropic-version": "2023-06-01"},
+                json={
+                    "model": "claude-sonnet-4-6",
+                    "max_tokens": 64,
+                    "system": "base system",
+                    "messages": [
+                        {"role": "user", "content": "frozen prefix"},
+                        {"role": "assistant", "content": "ack"},
+                        {"role": "user", "content": "latest user"},
+                    ],
+                },
+            )
+            assert response.status_code == 200
+            return response.headers.get("x-headroom-transforms", "")
+
+        # Turn 1: shaped and pinned, establishing the L2 system bytes.
+        transforms = _post()
+        established = captured["body"]["system"]
+        assert STEERING_SENTINEL in str(established)
+        assert fake_tracker.output_shaping_level == 2
+        assert _STRATUM_LABEL in transforms
+
+        # Turn N: the prefix is frozen and shaping stops for this conversation.
+        fake_tracker._frozen_count = 1
+        if drift == "shaper_disabled":
+            # Enablement is a resolved snapshot since #1490, so a live disable
+            # reaches the handler as a re-resolved gate with the legacy alias
+            # falsey — the same route a rollout-channel demotion takes.
+            proxy.config.rollout = resolve_rollout(
+                {"HEADROOM_ROLLOUT_CHANNEL": "beta", "HEADROOM_OUTPUT_SHAPER": "0"}
+            )
+        else:
+            monkeypatch.setenv("HEADROOM_OUTPUT_HOLDOUT", "1")
+        transforms = _post()
+        assert captured["body"]["system"] == established
+        assert fake_tracker.output_shaping_level == 2
+
+        # The treatment path is unreachable in both drifts, so surviving bytes
+        # can only have come from the replay. The missing treatment label is
+        # what proves the drift actually took effect — and that a cache replay
+        # is never counted as a treatment observation.
+        assert _STRATUM_LABEL not in transforms
+
+
 def test_ccr_system_instruction_injection_disabled_when_prefix_frozen(monkeypatch) -> None:
     captured = {"inject_system": None}
     with _make_proxy_client() as client:

tests/test_release_workflows.py

@@ -68,6 +68,40 @@ def test_docker_latest_promotion_is_owned_by_root_manifest_cell() -> None:
     assert guard_start < manifest_script.index("exit 1", guard_start) < create_start
 
 
+def test_docker_manifest_downloads_exactly_one_artifact_per_architecture() -> None:
+    """Each manifest cell must download exactly its two architecture digests.
+
+    Keeping the variant before a trailing wildcard makes prefix-related names
+    overlap: ``digests-code-*`` also selects code-nonroot, code-slim, and
+    code-slim-nonroot. The 0.35.0 Docker release exposed this by downloading
+    eight markers into the code manifest job instead of two.
+    """
+    workflow = yaml.safe_load((ROOT / ".github" / "workflows" / "docker.yml").read_text())
+    jobs = workflow["jobs"]
+    build = jobs["docker-build"]
+    manifest = jobs["docker-manifest"]
+    upload = next(step for step in build["steps"] if step.get("name") == "Upload digest marker")
+    downloads = [
+        step
+        for step in manifest["steps"]
+        if step.get("name")
+        in {
+            "Download amd64 digest for this variant",
+            "Download arm64 digest for this variant",
+        }
+    ]
+
+    assert upload["with"]["name"] == (
+        "digests-${{ matrix.variant.name || 'root' }}-${{ matrix.arch.name }}"
+    )
+    assert [step["with"]["name"] for step in downloads] == [
+        "digests-${{ matrix.variant.name || 'root' }}-amd64",
+        "digests-${{ matrix.variant.name || 'root' }}-arm64",
+    ]
+    assert all("pattern" not in step["with"] for step in downloads)
+    assert all(step["with"]["path"] == "${{ runner.temp }}/digests" for step in downloads)
+
+
 def test_release_workflow_publishes_both_node_packages_to_github_packages() -> None:
     content = (ROOT / ".github" / "workflows" / "release.yml").read_text(encoding="utf-8")
 
@@ -1287,6 +1321,22 @@ def test_release_please_config_and_manifest_are_present_and_consistent() -> None
         "changelog because the bot can't find its baseline."
     )
 
+    # This manifest has one package. Sending it through the merge plugin
+    # produces the group title `chore: release main`, which contains neither
+    # the package component nor its version. On merge, release-please cannot
+    # associate that title with `headroom-ai`, leaves the PR tagged
+    # `autorelease: pending`, and never emits the release event that publishes
+    # to PyPI. Keep the single package on the normal, versioned PR path and
+    # preserve the component in the title used to match the merged PR.
+    assert config.get("separate-pull-requests") is True, (
+        "The single root package must bypass release-please's merge plugin; "
+        "its grouped PR title is `chore: release main` and cannot be tagged."
+    )
+    assert config.get("pull-request-title-pattern") == ("chore: release${component} ${version}"), (
+        "Release PR titles must include both component and version so "
+        "release-please can match the merged PR back to headroom-ai."
+    )
+
     # extra-files: TypeScript SDK and npm plugin package.json files
     # files must be in lockstep with pyproject.toml.
     extra_paths = {ef["path"] for ef in root_pkg.get("extra-files", [])}