diff --git a/src/plyngent/agent/tools.py b/src/plyngent/agent/tools.py index 7a9d10e..94b1002 100644 --- a/src/plyngent/agent/tools.py +++ b/src/plyngent/agent/tools.py @@ -4,7 +4,9 @@ import asyncio import inspect import types from collections.abc import Awaitable, Callable, Mapping +from contextlib import AsyncExitStack from enum import Flag, auto +from pathlib import Path from typing import Any, cast, get_args, get_origin, get_type_hints, overload import msgspec @@ -46,13 +48,19 @@ class ToolTag(Flag): 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 description: str parameters: JSONSchema handler: ToolHandler tags: ToolTag + writes: tuple[str, ...] def __init__( self, @@ -62,6 +70,7 @@ class ToolDefinition: handler: ToolHandler, *, tags: ToolTag = ToolTag.LOCAL, + writes: tuple[str, ...] = (), ) -> None: if not (tags & (ToolTag.LOCAL | ToolTag.PUBLIC)): msg = f"tool {name!r} tags must include LOCAL and/or PUBLIC" @@ -71,6 +80,7 @@ class ToolDefinition: self.parameters = parameters self.handler = handler self.tags = tags + self.writes = writes def to_tool_item(self) -> ToolFunctionItem: return ToolFunctionItem( @@ -130,6 +140,7 @@ def _build_definition( name: str | None, description: str | None, tags: ToolTag, + writes: tuple[str, ...] = (), ) -> ToolDefinition: tool_name = name or func.__name__ 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), handler=func, tags=tags, + writes=writes, ) @@ -162,6 +174,7 @@ def tool[**PS, R]( name: str | None = None, description: str | None = None, tags: ToolTag = ToolTag.LOCAL, + writes: tuple[str, ...] = (), register: bool = True, ) -> Callable[[Callable[PS, R]], ToolDefinition]: ... @@ -173,16 +186,19 @@ def tool[**PS, R]( name: str | None = None, description: str | None = None, tags: ToolTag = ToolTag.LOCAL, + writes: tuple[str, ...] = (), register: bool = True, ) -> ToolDefinition | Callable[[Callable[PS, R]], ToolDefinition]: """Define (and by default catalog-register) a function as an agent tool. Schema is inferred from type hints; description defaults to the docstring. - Default ``tags`` is :attr:`ToolTag.LOCAL`. When ``register`` is true, the - definition is added to the process :class:`~plyngent.tools.catalog.ToolCatalog` - with the current registration 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`. + Default ``tags`` is :attr:`ToolTag.LOCAL`. ``writes`` names the arguments + that carry the files the handler writes (:attr:`ToolDefinition.writes`). + When ``register`` is true, the definition is added to the process + :class:`~plyngent.tools.catalog.ToolCatalog` with the current registration + 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: @@ -191,6 +207,7 @@ def tool[**PS, R]( name=name, description=description, tags=tags, + writes=writes, ) if register: _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}" +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: """Name → tool definition map with execution helpers.""" @@ -232,6 +258,9 @@ class ToolRegistry: _danger: DangerClassifier | None _on_confirm: ToolConfirmHook | None _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 _auto_bind_state: bool _instance: object | None @@ -252,6 +281,7 @@ class ToolRegistry: self._danger = danger self._on_confirm = on_confirm self._confirm_lock = asyncio.Lock() + self._write_locks = {} self._yolo = yolo self._auto_bind_state = auto_bind_state self._instance = instance_state @@ -381,6 +411,25 @@ class ToolRegistry: return cast("Any", self._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: if not (tags & ToolTag.TRUSTABLE): return False @@ -468,10 +517,15 @@ class ToolRegistry: missing = self._check_state_tags(definition) if missing is not None: return missing - denied = await self._maybe_confirm(name, args, tags=definition.tags) - if denied is not None: - return denied - return await self._invoke(definition, args) + # 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) + if denied is not None: + return denied + return await self._invoke(definition, args) # Bind host-provided state for tools that read contextvars (todo, workspace). # Tag enforcement remains gated by ``auto_bind_state``; binding is always useful diff --git a/src/plyngent/tools/file/edit_lineno.py b/src/plyngent/tools/file/edit_lineno.py index 6860a77..ccbedec 100644 --- a/src/plyngent/tools/file/edit_lineno.py +++ b/src/plyngent/tools/file/edit_lineno.py @@ -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: """Replace lines ``start_line``..``end_line`` (1-based, inclusive) in a file. diff --git a/src/plyngent/tools/file/edit_replace.py b/src/plyngent/tools/file/edit_replace.py index d141460..598eb38 100644 --- a/src/plyngent/tools/file/edit_replace.py +++ b/src/plyngent/tools/file/edit_replace.py @@ -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: """Replace occurrences of ``old_string`` with ``new_string`` in a file. diff --git a/src/plyngent/tools/file/fs_ops.py b/src/plyngent/tools/file/fs_ops.py index 3ca3ad7..78c23d6 100644 --- a/src/plyngent/tools/file/fs_ops.py +++ b/src/plyngent/tools/file/fs_ops.py @@ -94,7 +94,7 @@ def _copy_or_move_validated( 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: """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") -@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: """Move/rename a file or directory under the workspace (``shutil.move``).""" 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}" -@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: """Delete a file or directory under the workspace. diff --git a/src/plyngent/tools/file/write.py b/src/plyngent/tools/file/write.py index 55f9e6c..78bb4ab 100644 --- a/src/plyngent/tools/file/write.py +++ b/src/plyngent/tools/file/write.py @@ -4,7 +4,7 @@ from plyngent.agent import ToolTag, tool 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: """Write text content to a file under the workspace (creates parents). diff --git a/tests/test_agent/test_parallel_tools.py b/tests/test_agent/test_parallel_tools.py index 938c4bf..bc98ba0 100644 --- a/tests/test_agent/test_parallel_tools.py +++ b/tests/test_agent/test_parallel_tools.py @@ -22,7 +22,7 @@ from plyngent.lmproto.openai_compatible.model import ( ) if TYPE_CHECKING: - from collections.abc import AsyncIterator + from collections.abc import AsyncIterator, Mapping 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)] assert {m.content for m in tool_msgs} == {"a", "b"} 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"] diff --git a/tests/test_tools/test_file.py b/tests/test_tools/test_file.py index fc873e4..7f0895d 100644 --- a/tests/test_tools/test_file.py +++ b/tests/test_tools/test_file.py @@ -1,11 +1,18 @@ from __future__ import annotations +import asyncio import importlib 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 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") @@ -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" +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: (workspace / "big.txt").write_text("a" * 1000, encoding="utf-8") out = call_sync(read_file, "big.txt", max_chars=200)