From f5b566b371e25401b3f07ba79fadd819db5f76ef Mon Sep 17 00:00:00 2001 From: worldmozara Date: Wed, 29 Jul 2026 01:40:31 +0800 Subject: [PATCH] core/tools: add per-turn guard requiring lineno read before edit_lineno --- src/plyngent/agent/__init__.py | 4 +++ src/plyngent/agent/chat.py | 40 ++++++++++++++++++++++++++ src/plyngent/tools/file/edit_lineno.py | 7 ++++- src/plyngent/tools/file/read.py | 3 +- tests/test_tools/test_edit_lineno.py | 19 ++++++++++++ 5 files changed, 71 insertions(+), 2 deletions(-) diff --git a/src/plyngent/agent/__init__.py b/src/plyngent/agent/__init__.py index 078f672..3dbb185 100644 --- a/src/plyngent/agent/__init__.py +++ b/src/plyngent/agent/__init__.py @@ -1,4 +1,8 @@ from .chat import ChatAgent as ChatAgent +from .chat import get_lineno_read_files as get_lineno_read_files +from .chat import mark_lineno_read as mark_lineno_read +from .chat import reset_lineno_tracker as reset_lineno_tracker +from .chat import was_lineno_read as was_lineno_read from .compact import build_compacted_seed_messages as build_compacted_seed_messages from .compact import summarize_messages as summarize_messages from .events import AgentEvent as AgentEvent diff --git a/src/plyngent/agent/chat.py b/src/plyngent/agent/chat.py index f47ffb2..e44aafc 100644 --- a/src/plyngent/agent/chat.py +++ b/src/plyngent/agent/chat.py @@ -1,5 +1,6 @@ from __future__ import annotations +from contextvars import ContextVar, Token from typing import TYPE_CHECKING from msgspec import UNSET @@ -46,6 +47,44 @@ if TYPE_CHECKING: type LimitContinueHook = Callable[[str], bool | Awaitable[bool]] +# Per-conversation tracking of files read with ``with_lineno=True``. +# Set to a fresh ``set`` at the start of each user turn via +# :func:`reset_lineno_tracker`; cleared after the assistant finishes. +# Tool handlers read/write via :func:`get_lineno_read_files`. +_lineno_read_files: ContextVar[set[str] | None] = ContextVar("_lineno_read_files", default=None) + + +def reset_lineno_tracker() -> Token[set[str] | None]: + """Replace the tracker with a fresh empty set; return the previous token. + + Call at the start of a new user turn. Restore with + ``_lineno_read_files.reset(token)`` when the turn completes. + """ + return _lineno_read_files.set(set()) + + +def get_lineno_read_files() -> set[str]: + """Return the current turn's set of line-numbered file paths. + + Returns an empty set when no turn is active (tools will reject edits). + """ + val = _lineno_read_files.get() + return val if val is not None else set() + + +def mark_lineno_read(path: str) -> None: + """Record that *path* was read with line numbers this turn.""" + val = _lineno_read_files.get() + if val is not None: + val.add(path) + + +def was_lineno_read(path: str) -> bool: + """Return True if *path* was read with line numbers this turn.""" + val = _lineno_read_files.get() + return val is not None and path in val + + def incomplete_turn_user_text(messages: Sequence[AnyChatMessage]) -> str | None: """User text for an incomplete turn, if history can be retried. @@ -484,6 +523,7 @@ class ChatAgent: async def run(self, user_text: str) -> AsyncIterator[AgentEvent]: """Append a user message (persist immediately), run the tool loop, yield events.""" + _ = reset_lineno_tracker() self._ensure_system_prompt() user_msg = UserChatMessage(content=user_text) self.messages.append(user_msg) diff --git a/src/plyngent/tools/file/edit_lineno.py b/src/plyngent/tools/file/edit_lineno.py index 6b91ac9..9738d88 100644 --- a/src/plyngent/tools/file/edit_lineno.py +++ b/src/plyngent/tools/file/edit_lineno.py @@ -2,7 +2,7 @@ from __future__ import annotations from typing import TYPE_CHECKING -from plyngent.agent import ToolTag, tool +from plyngent.agent import ToolTag, tool, was_lineno_read from plyngent.tools.workspace import resolve_path if TYPE_CHECKING: @@ -85,6 +85,11 @@ async def edit_lineno(path: str, start_line: int, end_line: int, new_content: st target = resolve_path(path) if not target.is_file(): return f"error: not a file: {path}" + if not was_lineno_read(str(target)): + return ( + f"error: read {path!r} first with read_file(path, with_lineno=true) " + "to get correct line numbers" + ) text = target.read_text(encoding="utf-8", errors="replace") lines = text.splitlines(keepends=True) diff --git a/src/plyngent/tools/file/read.py b/src/plyngent/tools/file/read.py index ed1de13..2b9cbd2 100644 --- a/src/plyngent/tools/file/read.py +++ b/src/plyngent/tools/file/read.py @@ -1,6 +1,6 @@ from __future__ import annotations -from plyngent.agent import ToolTag, tool +from plyngent.agent import ToolTag, mark_lineno_read, tool from plyngent.tools.workspace import resolve_path _LINENO_WIDTH = 6 @@ -44,5 +44,6 @@ async def read_file( return "" slice_lines = lines[start:end] if with_lineno: + mark_lineno_read(str(target)) return _format_with_lineno(slice_lines, start_lineno=start + 1) return "".join(slice_lines) diff --git a/tests/test_tools/test_edit_lineno.py b/tests/test_tools/test_edit_lineno.py index 4aa8d65..e13262d 100644 --- a/tests/test_tools/test_edit_lineno.py +++ b/tests/test_tools/test_edit_lineno.py @@ -2,13 +2,21 @@ from __future__ import annotations from pathlib import Path +from plyngent.agent import reset_lineno_tracker from plyngent.tools.file import edit_lineno, read_file, write_file from tests.test_tools.helpers import call_sync +def _read_lineno(path: str) -> str: + """Read file with line numbers so edit_lineno allows the edit.""" + _ = reset_lineno_tracker() + return call_sync(read_file, path, with_lineno=True) + + def test_edit_lineno_replace_middle(workspace: object) -> None: del workspace _ = call_sync(write_file, "a.txt", "one\ntwo\nthree\nfour\n") + _ = _read_lineno("a.txt") out = call_sync(edit_lineno, "a.txt", 2, 3, "TWO\nTHREE\n") assert "replaced lines 2-3" in out assert call_sync(read_file, "a.txt") == "one\nTWO\nTHREE\nfour\n" @@ -17,6 +25,7 @@ def test_edit_lineno_replace_middle(workspace: object) -> None: def test_edit_lineno_delete_range(workspace: object) -> None: del workspace _ = call_sync(write_file, "b.txt", "a\nb\nc\n") + _ = _read_lineno("b.txt") out = call_sync(edit_lineno, "b.txt", 2, 2, "") assert "replaced" in out assert call_sync(read_file, "b.txt") == "a\nc\n" @@ -25,6 +34,7 @@ def test_edit_lineno_delete_range(workspace: object) -> None: def test_edit_lineno_append(workspace: object) -> None: assert isinstance(workspace, Path) _ = call_sync(write_file, "c.txt", "only\n") + _ = _read_lineno("c.txt") out = call_sync(edit_lineno, "c.txt", 2, 2, "more\n") assert "appended" in out assert call_sync(read_file, "c.txt") == "only\nmore\n" @@ -33,5 +43,14 @@ def test_edit_lineno_append(workspace: object) -> None: def test_edit_lineno_invalid(workspace: object) -> None: del workspace _ = call_sync(write_file, "d.txt", "x\n") + _ = _read_lineno("d.txt") assert "start_line" in call_sync(edit_lineno, "d.txt", 0, 1, "y") assert "end_line" in call_sync(edit_lineno, "d.txt", 2, 1, "y") + + +def test_edit_lineno_rejects_without_prior_lineno_read(workspace: object) -> None: + del workspace + _ = call_sync(write_file, "e.txt", "hello\nworld\n") + _ = reset_lineno_tracker() + out = call_sync(edit_lineno, "e.txt", 1, 1, "hi\n") + assert "read" in out and "with_lineno=true" in out