Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license
This commit is contained in:
@@ -0,0 +1,253 @@
|
||||
"""Tests for the per-turn ``ContextEngine.select_context()`` hook.
|
||||
|
||||
``select_context()`` is the *selection / routing* verb — distinct from
|
||||
compression — that lets an external context engine replace which context
|
||||
enters the prompt for a single request, every turn, independent of
|
||||
``should_compress()``. It is additive and no-op by default, and the host
|
||||
call site (``_apply_context_engine_selection``) is fail-open: a missing hook,
|
||||
an exception, or an invalid return value must leave the assembled request
|
||||
untouched and must never mutate persisted history.
|
||||
|
||||
This pins the contract that engines such as retrieval-augmented, topic-routed,
|
||||
and role-switching engines rely on (RFC #36765), consolidating the per-turn
|
||||
request-assembly surface proposed across #41918, #24949, #47109, and #50053.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from agent.context_engine import ContextEngine
|
||||
from agent.conversation_loop import (
|
||||
_apply_context_engine_selection,
|
||||
_notify_context_engine_turn_complete,
|
||||
)
|
||||
|
||||
|
||||
class _MinimalEngine(ContextEngine):
|
||||
"""Concrete engine implementing only the abstract methods."""
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "minimal"
|
||||
|
||||
def update_from_response(self, usage: Dict[str, Any]) -> None:
|
||||
pass
|
||||
|
||||
def should_compress(self, prompt_tokens: int = None) -> bool:
|
||||
return False
|
||||
|
||||
def compress(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
current_tokens: int = None,
|
||||
focus_topic: str = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
return messages
|
||||
|
||||
|
||||
def _agent_with(engine) -> Any:
|
||||
agent = MagicMock()
|
||||
agent.session_id = "test-session"
|
||||
agent.context_compressor = engine
|
||||
return agent
|
||||
|
||||
|
||||
REQUEST = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "hello"},
|
||||
]
|
||||
HISTORY = [{"role": "user", "content": "hello"}]
|
||||
|
||||
|
||||
# -- ABC default -----------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
# -- Host call site: _apply_context_engine_selection -----------------------
|
||||
|
||||
|
||||
|
||||
def test_base_noop_select_context_is_short_circuited_not_called():
|
||||
"""Non-implementing engines skip the hook entirely (no call, no copies).
|
||||
|
||||
The built-in ContextCompressor — and any engine that merely inherits the
|
||||
ABC default — must keep the default request path byte-identical AND pay
|
||||
nothing per request. ``hasattr`` alone cannot distinguish "inherits the
|
||||
no-op default" from "implements the hook" because the ABC defines
|
||||
``select_context`` on every engine; the host therefore identity-checks the
|
||||
bound method against ``ContextEngine.select_context`` and short-circuits
|
||||
WITHOUT calling it or building the shallow reference copies. This pins
|
||||
that: even a base implementation patched to raise is never invoked.
|
||||
"""
|
||||
from unittest.mock import patch as _patch
|
||||
|
||||
def _explode(self, request_messages, **kwargs):
|
||||
raise AssertionError("base select_context must not be invoked")
|
||||
|
||||
engine = _MinimalEngine() # inherits the ABC default
|
||||
agent = _agent_with(engine)
|
||||
logger = MagicMock()
|
||||
with _patch.object(ContextEngine, "select_context", _explode):
|
||||
out = _apply_context_engine_selection(
|
||||
agent, REQUEST, HISTORY, HISTORY[-1], logger=logger
|
||||
)
|
||||
assert out is REQUEST
|
||||
assert not logger.warning.called
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def test_empty_list_keeps_original_request():
|
||||
"""An empty list must fall open to the original request.
|
||||
|
||||
``all([])`` is ``True``, so without an emptiness check a ``[]`` returned by
|
||||
a failing/buggy engine would replace a valid assembled request with an
|
||||
empty message list the downstream sanitizers cannot restore — reaching the
|
||||
provider as an invalid request instead of failing open. Guards the fail-open
|
||||
contract.
|
||||
"""
|
||||
|
||||
class _Engine(_MinimalEngine):
|
||||
def select_context(self, request_messages, **kwargs):
|
||||
return []
|
||||
|
||||
logger = MagicMock()
|
||||
agent = _agent_with(_Engine())
|
||||
out = _apply_context_engine_selection(
|
||||
agent, REQUEST, HISTORY, HISTORY[-1], logger=logger
|
||||
)
|
||||
assert out is REQUEST
|
||||
assert logger.warning.called
|
||||
|
||||
|
||||
def test_engine_mutating_inputs_cannot_corrupt_persisted_state():
|
||||
"""An engine that mutates its read-only inputs in place must not affect the
|
||||
persisted conversation history / incoming message.
|
||||
|
||||
``conversation_messages`` and ``incoming_message`` are reference-only
|
||||
context. The host passes shallow copies, so even a misbehaving engine that
|
||||
appends to / edits them in ``select_context()`` cannot alter the live
|
||||
persisted objects. Enforces the request-only contract (not just documents).
|
||||
"""
|
||||
history = [{"role": "user", "content": "hello"}]
|
||||
incoming = history[-1]
|
||||
history_snapshot = [dict(m) for m in history]
|
||||
incoming_snapshot = dict(incoming)
|
||||
|
||||
class _Engine(_MinimalEngine):
|
||||
def select_context(self, request_messages, *, conversation_messages=None,
|
||||
incoming_message=None, **kwargs):
|
||||
# Misbehaving engine: mutate the read-only inputs in place.
|
||||
if conversation_messages is not None:
|
||||
conversation_messages.append({"role": "user", "content": "INJECTED"})
|
||||
if conversation_messages and isinstance(conversation_messages[0], dict):
|
||||
conversation_messages[0]["content"] = "TAMPERED"
|
||||
if isinstance(incoming_message, dict):
|
||||
incoming_message["content"] = "TAMPERED"
|
||||
return None
|
||||
|
||||
agent = _agent_with(_Engine())
|
||||
_apply_context_engine_selection(
|
||||
agent, REQUEST, history, incoming, logger=MagicMock()
|
||||
)
|
||||
# Persisted history + incoming message are untouched despite the engine's
|
||||
# in-place mutation of the copies it received.
|
||||
assert history == history_snapshot
|
||||
assert incoming == incoming_snapshot
|
||||
|
||||
|
||||
def test_persisted_history_not_mutated():
|
||||
"""The hook must not mutate the persisted conversation history."""
|
||||
|
||||
class _Engine(_MinimalEngine):
|
||||
def select_context(self, request_messages, *, conversation_messages=None, **kwargs):
|
||||
# Even a misbehaving engine touching its inputs must not affect
|
||||
# what the host persists — the host passes the live list, so we
|
||||
# assert the host contract by checking the engine received it and
|
||||
# the canonical copy is unchanged after the call.
|
||||
return list(request_messages)
|
||||
|
||||
history_snapshot = [dict(m) for m in HISTORY]
|
||||
agent = _agent_with(_Engine())
|
||||
_apply_context_engine_selection(
|
||||
agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
|
||||
)
|
||||
assert HISTORY == history_snapshot
|
||||
|
||||
|
||||
# -- cache-stability + downstream-sanitizer contract -----------------------
|
||||
|
||||
|
||||
|
||||
def test_role_unusual_replacement_passed_through_for_downstream_sanitizers():
|
||||
"""The hook does structural validation only; role/tool normalization is
|
||||
deferred to the existing downstream sanitizers.
|
||||
|
||||
A `system -> user -> user` replacement (the exact shape flagged as a
|
||||
role-alternation risk on sibling PRs) is well-formed structurally, so the
|
||||
host returns it verbatim. Role-pairing/orphaned-tool cleanup runs *after*
|
||||
this hook in the request pipeline (`_sanitize_api_messages`,
|
||||
`_drop_thinking_only_and_merge_users`), so select_context cannot emit a
|
||||
malformed request that bypasses validation.
|
||||
"""
|
||||
role_unusual = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "a"},
|
||||
{"role": "user", "content": "b"},
|
||||
]
|
||||
|
||||
class _Engine(_MinimalEngine):
|
||||
def select_context(self, request_messages, **kwargs):
|
||||
return role_unusual
|
||||
|
||||
agent = _agent_with(_Engine())
|
||||
out = _apply_context_engine_selection(
|
||||
agent, REQUEST, HISTORY, HISTORY[-1], logger=MagicMock()
|
||||
)
|
||||
assert out is role_unusual # accepted structurally; downstream sanitizers normalize
|
||||
|
||||
|
||||
# -- on_turn_complete (post-turn observation) ------------------------------
|
||||
|
||||
|
||||
|
||||
def test_on_turn_complete_called_with_snapshot_and_meta():
|
||||
"""Host forwards a transcript copy + metadata; base no-op is skipped."""
|
||||
captured = {}
|
||||
|
||||
class _Engine(_MinimalEngine):
|
||||
def on_turn_complete(self, messages, usage=None, **kwargs):
|
||||
captured["messages"] = messages
|
||||
captured["usage"] = usage
|
||||
captured["kwargs"] = kwargs
|
||||
|
||||
agent = _agent_with(_Engine())
|
||||
_notify_context_engine_turn_complete(
|
||||
agent, HISTORY, usage={"total_tokens": 12}, logger=MagicMock(),
|
||||
turn_id="t1", api_call_count=1,
|
||||
)
|
||||
assert captured["messages"] == HISTORY
|
||||
assert captured["messages"] is not HISTORY # shallow copy
|
||||
assert captured["usage"] == {"total_tokens": 12}
|
||||
assert captured["kwargs"]["turn_id"] == "t1"
|
||||
assert captured["kwargs"]["api_call_count"] == 1
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user