mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix: A7 — memory tool injection session-sticky for both Anthropic and OpenAI
Closes the second half of P0-6: once memory injects memory_save / memory_search into body["tools"] for a session, every subsequent turn injects the byte-equal same definitions — even if memory is disabled mid-session. Toggling tool list mid-session busts Anthropic prefix cache per guide §6.3 #2. Adds in headroom/proxy/helpers.py: * SessionToolTracker — bounded LRU keyed by (provider, session_id) storing GOLDEN tool-definition bytes from the first injection. Tracker is provider-aware so the same session_id under Anthropic and OpenAI keeps independent state. Reentrant lock for concurrent access; LRU eviction at HEADROOM_TOOL_TRACKER_MAX_SESSIONS (default 1000). * apply_session_sticky_memory_tools — single coordination point with three paths: first-time inject (record golden bytes), sticky replay (always inject golden bytes regardless of inject_this_turn), and skip. Honors HEADROOM_TOOL_INJECTION_STICKY=disabled as a loud operator opt-in for rollback (NOT a fallback). * serialize_tool_definition_canonical — deterministic byte serialization via the same separators=(",",":")/ensure_ascii=False rules as serialize_body_canonical. * log_tool_injection_decision — structured per-decision log line; never logs the tool definition contents. Wires the helper into all four memory tool injection sites: * handlers/anthropic.py — /v1/messages * handlers/openai.py — /v1/chat/completions * handlers/openai.py — /v1/responses * handlers/openai.py — Codex WS path memory_handler.MemoryHandler gains compute_memory_tool_definitions(provider) — a pure builder that returns the tool definitions without mutating a tools list, so the proxy can route through the sticky tracker. The legacy inject_tools(...) is preserved for callers without a session_id. Tests: tests/test_memory_tool_session_sticky.py — 29 unit + integration cases covering: turn-1→turn-2 byte-equality (Anthropic + OpenAI), sticky replay after memory disabled, golden-fixture pin, LRU eviction, provider isolation under shared session_id, thread-safe concurrent access, env-var contract, disabled-mode passthrough, dedupe with client tools. Golden fixtures pin canonical bytes: * tests/fixtures/memory_tool_definitions/anthropic.json * tests/fixtures/memory_tool_definitions/openai.json No regex. No hardcodes (env-configurable: HEADROOM_TOOL_INJECTION_STICKY, HEADROOM_TOOL_TRACKER_MAX_SESSIONS). No silent fallbacks. Per-decision structured logging. Realignment build constraints satisfied.
This commit is contained in:
chopratejas
parent
aec5ba3253
commit
8dcd474aca
7 changed files with 1813 additions and 135 deletions
|
|
@ -1268,69 +1268,87 @@ class AnthropicHandlerMixin:
|
|||
except Exception as e:
|
||||
logger.warning(f"[{request_id}] Memory: Context injection failed: {e}")
|
||||
|
||||
# Inject memory tools
|
||||
if self.memory_handler.config.inject_tools:
|
||||
tools, mem_tools_injected = self.memory_handler.inject_tools(tools, "anthropic")
|
||||
if mem_tools_injected:
|
||||
memory_tools_injected = True
|
||||
tool_names = [
|
||||
t.get("name") or t.get("type", "")
|
||||
for t in tools
|
||||
if t.get("name", "").startswith("memory")
|
||||
or t.get("type", "").startswith("memory")
|
||||
]
|
||||
logger.info(f"[{request_id}] Memory: Injected tools: {tool_names}")
|
||||
# Inject memory tools — PR-A7 (P0-6) routes through
|
||||
# `apply_session_sticky_memory_tools` so tool list bytes
|
||||
# stay byte-stable across turns: once a session injects,
|
||||
# every subsequent turn replays the same canonical bytes.
|
||||
# `inject_this_turn` is True iff memory is enabled this
|
||||
# turn (i.e. memory_handler.config.inject_tools and we
|
||||
# have a memory_user_id, which the outer guard at line
|
||||
# 1192 already enforces).
|
||||
from headroom.proxy.helpers import (
|
||||
apply_session_sticky_memory_tools,
|
||||
)
|
||||
|
||||
# Add beta headers for native memory tool. PR-A6
|
||||
# (P5-50): use the deterministic `merge_anthropic_beta`
|
||||
# helper instead of ad-hoc string concat. Order:
|
||||
# client tokens first (preserved from session-sticky
|
||||
# baseline above), then Headroom-required tokens.
|
||||
# The session tracker already recorded the client
|
||||
# value; we append Headroom-required tokens here so
|
||||
# the next turn re-applies them deterministically.
|
||||
beta_headers = self.memory_handler.get_beta_headers()
|
||||
if beta_headers:
|
||||
from headroom.proxy.helpers import (
|
||||
log_beta_header_merge as _log_beta_header_merge_mem,
|
||||
)
|
||||
from headroom.proxy.helpers import (
|
||||
merge_anthropic_beta,
|
||||
)
|
||||
memory_tool_defs = (
|
||||
self.memory_handler.compute_memory_tool_definitions("anthropic")
|
||||
if self.memory_handler.config.inject_tools
|
||||
else []
|
||||
)
|
||||
tools, mem_tools_injected = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id=session_id,
|
||||
request_id=request_id,
|
||||
existing_tools=tools,
|
||||
memory_tools_to_inject=memory_tool_defs,
|
||||
inject_this_turn=bool(self.memory_handler.config.inject_tools),
|
||||
)
|
||||
if mem_tools_injected:
|
||||
memory_tools_injected = True
|
||||
tool_names = [
|
||||
t.get("name") or t.get("type", "")
|
||||
for t in tools
|
||||
if t.get("name", "").startswith("memory")
|
||||
or t.get("type", "").startswith("memory")
|
||||
]
|
||||
logger.info(f"[{request_id}] Memory: Injected tools: {tool_names}")
|
||||
|
||||
for key, value in beta_headers.items():
|
||||
if key.lower() != "anthropic-beta":
|
||||
# Defensive: memory handler currently
|
||||
# only emits anthropic-beta. Any future
|
||||
# provider-specific beta header would
|
||||
# need its own merge helper.
|
||||
headers[key] = value
|
||||
continue
|
||||
existing_value = headers.get(key, "")
|
||||
required_tokens = [t.strip() for t in value.split(",") if t.strip()]
|
||||
merged = merge_anthropic_beta(existing_value, required_tokens)
|
||||
_existing_count = (
|
||||
len([t for t in existing_value.split(",") if t.strip()])
|
||||
if existing_value
|
||||
else 0
|
||||
)
|
||||
_merged_count = (
|
||||
len([t for t in merged.split(",") if t.strip()])
|
||||
if merged
|
||||
else 0
|
||||
)
|
||||
headers[key] = merged
|
||||
_log_beta_header_merge_mem(
|
||||
provider="anthropic",
|
||||
session_id=session_id,
|
||||
client_betas_count=_existing_count,
|
||||
sticky_betas_count=_merged_count,
|
||||
headroom_added=required_tokens,
|
||||
request_id=request_id,
|
||||
)
|
||||
logger.info(
|
||||
f"[{request_id}] Memory: Added beta header: {key}={merged}"
|
||||
)
|
||||
# Add beta headers for native memory tool. PR-A6
|
||||
# (P5-50): use the deterministic `merge_anthropic_beta`
|
||||
# helper instead of ad-hoc string concat. Order:
|
||||
# client tokens first (preserved from session-sticky
|
||||
# baseline above), then Headroom-required tokens.
|
||||
# The session tracker already recorded the client
|
||||
# value; we append Headroom-required tokens here so
|
||||
# the next turn re-applies them deterministically.
|
||||
beta_headers = self.memory_handler.get_beta_headers()
|
||||
if beta_headers:
|
||||
from headroom.proxy.helpers import (
|
||||
log_beta_header_merge as _log_beta_header_merge_mem,
|
||||
)
|
||||
from headroom.proxy.helpers import (
|
||||
merge_anthropic_beta,
|
||||
)
|
||||
|
||||
for key, value in beta_headers.items():
|
||||
if key.lower() != "anthropic-beta":
|
||||
# Defensive: memory handler currently
|
||||
# only emits anthropic-beta. Any future
|
||||
# provider-specific beta header would
|
||||
# need its own merge helper.
|
||||
headers[key] = value
|
||||
continue
|
||||
existing_value = headers.get(key, "")
|
||||
required_tokens = [t.strip() for t in value.split(",") if t.strip()]
|
||||
merged = merge_anthropic_beta(existing_value, required_tokens)
|
||||
_existing_count = (
|
||||
len([t for t in existing_value.split(",") if t.strip()])
|
||||
if existing_value
|
||||
else 0
|
||||
)
|
||||
_merged_count = (
|
||||
len([t for t in merged.split(",") if t.strip()]) if merged else 0
|
||||
)
|
||||
headers[key] = merged
|
||||
_log_beta_header_merge_mem(
|
||||
provider="anthropic",
|
||||
session_id=session_id,
|
||||
client_betas_count=_existing_count,
|
||||
sticky_betas_count=_merged_count,
|
||||
headroom_added=required_tokens,
|
||||
request_id=request_id,
|
||||
)
|
||||
logger.info(f"[{request_id}] Memory: Added beta header: {key}={merged}")
|
||||
|
||||
if memory_context_injected or memory_tools_injected:
|
||||
remembered_event = self.pipeline_extensions.emit(
|
||||
|
|
|
|||
|
|
@ -636,12 +636,28 @@ class OpenAIHandlerMixin:
|
|||
query=None,
|
||||
)
|
||||
|
||||
# Inject memory tools
|
||||
if self.memory_handler.config.inject_tools:
|
||||
tools, mem_tools_injected = self.memory_handler.inject_tools(tools, "openai")
|
||||
if mem_tools_injected:
|
||||
memory_tools_injected = True
|
||||
logger.info(f"[{request_id}] Memory: Injected memory tools (openai)")
|
||||
# Inject memory tools — PR-A7 (P0-6) routes through
|
||||
# `apply_session_sticky_memory_tools` so byte-stable across turns.
|
||||
from headroom.proxy.helpers import (
|
||||
apply_session_sticky_memory_tools as _apply_sticky_mem_tools,
|
||||
)
|
||||
|
||||
memory_tool_defs = (
|
||||
self.memory_handler.compute_memory_tool_definitions("openai")
|
||||
if self.memory_handler.config.inject_tools
|
||||
else []
|
||||
)
|
||||
tools, mem_tools_injected = _apply_sticky_mem_tools(
|
||||
provider="openai",
|
||||
session_id=openai_session_id,
|
||||
request_id=request_id,
|
||||
existing_tools=tools,
|
||||
memory_tools_to_inject=memory_tool_defs,
|
||||
inject_this_turn=bool(self.memory_handler.config.inject_tools),
|
||||
)
|
||||
if mem_tools_injected:
|
||||
memory_tools_injected = True
|
||||
logger.info(f"[{request_id}] Memory: Injected memory tools (openai)")
|
||||
except Exception as e:
|
||||
logger.warning(f"[{request_id}] Memory injection failed: {e}")
|
||||
|
||||
|
|
@ -1423,32 +1439,47 @@ class OpenAIHandlerMixin:
|
|||
query=user_query,
|
||||
)
|
||||
|
||||
# Inject memory tools (Responses API format)
|
||||
if self.memory_handler.config.inject_tools:
|
||||
resp_tools = body.get("tools") or []
|
||||
resp_tools, mem_tools_injected = self.memory_handler.inject_tools(
|
||||
resp_tools, "openai"
|
||||
)
|
||||
if mem_tools_injected:
|
||||
# Convert Chat Completions format to Responses API format
|
||||
converted_tools = []
|
||||
for t in resp_tools:
|
||||
if t.get("type") == "function" and "function" in t:
|
||||
fn = t["function"]
|
||||
converted_tools.append(
|
||||
{
|
||||
"type": "function",
|
||||
"name": fn.get("name"),
|
||||
"description": fn.get("description", ""),
|
||||
"parameters": fn.get("parameters", {}),
|
||||
}
|
||||
)
|
||||
else:
|
||||
converted_tools.append(t)
|
||||
body["tools"] = converted_tools
|
||||
logger.info(
|
||||
f"[{request_id}] Memory: Injected memory tools (openai/responses)"
|
||||
# Inject memory tools (Responses API format) — PR-A7 (P0-6).
|
||||
# Pre-convert the Chat-Completions schema to Responses API
|
||||
# format BEFORE handing to the sticky tracker so the
|
||||
# canonical bytes pinned in turn 1 already reflect the
|
||||
# exact bytes that will hit the wire.
|
||||
from headroom.proxy.helpers import (
|
||||
apply_session_sticky_memory_tools as _apply_sticky_mem_tools_resp,
|
||||
)
|
||||
|
||||
memory_tool_defs_chat = (
|
||||
self.memory_handler.compute_memory_tool_definitions("openai")
|
||||
if self.memory_handler.config.inject_tools
|
||||
else []
|
||||
)
|
||||
memory_tool_defs_responses: list[dict[str, Any]] = []
|
||||
for t in memory_tool_defs_chat:
|
||||
if t.get("type") == "function" and "function" in t:
|
||||
fn = t["function"]
|
||||
memory_tool_defs_responses.append(
|
||||
{
|
||||
"type": "function",
|
||||
"name": fn.get("name"),
|
||||
"description": fn.get("description", ""),
|
||||
"parameters": fn.get("parameters", {}),
|
||||
}
|
||||
)
|
||||
else:
|
||||
memory_tool_defs_responses.append(t)
|
||||
|
||||
resp_tools = body.get("tools") or []
|
||||
resp_tools, mem_tools_injected = _apply_sticky_mem_tools_resp(
|
||||
provider="openai",
|
||||
session_id=_responses_session_id,
|
||||
request_id=request_id,
|
||||
existing_tools=resp_tools,
|
||||
memory_tools_to_inject=memory_tool_defs_responses,
|
||||
inject_this_turn=bool(self.memory_handler.config.inject_tools),
|
||||
)
|
||||
if mem_tools_injected:
|
||||
body["tools"] = resp_tools
|
||||
logger.info(f"[{request_id}] Memory: Injected memory tools (openai/responses)")
|
||||
except Exception as e:
|
||||
logger.warning(f"[{request_id}] Memory injection failed (responses): {e}")
|
||||
|
||||
|
|
@ -2024,51 +2055,68 @@ class OpenAIHandlerMixin:
|
|||
f"of context into instructions"
|
||||
)
|
||||
|
||||
# Inject memory tools (Responses API format)
|
||||
if self.memory_handler.config.inject_tools:
|
||||
ws_tools = ws_response_body.get("tools") or []
|
||||
ws_tools, mem_injected = self.memory_handler.inject_tools(
|
||||
ws_tools, "openai"
|
||||
)
|
||||
if mem_injected:
|
||||
converted_tools = []
|
||||
for t in ws_tools:
|
||||
if t.get("type") == "function" and "function" in t:
|
||||
fn = t["function"]
|
||||
converted_tools.append(
|
||||
{
|
||||
"type": "function",
|
||||
"name": fn.get("name"),
|
||||
"description": fn.get("description", ""),
|
||||
"parameters": fn.get("parameters", {}),
|
||||
}
|
||||
)
|
||||
else:
|
||||
converted_tools.append(t)
|
||||
ws_response_body["tools"] = converted_tools
|
||||
# Inject memory tools (Responses API format) — PR-A7 (P0-6).
|
||||
# WS path uses a per-connection UUID; tracker scope is
|
||||
# the WS session (short-lived). Pre-convert to Responses
|
||||
# API format so canonical bytes match the wire format.
|
||||
from headroom.proxy.helpers import (
|
||||
apply_session_sticky_memory_tools as _apply_sticky_mem_tools_ws,
|
||||
)
|
||||
|
||||
# Add memory instruction so the model uses
|
||||
# memory tools as persistent cross-session knowledge.
|
||||
mem_instruction = (
|
||||
"\n\n## Memory\n"
|
||||
"You have persistent memory via memory_search and "
|
||||
"memory_save tools. Memory stores knowledge across "
|
||||
"sessions — user info, project details, org context, "
|
||||
"decisions, architecture, conventions, anything worth "
|
||||
"remembering.\n\n"
|
||||
"- ALWAYS call memory_search BEFORE searching files "
|
||||
"when the user asks a question that could be answered "
|
||||
"from prior knowledge.\n"
|
||||
"- Call memory_save to store important facts, decisions, "
|
||||
"or context that would be useful in future sessions.\n"
|
||||
"- Memory is your first source of truth for anything "
|
||||
"not visible in the current conversation."
|
||||
)
|
||||
existing_instr = ws_response_body.get("instructions") or ""
|
||||
ws_response_body["instructions"] = existing_instr + mem_instruction
|
||||
logger.info(
|
||||
f"[{request_id}] WS Memory: Injected memory tools + instruction"
|
||||
ws_mem_defs_chat = (
|
||||
self.memory_handler.compute_memory_tool_definitions("openai")
|
||||
if self.memory_handler.config.inject_tools
|
||||
else []
|
||||
)
|
||||
ws_mem_defs_responses: list[dict[str, Any]] = []
|
||||
for t in ws_mem_defs_chat:
|
||||
if t.get("type") == "function" and "function" in t:
|
||||
fn = t["function"]
|
||||
ws_mem_defs_responses.append(
|
||||
{
|
||||
"type": "function",
|
||||
"name": fn.get("name"),
|
||||
"description": fn.get("description", ""),
|
||||
"parameters": fn.get("parameters", {}),
|
||||
}
|
||||
)
|
||||
else:
|
||||
ws_mem_defs_responses.append(t)
|
||||
|
||||
ws_tools = ws_response_body.get("tools") or []
|
||||
ws_tools, mem_injected = _apply_sticky_mem_tools_ws(
|
||||
provider="openai",
|
||||
session_id=session_id,
|
||||
request_id=request_id,
|
||||
existing_tools=ws_tools,
|
||||
memory_tools_to_inject=ws_mem_defs_responses,
|
||||
inject_this_turn=bool(self.memory_handler.config.inject_tools),
|
||||
)
|
||||
if mem_injected:
|
||||
ws_response_body["tools"] = ws_tools
|
||||
|
||||
# Add memory instruction so the model uses
|
||||
# memory tools as persistent cross-session knowledge.
|
||||
mem_instruction = (
|
||||
"\n\n## Memory\n"
|
||||
"You have persistent memory via memory_search and "
|
||||
"memory_save tools. Memory stores knowledge across "
|
||||
"sessions — user info, project details, org context, "
|
||||
"decisions, architecture, conventions, anything worth "
|
||||
"remembering.\n\n"
|
||||
"- ALWAYS call memory_search BEFORE searching files "
|
||||
"when the user asks a question that could be answered "
|
||||
"from prior knowledge.\n"
|
||||
"- Call memory_save to store important facts, decisions, "
|
||||
"or context that would be useful in future sessions.\n"
|
||||
"- Memory is your first source of truth for anything "
|
||||
"not visible in the current conversation."
|
||||
)
|
||||
existing_instr = ws_response_body.get("instructions") or ""
|
||||
ws_response_body["instructions"] = existing_instr + mem_instruction
|
||||
logger.info(
|
||||
f"[{request_id}] WS Memory: Injected memory tools + instruction"
|
||||
)
|
||||
|
||||
# Write back into envelope if it was wrapped
|
||||
if "response" in body and isinstance(body["response"], dict):
|
||||
|
|
|
|||
|
|
@ -975,6 +975,495 @@ def log_beta_header_merge(
|
|||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Memory-tool injection session-stickiness (PR-A7 — closes P0-6).
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Memory adds `memory_save` / `memory_search` tool definitions to
|
||||
# `body["tools"]` when memory is enabled for a request. The cache-killer
|
||||
# pattern motivated by guide §6.3 #2 ("tool list change → cache bust"):
|
||||
#
|
||||
# * Mid-session toggle: memory is enabled in turn N (tool definitions
|
||||
# injected) and disabled in turn N+1 (tool list shrinks). The next
|
||||
# turn's prefix bytes hash differently, prefix-cache misses, and the
|
||||
# full prompt re-runs at provider cost.
|
||||
#
|
||||
# * Tool definition drift: memory adds the SAME logical tool but the
|
||||
# bytes differ across turns (insertion order, dict key order, schema
|
||||
# drift between deploys, etc.). Even with the tool list intact the
|
||||
# prefix bytes change.
|
||||
#
|
||||
# PR-A7 introduces:
|
||||
#
|
||||
# * `SessionToolTracker`: bounded LRU keyed by (provider, session_id)
|
||||
# storing the GOLDEN tool-definition bytes injected on the first
|
||||
# turn. Subsequent turns of that session always inject the same
|
||||
# bytes — even if memory is disabled mid-session (sticky-on per
|
||||
# guide §6.3 #2). Provider-aware so the same `session_id` under
|
||||
# two providers keeps independent state.
|
||||
#
|
||||
# The golden bytes are produced by `serialize_body_canonical` of the
|
||||
# tool definition object so they are deterministic across deploys
|
||||
# regardless of dict insertion ordering quirks.
|
||||
#
|
||||
# Operator opt-in `HEADROOM_TOOL_INJECTION_STICKY=disabled` short-
|
||||
# circuits the tracker; per-turn decision flows through unchanged. That
|
||||
# mode is loud and explicit per realignment build constraint #4 — NOT a
|
||||
# silent fallback. It exists for diagnostic shadow tracing / emergency
|
||||
# rollback only.
|
||||
|
||||
_TOOL_INJECTION_STICKY_ENV = "HEADROOM_TOOL_INJECTION_STICKY"
|
||||
ToolInjectionStickyMode = Literal["enabled", "disabled"]
|
||||
_TOOL_INJECTION_STICKY_DEFAULT: ToolInjectionStickyMode = "enabled"
|
||||
|
||||
_TOOL_TRACKER_MAX_SESSIONS_ENV = "HEADROOM_TOOL_TRACKER_MAX_SESSIONS"
|
||||
_TOOL_TRACKER_MAX_SESSIONS_DEFAULT = 1000
|
||||
|
||||
|
||||
def get_tool_injection_sticky_mode() -> ToolInjectionStickyMode:
|
||||
"""Return the active memory-tool stickiness mode.
|
||||
|
||||
Read at request time so operators can flip behaviour without a
|
||||
restart. Unknown values raise loudly per the no-silent-fallback
|
||||
build constraint.
|
||||
"""
|
||||
raw = os.environ.get(_TOOL_INJECTION_STICKY_ENV, "").strip().lower()
|
||||
if not raw:
|
||||
return _TOOL_INJECTION_STICKY_DEFAULT
|
||||
if raw in ("enabled", "disabled"):
|
||||
return cast(ToolInjectionStickyMode, raw)
|
||||
raise ValueError(
|
||||
f"Invalid {_TOOL_INJECTION_STICKY_ENV}={raw!r}; expected 'enabled' or 'disabled'"
|
||||
)
|
||||
|
||||
|
||||
def get_tool_tracker_max_sessions() -> int:
|
||||
"""Return the LRU bound for `SessionToolTracker` (sessions cap)."""
|
||||
raw = os.environ.get(_TOOL_TRACKER_MAX_SESSIONS_ENV, "").strip()
|
||||
if not raw:
|
||||
return _TOOL_TRACKER_MAX_SESSIONS_DEFAULT
|
||||
try:
|
||||
value = int(raw)
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"Invalid {_TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int"
|
||||
) from exc
|
||||
if value <= 0:
|
||||
raise ValueError(f"Invalid {_TOOL_TRACKER_MAX_SESSIONS_ENV}={raw!r}; expected positive int")
|
||||
return value
|
||||
|
||||
|
||||
def serialize_tool_definition_canonical(tool_definition: dict[str, Any]) -> bytes:
|
||||
"""Deterministic byte serialization of a single memory tool definition.
|
||||
|
||||
Uses ``serialize_body_canonical`` semantics (compact separators, UTF-8,
|
||||
no ASCII escaping). Python 3.7+ dict insertion order is preserved by
|
||||
``json.dumps`` so callers must construct the tool definition with a
|
||||
stable key order — which the static schemas in
|
||||
``headroom/proxy/memory_handler.py`` and
|
||||
``headroom/proxy/memory_tool_adapter.py`` already do.
|
||||
|
||||
Returned bytes pin the golden tool definition for a session: every
|
||||
follow-up turn must inject byte-equal output to keep the prefix
|
||||
cache hot.
|
||||
"""
|
||||
return serialize_body_canonical(tool_definition)
|
||||
|
||||
|
||||
class SessionToolTracker:
|
||||
"""Bounded LRU tracker recording per-session memory-tool injection state.
|
||||
|
||||
Once memory injects tool definitions into a session, future requests
|
||||
in that session always inject the byte-equal same definitions —
|
||||
never toggling on/off mid-session (guide §6.3 #2). The first turn's
|
||||
canonical bytes are stored as the golden definition; subsequent
|
||||
turns reuse those bytes verbatim.
|
||||
|
||||
State per session: ordered list of (tool_name → golden_bytes)
|
||||
pairs. Order is preserved so the rebuilt tool list matches the
|
||||
original injection order.
|
||||
|
||||
Bounded by ``max_sessions`` (default 1000) via ``OrderedDict`` LRU
|
||||
eviction: hits move-to-end; overflow pops oldest. Reentrant lock so
|
||||
future callers from inside another locked method don't self-deadlock
|
||||
(mirrors `SessionBetaTracker` / `CompressionCache` pattern).
|
||||
|
||||
The tracker is provider-aware: the same ``session_id`` for Anthropic
|
||||
and OpenAI keeps independent state (the tool schemas differ in
|
||||
format).
|
||||
"""
|
||||
|
||||
def __init__(self, max_sessions: int | None = None) -> None:
|
||||
if max_sessions is None:
|
||||
max_sessions = get_tool_tracker_max_sessions()
|
||||
if max_sessions <= 0:
|
||||
raise ValueError("max_sessions must be > 0")
|
||||
self._max_sessions: int = max_sessions
|
||||
self._lock = threading.RLock()
|
||||
# Value is an OrderedDict[tool_name -> golden_definition_bytes].
|
||||
# Storing per-tool bytes (not the entire tools list) keeps the
|
||||
# tracker resilient to non-memory tool list changes by the client
|
||||
# (which are the client's responsibility, not ours to gate).
|
||||
self._sessions: OrderedDict[tuple[str, str], OrderedDict[str, bytes]] = OrderedDict()
|
||||
|
||||
@property
|
||||
def active_sessions(self) -> int:
|
||||
with self._lock:
|
||||
return len(self._sessions)
|
||||
|
||||
def _key(self, provider: str, session_id: str) -> tuple[str, str]:
|
||||
return (provider, session_id)
|
||||
|
||||
def should_inject(self, provider: str, session_id: str) -> bool:
|
||||
"""Return True iff this session has previously injected memory tools.
|
||||
|
||||
Used by the sticky-on path: when memory is disabled this turn but
|
||||
the session previously injected, we still inject the golden bytes.
|
||||
"""
|
||||
if not provider:
|
||||
raise ValueError("provider must be non-empty")
|
||||
if not session_id:
|
||||
raise ValueError("session_id must be non-empty")
|
||||
with self._lock:
|
||||
entry = self._sessions.get(self._key(provider, session_id))
|
||||
if entry is None:
|
||||
return False
|
||||
# LRU touch on read so the carry-over decision keeps the
|
||||
# session in the hot set.
|
||||
self._sessions.move_to_end(self._key(provider, session_id))
|
||||
return len(entry) > 0
|
||||
|
||||
def get_golden_definitions(
|
||||
self, provider: str, session_id: str
|
||||
) -> list[tuple[str, bytes]] | None:
|
||||
"""Return the previously-recorded (name, bytes) pairs for the session.
|
||||
|
||||
Returns ``None`` when the session has never injected memory tools.
|
||||
Callers replay the bytes verbatim into ``body["tools"]``.
|
||||
"""
|
||||
if not provider:
|
||||
raise ValueError("provider must be non-empty")
|
||||
if not session_id:
|
||||
raise ValueError("session_id must be non-empty")
|
||||
with self._lock:
|
||||
entry = self._sessions.get(self._key(provider, session_id))
|
||||
if entry is None:
|
||||
return None
|
||||
self._sessions.move_to_end(self._key(provider, session_id))
|
||||
# Snapshot — never expose internal storage directly.
|
||||
return [(name, golden_bytes) for name, golden_bytes in entry.items()]
|
||||
|
||||
def record_injection(
|
||||
self,
|
||||
provider: str,
|
||||
session_id: str,
|
||||
tool_name: str,
|
||||
tool_definition_bytes: bytes,
|
||||
) -> None:
|
||||
"""Record the golden bytes for a single memory tool in this session.
|
||||
|
||||
First-write wins: re-recording the same ``tool_name`` for an
|
||||
existing session is a no-op (prevents drift if the canonical
|
||||
serialization output changed between deploys mid-session). For
|
||||
a *new* session, record fresh. LRU bound applies on every write.
|
||||
"""
|
||||
if not provider:
|
||||
raise ValueError("provider must be non-empty")
|
||||
if not session_id:
|
||||
raise ValueError("session_id must be non-empty")
|
||||
if not tool_name:
|
||||
raise ValueError("tool_name must be non-empty")
|
||||
if not tool_definition_bytes:
|
||||
raise ValueError("tool_definition_bytes must be non-empty")
|
||||
|
||||
key = self._key(provider, session_id)
|
||||
|
||||
with self._lock:
|
||||
entry = self._sessions.get(key)
|
||||
if entry is None:
|
||||
entry = OrderedDict()
|
||||
self._sessions[key] = entry
|
||||
|
||||
# First-write wins: only record if not already pinned.
|
||||
if tool_name not in entry:
|
||||
entry[tool_name] = tool_definition_bytes
|
||||
|
||||
# LRU touch + bound enforcement.
|
||||
self._sessions.move_to_end(key)
|
||||
while len(self._sessions) > self._max_sessions:
|
||||
self._sessions.popitem(last=False)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Clear all session state (test helper)."""
|
||||
with self._lock:
|
||||
self._sessions.clear()
|
||||
|
||||
|
||||
# Process-wide singleton. Lazily replaced by tests via
|
||||
# `_reset_session_tool_tracker_for_test`.
|
||||
_session_tool_tracker_lock = threading.Lock()
|
||||
_session_tool_tracker: SessionToolTracker | None = None
|
||||
|
||||
|
||||
def get_session_tool_tracker() -> SessionToolTracker:
|
||||
"""Return the process-wide `SessionToolTracker` singleton.
|
||||
|
||||
Lazily constructed so the env-var bound
|
||||
(`HEADROOM_TOOL_TRACKER_MAX_SESSIONS`) is honored at first use.
|
||||
Tests use ``_reset_session_tool_tracker_for_test``.
|
||||
"""
|
||||
global _session_tool_tracker
|
||||
with _session_tool_tracker_lock:
|
||||
if _session_tool_tracker is None:
|
||||
_session_tool_tracker = SessionToolTracker()
|
||||
return _session_tool_tracker
|
||||
|
||||
|
||||
def _reset_session_tool_tracker_for_test() -> None:
|
||||
"""Clear the process-wide tracker (test-only)."""
|
||||
global _session_tool_tracker
|
||||
with _session_tool_tracker_lock:
|
||||
_session_tool_tracker = None
|
||||
|
||||
|
||||
def log_tool_injection_decision(
|
||||
*,
|
||||
provider: str,
|
||||
session_id: str | None,
|
||||
decision: Literal[
|
||||
"inject_first_time",
|
||||
"inject_sticky_replay",
|
||||
"skip",
|
||||
"skip_disabled_via_env",
|
||||
],
|
||||
tool_definition_bytes_count: int,
|
||||
request_id: str | None,
|
||||
) -> None:
|
||||
"""Structured log for every cache-affecting tool-injection decision.
|
||||
|
||||
Per realignment build constraint #8 we log every cache-affecting
|
||||
decision. ``tool_definition_bytes_count`` is the per-tool byte count
|
||||
summed across all memory tools injected this turn. We do NOT log the
|
||||
tool definition contents (might contain user-specific schemas) per
|
||||
constraint #11.
|
||||
"""
|
||||
logger.info(
|
||||
"event=tool_injection_decision provider=%s session_id=%s "
|
||||
"decision=%s tool_definition_bytes_count=%d request_id=%s",
|
||||
provider,
|
||||
session_id or "",
|
||||
decision,
|
||||
tool_definition_bytes_count,
|
||||
request_id or "",
|
||||
)
|
||||
|
||||
|
||||
def _extract_tool_name(tool_definition: dict[str, Any]) -> str | None:
|
||||
"""Extract a stable tool name from a memory tool definition.
|
||||
|
||||
Handles three formats:
|
||||
* Anthropic custom: ``{"name": "memory_save", ...}``
|
||||
* Anthropic native: ``{"type": "memory_20250818", "name": "memory"}``
|
||||
* OpenAI function: ``{"type": "function", "function": {"name": "memory_save", ...}}``
|
||||
"""
|
||||
name = tool_definition.get("name")
|
||||
if isinstance(name, str) and name:
|
||||
return name
|
||||
fn = tool_definition.get("function")
|
||||
if isinstance(fn, dict):
|
||||
fn_name = fn.get("name")
|
||||
if isinstance(fn_name, str) and fn_name:
|
||||
return fn_name
|
||||
# Native memory tool with no explicit name uses ``type`` as its identifier.
|
||||
type_val = tool_definition.get("type")
|
||||
if isinstance(type_val, str) and type_val:
|
||||
return type_val
|
||||
return None
|
||||
|
||||
|
||||
def apply_session_sticky_memory_tools(
|
||||
*,
|
||||
provider: Literal["anthropic", "openai"],
|
||||
session_id: str | None,
|
||||
request_id: str | None,
|
||||
existing_tools: list[dict[str, Any]] | None,
|
||||
memory_tools_to_inject: list[dict[str, Any]],
|
||||
inject_this_turn: bool,
|
||||
) -> tuple[list[dict[str, Any]], bool]:
|
||||
"""Apply sticky-on memory tool injection per `SessionToolTracker`.
|
||||
|
||||
The single coordination point for all memory-tool injection sites
|
||||
(Anthropic custom tools, Anthropic native tool, OpenAI function tools).
|
||||
|
||||
Logic (guide §6.3 #2):
|
||||
|
||||
* If ``HEADROOM_TOOL_INJECTION_STICKY=disabled``: bypass tracker,
|
||||
inject only when ``inject_this_turn`` is True. Diagnostic mode.
|
||||
|
||||
* If session previously injected and tracker has golden bytes:
|
||||
ALWAYS inject the golden bytes verbatim (sticky-on). Memory-this-
|
||||
turn flag is irrelevant — once injected, always injected.
|
||||
|
||||
* If session has NOT previously injected:
|
||||
- ``inject_this_turn=True``: serialize ``memory_tools_to_inject``,
|
||||
record golden bytes, append to tools list.
|
||||
- ``inject_this_turn=False``: skip; no future replay obligation.
|
||||
|
||||
Memory tools whose names already appear in ``existing_tools`` are
|
||||
NOT re-appended (the client owns the canonical definition then).
|
||||
|
||||
``session_id`` may be ``None`` (e.g. WS path with no per-turn
|
||||
session); in that case the tracker is bypassed and the caller's
|
||||
``inject_this_turn`` flag drives the decision verbatim. We log the
|
||||
bypass once so operators can see it.
|
||||
|
||||
Returns ``(updated_tools, was_injected)``. The returned list is a
|
||||
fresh list (caller-safe). ``was_injected`` is True iff at least one
|
||||
memory tool was added to the list.
|
||||
"""
|
||||
if provider not in ("anthropic", "openai"):
|
||||
raise ValueError(f"unsupported provider: {provider!r}")
|
||||
|
||||
tools_out: list[dict[str, Any]] = list(existing_tools) if existing_tools else []
|
||||
existing_names: set[str] = set()
|
||||
for t in tools_out:
|
||||
n = _extract_tool_name(t)
|
||||
if n:
|
||||
existing_names.add(n)
|
||||
|
||||
# Diagnostic / rollback path.
|
||||
if get_tool_injection_sticky_mode() == "disabled":
|
||||
if not inject_this_turn:
|
||||
log_tool_injection_decision(
|
||||
provider=provider,
|
||||
session_id=session_id,
|
||||
decision="skip_disabled_via_env",
|
||||
tool_definition_bytes_count=0,
|
||||
request_id=request_id,
|
||||
)
|
||||
return tools_out, False
|
||||
# Disabled mode + inject_this_turn=True: append the definitions
|
||||
# verbatim without recording golden bytes (per-turn decision
|
||||
# passes through as the broken behavior — explicit operator
|
||||
# opt-in only). Skip names already in the list.
|
||||
added_bytes = 0
|
||||
for tool_def in memory_tools_to_inject:
|
||||
tn = _extract_tool_name(tool_def)
|
||||
if tn is None or tn in existing_names:
|
||||
continue
|
||||
tools_out.append(tool_def)
|
||||
existing_names.add(tn)
|
||||
added_bytes += len(serialize_tool_definition_canonical(tool_def))
|
||||
log_tool_injection_decision(
|
||||
provider=provider,
|
||||
session_id=session_id,
|
||||
decision="skip_disabled_via_env",
|
||||
tool_definition_bytes_count=added_bytes,
|
||||
request_id=request_id,
|
||||
)
|
||||
return tools_out, added_bytes > 0
|
||||
|
||||
# Sticky path requires a session_id. None means we cannot track —
|
||||
# fall back to the caller's per-turn decision (loud, single log line)
|
||||
# so WS handlers / pre-session paths remain functional.
|
||||
if not session_id:
|
||||
if not inject_this_turn:
|
||||
log_tool_injection_decision(
|
||||
provider=provider,
|
||||
session_id=None,
|
||||
decision="skip",
|
||||
tool_definition_bytes_count=0,
|
||||
request_id=request_id,
|
||||
)
|
||||
return tools_out, False
|
||||
added_bytes = 0
|
||||
for tool_def in memory_tools_to_inject:
|
||||
tn = _extract_tool_name(tool_def)
|
||||
if tn is None or tn in existing_names:
|
||||
continue
|
||||
tools_out.append(tool_def)
|
||||
existing_names.add(tn)
|
||||
added_bytes += len(serialize_tool_definition_canonical(tool_def))
|
||||
log_tool_injection_decision(
|
||||
provider=provider,
|
||||
session_id=None,
|
||||
decision="inject_first_time",
|
||||
tool_definition_bytes_count=added_bytes,
|
||||
request_id=request_id,
|
||||
)
|
||||
return tools_out, added_bytes > 0
|
||||
|
||||
tracker = get_session_tool_tracker()
|
||||
previously_injected = tracker.should_inject(provider, session_id)
|
||||
|
||||
if previously_injected:
|
||||
# Sticky replay: always inject the golden bytes. inject_this_turn
|
||||
# flag is intentionally ignored (memory may be disabled this turn
|
||||
# but the cache prefix demands the same tool list as before).
|
||||
golden = tracker.get_golden_definitions(provider, session_id) or []
|
||||
replay_bytes = 0
|
||||
for tool_name, golden_bytes in golden:
|
||||
if tool_name in existing_names:
|
||||
# Client also has a tool by this name — don't double up.
|
||||
# Their bytes win (the client's choice, not ours to gate).
|
||||
continue
|
||||
try:
|
||||
tool_def = json.loads(golden_bytes.decode("utf-8"))
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
||||
# Should never happen — golden bytes were produced by us.
|
||||
# Loud failure per build constraint #4.
|
||||
raise RuntimeError(
|
||||
f"corrupt golden tool bytes for session {session_id} tool {tool_name}: {exc}"
|
||||
) from exc
|
||||
tools_out.append(tool_def)
|
||||
existing_names.add(tool_name)
|
||||
replay_bytes += len(golden_bytes)
|
||||
log_tool_injection_decision(
|
||||
provider=provider,
|
||||
session_id=session_id,
|
||||
decision="inject_sticky_replay",
|
||||
tool_definition_bytes_count=replay_bytes,
|
||||
request_id=request_id,
|
||||
)
|
||||
return tools_out, replay_bytes > 0
|
||||
|
||||
# Fresh session.
|
||||
if not inject_this_turn:
|
||||
log_tool_injection_decision(
|
||||
provider=provider,
|
||||
session_id=session_id,
|
||||
decision="skip",
|
||||
tool_definition_bytes_count=0,
|
||||
request_id=request_id,
|
||||
)
|
||||
return tools_out, False
|
||||
|
||||
# First-time inject: serialize, record, append.
|
||||
added_bytes = 0
|
||||
for tool_def in memory_tools_to_inject:
|
||||
tn = _extract_tool_name(tool_def)
|
||||
if tn is None or tn in existing_names:
|
||||
continue
|
||||
golden_bytes = serialize_tool_definition_canonical(tool_def)
|
||||
tracker.record_injection(
|
||||
provider=provider,
|
||||
session_id=session_id,
|
||||
tool_name=tn,
|
||||
tool_definition_bytes=golden_bytes,
|
||||
)
|
||||
tools_out.append(tool_def)
|
||||
existing_names.add(tn)
|
||||
added_bytes += len(golden_bytes)
|
||||
log_tool_injection_decision(
|
||||
provider=provider,
|
||||
session_id=session_id,
|
||||
decision="inject_first_time",
|
||||
tool_definition_bytes_count=added_bytes,
|
||||
request_id=request_id,
|
||||
)
|
||||
return tools_out, added_bytes > 0
|
||||
|
||||
|
||||
async def _read_request_body_bytes(request: Request) -> bytes:
|
||||
"""Read and (if needed) decompress the request body, returning raw UTF-8 bytes.
|
||||
|
||||
|
|
|
|||
|
|
@ -347,6 +347,51 @@ class MemoryHandler:
|
|||
self._memory_tools = get_memory_tools_optimized()
|
||||
return self._memory_tools
|
||||
|
||||
def compute_memory_tool_definitions(
|
||||
self,
|
||||
provider: str = "anthropic",
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return the memory tool definitions for ``provider`` (pure, no I/O).
|
||||
|
||||
Replaces the building half of ``inject_tools`` so the proxy
|
||||
injection path can route through ``SessionToolTracker`` (PR-A7).
|
||||
Honors ``self.config.use_native_tool`` for Anthropic so the
|
||||
native ``memory_20250818`` tool flows through the same sticky
|
||||
codepath as the custom ``memory_save`` / ``memory_search`` set.
|
||||
|
||||
The returned list is a fresh list of dicts. Order is stable
|
||||
(matches ``_get_memory_tools()`` order) so the canonical bytes
|
||||
are deterministic across calls.
|
||||
"""
|
||||
if not self.config.inject_tools:
|
||||
return []
|
||||
|
||||
if self.config.use_native_tool and provider == "anthropic":
|
||||
return [
|
||||
{
|
||||
"type": NATIVE_MEMORY_TOOL_TYPE,
|
||||
"name": NATIVE_MEMORY_TOOL_NAME,
|
||||
}
|
||||
]
|
||||
|
||||
out: list[dict[str, Any]] = []
|
||||
for memory_tool in self._get_memory_tools():
|
||||
tool_name = memory_tool["function"]["name"]
|
||||
if provider == "anthropic":
|
||||
out.append(
|
||||
{
|
||||
"name": tool_name,
|
||||
"description": memory_tool["function"]["description"],
|
||||
"input_schema": memory_tool["function"]["parameters"],
|
||||
}
|
||||
)
|
||||
else:
|
||||
# OpenAI format — return a fresh shallow copy so callers
|
||||
# can mutate without surprise. dict() is sufficient: the
|
||||
# nested schema is treated as immutable downstream.
|
||||
out.append(dict(memory_tool))
|
||||
return out
|
||||
|
||||
def inject_tools(
|
||||
self,
|
||||
tools: list[dict[str, Any]] | None,
|
||||
|
|
@ -360,6 +405,12 @@ class MemoryHandler:
|
|||
|
||||
Returns:
|
||||
Tuple of (updated_tools, was_injected).
|
||||
|
||||
NOTE (PR-A7): The proxy now wires injection through
|
||||
``apply_session_sticky_memory_tools`` so tool list bytes stay
|
||||
cache-stable across turns. This method remains as the
|
||||
non-session-aware fallback for tests / callers that don't have
|
||||
a session_id (e.g. diagnostic shadow runs).
|
||||
"""
|
||||
if not self.config.inject_tools:
|
||||
return tools or [], False
|
||||
|
|
|
|||
194
tests/fixtures/memory_tool_definitions/anthropic.json
vendored
Normal file
194
tests/fixtures/memory_tool_definitions/anthropic.json
vendored
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
{
|
||||
"provider": "anthropic",
|
||||
"tools": [
|
||||
{
|
||||
"name": "memory_save",
|
||||
"description": "Save important information to long-term memory with optional pre-extraction.\n\nIMPORTANT: For efficiency, extract facts, entities, and relationships yourself when calling this tool.\nThis avoids redundant LLM calls in the storage backend.\n\nUse this tool when you encounter information that should be remembered:\n- User preferences, personal facts, project context, decisions, relationships\n\nPRE-EXTRACTION (recommended for efficiency):\n- facts: List of discrete, self-contained fact strings\n Example: [\"Prefers Python over JavaScript\", \"Works at Acme Corp\"]\n- extracted_entities: List of entities with types\n Example: [{\"entity\": \"Python\", \"entity_type\": \"technology\"}]\n- extracted_relationships: List of entity relationships\n Example: [{\"source\": \"user\", \"relationship\": \"works_at\", \"destination\": \"Acme Corp\"}]\n\nASYNC/BACKGROUND MODE (for zero latency):\n- Set background=true to return immediately while saving happens in background\n- Returns a task_id that can be used to check save status\n- Ideal for real-time conversations where response speed is critical\n\nThe importance score (0.0-1.0) helps prioritize memories:\n- 0.9-1.0: Critical facts\n- 0.7-0.8: Important preferences\n- 0.5-0.6: Useful information\n- 0.3-0.4: Background context\n\nDO NOT save: transient information, sensitive data (passwords, keys), redundant info",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "The original information to remember. Used as context and fallback if no facts provided."
|
||||
},
|
||||
"importance": {
|
||||
"type": "number",
|
||||
"minimum": 0.0,
|
||||
"maximum": 1.0,
|
||||
"description": "Importance score from 0.0 (low) to 1.0 (critical)."
|
||||
},
|
||||
"facts": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "Pre-extracted discrete facts. Each should be self-contained and specific. Example: ['Uses PyTorch for deep learning', 'Prefers dark mode']"
|
||||
},
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "List of entity names referenced (simple format for backwards compatibility)."
|
||||
},
|
||||
"extracted_entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entity": {
|
||||
"type": "string",
|
||||
"description": "Entity name"
|
||||
},
|
||||
"entity_type": {
|
||||
"type": "string",
|
||||
"description": "Type: person, organization, technology, location, project, concept"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"entity",
|
||||
"entity_type"
|
||||
]
|
||||
},
|
||||
"description": "Pre-extracted entities with types for graph storage."
|
||||
},
|
||||
"relationships": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string"
|
||||
},
|
||||
"relation": {
|
||||
"type": "string"
|
||||
},
|
||||
"target": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relation",
|
||||
"target"
|
||||
]
|
||||
},
|
||||
"description": "Simple relationship format (backwards compatible)."
|
||||
},
|
||||
"extracted_relationships": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "Source entity"
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "Relationship type: works_at, uses, knows, manages, depends_on, etc."
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "Destination entity"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination"
|
||||
]
|
||||
},
|
||||
"description": "Pre-extracted relationships for graph storage."
|
||||
},
|
||||
"background": {
|
||||
"type": "boolean",
|
||||
"description": "If true, save in background and return immediately with task_id. Use for zero-latency responses. The save will complete asynchronously. Check status via memory system's get_task_status(task_id)."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"content",
|
||||
"importance"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "memory_search",
|
||||
"description": "Search stored memories to recall relevant information.\n\nUse this tool to retrieve previously saved information before responding to questions about:\n- User preferences or past decisions\n- Personal or professional context\n- Previously discussed topics or projects\n- Relationships between people, systems, or concepts\n- Historical context from past conversations\n\nSearch strategies:\n1. Semantic search (default): Use natural language queries that describe what you're looking for\n - \"user's programming language preferences\"\n - \"information about the current project\"\n - \"past decisions about database choices\"\n\n2. Entity-based search: Specify entities to find memories mentioning specific people/things\n - entities=[\"Alice\", \"Project X\"] finds memories involving Alice or Project X\n\n3. Related memories: Set include_related=true to also retrieve connected memories\n - Finds memories linked by shared entities or explicit relationships\n\nBest practices:\n- Search BEFORE saving to avoid duplicates\n- Search when answering questions that might rely on remembered information\n- Use specific queries for better precision\n- Combine entity filters with semantic queries for targeted retrieval",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Natural language search query describing what information you're looking for. Be specific but not too narrow."
|
||||
},
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "Filter to memories mentioning any of these entities. Useful for finding information about specific people, projects, or systems."
|
||||
},
|
||||
"include_related": {
|
||||
"type": "boolean",
|
||||
"description": "If true, also retrieve memories connected to the results via entity relationships. Helps build fuller context around a topic."
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"minimum": 1,
|
||||
"maximum": 50,
|
||||
"description": "Maximum number of memories to retrieve. Default is 10. Use higher values when you need comprehensive context."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"query"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "memory_update",
|
||||
"description": "Update an existing memory with corrected or evolved information.\n\nUse this tool when:\n- The user provides a correction to previously stored information\n - \"Actually, I prefer TypeScript now, not JavaScript\"\n - \"My project is called ProjectX, not Project Y\"\n\n- Information has changed over time\n - \"I've switched teams from Engineering to Product\"\n - \"We migrated from MySQL to PostgreSQL\"\n\n- You need to add detail or clarification to an existing memory\n - Original: \"Uses React\" -> Updated: \"Uses React 18 with TypeScript and Vite\"\n\n- Consolidating multiple related memories into one clearer entry\n\nDO NOT use this to:\n- Add completely new information (use memory_save instead)\n- Delete memories (use memory_delete instead)\n- Update memories with unrelated content\n\nThe update creates a new version while preserving history, allowing point-in-time queries of past states. Always provide a clear reason for the update to maintain an audit trail.",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to update. Obtain this from a memory_search result."
|
||||
},
|
||||
"new_content": {
|
||||
"type": "string",
|
||||
"description": "The updated content that will replace the existing memory content. Should be complete and self-contained."
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
"description": "Explanation for why this memory is being updated (e.g., 'user correction', 'information changed', 'adding detail'). Stored for audit trail."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"memory_id",
|
||||
"new_content"
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "memory_delete",
|
||||
"description": "Delete a memory that is no longer relevant or was stored in error.\n\nUse this tool when:\n- The user explicitly asks to forget something\n - \"Please forget that I mentioned working at Acme\"\n - \"Delete what you remember about Project X\"\n\n- Information is outdated and no longer applicable (not just changed - use update for that)\n - A completed project that's no longer relevant\n - A temporary context that has expired\n\n- A memory was saved in error\n - Duplicate information\n - Misunderstood or incorrect context\n\n- Privacy or data hygiene reasons\n - User requests removal of personal information\n - Cleaning up test or debug memories\n\nBefore deleting:\n1. Search to find the specific memory and confirm its ID\n2. Verify with the user if the deletion intent is ambiguous\n3. Consider if update would be more appropriate (for changed vs. obsolete info)\n\nDeletions are soft by default - the memory history is preserved but marked as deleted.\nAlways provide a reason for deletion to maintain an audit trail.",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to delete. Obtain this from a memory_search result."
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
"description": "Explanation for why this memory is being deleted (e.g., 'user request', 'outdated', 'stored in error'). Required for audit trail."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"memory_id"
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
206
tests/fixtures/memory_tool_definitions/openai.json
vendored
Normal file
206
tests/fixtures/memory_tool_definitions/openai.json
vendored
Normal file
|
|
@ -0,0 +1,206 @@
|
|||
{
|
||||
"provider": "openai",
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_save",
|
||||
"description": "Save important information to long-term memory with optional pre-extraction.\n\nIMPORTANT: For efficiency, extract facts, entities, and relationships yourself when calling this tool.\nThis avoids redundant LLM calls in the storage backend.\n\nUse this tool when you encounter information that should be remembered:\n- User preferences, personal facts, project context, decisions, relationships\n\nPRE-EXTRACTION (recommended for efficiency):\n- facts: List of discrete, self-contained fact strings\n Example: [\"Prefers Python over JavaScript\", \"Works at Acme Corp\"]\n- extracted_entities: List of entities with types\n Example: [{\"entity\": \"Python\", \"entity_type\": \"technology\"}]\n- extracted_relationships: List of entity relationships\n Example: [{\"source\": \"user\", \"relationship\": \"works_at\", \"destination\": \"Acme Corp\"}]\n\nASYNC/BACKGROUND MODE (for zero latency):\n- Set background=true to return immediately while saving happens in background\n- Returns a task_id that can be used to check save status\n- Ideal for real-time conversations where response speed is critical\n\nThe importance score (0.0-1.0) helps prioritize memories:\n- 0.9-1.0: Critical facts\n- 0.7-0.8: Important preferences\n- 0.5-0.6: Useful information\n- 0.3-0.4: Background context\n\nDO NOT save: transient information, sensitive data (passwords, keys), redundant info",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"content": {
|
||||
"type": "string",
|
||||
"description": "The original information to remember. Used as context and fallback if no facts provided."
|
||||
},
|
||||
"importance": {
|
||||
"type": "number",
|
||||
"minimum": 0.0,
|
||||
"maximum": 1.0,
|
||||
"description": "Importance score from 0.0 (low) to 1.0 (critical)."
|
||||
},
|
||||
"facts": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "Pre-extracted discrete facts. Each should be self-contained and specific. Example: ['Uses PyTorch for deep learning', 'Prefers dark mode']"
|
||||
},
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "List of entity names referenced (simple format for backwards compatibility)."
|
||||
},
|
||||
"extracted_entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"entity": {
|
||||
"type": "string",
|
||||
"description": "Entity name"
|
||||
},
|
||||
"entity_type": {
|
||||
"type": "string",
|
||||
"description": "Type: person, organization, technology, location, project, concept"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"entity",
|
||||
"entity_type"
|
||||
]
|
||||
},
|
||||
"description": "Pre-extracted entities with types for graph storage."
|
||||
},
|
||||
"relationships": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string"
|
||||
},
|
||||
"relation": {
|
||||
"type": "string"
|
||||
},
|
||||
"target": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relation",
|
||||
"target"
|
||||
]
|
||||
},
|
||||
"description": "Simple relationship format (backwards compatible)."
|
||||
},
|
||||
"extracted_relationships": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "Source entity"
|
||||
},
|
||||
"relationship": {
|
||||
"type": "string",
|
||||
"description": "Relationship type: works_at, uses, knows, manages, depends_on, etc."
|
||||
},
|
||||
"destination": {
|
||||
"type": "string",
|
||||
"description": "Destination entity"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"source",
|
||||
"relationship",
|
||||
"destination"
|
||||
]
|
||||
},
|
||||
"description": "Pre-extracted relationships for graph storage."
|
||||
},
|
||||
"background": {
|
||||
"type": "boolean",
|
||||
"description": "If true, save in background and return immediately with task_id. Use for zero-latency responses. The save will complete asynchronously. Check status via memory system's get_task_status(task_id)."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"content",
|
||||
"importance"
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_search",
|
||||
"description": "Search stored memories to recall relevant information.\n\nUse this tool to retrieve previously saved information before responding to questions about:\n- User preferences or past decisions\n- Personal or professional context\n- Previously discussed topics or projects\n- Relationships between people, systems, or concepts\n- Historical context from past conversations\n\nSearch strategies:\n1. Semantic search (default): Use natural language queries that describe what you're looking for\n - \"user's programming language preferences\"\n - \"information about the current project\"\n - \"past decisions about database choices\"\n\n2. Entity-based search: Specify entities to find memories mentioning specific people/things\n - entities=[\"Alice\", \"Project X\"] finds memories involving Alice or Project X\n\n3. Related memories: Set include_related=true to also retrieve connected memories\n - Finds memories linked by shared entities or explicit relationships\n\nBest practices:\n- Search BEFORE saving to avoid duplicates\n- Search when answering questions that might rely on remembered information\n- Use specific queries for better precision\n- Combine entity filters with semantic queries for targeted retrieval",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Natural language search query describing what information you're looking for. Be specific but not too narrow."
|
||||
},
|
||||
"entities": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"description": "Filter to memories mentioning any of these entities. Useful for finding information about specific people, projects, or systems."
|
||||
},
|
||||
"include_related": {
|
||||
"type": "boolean",
|
||||
"description": "If true, also retrieve memories connected to the results via entity relationships. Helps build fuller context around a topic."
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"minimum": 1,
|
||||
"maximum": 50,
|
||||
"description": "Maximum number of memories to retrieve. Default is 10. Use higher values when you need comprehensive context."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"query"
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_update",
|
||||
"description": "Update an existing memory with corrected or evolved information.\n\nUse this tool when:\n- The user provides a correction to previously stored information\n - \"Actually, I prefer TypeScript now, not JavaScript\"\n - \"My project is called ProjectX, not Project Y\"\n\n- Information has changed over time\n - \"I've switched teams from Engineering to Product\"\n - \"We migrated from MySQL to PostgreSQL\"\n\n- You need to add detail or clarification to an existing memory\n - Original: \"Uses React\" -> Updated: \"Uses React 18 with TypeScript and Vite\"\n\n- Consolidating multiple related memories into one clearer entry\n\nDO NOT use this to:\n- Add completely new information (use memory_save instead)\n- Delete memories (use memory_delete instead)\n- Update memories with unrelated content\n\nThe update creates a new version while preserving history, allowing point-in-time queries of past states. Always provide a clear reason for the update to maintain an audit trail.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to update. Obtain this from a memory_search result."
|
||||
},
|
||||
"new_content": {
|
||||
"type": "string",
|
||||
"description": "The updated content that will replace the existing memory content. Should be complete and self-contained."
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
"description": "Explanation for why this memory is being updated (e.g., 'user correction', 'information changed', 'adding detail'). Stored for audit trail."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"memory_id",
|
||||
"new_content"
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_delete",
|
||||
"description": "Delete a memory that is no longer relevant or was stored in error.\n\nUse this tool when:\n- The user explicitly asks to forget something\n - \"Please forget that I mentioned working at Acme\"\n - \"Delete what you remember about Project X\"\n\n- Information is outdated and no longer applicable (not just changed - use update for that)\n - A completed project that's no longer relevant\n - A temporary context that has expired\n\n- A memory was saved in error\n - Duplicate information\n - Misunderstood or incorrect context\n\n- Privacy or data hygiene reasons\n - User requests removal of personal information\n - Cleaning up test or debug memories\n\nBefore deleting:\n1. Search to find the specific memory and confirm its ID\n2. Verify with the user if the deletion intent is ambiguous\n3. Consider if update would be more appropriate (for changed vs. obsolete info)\n\nDeletions are soft by default - the memory history is preserved but marked as deleted.\nAlways provide a reason for deletion to maintain an audit trail.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"memory_id": {
|
||||
"type": "string",
|
||||
"description": "The unique ID of the memory to delete. Obtain this from a memory_search result."
|
||||
},
|
||||
"reason": {
|
||||
"type": "string",
|
||||
"description": "Explanation for why this memory is being deleted (e.g., 'user request', 'outdated', 'stored in error'). Required for audit trail."
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"memory_id"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
672
tests/test_memory_tool_session_sticky.py
Normal file
672
tests/test_memory_tool_session_sticky.py
Normal file
|
|
@ -0,0 +1,672 @@
|
|||
"""Session-sticky memory tool injection tests for PR-A7 (closes P0-6).
|
||||
|
||||
The cache-killer pattern this guards against (guide §6.3 #2):
|
||||
|
||||
* Mid-session toggle: memory enabled in turn N injects `memory_save` /
|
||||
`memory_search` tool definitions into `body["tools"]`. Turn N+1
|
||||
disables memory; tool list shrinks; prefix bytes hash differently;
|
||||
prefix-cache misses; full prompt re-runs at provider cost.
|
||||
|
||||
* Tool definition drift across deploys: the same logical tool is
|
||||
injected but the bytes differ (key insertion order, schema bump,
|
||||
description tweak). Even with the tool list intact, prefix bytes
|
||||
change.
|
||||
|
||||
The fix:
|
||||
|
||||
* `SessionToolTracker`: bounded LRU keyed by (provider, session_id)
|
||||
storing GOLDEN tool-definition bytes from the first injection.
|
||||
Subsequent turns of that session always replay those bytes — even
|
||||
when memory is disabled mid-session (sticky-on per §6.3 #2).
|
||||
|
||||
* `apply_session_sticky_memory_tools`: single coordination point
|
||||
used at every memory injection site (Anthropic custom + native,
|
||||
OpenAI Chat-Completions + Responses + WS).
|
||||
|
||||
Operator opt-in `HEADROOM_TOOL_INJECTION_STICKY=disabled` short-
|
||||
circuits the tracker (per-turn decision flows through verbatim — the
|
||||
broken behavior). That mode is loud and explicit per realignment build
|
||||
constraint #4 — NOT a silent fallback. It exists for diagnostic shadow
|
||||
tracing and emergency rollback only.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.proxy.helpers import (
|
||||
SessionToolTracker,
|
||||
_reset_session_tool_tracker_for_test,
|
||||
apply_session_sticky_memory_tools,
|
||||
get_session_tool_tracker,
|
||||
get_tool_injection_sticky_mode,
|
||||
get_tool_tracker_max_sessions,
|
||||
serialize_tool_definition_canonical,
|
||||
)
|
||||
from headroom.proxy.memory_handler import MemoryConfig, MemoryHandler
|
||||
|
||||
FIXTURES_DIR = Path(__file__).parent / "fixtures" / "memory_tool_definitions"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test isolation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_tracker(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Reset env + tracker singleton between tests."""
|
||||
monkeypatch.delenv("HEADROOM_TOOL_INJECTION_STICKY", raising=False)
|
||||
monkeypatch.delenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", raising=False)
|
||||
_reset_session_tool_tracker_for_test()
|
||||
yield
|
||||
_reset_session_tool_tracker_for_test()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _anthropic_memory_defs() -> list[dict[str, Any]]:
|
||||
h = MemoryHandler(MemoryConfig(enabled=True, inject_tools=True))
|
||||
return h.compute_memory_tool_definitions("anthropic")
|
||||
|
||||
|
||||
def _openai_memory_defs() -> list[dict[str, Any]]:
|
||||
h = MemoryHandler(MemoryConfig(enabled=True, inject_tools=True))
|
||||
return h.compute_memory_tool_definitions("openai")
|
||||
|
||||
|
||||
def _names_in(tools: list[dict[str, Any]]) -> set[str]:
|
||||
out: set[str] = set()
|
||||
for t in tools:
|
||||
n = t.get("name") or (t.get("function") or {}).get("name") or t.get("type")
|
||||
if n:
|
||||
out.add(n)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# `SessionToolTracker` direct unit tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_should_inject_false_for_unknown_session() -> None:
|
||||
tracker = SessionToolTracker(max_sessions=10)
|
||||
assert tracker.should_inject("anthropic", "s-1") is False
|
||||
|
||||
|
||||
def test_record_then_should_inject_true() -> None:
|
||||
tracker = SessionToolTracker(max_sessions=10)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b'{"name":"memory_save"}',
|
||||
)
|
||||
assert tracker.should_inject("anthropic", "s-1") is True
|
||||
|
||||
|
||||
def test_get_golden_definitions_returns_recorded_bytes() -> None:
|
||||
tracker = SessionToolTracker(max_sessions=10)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b'{"name":"memory_save","x":1}',
|
||||
)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
tool_name="memory_search",
|
||||
tool_definition_bytes=b'{"name":"memory_search","x":2}',
|
||||
)
|
||||
golden = tracker.get_golden_definitions("anthropic", "s-1")
|
||||
assert golden is not None
|
||||
assert [name for name, _ in golden] == ["memory_save", "memory_search"]
|
||||
assert golden[0][1] == b'{"name":"memory_save","x":1}'
|
||||
assert golden[1][1] == b'{"name":"memory_search","x":2}'
|
||||
|
||||
|
||||
def test_record_first_write_wins_on_duplicate_name() -> None:
|
||||
"""Re-recording the same tool name is a no-op (prevents drift mid-session)."""
|
||||
tracker = SessionToolTracker(max_sessions=10)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b"original",
|
||||
)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b"new-bytes",
|
||||
)
|
||||
golden = tracker.get_golden_definitions("anthropic", "s-1") or []
|
||||
assert golden == [("memory_save", b"original")]
|
||||
|
||||
|
||||
def test_provider_isolation_anthropic_vs_openai_same_session_id() -> None:
|
||||
"""Same session_id under two providers keeps independent state."""
|
||||
tracker = SessionToolTracker(max_sessions=10)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="shared",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b"anthropic-bytes",
|
||||
)
|
||||
tracker.record_injection(
|
||||
provider="openai",
|
||||
session_id="shared",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b"openai-bytes",
|
||||
)
|
||||
a_golden = tracker.get_golden_definitions("anthropic", "shared") or []
|
||||
o_golden = tracker.get_golden_definitions("openai", "shared") or []
|
||||
assert a_golden == [("memory_save", b"anthropic-bytes")]
|
||||
assert o_golden == [("memory_save", b"openai-bytes")]
|
||||
|
||||
|
||||
def test_lru_eviction_at_max_sessions() -> None:
|
||||
"""Bounded LRU pops oldest session when overflowing."""
|
||||
tracker = SessionToolTracker(max_sessions=2)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b"a",
|
||||
)
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-2",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b"b",
|
||||
)
|
||||
assert tracker.active_sessions == 2
|
||||
|
||||
# Touch s-1 so s-2 becomes the LRU.
|
||||
assert tracker.should_inject("anthropic", "s-1") is True
|
||||
|
||||
# Add s-3: pops s-2.
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s-3",
|
||||
tool_name="memory_save",
|
||||
tool_definition_bytes=b"c",
|
||||
)
|
||||
assert tracker.active_sessions == 2
|
||||
assert tracker.should_inject("anthropic", "s-2") is False
|
||||
assert tracker.should_inject("anthropic", "s-1") is True
|
||||
assert tracker.should_inject("anthropic", "s-3") is True
|
||||
|
||||
|
||||
def test_max_sessions_invalid_raises() -> None:
|
||||
with pytest.raises(ValueError):
|
||||
SessionToolTracker(max_sessions=0)
|
||||
with pytest.raises(ValueError):
|
||||
SessionToolTracker(max_sessions=-1)
|
||||
|
||||
|
||||
def test_blank_provider_or_session_raises() -> None:
|
||||
tracker = SessionToolTracker(max_sessions=10)
|
||||
with pytest.raises(ValueError):
|
||||
tracker.should_inject("", "s")
|
||||
with pytest.raises(ValueError):
|
||||
tracker.should_inject("anthropic", "")
|
||||
with pytest.raises(ValueError):
|
||||
tracker.record_injection(
|
||||
provider="",
|
||||
session_id="s",
|
||||
tool_name="x",
|
||||
tool_definition_bytes=b"y",
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s",
|
||||
tool_name="",
|
||||
tool_definition_bytes=b"y",
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="s",
|
||||
tool_name="x",
|
||||
tool_definition_bytes=b"",
|
||||
)
|
||||
|
||||
|
||||
def test_thread_safe_concurrent_access() -> None:
|
||||
"""N threads on same session: no exceptions, all pinned bytes survive."""
|
||||
tracker = SessionToolTracker(max_sessions=10)
|
||||
n_threads = 16
|
||||
iterations = 50
|
||||
errors: list[BaseException] = []
|
||||
|
||||
def worker(thread_idx: int) -> None:
|
||||
try:
|
||||
for i in range(iterations):
|
||||
tracker.record_injection(
|
||||
provider="anthropic",
|
||||
session_id="shared",
|
||||
tool_name=f"t{thread_idx}-i{i}",
|
||||
tool_definition_bytes=f"bytes-{thread_idx}-{i}".encode(),
|
||||
)
|
||||
# Concurrent reads.
|
||||
tracker.should_inject("anthropic", "shared")
|
||||
tracker.get_golden_definitions("anthropic", "shared")
|
||||
except BaseException as e: # noqa: BLE001
|
||||
errors.append(e)
|
||||
|
||||
threads = [threading.Thread(target=worker, args=(idx,)) for idx in range(n_threads)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
assert errors == []
|
||||
golden = tracker.get_golden_definitions("anthropic", "shared") or []
|
||||
names = {name for name, _ in golden}
|
||||
expected = {f"t{idx}-i{i}" for idx in range(n_threads) for i in range(iterations)}
|
||||
assert expected.issubset(names)
|
||||
|
||||
|
||||
def test_singleton_returns_same_instance() -> None:
|
||||
a = get_session_tool_tracker()
|
||||
b = get_session_tool_tracker()
|
||||
assert a is b
|
||||
|
||||
|
||||
def test_singleton_reset_replaces_instance() -> None:
|
||||
a = get_session_tool_tracker()
|
||||
_reset_session_tool_tracker_for_test()
|
||||
b = get_session_tool_tracker()
|
||||
assert a is not b
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Env vars
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_max_sessions_env_var_default() -> None:
|
||||
assert get_tool_tracker_max_sessions() == 1000
|
||||
|
||||
|
||||
def test_max_sessions_env_var_custom(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "42")
|
||||
assert get_tool_tracker_max_sessions() == 42
|
||||
|
||||
|
||||
def test_max_sessions_env_var_invalid_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "0")
|
||||
with pytest.raises(ValueError):
|
||||
get_tool_tracker_max_sessions()
|
||||
monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "-3")
|
||||
with pytest.raises(ValueError):
|
||||
get_tool_tracker_max_sessions()
|
||||
monkeypatch.setenv("HEADROOM_TOOL_TRACKER_MAX_SESSIONS", "not-int")
|
||||
with pytest.raises(ValueError):
|
||||
get_tool_tracker_max_sessions()
|
||||
|
||||
|
||||
def test_sticky_mode_default_enabled() -> None:
|
||||
assert get_tool_injection_sticky_mode() == "enabled"
|
||||
|
||||
|
||||
def test_sticky_mode_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "disabled")
|
||||
assert get_tool_injection_sticky_mode() == "disabled"
|
||||
|
||||
|
||||
def test_sticky_mode_invalid_raises(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "yolo")
|
||||
with pytest.raises(ValueError, match="HEADROOM_TOOL_INJECTION_STICKY"):
|
||||
get_tool_injection_sticky_mode()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# `apply_session_sticky_memory_tools` integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_injection_in_turn_1_repeats_in_turn_2_same_session_anthropic() -> None:
|
||||
"""Core sticky-on guarantee for Anthropic provider."""
|
||||
defs = _anthropic_memory_defs()
|
||||
assert len(defs) >= 2
|
||||
|
||||
# Turn 1: memory enabled — first-time injection.
|
||||
tools1, was1 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was1 is True
|
||||
names1 = _names_in(tools1)
|
||||
assert "memory_save" in names1
|
||||
assert "memory_search" in names1
|
||||
|
||||
# Turn 2: memory STILL enabled — bytes match turn 1.
|
||||
tools2, was2 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="s-1",
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was2 is True
|
||||
# Same set of memory tools.
|
||||
assert _names_in(tools2) == names1
|
||||
|
||||
|
||||
def test_injection_in_turn_1_repeats_in_turn_2_same_session_openai() -> None:
|
||||
defs = _openai_memory_defs()
|
||||
assert len(defs) >= 2
|
||||
|
||||
tools1, was1 = apply_session_sticky_memory_tools(
|
||||
provider="openai",
|
||||
session_id="o-1",
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was1 is True
|
||||
names1 = _names_in(tools1)
|
||||
assert "memory_save" in names1
|
||||
assert "memory_search" in names1
|
||||
|
||||
tools2, was2 = apply_session_sticky_memory_tools(
|
||||
provider="openai",
|
||||
session_id="o-1",
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was2 is True
|
||||
assert _names_in(tools2) == names1
|
||||
|
||||
|
||||
def test_byte_equal_tool_definition_across_turns() -> None:
|
||||
"""The injected tool list serialization is BYTE-equal turn 1 vs turn 2.
|
||||
|
||||
Pin the bytes via the golden snapshot fixture.
|
||||
"""
|
||||
defs = _anthropic_memory_defs()
|
||||
|
||||
tools1, _ = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="bytestable-1",
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
tools2, _ = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="bytestable-1",
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
|
||||
# Byte-equality (the cache-stable invariant).
|
||||
bytes1 = b"".join(serialize_tool_definition_canonical(t) for t in tools1)
|
||||
bytes2 = b"".join(serialize_tool_definition_canonical(t) for t in tools2)
|
||||
assert bytes1 == bytes2
|
||||
|
||||
# Match the golden fixture (computed via the pinned helper).
|
||||
fixture = json.loads((FIXTURES_DIR / "anthropic.json").read_text())
|
||||
fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
|
||||
assert bytes1 == fixture_bytes
|
||||
|
||||
|
||||
def test_byte_equal_tool_definition_across_turns_openai() -> None:
|
||||
defs = _openai_memory_defs()
|
||||
tools1, _ = apply_session_sticky_memory_tools(
|
||||
provider="openai",
|
||||
session_id="bytestable-2",
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
tools2, _ = apply_session_sticky_memory_tools(
|
||||
provider="openai",
|
||||
session_id="bytestable-2",
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
bytes1 = b"".join(serialize_tool_definition_canonical(t) for t in tools1)
|
||||
bytes2 = b"".join(serialize_tool_definition_canonical(t) for t in tools2)
|
||||
assert bytes1 == bytes2
|
||||
|
||||
fixture = json.loads((FIXTURES_DIR / "openai.json").read_text())
|
||||
fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
|
||||
assert bytes1 == fixture_bytes
|
||||
|
||||
|
||||
def test_memory_disabled_after_inject_still_injects() -> None:
|
||||
"""Turn 1 injects; turn 2 has memory disabled; turn 2 still injects golden bytes."""
|
||||
defs = _anthropic_memory_defs()
|
||||
|
||||
# Turn 1: memory enabled.
|
||||
tools1, was1 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="s-cancel-1",
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was1 is True
|
||||
names1 = _names_in(tools1)
|
||||
|
||||
# Turn 2: memory DISABLED for this turn (e.g. inject_tools flag flipped).
|
||||
# `inject_this_turn=False` AND `memory_tools_to_inject=[]` mimic the
|
||||
# caller's behavior under disabled-memory: nothing fresh to inject.
|
||||
tools2, was2 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="s-cancel-1",
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=[],
|
||||
inject_this_turn=False,
|
||||
)
|
||||
# Sticky-on: the golden bytes are still injected even though caller
|
||||
# passed nothing this turn.
|
||||
assert was2 is True
|
||||
assert _names_in(tools2) == names1
|
||||
|
||||
# Bytes match.
|
||||
bytes1 = b"".join(serialize_tool_definition_canonical(t) for t in tools1)
|
||||
bytes2 = b"".join(serialize_tool_definition_canonical(t) for t in tools2)
|
||||
assert bytes1 == bytes2
|
||||
|
||||
|
||||
def test_different_sessions_independent() -> None:
|
||||
"""Session A injects; session B doesn't; verify isolation."""
|
||||
defs = _anthropic_memory_defs()
|
||||
|
||||
# Session A: inject.
|
||||
_, was_a = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="A",
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was_a is True
|
||||
|
||||
# Session B: NO inject this turn, no prior history.
|
||||
tools_b, was_b = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="B",
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=[],
|
||||
inject_this_turn=False,
|
||||
)
|
||||
assert was_b is False
|
||||
assert _names_in(tools_b) == set()
|
||||
|
||||
# Session A still has its golden state.
|
||||
tools_a2, was_a2 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="A",
|
||||
request_id="r-3",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=[],
|
||||
inject_this_turn=False,
|
||||
)
|
||||
assert was_a2 is True
|
||||
assert "memory_save" in _names_in(tools_a2)
|
||||
|
||||
|
||||
def test_disabled_mode_passes_through_per_turn_decision(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""`HEADROOM_TOOL_INJECTION_STICKY=disabled` reverts to per-turn behavior.
|
||||
|
||||
This is the broken behavior — explicit operator opt-in only. Turn 1
|
||||
injects; turn 2 with `inject_this_turn=False` does NOT replay (the
|
||||
sticky guarantee is bypassed).
|
||||
"""
|
||||
monkeypatch.setenv("HEADROOM_TOOL_INJECTION_STICKY", "disabled")
|
||||
defs = _anthropic_memory_defs()
|
||||
|
||||
tools1, was1 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="s-disabled",
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was1 is True
|
||||
assert "memory_save" in _names_in(tools1)
|
||||
|
||||
# Turn 2: caller says don't inject. Disabled mode → tracker bypassed.
|
||||
tools2, was2 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="s-disabled",
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=[],
|
||||
inject_this_turn=False,
|
||||
)
|
||||
assert was2 is False
|
||||
assert _names_in(tools2) == set()
|
||||
|
||||
|
||||
def test_existing_tool_with_memory_name_not_double_injected() -> None:
|
||||
"""If client already has a tool by the same name, do not re-append it."""
|
||||
defs = _anthropic_memory_defs()
|
||||
client_tools: list[dict[str, Any]] = [
|
||||
{"name": "memory_save", "description": "client's own", "input_schema": {}}
|
||||
]
|
||||
|
||||
tools, _ = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id="s-dedup",
|
||||
request_id="r-1",
|
||||
existing_tools=client_tools,
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
# Exactly one tool named "memory_save".
|
||||
save_count = sum(1 for t in tools if t.get("name") == "memory_save")
|
||||
assert save_count == 1
|
||||
|
||||
|
||||
def test_no_session_id_falls_back_to_per_turn(caplog: pytest.LogCaptureFixture) -> None:
|
||||
"""`session_id=None` (e.g. WS pre-session) bypasses the tracker."""
|
||||
defs = _anthropic_memory_defs()
|
||||
tools1, was1 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id=None,
|
||||
request_id="r-1",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=defs,
|
||||
inject_this_turn=True,
|
||||
)
|
||||
assert was1 is True
|
||||
assert "memory_save" in _names_in(tools1)
|
||||
|
||||
# Without session_id we can't replay across turns.
|
||||
tools2, was2 = apply_session_sticky_memory_tools(
|
||||
provider="anthropic",
|
||||
session_id=None,
|
||||
request_id="r-2",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=[],
|
||||
inject_this_turn=False,
|
||||
)
|
||||
assert was2 is False
|
||||
|
||||
|
||||
def test_unknown_provider_raises() -> None:
|
||||
with pytest.raises(ValueError, match="unsupported provider"):
|
||||
apply_session_sticky_memory_tools(
|
||||
provider="gemini", # type: ignore[arg-type]
|
||||
session_id="s",
|
||||
request_id="r",
|
||||
existing_tools=[],
|
||||
memory_tools_to_inject=[],
|
||||
inject_this_turn=True,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Golden fixture pinning
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_anthropic_fixture_matches_helper_output() -> None:
|
||||
"""Fixture file pins the canonical bytes — regenerate if this fails."""
|
||||
fixture = json.loads((FIXTURES_DIR / "anthropic.json").read_text())
|
||||
assert fixture["provider"] == "anthropic"
|
||||
|
||||
helper_defs = _anthropic_memory_defs()
|
||||
helper_bytes = b"".join(serialize_tool_definition_canonical(t) for t in helper_defs)
|
||||
fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
|
||||
assert helper_bytes == fixture_bytes, (
|
||||
"Anthropic memory tool definitions drifted from golden fixture. "
|
||||
"If intentional, regenerate "
|
||||
"tests/fixtures/memory_tool_definitions/anthropic.json. "
|
||||
f"Helper SHA-256: {hashlib.sha256(helper_bytes).hexdigest()} "
|
||||
f"Fixture SHA-256: {hashlib.sha256(fixture_bytes).hexdigest()}"
|
||||
)
|
||||
|
||||
|
||||
def test_openai_fixture_matches_helper_output() -> None:
|
||||
fixture = json.loads((FIXTURES_DIR / "openai.json").read_text())
|
||||
assert fixture["provider"] == "openai"
|
||||
|
||||
helper_defs = _openai_memory_defs()
|
||||
helper_bytes = b"".join(serialize_tool_definition_canonical(t) for t in helper_defs)
|
||||
fixture_bytes = b"".join(serialize_tool_definition_canonical(t) for t in fixture["tools"])
|
||||
assert helper_bytes == fixture_bytes, (
|
||||
"OpenAI memory tool definitions drifted from golden fixture. "
|
||||
"If intentional, regenerate "
|
||||
"tests/fixtures/memory_tool_definitions/openai.json. "
|
||||
f"Helper SHA-256: {hashlib.sha256(helper_bytes).hexdigest()} "
|
||||
f"Fixture SHA-256: {hashlib.sha256(fixture_bytes).hexdigest()}"
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue