fix(nous-portal): give auxiliary calls the Portal sticky routing key

The Nous Portal profile publishes a top-level `session_id` that the Portal
uses as its sticky routing key, pinning a conversation to one upstream
endpoint so explicit Anthropic `cache_control` breakpoints stay warm
(those caches are instance-local, so a reroute cold-writes instead of reads).

`NousProfile.build_extra_body` sourced that key only from the explicit
`session_id` argument. Auxiliary call sites — compression, title generation,
vision, web_extract, session_search, MoA slots — funnel through
`agent.auxiliary_client`, which has no session handle and never passes one.
They carried the `conversation=` tag but NO sticky key at all, so each one
routed independently of the conversation it belonged to.

Resolve the key the way `nous_portal_tags` already resolves that tag:
ambient lineage-ROOT contextvar first, explicit argument as fallback. Fixes
every aux call site at once with no per-call-site plumbing.

Also publish the root in the `_compress_context` forwarder when nothing is
ambient. Out-of-turn entry points (`/compact`, the gateway `/compress`
command and its hygiene sweep, partial head compression) call it outside
`run_conversation`'s scope, so the summarizer's aux call ran untagged and
unkeyed; the caller's value is restored in `finally` so a compaction never
leaks its tag.

Scope note: this does NOT keep the compaction turn's own prompt cache warm.
Compaction replaces the history with a summary and rebuilds the system
prompt, so that request is a cold write on any endpoint — what a stable key
buys is the turns AFTER compaction reading the cache it wrote. Under the
default `compression.in_place: true` (#38763) the session id does not rotate
at compaction, so for the main loop the ambient root and the explicit
argument already agree; the ambient resolution additionally holds the key
stable for installs that opt back into rotating compaction and across
delegate-subagent trees.

Salvaged from #71747. Subsumes the aux-call half of #70883 (@webtecnica),
which threaded `session_id` through `call_llm` to close the same gap.

Tests: 4 cases in tests/agent/test_portal_tags.py; both fix halves verified
to fail against the pre-fix behavior via sabotage runs.
This commit is contained in:
joaomarcos 2026-07-25 23:42:00 -03:00 committed by Teknium
parent a1c4d99953
commit f2f4df064d
3 changed files with 175 additions and 16 deletions

View file

@ -2,7 +2,7 @@
from typing import Any
from agent.portal_tags import nous_portal_tags
from agent.portal_tags import get_conversation_context, nous_portal_tags
from providers import register_provider
from providers.base import ProviderProfile
@ -14,15 +14,35 @@ class NousProfile(ProviderProfile):
self, *, session_id: str | None = None, **context
) -> dict[str, Any]:
body: dict[str, Any] = {"tags": nous_portal_tags(session_id=session_id)}
if session_id:
# Top-level session_id → provider sticky routing key. Pins every
# turn of a session to the same upstream endpoint so explicit
# Anthropic cache_control breakpoints stay warm instead of
# cold-writing a fresh cache on each reroute (Anthropic/Vertex/
# Bedrock caches are instance-local). Mirrors the OpenRouter
# profile; without it the portal falls back to hashing the opening
# messages, which breaks pinning whenever those shift.
body["session_id"] = session_id
# Top-level session_id → provider sticky routing key. Pins every
# turn of a session to the same upstream endpoint so explicit
# Anthropic cache_control breakpoints stay warm instead of
# cold-writing a fresh cache on each reroute (Anthropic/Vertex/
# Bedrock caches are instance-local). Mirrors the OpenRouter
# profile; without it the portal falls back to hashing the opening
# messages, which breaks pinning whenever those shift.
#
# Resolve it exactly like ``nous_portal_tags`` resolves the
# ``conversation=`` tag: ambient context first (the lineage ROOT id
# published by the agent loop), explicit argument as fallback.
#
# The gap this closes is the auxiliary call sites — compression,
# title generation, vision, web_extract, session_search, MoA slots.
# They funnel through ``agent.auxiliary_client`` which has no session
# handle, so they never pass ``session_id``: they carried the
# ``conversation=`` tag but NO sticky key at all, and each one routed
# independently of the conversation it belongs to. Reading the same
# ambient contextvar the tag already uses fixes that with zero
# per-call-site plumbing.
#
# For the main loop the two agree anyway under the default
# ``compression.in_place: true`` (#38763), where compaction keeps the
# session id; the ambient root additionally keeps the key stable for
# installs that opt back into rotating compaction, and across
# delegate-subagent trees.
sticky_key = get_conversation_context() or session_id
if sticky_key:
body["session_id"] = sticky_key
provider_preferences = context.get("provider_preferences")
if provider_preferences:
body["provider"] = provider_preferences

View file

@ -6462,13 +6462,43 @@ class AIAgent:
``force=False``.
"""
from agent.conversation_compression import compress_context
return compress_context(
self, messages, system_message,
approx_tokens=approx_tokens, task_id=task_id, focus_topic=focus_topic,
force=force,
defer_context_engine_notification=defer_context_engine_notification,
commit_fence=commit_fence,
from agent.portal_tags import (
get_conversation_context,
reset_conversation_context,
set_conversation_context,
)
# Out-of-turn compaction entry points — ``/compact`` (cli.py), the
# gateway ``/compress`` command and its hygiene sweep (both of which
# build a throwaway agent), and partial head compression — call this
# forwarder directly, outside ``run_conversation``'s ambient scope.
# With nothing ambient the summarizer's auxiliary call carries no
# conversation tag and no Portal sticky key, so it routes independently
# of the conversation it belongs to. Publish the root here as a
# fallback; in-turn callers already have it set to the same value, so
# this is a no-op for them.
#
# Note this does NOT keep the compaction turn's own prompt cache warm:
# compaction replaces the history with a summary and rebuilds the
# system prompt, so that request is a cold write on any endpoint. What
# it buys is the turns AFTER compaction reading the cache it wrote.
token = None
if get_conversation_context() is None:
root = self._conversation_root_id()
if root:
token = set_conversation_context(root)
try:
return compress_context(
self, messages, system_message,
approx_tokens=approx_tokens, task_id=task_id, focus_topic=focus_topic,
force=force,
defer_context_engine_notification=defer_context_engine_notification,
commit_fence=commit_fence,
)
finally:
# Restore whatever the caller had, so a compaction never leaks its
# tag into the surrounding scope.
if token is not None:
reset_conversation_context(token)
def _set_tool_guardrail_halt(self, decision: ToolGuardrailDecision) -> None:
"""Record the first guardrail decision that should stop this turn."""

View file

@ -217,3 +217,112 @@ def test_nous_provider_profile_uses_helper():
assert profile is not None
body = profile.build_extra_body()
assert body["tags"] == nous_portal_tags()
def test_nous_sticky_key_matches_conversation_tag():
"""Sticky routing key must resolve like the ``conversation=`` tag does.
The load-bearing case is the auxiliary call sites (compression, titles,
vision, MoA slots): they pass no ``session_id`` at all, so before this
resolution they carried the conversation tag but NO Portal sticky key and
routed independently of their conversation.
The explicit-argument case matters for installs that opt out of the
default ``compression.in_place: true`` (#38763) and therefore still rotate
``agent.session_id`` at compaction, and for delegate-subagent trees that
should tag as the parent conversation.
"""
from agent.portal_tags import (
conversation_tag,
reset_conversation_context,
set_conversation_context,
)
from providers import get_provider_profile
profile = get_provider_profile("nous")
token = set_conversation_context("root-conversation")
try:
# Rotated segment id passed explicitly — root still wins, both places.
body = profile.build_extra_body(session_id="segment-after-compaction")
assert body["session_id"] == "root-conversation"
assert conversation_tag("root-conversation") in body["tags"]
# Auxiliary call sites pass no session_id but inherit the context.
aux = profile.build_extra_body()
assert aux["session_id"] == "root-conversation"
finally:
reset_conversation_context(token)
def test_nous_sticky_key_falls_back_to_explicit_session_id():
"""Outside any agent turn the explicit session_id remains the sticky key."""
from providers import get_provider_profile
profile = get_provider_profile("nous")
body = profile.build_extra_body(session_id="explicit-only")
assert body["session_id"] == "explicit-only"
assert profile.build_extra_body().get("session_id") is None
def test_compress_context_publishes_root_when_called_out_of_turn(monkeypatch):
"""Out-of-turn compaction must still carry the conversation tag.
``/compact``, the gateway ``/compress`` command and its hygiene sweep call
``_compress_context`` directly, outside ``run_conversation``'s ambient
scope. That call ships the full uncompressed history the largest prompt
of the session so losing the sticky key there reroutes it to a cold
endpoint.
"""
import agent.conversation_compression as cc
from agent.portal_tags import get_conversation_context
from run_agent import AIAgent
seen = {}
def _fake_compress(agent, messages, system_message, **kwargs):
seen["conversation"] = get_conversation_context()
return ([], "")
monkeypatch.setattr(cc, "compress_context", _fake_compress)
class _Agent:
def _conversation_root_id(self):
return "root-abc"
AIAgent._compress_context(_Agent(), [], "sys")
assert seen["conversation"] == "root-abc"
# The scope is local: nothing leaks into the caller's context.
assert get_conversation_context() is None
def test_compress_context_preserves_ambient_context(monkeypatch):
"""In-turn compaction inherits the turn's root and restores it untouched."""
import agent.conversation_compression as cc
from agent.portal_tags import (
get_conversation_context,
reset_conversation_context,
set_conversation_context,
)
from run_agent import AIAgent
seen = {}
def _fake_compress(agent, messages, system_message, **kwargs):
seen["conversation"] = get_conversation_context()
return ([], "")
monkeypatch.setattr(cc, "compress_context", _fake_compress)
class _Agent:
def _conversation_root_id(self):
# A rotated segment id must never win over the ambient root.
return "segment-after-compaction"
token = set_conversation_context("outer-root")
try:
AIAgent._compress_context(_Agent(), [], "sys")
assert seen["conversation"] == "outer-root"
assert get_conversation_context() == "outer-root"
finally:
reset_conversation_context(token)