mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 23:47:54 +08:00
core/tools: add per-turn guard requiring lineno read before edit_lineno
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user