From 35701ce8098dddad955bd3d51bd632d93ad05a73 Mon Sep 17 00:00:00 2001
From: thejesh23 <35212698+thejesh23@users.noreply.github.com>
Date: Mon, 13 Jul 2026 16:57:32 -0700
Subject: [PATCH] fix(shared_context): don't evict an unrelated entry on an
 update at capacity (#2136)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes #2135.

## Summary

`SharedContext.put` ran `_evict_if_needed` before writing, and the
eviction loop only checked `len(self._entries) >= self._max_entries`.
When a caller updated a key that was already cached at capacity, the put
would not have grown the map — but the loop still dropped the oldest
unrelated entry.

Same defect class as fixed for `SemanticCache` in #2094: the eviction
path must know the incoming key so an update is not treated as an
insert. This mirrors that fix over to `SharedContext`.

Threads the incoming key through `_evict_if_needed` and skips capacity
eviction when it names an entry that already exists. Expired-entry
cleanup still runs unconditionally.

Issue #2135 has the reproduction and impact writeup.

## Test plan

- [x] `uv run pytest tests/test_shared_context.py` — 16 passed (added
`test_updating_existing_key_at_capacity_does_not_evict`).
- [x] `uv run ruff check headroom/shared_context.py
tests/test_shared_context.py` — clean.
- [x] `uv run ruff format --check headroom/shared_context.py
tests/test_shared_context.py` — already formatted.

## Real behavior proof

**Setup:** macOS 25.4 (Darwin arm64), Python 3.12.13, `uv 0.11.28`, this
branch (`fix/shared-context-evict-on-update`).

**Before the patch (unpatched `main`)**

\`\`\`
before update: ['a', 'b', 'c']
after update: ['b', 'c'] # <-- 'a' evicted, even though 'c' was an
update
\`\`\`

**After the patch (this branch)**

\`\`\`
\$ uv run python <<'PY'
from headroom.shared_context import SharedContext
ctx = SharedContext(ttl=3600, max_entries=3)
ctx.put(\"a\", \"x\"*400)
ctx.put(\"b\", \"x\"*400)
ctx.put(\"c\", \"x\"*400)
print(\"before update:\", sorted(ctx.keys()))
ctx.put(\"c\", \"y\"*400)   # update existing at capacity
print(\"after update: \", sorted(ctx.keys()))
print(\"c value:\", ctx.get(\"c\", full=True)[:12] + \"...\")
PY
before update: ['a', 'b', 'c']
after update:  ['a', 'b', 'c']
c value: yyyyyyyyyyyy...
\`\`\`

**Test output**

\`\`\`
\$ uv run pytest tests/test_shared_context.py -q
................ [100%]
16 passed in 2.17s
\`\`\`

**What I did NOT test**

- Multi-thread test — the fix is inside the existing `self._lock`, so
serialization semantics are unchanged; I did not add a concurrent-put
stress test.
- Interaction with TTL expiry AND capacity in one call — the existing
`test_evicts_oldest_at_capacity` and `test_expired_entry_returns_none`
still pass, but I did not add a combined case.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
---
 CHANGELOG.md                 |  1 +
 headroom/shared_context.py   | 16 +++++++++++++---
 tests/test_shared_context.py | 21 +++++++++++++++++++++
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 395b3d338..636da84a7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 
 ### Fixed
+- **shared_context:** `SharedContext.put` no longer evicts an unrelated entry when it merely updates a key that is already cached at capacity — same defect class fixed for `SemanticCache` in [#2094](https://github.com/headroomlabs-ai/headroom/pull/2094).
 - **compress:** stop mutating the caller's `CompressConfig`. `compress(config=my_cfg, protect_recent=0, target_ratio=0.2)` used to write those kwargs onto `my_cfg`, so a shared per-agent config was silently rewritten by every request that overrode a single option.
 - **paths:** reject `.`, `..`, and NUL as plugin names so `plugin_config_dir` / `plugin_workspace_dir` cannot resolve outside the `plugins/` sandbox. Previously `plugin_config_dir("..")` returned the entire config root and `plugin_workspace_dir("..")` returned the workspace root (savings ledger, memory DB, license cache, logs).
 - **backends/litellm:** drop tool names over 64 chars before calling Bedrock Converse (`send_message` and `stream_message`), instead of letting the whole request 401. The Bedrock Converse API hard-rejects any tool name past that length, and Claude Code includes every globally-added claude.ai MCP connector tool in every request, even ones the user hasn't enabled locally, so a single oversized connector name broke every call through this backend. Only the `bedrock` provider filters; other providers forward tool names unfiltered.
diff --git a/headroom/shared_context.py b/headroom/shared_context.py
index fd4334a6a..d161c5a4d 100644
--- a/headroom/shared_context.py
+++ b/headroom/shared_context.py
@@ -128,7 +128,7 @@ class SharedContext:
         )
 
         with self._lock:
-            self._evict_if_needed()
+            self._evict_if_needed(incoming_key=key)
             self._entries[key] = entry
 
         logger.debug(
@@ -206,13 +206,23 @@ class SharedContext:
         with self._lock:
             self._entries.clear()
 
-    def _evict_if_needed(self) -> None:
-        """Evict expired and oldest entries if at capacity. Lock must be held."""
+    def _evict_if_needed(self, *, incoming_key: str | None = None) -> None:
+        """Evict expired and oldest entries if at capacity. Lock must be held.
+
+        ``incoming_key`` is the key about to be written. When it names an entry
+        that already exists, this ``put`` is an update — the map size will not
+        grow — so no eviction is required. Without this guard, updating a key
+        at capacity dropped an unrelated entry (same defect fixed for
+        ``SemanticCache`` in #2094).
+        """
         now = time.time()
         expired = [k for k, e in self._entries.items() if now - e.timestamp > self._ttl]
         for k in expired:
             del self._entries[k]
 
+        if incoming_key is not None and incoming_key in self._entries:
+            return
+
         while len(self._entries) >= self._max_entries:
             oldest_key = min(self._entries, key=lambda k: self._entries[k].timestamp)
             del self._entries[oldest_key]
diff --git a/tests/test_shared_context.py b/tests/test_shared_context.py
index 80a13b21f..718e5b46d 100644
--- a/tests/test_shared_context.py
+++ b/tests/test_shared_context.py
@@ -110,6 +110,27 @@ class TestEviction:
         assert ctx.get("second") is not None
         assert ctx.get("third") is not None
 
+    def test_updating_existing_key_at_capacity_does_not_evict(self) -> None:
+        """Overwriting an existing key at capacity must not evict an unrelated one.
+
+        Regression: ``_evict_if_needed`` runs before the assignment, so it
+        drops the oldest entry even when the ``put`` was going to overwrite
+        (not grow) the map — the size stays inside the cap without eviction.
+        Same class of bug as fixed for ``SemanticCache`` in #2094.
+        """
+
+        ctx = SharedContext(max_entries=3)
+        ctx.put("a", "data a")
+        ctx.put("b", "data b")
+        ctx.put("c", "data c")
+        assert set(ctx.keys()) == {"a", "b", "c"}
+
+        # Update an existing key at capacity — must be a no-op for the others.
+        ctx.put("c", "data c UPDATED")
+
+        assert set(ctx.keys()) == {"a", "b", "c"}
+        assert ctx.get("c", full=True) == "data c UPDATED"
+
 
 class TestClear:
     def test_clear_removes_all(self) -> None: