fix: better snapshot and gate

This commit is contained in:
Alex committed 2026-08-22 14:00:15 +01:00
1 parent 29f661f3c7
commit 3f88dc7b57
12 files changed
+449 -22

No files matched your search

+26 -3
View File
@@ -886,6 +886,31 @@ class ToolExecutor:
)
return True
@staticmethod
def _advertisable_action_names(tools_dict: Dict) -> set:
"""Action names a ``tools_dict`` would expose, mirroring pass 1 of
:meth:`prepare_tools_for_llm`.
The model calls ACTION names (``run_code``), never tool names
(``code_executor``), so a fallback built from ``tool["name"]`` names
strings that cannot resolve — feeding the very retry loop the
correctable error exists to break.
"""
names = set()
for tool in (tools_dict or {}).values():
tool = tool or {}
if tool.get("name") == "api_tool":
actions = (tool.get("config") or {}).get("actions") or {}
actions = actions.values() if isinstance(actions, dict) else actions
else:
actions = tool.get("actions") or []
for action in actions:
if not isinstance(action, dict) or not action.get("active", True):
continue
if action.get("name"):
names.add(str(action["name"]))
return names
def _available_tool_names(self, tools_dict: Dict, exclude: Optional[str] = None) -> str:
"""Render the names the model can actually call, for a correctable error.
@@ -909,9 +934,7 @@ class ToolExecutor:
if not scope or tool_id in scope
)
if not names:
names = sorted(
{(tool or {}).get("name") or tool_id for tool_id, tool in (tools_dict or {}).items()}
)
names = sorted(self._advertisable_action_names(tools_dict or {}))
names = [str(name) for name in names if name != exclude]
if len(names) > self.MAX_ADVERTISED_TOOL_NAMES:
hidden = len(names) - self.MAX_ADVERTISED_TOOL_NAMES
+11 -7
View File
@@ -27,13 +27,17 @@ config = context.config
if settings.POSTGRES_URI:
config.set_main_option("sqlalchemy.url", settings.POSTGRES_URI)
if config.config_file_name is not None:
# ``disable_existing_loggers`` defaults to True, which would switch off
# every ``application.*`` logger already imported by the time migrations
# run. ``app.py`` calls ``setup_logging()`` (line 16) BEFORE
# ``ensure_database_ready()`` (line 58), so the web tier never
# re-enables them and loses all application logging for the life of the
# process — on exactly the boot where a schema upgrade happened.
if config.config_file_name is not None and config.attributes.get("configure_logger", True):
# ``fileConfig`` reconfigures the ROOT logger from ``[logger_root]``
# (level WARNING, a stderr handler), which drops every ``application.*``
# INFO record, detaches the OTLP handler and the request/trace filter that
# ``setup_logging()`` installed, and cannot be undone by
# ``disable_existing_loggers=False`` — that flag only governs whether
# existing loggers are switched off. ``app.py`` calls ``setup_logging()``
# before ``ensure_database_ready()``, so an in-process migration would
# silence the web tier for the life of the process, on exactly the boot
# where a schema upgrade happened. In-process callers opt out via
# ``cfg.attributes``; the alembic CLI still gets its configured logging.
fileConfig(config.config_file_name, disable_existing_loggers=False)
+24 -2
View File
@@ -1437,6 +1437,25 @@ class LLMHandler(ABC):
tool_calls[call.index] = call
else:
existing = tool_calls[call.index]
# Decide BEFORE the id/name overwrite below: a gateway
# that reuses one index for sequential DISTINCT calls
# is indistinguishable from a restatement once the
# identity is gone. Restatement frames carry neither
# id nor name, so a differing one means a new call.
is_new_call = bool(
(call.id and existing.id and call.id != existing.id)
or (call.name and existing.name and call.name != existing.name)
)
if is_new_call:
logger.warning(
"tool_call_index_reused_by_distinct_call",
extra={
"index": call.index,
"previous_call": f"{existing.id}/{existing.name}",
"incoming_call": f"{call.id}/{call.name}",
"dropped_prefix": str(existing.arguments)[:120],
},
)
if call.id:
existing.id = call.id
if call.name:
@@ -1447,8 +1466,11 @@ class LLMHandler(ABC):
elif isinstance(existing.arguments, str) and isinstance(
call.arguments, str
):
if _is_restated_payload(
existing.arguments, call.arguments
if (
_is_restated_payload(
existing.arguments, call.arguments
)
and not is_new_call
):
# Some OpenAI-compatible gateways restate
# the WHOLE argument payload on the finish
+4
View File
@@ -400,6 +400,10 @@ def _run_migrations(log: logging.Logger) -> None:
cfg = Config(str(alembic_ini))
cfg.set_main_option("script_location", str(alembic_ini.parent / "alembic"))
# We migrate in-process, after setup_logging() has configured the root
# logger. env.py honours this by skipping fileConfig, which would
# otherwise replace root's handlers and level for the rest of the process.
cfg.attributes["configure_logger"] = False
# Cheap pre-check: if we're already at head, say so explicitly.
try:
@@ -176,7 +176,12 @@ class PendingToolStateRepository:
)
return result.rowcount > 0
def release_claim(self, conversation_id: str, user_id: str) -> bool:
def release_claim(
self,
conversation_id: str,
user_id: str,
ttl_extension_seconds: int = PENDING_STATE_TTL_SECONDS,
) -> bool:
"""Flip a ``resuming`` row back to ``pending`` so a retry can claim it.
The inverse of :meth:`mark_resuming`, for the process that TOOK the
@@ -187,8 +192,16 @@ class PendingToolStateRepository:
:meth:`load_state` only sees ``pending`` rows.
Deliberately narrow: it matches on ``status = 'resuming'`` only, so it
can never resurrect a row another request has since deleted, and it
does not touch ``expires_at`` (the TTL is unrelated to the claim).
can never resurrect a row another request has since deleted.
Extends ``expires_at`` like :meth:`revert_stale_resuming` does. A
resume claimed near the end of the TTL that fails minutes later would
otherwise be handed back already expired: ``load_state`` and
``mark_resuming`` both gate on ``expires_at > clock_timestamp()``, and
flipping to ``pending`` also hides the row from
:meth:`revert_stale_resuming`, which matches ``resuming`` only — so
releasing the claim would REMOVE the rescue instead of speeding it up.
``GREATEST`` so a failed resume can never shorten a healthy TTL.
Returns:
True when a claim was released.
@@ -197,13 +210,22 @@ class PendingToolStateRepository:
text(
"""
UPDATE pending_tool_state
SET status = 'pending', resumed_at = NULL
SET status = 'pending',
resumed_at = NULL,
expires_at = GREATEST(
expires_at,
clock_timestamp() + make_interval(secs => :ttl)
)
WHERE conversation_id = CAST(:conv_id AS uuid)
AND user_id = :user_id
AND status = 'resuming'
"""
),
{"conv_id": conversation_id, "user_id": user_id},
{
"conv_id": conversation_id,
"user_id": user_id,
"ttl": ttl_extension_seconds,
},
)
return result.rowcount > 0
@@ -297,6 +297,58 @@ describe('scheme-less and file: hrefs', () => {
});
});
// A chip REPLACES the anchor, so a false positive here is worse than a miss:
// the user clicks an in-app link, gets a viewer onto an unrelated file, and has
// no href left to fall back to. In a shared conversation, where no
// `onOpenArtifact` is wired, the link degrades to inert plain text instead.
// The matchers are deliberately loose because a `sandbox:` URL is already known
// to be ours; against an ordinary relative link they are far too eager, so a
// scheme-less href must name a FILE before any of them may run.
describe('an ordinary relative link is not an artifact reference', () => {
// `deriveArtifactChips` labels an artifact `formatToolName(tool_name)` when
// no filename was reported, so 'Artifact' and 'Code Executor' are real
// labels — and every conversation with one artifact has a ref of `A1`.
const available = [
{ id: 'id-one', label: 'Artifact', ref: 'A1' },
{ id: 'id-two', label: 'Code Executor', ref: 'A2' },
];
it.each([
'/artifact',
'/artifact/',
'/Artifact?x=1',
'/code%20executor',
'/a1',
'/A2',
'/id-one',
'/settings/general',
'about',
])('leaves %s as an ordinary link', (href) => {
expect(resolveSandboxLink(href, available)).toEqual({ kind: 'external' });
});
it('still resolves a scheme-less href that names a file', () => {
const withFile = [{ id: 'id-three', label: 'report.pdf', ref: 'A3' }];
expect(resolveSandboxLink('/mnt/data/report.pdf', withFile)).toEqual({
kind: 'artifact',
artifact: withFile[0],
});
});
it('still resolves those same forms under an explicit scheme', () => {
// The gate is scoped to scheme-less hrefs: `sandbox:` and `artifact:` are
// unambiguously ours, so a bare ref or label stays resolvable there.
expect(resolveSandboxLink('artifact:A1', available)).toEqual({
kind: 'artifact',
artifact: available[0],
});
expect(resolveSandboxLink('sandbox:/artifact/id-two', available)).toEqual({
kind: 'artifact',
artifact: available[1],
});
});
});
// Two turns each produced a file called `report.pdf`. That is not exotic:
// `_filename()` in artifact_generator.py hands EVERY html artifact the name
// `report.html`, and re-running a script that writes `chart.png` collides too.
+13
View File
@@ -91,6 +91,10 @@ function matchedScheme(href: string | undefined | null): string | null {
/** Short refs the artifact tool hands back, e.g. ``A1`` is the first artifact. */
const SHORT_REF = /^a(\d+)$/i;
// A trailing extension is what separates "this segment names a file" from
// "this is an ordinary route". Only consulted for scheme-less hrefs.
const HAS_EXTENSION = /\.[a-z0-9]{1,8}$/i;
/** Any URL scheme, so ``https:`` and ``mailto:`` are left alone. */
const ANY_SCHEME = /^[a-z][a-z0-9+.-]*:/i;
@@ -162,6 +166,15 @@ export function resolveSandboxLink(
const last = segments[segments.length - 1];
// A scheme-less href stays a link unless the segment names a FILE. The
// matchers below are deliberately loose because a `sandbox:` URL is already
// known to be ours; against an ordinary relative link they are far too
// eager. `[see the artifact](/artifact)` hits the `formatToolName` label an
// artifact carries when no filename was reported, and `[notes](/a1)` hits
// the ref pattern — and a chip REPLACES the anchor, so the user is left
// with a button onto an unrelated file and no link to fall back to.
if (!scheme && !HAS_EXTENSION.test(last)) return miss;
// Try every identifier the segment could be, in order of certainty. A ref is
// matched against the artifact's own ``ref``, never by position: ``A2`` is the
// conversation's second artifact, which on a later turn is not this turn's
+11 -3
View File
@@ -511,15 +511,23 @@ class TestBaseAgentToolExecution:
call.name = "action_999"
call.arguments = "{}"
tools_dict = {"1": {"name": "tool1", "config": {}, "actions": []}}
tools_dict = {
"1": {
"name": "tool1",
"config": {},
"actions": [{"name": "tool1_read", "description": "D"}],
}
}
results = list(agent._execute_tool_action(tools_dict, call))
assert results[0]["type"] == "tool_call"
assert results[0]["data"]["status"] == "error"
assert "no such tool" in results[0]["data"]["result"]
# Tool *names*, not internal ids — the model never sees the ids.
assert "tool1" in results[0]["data"]["result"]
# LLM-visible ACTION names. Not the tool name and never the internal
# id: both name a string the model cannot actually call, which just
# buys another failed round.
assert "tool1_read" in results[0]["data"]["result"]
def test_execute_tool_action_with_parameters(
self,
@@ -17,10 +17,16 @@ import pytest
from application.agents.tool_executor import ToolExecutor
def _action(name):
return {"name": name, "description": "D", "active": True, "parameters": {"properties": {}}}
def _tools_dict():
# Rows carry actions, as every production row does: the model calls the
# ACTION name, so that is what an error may advertise.
return {
"t1": {"name": "memory", "actions": [], "config": {}},
"t2": {"name": "read_webpage", "actions": [], "config": {}},
"t1": {"name": "memory", "actions": [_action("memory_view")], "config": {}},
"t2": {"name": "read_webpage", "actions": [_action("read_webpage")], "config": {}},
}
@@ -119,6 +125,39 @@ class TestErrorNamesWhatTheModelCanCall:
_events, (result, _call_id) = _drain(executor, _call("make_a_pdf"), tools=tools_dict)
assert "create_artifact" in result
def test_the_fallback_advertises_action_names_not_tool_names(self):
"""With no name mapping built, the fallback must still name callables.
``_tool_to_name`` is empty whenever a turn produced zero LLM schemas —
every action toggled off, an unsynced MCP row, an ``api_tool`` with no
``config.actions``. Advertising ``code_executor`` there invites a call
named ``code_executor``, which cannot resolve: the error feeds the very
loop it exists to break.
"""
executor = ToolExecutor()
tools_dict = {
"t1": {"name": "code_executor", "actions": [_action("run_code")]},
"t2": {
"name": "api_tool",
"config": {"actions": {"a": _action("fetch_invoice")}},
},
}
assert executor._tool_to_name == {}
_events, (result, _call_id) = _drain(executor, _call("bash"), tools=tools_dict)
assert "run_code" in result and "fetch_invoice" in result, result
assert "code_executor" not in result, result
assert "api_tool" not in result, result
def test_the_fallback_skips_inactive_actions(self):
"""An action the user switched off is not callable, so it is not offered."""
executor = ToolExecutor()
off = _action("run_code")
off["active"] = False
tools_dict = {"t1": {"name": "code_executor", "actions": [off]}}
_events, (result, _call_id) = _drain(executor, _call("bash"), tools=tools_dict)
assert "(none available)" in result, result
def test_only_advertises_tools_in_scope_for_this_call(self):
"""A narrowed ``tools_dict`` must not be told about out-of-scope tools.
+65
View File
@@ -10,8 +10,10 @@ when OTLP log export is enabled.
from __future__ import annotations
import logging
import pathlib
import sys
import types
from unittest.mock import Mock
import pytest
@@ -164,3 +166,66 @@ class TestAlembicDoesNotSilenceApplicationLoggers:
).stdout
assert "default_disables True" in out
assert "flag_preserves True" in out
@pytest.mark.unit
class TestMigrationsDoNotClobberAppLogging:
"""An in-process ``alembic upgrade`` must not take application logging down.
``app.py`` calls ``setup_logging()`` before ``ensure_database_ready()``, so
``env.py``'s ``fileConfig`` runs against an already-configured root logger.
``fileConfig`` rebuilds root from ``[logger_root]`` (level WARNING, a stderr
handler), which drops every ``application.*`` INFO record and detaches the
OTEL handler and context filter for the life of the process —
``disable_existing_loggers=False`` does not prevent it, since that flag only
governs whether existing loggers are switched off.
"""
def test_bootstrap_tells_env_py_to_leave_logging_alone(self, monkeypatch):
from logging.config import fileConfig
from application.storage.db import bootstrap
# Skip the best-effort revision precheck; it needs a live DB. It is
# wrapped in try/except, so raising here lands on the "upgrade anyway"
# path we actually want to exercise.
monkeypatch.setattr(
"alembic.script.ScriptDirectory.from_config",
Mock(side_effect=RuntimeError("no db")),
)
root = logging.getLogger()
root.handlers = []
setup_logging()
marker = logging.NullHandler()
root.addHandler(marker)
before_level = root.level
seen = {}
def fake_upgrade(cfg, revision):
# Stand in for env.py, which guards its fileConfig on this key.
seen["configure_logger"] = cfg.attributes.get("configure_logger", True)
if seen["configure_logger"]:
fileConfig(cfg.config_file_name, disable_existing_loggers=False)
monkeypatch.setattr(
"alembic.command.upgrade", fake_upgrade, raising=False
)
bootstrap._run_migrations(logging.getLogger("test"))
assert seen["configure_logger"] is False
assert root.level == before_level
assert marker in root.handlers
assert logging.getLogger("application.probe").isEnabledFor(logging.INFO)
def test_env_py_honours_the_opt_out(self):
"""The other half of the contract lives in env.py's module-level guard."""
env_py = (
pathlib.Path(__file__).resolve().parents[2]
/ "application"
/ "alembic"
/ "env.py"
)
source = env_py.read_text()
assert 'config.attributes.get("configure_logger", True)' in source
+128
View File
@@ -1367,6 +1367,134 @@ class TestHandleStreaming:
# The whole point: the merged payload parses.
_json.loads(executed_call.arguments)
def test_a_distinct_call_reusing_an_index_is_not_treated_as_a_restatement(
self, caplog
):
"""Two complete payloads under one index may be two different calls.
Some OpenAI-compatible gateways tag sequential parallel calls with the
same ``index``. That is indistinguishable from a restatement by payload
shape alone, so the identity decides — and it must be read BEFORE the
id/name overwrite, which would otherwise erase the evidence. Falling
back to concatenation keeps the failure loud: ``tool_action_parser``
logs the unparseable prefix, where a silent replace would run one tool
with the other's arguments and report the turn complete.
"""
import logging as _logging
handler = ConcreteHandler()
agent = Mock()
agent.llm = Mock()
agent.model_id = "test"
agent.tools = []
agent._check_context_limit = Mock(return_value=False)
agent.context_limit_reached = False
agent.llm.__class__.__name__ = "MockLLM"
agent.tool_executor.check_pause = Mock(return_value=None)
chunk1 = LLMResponse(
content="",
tool_calls=[
ToolCall(id="call_A", name="memory_view", arguments='{"a":1}', index=0)
],
finish_reason="",
raw_response={},
)
chunk2 = LLMResponse(
content="",
tool_calls=[
ToolCall(id="call_B", name="memory_write", arguments='{"b":2}', index=0)
],
finish_reason="tool_calls",
raw_response={},
)
handler.parse_response = lambda c: c
def fake_iterate(response):
yield from response
handler._iterate_stream = fake_iterate
def fake_execute(tools_dict, call):
yield {"type": "tool_call", "data": {"status": "pending"}}
return ("result", call.id)
agent._execute_tool_action = Mock(side_effect=fake_execute)
agent.llm.gen_stream = Mock(
return_value=[
LLMResponse(
content="done", tool_calls=[], finish_reason="stop", raw_response={}
)
]
)
with caplog.at_level(_logging.WARNING, logger="application.llm.handlers.base"):
list(
handler.handle_streaming(agent, [chunk1, chunk2], {"1": {"name": "t"}}, [])
)
executed_call = agent._execute_tool_action.call_args[0][1]
# Concatenated, not silently replaced: call A's payload survives in a
# form the parser reports rather than vanishing.
assert executed_call.arguments == '{"a":1}{"b":2}'
assert "tool_call_index_reused_by_distinct_call" in caplog.text
def test_a_restatement_carrying_the_same_identity_still_replaces(self):
"""The guard must not undo the fix: same id/name means one call."""
handler = ConcreteHandler()
agent = Mock()
agent.llm = Mock()
agent.model_id = "test"
agent.tools = []
agent._check_context_limit = Mock(return_value=False)
agent.context_limit_reached = False
agent.llm.__class__.__name__ = "MockLLM"
agent.tool_executor.check_pause = Mock(return_value=None)
chunk1 = LLMResponse(
content="",
tool_calls=[
ToolCall(id="c1", name="search", arguments='{"q":"x"}', index=0)
],
finish_reason="",
raw_response={},
)
# The finish frame repeats the identity as well as the payload.
chunk2 = LLMResponse(
content="",
tool_calls=[
ToolCall(id="c1", name="search", arguments='{"q":"x"}', index=0)
],
finish_reason="tool_calls",
raw_response={},
)
handler.parse_response = lambda c: c
def fake_iterate(response):
yield from response
handler._iterate_stream = fake_iterate
def fake_execute(tools_dict, call):
yield {"type": "tool_call", "data": {"status": "pending"}}
return ("result", call.id)
agent._execute_tool_action = Mock(side_effect=fake_execute)
agent.llm.gen_stream = Mock(
return_value=[
LLMResponse(
content="done", tool_calls=[], finish_reason="stop", raw_response={}
)
]
)
list(handler.handle_streaming(agent, [chunk1, chunk2], {"1": {"name": "t"}}, []))
executed_call = agent._execute_tool_action.call_args[0][1]
assert executed_call.arguments == '{"q":"x"}'
def test_context_limit_adds_system_message(self):
handler = ConcreteHandler()
agent = Mock()
@@ -248,6 +248,53 @@ class TestReleaseClaim:
# ...and the next resume attempt can take it straight away.
assert repo.claim_state(conv["id"], "user-1") is not None
def test_bumps_expires_at_to_protect_from_the_ttl_sweep(self, pg_conn):
"""A claim released after a long resume must not come back expired.
``claim_state``/``mark_resuming`` gate on ``expires_at``, and flipping
to ``pending`` also hides the row from ``revert_stale_resuming``, which
matches ``resuming`` only. Without the bump, releasing a claim on a row
whose TTL lapsed mid-resume DELETES the user's rescue instead of
speeding it up — the turn is reported retryable but can never be
re-claimed. This is the sibling of
``TestRevertStaleResuming::test_bumps_expires_at_to_protect_from_ttl_sweep``.
"""
conv = _conv(pg_conn)
repo = _repo(pg_conn)
repo.save_state(conv["id"], "user-1", **_sample_state())
repo.claim_state(conv["id"], "user-1")
# The resume outlived the row's TTL before it failed.
pg_conn.execute(
text(
"""
UPDATE pending_tool_state
SET expires_at = clock_timestamp() - interval '1 minute'
WHERE conversation_id = CAST(:conv_id AS uuid)
"""
),
{"conv_id": conv["id"]},
)
assert repo.release_claim(conv["id"], "user-1") is True
reclaimed = repo.load_state(conv["id"], "user-1")
assert reclaimed is not None, "released claim came back already expired"
assert repo.claim_state(conv["id"], "user-1") is not None
assert repo.cleanup_expired() == []
def test_does_not_shorten_a_healthy_ttl(self, pg_conn):
"""A failed resume must never cut short a row with time left."""
conv = _conv(pg_conn)
repo = _repo(pg_conn)
repo.save_state(conv["id"], "user-1", **_sample_state())
repo.claim_state(conv["id"], "user-1")
before = repo.load_state_any(conv["id"], "user-1")["expires_at"]
repo.release_claim(conv["id"], "user-1", ttl_extension_seconds=1)
after = repo.load_state_any(conv["id"], "user-1")["expires_at"]
assert after == before
def test_noop_on_a_pending_row(self, pg_conn):
"""Only a claim can be released; a live pending row is left alone."""
conv = _conv(pg_conn)