mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 16:15:14 +00:00
fix: better snapshot and gate
This commit is contained in:
1 parent
29f661f3c7
commit
3f88dc7b57
12 files changed
+449
-22
No files matched your search
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user