core/agent+tools: apply a batch's same-file writes in call order

This commit is contained in:
2026-10-16 19:37:29 +08:00
parent ab6b5d590c
commit 85b0318b77
7 changed files with 208 additions and 17 deletions
+60 -6
View File
@@ -4,7 +4,9 @@ import asyncio
import inspect import inspect
import types import types
from collections.abc import Awaitable, Callable, Mapping from collections.abc import Awaitable, Callable, Mapping
from contextlib import AsyncExitStack
from enum import Flag, auto from enum import Flag, auto
from pathlib import Path
from typing import Any, cast, get_args, get_origin, get_type_hints, overload from typing import Any, cast, get_args, get_origin, get_type_hints, overload
import msgspec import msgspec
@@ -46,13 +48,19 @@ class ToolTag(Flag):
class ToolDefinition: class ToolDefinition:
"""A registered tool: schema for the model plus a callable handler.""" """A registered tool: schema for the model plus a callable handler.
``writes`` names the arguments whose values name files the handler writes.
Calls of one batch that name the same file apply one at a time, in call
order, so the later one sees what the earlier one wrote.
"""
name: str name: str
description: str description: str
parameters: JSONSchema parameters: JSONSchema
handler: ToolHandler handler: ToolHandler
tags: ToolTag tags: ToolTag
writes: tuple[str, ...]
def __init__( def __init__(
self, self,
@@ -62,6 +70,7 @@ class ToolDefinition:
handler: ToolHandler, handler: ToolHandler,
*, *,
tags: ToolTag = ToolTag.LOCAL, tags: ToolTag = ToolTag.LOCAL,
writes: tuple[str, ...] = (),
) -> None: ) -> None:
if not (tags & (ToolTag.LOCAL | ToolTag.PUBLIC)): if not (tags & (ToolTag.LOCAL | ToolTag.PUBLIC)):
msg = f"tool {name!r} tags must include LOCAL and/or PUBLIC" msg = f"tool {name!r} tags must include LOCAL and/or PUBLIC"
@@ -71,6 +80,7 @@ class ToolDefinition:
self.parameters = parameters self.parameters = parameters
self.handler = handler self.handler = handler
self.tags = tags self.tags = tags
self.writes = writes
def to_tool_item(self) -> ToolFunctionItem: def to_tool_item(self) -> ToolFunctionItem:
return ToolFunctionItem( return ToolFunctionItem(
@@ -130,6 +140,7 @@ def _build_definition(
name: str | None, name: str | None,
description: str | None, description: str | None,
tags: ToolTag, tags: ToolTag,
writes: tuple[str, ...] = (),
) -> ToolDefinition: ) -> ToolDefinition:
tool_name = name or func.__name__ tool_name = name or func.__name__
tool_description = description if description is not None else (inspect.getdoc(func) or "") tool_description = description if description is not None else (inspect.getdoc(func) or "")
@@ -139,6 +150,7 @@ def _build_definition(
parameters=schema_from_callable(func), parameters=schema_from_callable(func),
handler=func, handler=func,
tags=tags, tags=tags,
writes=writes,
) )
@@ -162,6 +174,7 @@ def tool[**PS, R](
name: str | None = None, name: str | None = None,
description: str | None = None, description: str | None = None,
tags: ToolTag = ToolTag.LOCAL, tags: ToolTag = ToolTag.LOCAL,
writes: tuple[str, ...] = (),
register: bool = True, register: bool = True,
) -> Callable[[Callable[PS, R]], ToolDefinition]: ... ) -> Callable[[Callable[PS, R]], ToolDefinition]: ...
@@ -173,16 +186,19 @@ def tool[**PS, R](
name: str | None = None, name: str | None = None,
description: str | None = None, description: str | None = None,
tags: ToolTag = ToolTag.LOCAL, tags: ToolTag = ToolTag.LOCAL,
writes: tuple[str, ...] = (),
register: bool = True, register: bool = True,
) -> ToolDefinition | Callable[[Callable[PS, R]], ToolDefinition]: ) -> ToolDefinition | Callable[[Callable[PS, R]], ToolDefinition]:
"""Define (and by default catalog-register) a function as an agent tool. """Define (and by default catalog-register) a function as an agent tool.
Schema is inferred from type hints; description defaults to the docstring. Schema is inferred from type hints; description defaults to the docstring.
Default ``tags`` is :attr:`ToolTag.LOCAL`. When ``register`` is true, the Default ``tags`` is :attr:`ToolTag.LOCAL`. ``writes`` names the arguments
definition is added to the process :class:`~plyngent.tools.catalog.ToolCatalog` that carry the files the handler writes (:attr:`ToolDefinition.writes`).
with the current registration source (builtin unless a plugin context is set). When ``register`` is true, the definition is added to the process
Catalog registration does **not** alone make the tool model-visible — hosts :class:`~plyngent.tools.catalog.ToolCatalog` with the current registration
still **select** tools into a :class:`ToolRegistry`. source (builtin unless a plugin context is set). Catalog registration does
**not** alone make the tool model-visible — hosts still **select** tools
into a :class:`ToolRegistry`.
""" """
def decorator(fn: Callable[PS, R]) -> ToolDefinition: def decorator(fn: Callable[PS, R]) -> ToolDefinition:
@@ -191,6 +207,7 @@ def tool[**PS, R](
name=name, name=name,
description=description, description=description,
tags=tags, tags=tags,
writes=writes,
) )
if register: if register:
_register_definition(definition) _register_definition(definition)
@@ -225,6 +242,15 @@ def _missing_required_error(definition: ToolDefinition, args: Mapping[str, objec
return f"error: `{definition.name}` is missing required argument(s): {', '.join(missing)}{hint}" return f"error: `{definition.name}` is missing required argument(s): {', '.join(missing)}{hint}"
def _write_path_key(value: str) -> str:
"""Comparable key for one path argument of a writing tool.
Relative arguments all mean the same workspace, so the text itself separates
two spellings of one file (``a.txt`` / ``./a.txt``) from two files.
"""
return str(Path(value).expanduser())
class ToolRegistry: class ToolRegistry:
"""Name → tool definition map with execution helpers.""" """Name → tool definition map with execution helpers."""
@@ -232,6 +258,9 @@ class ToolRegistry:
_danger: DangerClassifier | None _danger: DangerClassifier | None
_on_confirm: ToolConfirmHook | None _on_confirm: ToolConfirmHook | None
_confirm_lock: asyncio.Lock _confirm_lock: asyncio.Lock
# One lock per written file, kept for the registry's life (a session's edits
# bound the map), so a batch's writes to one file apply one at a time.
_write_locks: dict[str, asyncio.Lock]
_yolo: bool _yolo: bool
_auto_bind_state: bool _auto_bind_state: bool
_instance: object | None _instance: object | None
@@ -252,6 +281,7 @@ class ToolRegistry:
self._danger = danger self._danger = danger
self._on_confirm = on_confirm self._on_confirm = on_confirm
self._confirm_lock = asyncio.Lock() self._confirm_lock = asyncio.Lock()
self._write_locks = {}
self._yolo = yolo self._yolo = yolo
self._auto_bind_state = auto_bind_state self._auto_bind_state = auto_bind_state
self._instance = instance_state self._instance = instance_state
@@ -381,6 +411,25 @@ class ToolRegistry:
return cast("Any", self._session) return cast("Any", self._session)
return session return session
def _write_locks_for(self, definition: ToolDefinition, args: Mapping[str, object]) -> list[asyncio.Lock]:
"""Locks for the files this call writes, in a stable order (no deadlock).
Empty for a tool that declares no :attr:`ToolDefinition.writes`.
"""
keys: set[str] = set()
for name in definition.writes:
value = args.get(name)
if isinstance(value, str) and value.strip():
keys.add(_write_path_key(value))
locks: list[asyncio.Lock] = []
for key in sorted(keys):
lock = self._write_locks.get(key)
if lock is None:
lock = asyncio.Lock()
self._write_locks[key] = lock
locks.append(lock)
return locks
def _grant_allows(self, name: str, *, tags: ToolTag) -> bool: def _grant_allows(self, name: str, *, tags: ToolTag) -> bool:
if not (tags & ToolTag.TRUSTABLE): if not (tags & ToolTag.TRUSTABLE):
return False return False
@@ -468,6 +517,11 @@ class ToolRegistry:
missing = self._check_state_tags(definition) missing = self._check_state_tags(definition)
if missing is not None: if missing is not None:
return missing return missing
# Held before the confirm prompt, so a sibling call writing the same
# file waits for this one instead of overtaking its prompt.
async with AsyncExitStack() as held:
for lock in self._write_locks_for(definition, args):
await held.enter_async_context(lock)
denied = await self._maybe_confirm(name, args, tags=definition.tags) denied = await self._maybe_confirm(name, args, tags=definition.tags)
if denied is not None: if denied is not None:
return denied return denied
+1 -1
View File
@@ -124,7 +124,7 @@ def _replace_range(
) )
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO) @tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO, writes=("path",))
async def edit_lineno(path: str, start_line: int, end_line: int, new_content: str) -> str: async def edit_lineno(path: str, start_line: int, end_line: int, new_content: str) -> str:
"""Replace lines ``start_line``..``end_line`` (1-based, inclusive) in a file. """Replace lines ``start_line``..``end_line`` (1-based, inclusive) in a file.
+1 -1
View File
@@ -48,7 +48,7 @@ def _success_message(
) )
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO) @tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO, writes=("path",))
async def edit_replace(path: str, old_string: str, new_string: str, max_replaces: int = 1) -> str: async def edit_replace(path: str, old_string: str, new_string: str, max_replaces: int = 1) -> str:
"""Replace occurrences of ``old_string`` with ``new_string`` in a file. """Replace occurrences of ``old_string`` with ``new_string`` in a file.
+3 -3
View File
@@ -94,7 +94,7 @@ def _copy_or_move_validated(
return f"error: {action} failed: {exc}" return f"error: {action} failed: {exc}"
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO) @tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO, writes=("dst",))
async def copy_path(src: str, dst: str, *, overwrite: bool = False) -> str: async def copy_path(src: str, dst: str, *, overwrite: bool = False) -> str:
"""Copy a file or directory under the workspace (``shutil.copy2`` / ``copytree``). """Copy a file or directory under the workspace (``shutil.copy2`` / ``copytree``).
@@ -110,7 +110,7 @@ async def copy_path(src: str, dst: str, *, overwrite: bool = False) -> str:
return _copy_or_move_validated(source, dest, root, src=src, dst=dst, overwrite=overwrite, action="copy") return _copy_or_move_validated(source, dest, root, src=src, dst=dst, overwrite=overwrite, action="copy")
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO) @tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO, writes=("src", "dst"))
async def move_path(src: str, dst: str, *, overwrite: bool = False) -> str: async def move_path(src: str, dst: str, *, overwrite: bool = False) -> str:
"""Move/rename a file or directory under the workspace (``shutil.move``).""" """Move/rename a file or directory under the workspace (``shutil.move``)."""
pair = _resolve_pair(src, dst) pair = _resolve_pair(src, dst)
@@ -149,7 +149,7 @@ def _delete_target(target: Path, path: str, *, recursive: bool) -> str:
return f"error: unsupported path type: {path}" return f"error: unsupported path type: {path}"
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO) @tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO, writes=("path",))
async def delete_path(path: str, *, recursive: bool = False) -> str: async def delete_path(path: str, *, recursive: bool = False) -> str:
"""Delete a file or directory under the workspace. """Delete a file or directory under the workspace.
+1 -1
View File
@@ -4,7 +4,7 @@ from plyngent.agent import ToolTag, tool
from plyngent.tools.workspace import AccessMode, resolve_path from plyngent.tools.workspace import AccessMode, resolve_path
@tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO) @tool(tags=ToolTag.LOCAL | ToolTag.INSTANCE_STATE | ToolTag.YOLO, writes=("path",))
async def write_file(path: str, content: str, *, overwrite: bool = False) -> str: async def write_file(path: str, content: str, *, overwrite: bool = False) -> str:
"""Write text content to a file under the workspace (creates parents). """Write text content to a file under the workspace (creates parents).
+100 -1
View File
@@ -22,7 +22,7 @@ from plyngent.lmproto.openai_compatible.model import (
) )
if TYPE_CHECKING: if TYPE_CHECKING:
from collections.abc import AsyncIterator from collections.abc import AsyncIterator, Mapping
def _response(message: AssistantChatMessage) -> ChatCompletionResponse: def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
@@ -156,3 +156,102 @@ async def test_parallel_tool_calls_run_concurrently() -> None:
tool_msgs = [e for e in agent.messages if isinstance(e, ToolChatMessage)] tool_msgs = [e for e in agent.messages if isinstance(e, ToolChatMessage)]
assert {m.content for m in tool_msgs} == {"a", "b"} assert {m.content for m in tool_msgs} == {"a", "b"}
assert any(getattr(e, "content", None) == "done" for e in events) assert any(getattr(e, "content", None) == "done" for e in events)
async def test_same_file_writes_apply_in_call_order() -> None:
"""A batch's writes to one file run one at a time, in call order."""
order: list[str] = []
@tool(register=False, writes=("path",))
async def write_a(path: str) -> str:
order.append(f"a start {path}")
await asyncio.sleep(0) # an unserialized sibling would interleave here
order.append(f"a end {path}")
return "a"
@tool(register=False, writes=("path",))
async def write_b(path: str) -> str:
order.append(f"b start {path}")
await asyncio.sleep(0)
order.append(f"b end {path}")
return "b"
registry = ToolRegistry([write_a, write_b])
results = await asyncio.gather(
registry.execute("write_a", '{"path": "a.txt"}'),
registry.execute("write_b", '{"path": "./a.txt"}'),
)
assert results == ["a", "b"]
assert order == ["a start a.txt", "a end a.txt", "b start ./a.txt", "b end ./a.txt"]
async def test_different_files_still_write_concurrently() -> None:
"""Serialization is per file: two files of one batch still overlap."""
overlaps = 0
gate = asyncio.Event()
async def _rendezvous() -> None:
nonlocal overlaps
overlaps += 1
if overlaps == 2:
_ = gate.set()
_ = await asyncio.wait_for(gate.wait(), timeout=2.0)
@tool(register=False, writes=("path",))
async def write_a(path: str) -> str:
del path
await _rendezvous()
return "a"
@tool(register=False, writes=("path",))
async def write_b(path: str) -> str:
del path
await _rendezvous()
return "b"
registry = ToolRegistry([write_a, write_b])
results = await asyncio.gather(
registry.execute("write_a", '{"path": "a.txt"}'),
registry.execute("write_b", '{"path": "b.txt"}'),
)
assert results == ["a", "b"]
async def test_a_write_holds_its_file_across_the_confirm() -> None:
"""The file is locked before the prompt, so a sibling write cannot overtake it."""
order: list[str] = []
prompt_shown = asyncio.Event()
answered = asyncio.Event()
@tool(register=False, writes=("path",))
async def write_a(path: str) -> str:
del path
order.append("a")
return "a"
@tool(register=False, writes=("path",))
async def write_b(path: str) -> str:
del path
order.append("b")
return "b"
async def on_confirm(name: str, args: Mapping[str, object], reason: str) -> bool:
del args, reason
order.append(f"confirm {name}")
prompt_shown.set()
await answered.wait()
return True
def danger(name: str, args: Mapping[str, object]) -> str | None:
del args
return "overwrite" if name == "write_a" else None
registry = ToolRegistry([write_a, write_b], danger=danger, on_confirm=on_confirm)
first = asyncio.create_task(registry.execute("write_a", '{"path": "a.txt"}'))
await prompt_shown.wait()
second = asyncio.create_task(registry.execute("write_b", '{"path": "a.txt"}'))
await asyncio.sleep(0) # the second write waits for the file, not for the human
assert order == ["confirm write_a"]
answered.set()
assert await asyncio.gather(first, second) == ["a", "b"]
assert order == ["confirm write_a", "a", "b"]
+38
View File
@@ -1,11 +1,18 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import importlib import importlib
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING
from plyngent.agent import ToolRegistry
from plyngent.tools.danger import classify_danger
from plyngent.tools.file import edit_replace, listdir, read_file, write_file from plyngent.tools.file import edit_replace, listdir, read_file, write_file
from tests.test_tools.helpers import call_sync from tests.test_tools.helpers import call_sync
if TYPE_CHECKING:
from collections.abc import Mapping
get_truncated_module = importlib.import_module("plyngent.tools.file.get_truncated") get_truncated_module = importlib.import_module("plyngent.tools.file.get_truncated")
@@ -37,6 +44,37 @@ def test_write_file_refuses_existing_file_without_overwrite(workspace: Path) ->
assert (workspace / "keep.txt").read_text(encoding="utf-8") == "second" assert (workspace / "keep.txt").read_text(encoding="utf-8") == "second"
async def test_same_file_writes_in_one_batch_keep_call_order(workspace: Path) -> None:
"""The second call writes onto the first one's text, not onto the stale text."""
target = workspace / "notes.txt"
_ = target.write_text("one", encoding="utf-8")
confirmed = asyncio.Event()
async def on_confirm(name: str, args: Mapping[str, object], reason: str) -> bool:
del name, args, reason
await confirmed.wait()
return True
registry = ToolRegistry([write_file, edit_replace], danger=classify_danger, on_confirm=on_confirm)
rewrite = asyncio.create_task(
registry.execute("write_file", '{"path": "notes.txt", "content": "two", "overwrite": true}')
)
await asyncio.sleep(0) # the rewrite is waiting on its overwrite confirm
edit = asyncio.create_task(
registry.execute(
"edit_replace",
'{"path": "notes.txt", "old_string": "two", "new_string": "three"}',
)
)
await asyncio.sleep(0) # the edit waits for the file instead of editing stale text
assert target.read_text(encoding="utf-8") == "one"
confirmed.set()
rewrite_out, edit_out = await asyncio.gather(rewrite, edit)
assert "wrote" in rewrite_out
assert "replaced" in edit_out
assert target.read_text(encoding="utf-8") == "three"
def test_read_file_max_chars_embeds_token(workspace: Path) -> None: def test_read_file_max_chars_embeds_token(workspace: Path) -> None:
(workspace / "big.txt").write_text("a" * 1000, encoding="utf-8") (workspace / "big.txt").write_text("a" * 1000, encoding="utf-8")
out = call_sync(read_file, "big.txt", max_chars=200) out = call_sync(read_file, "big.txt", max_chars=200)