feat(engine): the turn runner — prompt builder, sandbox tools, agent loop

Phase 2 of TODO-agent-2: one employee, one turn, real tools in the sandbox.

- engine/prompt.py: the stable, cache-friendly prefix in the plan's order
  (rules → role protocol → toolkit → app profile → memory), a task brief
  capped at 1200 chars, role-specific stage context (spec for the developer,
  the diff only for the reviewer, test output for the tester) and a hard
  prompt cap that drops the oldest tool results first.
- engine/tools.py: READ/WRITE/EDIT/LIST_DIR/SEARCH/RUN executed inside the
  sandbox through SandboxProviderPort, rooted at the worktree's sandbox path.
  Path escapes, .. and .git are refused; git push/remote/config are refused
  with "the platform does git"; RUN clamps its timeout (120s default, 600s
  max) and keeps the tail of its output; every execution writes a tool_calls
  row so the approvals UI shows real work.
- engine/runner.py: the loop — model call, parse tags, execute tool tags,
  feed results back — until a terminal tag, the tool cap, the turn timeout or
  a cancellation. The employee's model is checked against ModelGrantPort
  before the first call (ungranted → error_kind=auth, no spend), errors are
  classified transient vs fatal, and every model call writes usage events
  through the one shared record_model_usage path (extracted from
  WorkforceService so manager chat counts too). Writes engine_turns, run_steps
  and tool_calls.
- engine/events.py: RunFeed batches run.log (≤1 event / 500 ms / stream, ≤200
  lines) and emits run.tool and task.gate through the generic envelope.
- tests: engine_fakes.py (ScriptedModel, FakeSandbox, FakeGrants,
  RecordingMemory) plus 54 new tests in test_engine_runner,
  test_engine_prompt and test_engine_tools.

Suite: 648 passed, ruff clean, mypy clean.
This commit is contained in:
mjalmoudy94
2026-10-14 16:32:52 +03:30
parent 9c0b911a00
commit c6488bfadf
13 changed files with 2712 additions and 19 deletions
+242
View File
@@ -0,0 +1,242 @@
"""Fakes for the engine tests: a scripted model and a scripted sandbox.
Two rules the engine tests insist on: the model never answers by calling a
real provider, and nothing ever executes on the host. Both fakes also *record*
what they were given, because half of what the engine must guarantee is about
prompts (no secret ever reaches one) and commands (git never runs).
``ScriptedModel.complete`` returns the queued responses in order and raises
the queued exceptions, so a test can script "read a file, run the tests, report"
exactly the way a real agent would behave — including the failures.
"""
from __future__ import annotations
from dataclasses import dataclass
from app.application.ports import ChatMessage, CommandResult, ModelHealth, ModelResponse
from app.application.ports.model_grants import GrantedModel
from app.application.ports.sandbox_provider import (
SessionInfo,
SnapshotInfo,
WorkspaceInfo,
)
@dataclass(slots=True, kw_only=True)
class ScriptedReply:
"""One queued model answer: text, or an exception to raise instead."""
content: str = ""
raises: BaseException | None = None
input_tokens: int = 10
output_tokens: int = 5
class ScriptedModel:
"""A ``ModelProviderPort`` that replays a script.
``seen`` holds every call it received (system prompt + messages), so a test
can assert what the model was actually told — the no-secrets rule is only
testable this way.
"""
def __init__(self, replies: list[ScriptedReply] | None = None) -> None:
self.replies: list[ScriptedReply] = list(replies or [])
self.seen: list[dict] = []
self.calls = 0
def queue(self, content: str, **kwargs) -> None:
self.replies.append(ScriptedReply(content=content, **kwargs))
async def complete(
self,
*,
model,
messages: list[ChatMessage],
system: str = "",
temperature: float = 0.2,
max_tokens: int = 2048,
credential=None,
credentials=None,
) -> ModelResponse:
self.calls += 1
self.seen.append(
{
"model": model,
"system": system,
"messages": list(messages),
"credential": credential,
}
)
if not self.replies:
return ModelResponse(
content="[[WORK_REPORT]]پاسخ پیش‌فرض[[/WORK_REPORT]]",
model_internal_name=getattr(model, "internal_name", ""),
input_tokens=1,
output_tokens=1,
)
reply = self.replies.pop(0)
if reply.raises is not None:
raise reply.raises
return ModelResponse(
content=reply.content,
model_internal_name=getattr(model, "internal_name", ""),
input_tokens=reply.input_tokens,
output_tokens=reply.output_tokens,
)
async def health(self) -> ModelHealth:
return ModelHealth(ok=True, detail="scripted")
@property
def prompts(self) -> list[str]:
"""Every system prompt the model was given, in order."""
return [call["system"] for call in self.seen]
def all_text(self) -> str:
"""Everything the model saw, for the "no secret" assertion."""
parts = [call["system"] for call in self.seen]
for call in self.seen:
parts.extend(message.content for message in call["messages"])
return "\n".join(parts)
class FakeSandbox:
"""A ``SandboxProviderPort`` over a dict filesystem and scripted commands.
``commands`` maps the command *prefix* to a queue of results, so a test can
make the tests red once and green afterwards: the engine must notice.
"""
def __init__(
self,
*,
files: dict[str, str] | None = None,
commands: dict[str, list[CommandResult]] | None = None,
) -> None:
self.files: dict[str, str] = dict(files or {})
self.commands: dict[str, list[CommandResult]] = {
key: list(value) for key, value in (commands or {}).items()
}
self.executed: list[dict] = []
self.writes: list[dict] = []
self.reads: list[str] = []
# -- SandboxProviderPort ------------------------------------------------
async def create_workspace(
self, *, organization_id: str, project_id: str, template_id: str, image: str
) -> WorkspaceInfo:
return WorkspaceInfo(external_id="ws-1", status="ready")
async def start_session(self, *, workspace_external_id: str) -> SessionInfo:
return SessionInfo(external_id="sess-1", status="running")
async def run_command(
self,
*,
session_external_id: str,
command: str,
timeout_seconds: int,
max_output_bytes: int,
env: dict[str, str] | None = None,
) -> CommandResult:
self.executed.append(
{
"command": command,
"timeout_seconds": timeout_seconds,
"max_output_bytes": max_output_bytes,
"session": session_external_id,
}
)
for prefix, queue in self.commands.items():
if command.startswith(prefix) and queue:
return queue.pop(0)
return CommandResult(exit_code=0, stdout="ok")
async def read_file(self, *, session_external_id: str, path: str) -> str:
self.reads.append(path)
if path not in self.files:
raise FileNotFoundError(path)
return self.files[path]
async def write_file(self, *, session_external_id: str, path: str, content: str) -> None:
self.files[path] = content
self.writes.append({"path": path, "content": content})
async def snapshot(self, *, session_external_id: str, name: str) -> SnapshotInfo:
return SnapshotInfo(external_id="snap-1", name=name)
async def pause(self, *, session_external_id: str) -> None:
return None
async def resume(self, *, session_external_id: str) -> SessionInfo:
return SessionInfo(external_id="sess-1", status="running")
async def destroy(self, *, session_external_id: str) -> None:
return None
async def get_preview_url(self, *, session_external_id: str, port: int) -> str:
return f"http://127.0.0.1:{port}"
class FakeGrants:
"""A ``ModelGrantPort`` with an explicit allow-list."""
def __init__(self, granted: list[str] | None = None) -> None:
self._granted = list(granted or ["model-strong", "model-fast"])
self.checks: list[str] = []
async def granted_models(self, *, organization_id: str) -> list[GrantedModel]:
return [
GrantedModel(ai_model_id=model_id, tier="balanced")
for model_id in self._granted
]
async def is_granted(self, *, organization_id: str, ai_model_id: str) -> bool:
self.checks.append(ai_model_id)
return ai_model_id in self._granted
def grant(self, model_id: str) -> None:
self._granted.append(model_id)
class RecordingMemory:
"""A ``MemoryPort`` that records its queries and returns fixed context."""
def __init__(self, context: str = "") -> None:
self.context = context
self.queries: list[dict] = []
async def build_context(
self, *, organization_id: str, role: str, query: str, budget_chars: int = 2500
) -> str:
self.queries.append(
{"organization_id": organization_id, "role": role, "query": query}
)
return self.context
async def record(
self,
*,
organization_id: str,
kind: str,
content: str,
source: str,
evidence_task_id: str | None = None,
key: str | None = None,
) -> str | None:
return None
async def on_task_finished(self, *, organization_id: str, task) -> None:
return None
__all__ = [
"FakeGrants",
"FakeSandbox",
"RecordingMemory",
"ScriptedModel",
"ScriptedReply",
]
+154
View File
@@ -0,0 +1,154 @@
"""The prompt builder: a stable prefix, a bounded tail.
What these tests defend:
* the prefix order is the plan's order (rules → protocol → toolkit → app
profile → memory), because that order is what makes prompt caching useful
* the task brief is capped, and the history truncation drops the **oldest tool
results first** — a model that forgets the rules is worse than one that
forgot what it read three calls ago
* nothing secret reaches the prompt: the builder only ever sees the strings it
is handed, and the runner never hands it settings
"""
from __future__ import annotations
from app.application.engine.prompt import (
PLATFORM_RULES,
TOOLKIT_DESCRIPTION,
TurnContext,
TurnMessage,
build_turn_prompt,
)
from app.core.config import settings
from app.domain.workforce.models import RoleKey
from app.domain.workforce.pipeline import Mode, Stage
def _ctx(**kwargs) -> TurnContext:
base = {
"organization_id": "org-1",
"task_id": "task-1",
"project_id": "proj-1",
"task_title": "ساخت صفحهٔ ورود",
"task_description": "توضیح کار",
"stage": Stage.PENDING_IMPLEMENTATION,
"role": RoleKey.DEVELOPER,
"mode": Mode.WORK,
"app_profile_text": "زبان: پایتون\nدستور تست: pytest -q",
"memory_context": "این سازمان تست‌ها را با pytest اجرا می‌کند.",
}
base.update(kwargs)
return TurnContext(**base)
def test_the_prefix_keeps_the_plan_order():
prompt = build_turn_prompt(_ctx())
system = prompt.system
positions = [
system.index(PLATFORM_RULES.splitlines()[0]), # 1. platform rules
system.index("تو برنامه‌نویس هستی"), # 2. the role protocol
system.index(TOOLKIT_DESCRIPTION.splitlines()[0]), # 3. the toolkit
system.index("دستور تست: pytest -q"), # 4. the app profile
system.index("این سازمان"), # 5. org memory
]
assert positions == sorted(positions)
assert prompt.messages[-1].role == "user" # 6/7. brief and history last
def test_the_toolkit_is_only_shown_to_roles_that_have_tools():
developer = build_turn_prompt(_ctx())
manager = build_turn_prompt(
_ctx(role=RoleKey.MANAGER, mode=Mode.TRIAGE, stage=Stage.PENDING_MANAGER_REVIEW)
)
assert "READ_FILE" in developer.system
assert "READ_FILE" not in manager.system
def test_the_task_brief_is_capped():
prompt = build_turn_prompt(_ctx(task_description="x" * 5000))
brief = prompt.messages[-1].content
# The description is cut, the title survives.
assert "ساخت صفحهٔ ورود" in brief
assert len(brief) < 2000
def test_stage_context_only_goes_to_the_role_that_needs_it():
"""A reviewer sees the diff; a developer sees the spec; nobody sees both."""
developer = build_turn_prompt(
_ctx(spec="برنامهٔ برنامه‌ریز", diff_text="diff --git a/x b/x")
)
assert "برنامهٔ برنامه‌ریز" in developer.messages[-1].content
assert "diff --git" not in developer.messages[-1].content
reviewer = build_turn_prompt(
_ctx(
role=RoleKey.REVIEWER,
mode=Mode.WORK,
stage=Stage.PENDING_REVIEW,
spec="برنامهٔ برنامه‌ریز",
diff_text="diff --git a/x b/x",
)
)
assert "diff --git" in reviewer.messages[-1].content
assert "برنامهٔ برنامه‌ریز" not in reviewer.messages[-1].content
def test_history_drops_the_oldest_results_first():
big = "y" * 4000
history = [
TurnMessage(kind="result", text=big), # oldest: dropped first
TurnMessage(kind="result", text=big),
TurnMessage(kind="result", text="نتیجهٔ تازه و مهم"),
]
settings.engine_prompt_char_cap = 9000 # keep the prefix small
try:
prompt = build_turn_prompt(_ctx(history=list(history)))
finally:
settings.engine_prompt_char_cap = 24000
joined = "\n".join(m.content for m in prompt.messages)
assert "نتیجهٔ تازه و مهم" in joined
assert joined.count("y" * 4000) <= 1
assert prompt.total_chars <= 9000 + 4000
def test_a_tiny_cap_still_keeps_the_rules_and_the_brief():
settings.engine_prompt_char_cap = 3000
try:
prompt = build_turn_prompt(
_ctx(history=[TurnMessage(kind="result", text="z" * 20000)])
)
finally:
settings.engine_prompt_char_cap = 24000
# The prefix and the brief survive; only history space is sacrificed.
assert "قوانین پلتفرم" in prompt.system
assert "ساخت صفحهٔ ورود" in prompt.messages[-1].content
def test_the_prompt_never_contains_the_settings(monkeypatch):
monkeypatch.setattr(settings, "jwt_secret", "top-secret")
monkeypatch.setattr(settings, "database_url", "sqlite:///private.db")
prompt = build_turn_prompt(
_ctx(history=[TurnMessage(kind="result", text="نتیجه")])
)
text = prompt.system + "\n" + "\n".join(m.content for m in prompt.messages)
assert "top-secret" not in text
assert "private.db" not in text
assert PLATFORM_RULES.splitlines()[0] in text
def test_manager_modes_get_their_own_protocol_hint():
triage = build_turn_prompt(
_ctx(role=RoleKey.MANAGER, mode=Mode.TRIAGE, stage=Stage.PENDING_MANAGER_REVIEW)
)
final = build_turn_prompt(
_ctx(role=RoleKey.MANAGER, mode=Mode.FINAL, stage=Stage.PENDING_FINAL_REVIEW)
)
assert "ASSIGN" in triage.system
assert "ASSIGN" not in final.system
assert "COMPLETE" in final.system
def test_the_toolkit_description_lists_every_tool_tag():
for tool in ("READ_FILE", "WRITE_FILE", "EDIT_FILE", "LIST_DIR", "SEARCH", "RUN"):
assert tool in TOOLKIT_DESCRIPTION
+551
View File
@@ -0,0 +1,551 @@
"""The turn runner: one employee turn, with a scripted model and a fake sandbox.
What these tests defend:
* the write → run → report loop really works: a file is written inside the
sandbox, a command runs there, and the report ends the turn
* the sandbox is the boundary: a path escape, a ``.git`` path and a git
command are refused *inside the tool*, with a message the model can read
* a turn is bounded: the tool-call cap and the turn timeout both stop it
* errors are classified, so a rate limit is retried later and an ungranted
model is not
* every model call writes usage events — including the ones a manager chat
makes, which agybot forgot
* no secret from the settings ever reaches a prompt
"""
from __future__ import annotations
from types import SimpleNamespace
import pytest
import pytest_asyncio
from sqlalchemy import select
from app.application.engine.runner import (
TurnRequest,
TurnRunner,
classify_error,
)
from app.application.ports import CommandResult
from app.core.config import settings
from app.core.errors import ProviderUnavailableError
from app.domain.ai_control.models import AiModel
from app.domain.base import new_id
from app.domain.workforce.engine import ErrorKind
from app.domain.workforce.models import RoleKey
from app.domain.workforce.pipeline import Mode, Stage
from app.infrastructure.db.models import (
AiProviderRow,
EmployeeRow,
OrganizationRow,
RunRow,
TaskRow,
UsageEventRow,
UserRow,
)
from app.infrastructure.repositories.ai_control import (
AiModelRepository,
ModelCredentialRepository,
RateCardRepository,
)
from app.infrastructure.repositories.audit import UsageLedgerRepository
from app.infrastructure.repositories.engine import (
EngineTurnRepository,
PipelineStateRepository,
)
from app.infrastructure.repositories.workforce import (
EmployeeRepository,
RunRepository,
RunStepRepository,
TaskRepository,
ToolCallRepository,
)
from tests.engine_fakes import (
FakeGrants,
FakeSandbox,
RecordingMemory,
ScriptedModel,
ScriptedReply,
)
DEVELOPER = {
"role": RoleKey.DEVELOPER,
"mode": Mode.WORK,
"stage": Stage.PENDING_IMPLEMENTATION,
}
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest_asyncio.fixture
async def _db(session_factory, seeded):
"""A tenant, a task, a run, a developer employee and a granted model."""
async with session_factory() as session:
user = UserRow(email=f"runner-{new_id()[:8]}@example.com", password_hash="x")
session.add(user)
await session.flush()
org = OrganizationRow(name="سازمان اجرا", slug=f"run-{user.id[:6]}", owner_user_id=user.id)
session.add(org)
await session.flush()
# The seeded provider row: the model must hang off a real one.
provider_id = (await session.execute(select(AiProviderRow.id).limit(1))).scalar_one()
model = await AiModelRepository(session).create(
AiModel(
provider_id=provider_id,
internal_name="scripted-model",
display_name_for_admin="مدل فیلم‌نویسی‌شده",
capabilities=["text", "code"],
is_active=True,
)
)
task = TaskRow(organization_id=org.id, title="ساخت صفحهٔ ورود")
session.add(task)
await session.flush()
run = RunRow(task_id=task.id)
session.add(run)
employee = EmployeeRow(organization_id=org.id, name="سارا", role_key="developer")
session.add(employee)
await session.commit()
yield SimpleNamespace(
session_factory=session_factory,
organization_id=org.id,
task_id=task.id,
run_id=run.id,
employee_id=employee.id,
model_id=model.id,
)
def _runner(session, *, model, sandbox, grants, feed=None) -> TurnRunner:
return TurnRunner(
provider=model,
models=AiModelRepository(session),
credentials=ModelCredentialRepository(session),
rate_cards=RateCardRepository(session),
ledger=UsageLedgerRepository(session),
grants=grants,
sandbox=sandbox,
memory=RecordingMemory(context="این سازمان تست‌ها را با pytest اجرا می‌کند."),
turns=EngineTurnRepository(session),
run_steps=RunStepRepository(session),
tool_calls=ToolCallRepository(session),
runs=RunRepository(session),
states=PipelineStateRepository(session),
feed=feed,
)
async def _request(db, *, sandbox_root="/workspace", session_id="sess-1", **kwargs) -> TurnRequest:
async with db.session_factory() as session:
task = await TaskRepository(session).get(db.task_id)
employee = await EmployeeRepository(session).get(db.employee_id)
return TurnRequest(
organization_id=db.organization_id,
task=task,
run_id=db.run_id,
employee=employee,
model_id=db.model_id,
sandbox_root=sandbox_root,
session_external_id=session_id,
**kwargs,
)
# ---------------------------------------------------------------------------
# The loop
# ---------------------------------------------------------------------------
async def test_write_then_run_then_report_ends_the_turn(_db, session_factory):
"""The whole agent loop: one file, one command, one report."""
scripted = ScriptedModel(
[
ScriptedReply(
content='[[WRITE_FILE path="app/login.py"]]print("login")[[/WRITE_FILE]]'
),
ScriptedReply(content='[[RUN cmd="pytest -q" timeout="30"/]]'),
ScriptedReply(content="[[WORK_REPORT]]صفحهٔ ورود نوشده و تست‌ها سبز شد.[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox(
commands={"pytest": [CommandResult(exit_code=0, stdout="3 passed")]},
)
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
request = await _request(_db, **DEVELOPER)
result = await runner.run_turn(request)
await session.commit()
assert result.ok is True
assert result.model_calls == 3
assert result.tool_calls == 2
assert result.tags[0].name.value == "WORK_REPORT"
# The file really was written inside the sandbox, rooted at the worktree.
assert sandbox.writes == [{"path": "/workspace/app/login.py", "content": 'print("login")'}]
assert [call["command"] for call in sandbox.executed] == ["pytest -q"]
assert scripted.calls == 3
async def test_tool_results_feed_back_into_the_next_prompt(_db, session_factory):
"""The model must see what its own command printed."""
scripted = ScriptedModel(
[
ScriptedReply(content='[[RUN cmd="pytest -q"/]]'),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox(
commands={"pytest": [CommandResult(exit_code=1, stdout="1 failed", stderr="boom")]}
)
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is True
last_prompt = scripted.seen[-1]
feedback = "\n".join(m.content for m in last_prompt["messages"])
assert "1 failed" in feedback and "boom" in feedback
async def test_each_tool_call_writes_a_tool_calls_row(_db, session_factory):
"""The approvals UI finally sees real executions, not intentions."""
scripted = ScriptedModel(
[
ScriptedReply(content='[[READ_FILE path="app/x.py"/]]'),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox(files={"/workspace/app/x.py": "print(1)"})
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
calls = await ToolCallRepository(session).list_for_run(_db.run_id)
assert len(calls) == 1
assert calls[0].tool_key == "sandbox.read_file"
assert calls[0].status.value == "succeeded"
assert calls[0].result_json["ok"] is True
async def test_a_failed_tool_is_recorded_as_failed(_db, session_factory):
scripted = ScriptedModel(
[
ScriptedReply(content='[[READ_FILE path="app/missing.py"/]]'),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox()
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
calls = await ToolCallRepository(session).list_for_run(_db.run_id)
assert result.ok is True
assert calls[0].status.value == "failed"
async def test_the_turn_leaves_an_engine_turns_row(_db, session_factory):
scripted = ScriptedModel([ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]")])
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=FakeSandbox(), grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
turns = await EngineTurnRepository(session).list_for_task(_db.task_id)
assert result.turn_id
assert len(turns) == 1
assert turns[0].ok is True
assert turns[0].role is RoleKey.DEVELOPER
assert turns[0].mode is Mode.WORK
assert turns[0].stage is Stage.PENDING_IMPLEMENTATION
assert turns[0].prompt_chars > 0
assert "[[WORK_REPORT]]" in turns[0].reply_excerpt
assert "تمام" in turns[0].reply_excerpt
assert turns[0].tags == {"WORK_REPORT": 1}
assert turns[0].finished_at is not None
# ---------------------------------------------------------------------------
# The sandbox boundary
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"path",
["../../etc/passwd", "/etc/passwd", ".git/config", "app/../../secret"],
)
async def test_a_path_escape_is_refused(_db, session_factory, path):
scripted = ScriptedModel(
[
ScriptedReply(content=f'[[READ_FILE path="{path}"/]]'),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox(files={"/etc/passwd": "root"})
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
# The refusal is fed back to the model, and nothing was read.
assert result.ok is True
assert sandbox.reads == []
assert "مجاز نیست" in scripted.all_text()
@pytest.mark.parametrize(
"command", ["git push origin dev", "git remote add x y", "git config user.email a@b.c"]
)
async def test_the_model_cannot_run_git(_db, session_factory, command):
scripted = ScriptedModel(
[
ScriptedReply(content=f'[[RUN cmd="{command}"/]]'),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox()
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is True
assert sandbox.executed == []
assert "پلتفرم" in scripted.all_text()
async def test_edit_file_applies_a_search_replace_block(_db, session_factory):
scripted = ScriptedModel(
[
ScriptedReply(
content=(
"[[EDIT_FILE path=\"app/x.py\"]]<<<OLD\nold line\n===\n"
"new line\nNEW>>>[[/EDIT_FILE]]"
)
),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox(files={"/workspace/app/x.py": "before\nold line\nafter"})
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is True
assert sandbox.files["/workspace/app/x.py"] == "before\nnew line\nafter"
async def test_edit_file_fails_when_the_old_text_is_not_unique(_db, session_factory):
scripted = ScriptedModel(
[
ScriptedReply(
content=(
"[[EDIT_FILE path=\"app/x.py\"]]<<<OLD\nline\n===\n"
"other\nNEW>>>[[/EDIT_FILE]]"
)
),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox(files={"/workspace/app/x.py": "line\nline"})
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is True
assert sandbox.files["/workspace/app/x.py"] == "line\nline"
assert "دقیقاً یک بار" in scripted.all_text()
# ---------------------------------------------------------------------------
# Bounds: the tool cap and the timeout
# ---------------------------------------------------------------------------
async def test_the_tool_call_cap_stops_a_runaway_turn(_db, session_factory, monkeypatch):
monkeypatch.setattr(settings, "engine_max_tool_calls_per_turn", 2)
scripted = ScriptedModel([ScriptedReply(content='[[LIST_DIR path="."/]]') for _ in range(5)])
sandbox = FakeSandbox()
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is False
assert result.incomplete is True
assert result.tool_calls == 2
assert result.error_kind == "protocol"
async def test_the_turn_timeout_stops_a_slow_loop(_db, session_factory, monkeypatch):
monkeypatch.setattr(settings, "engine_turn_timeout_seconds", 0)
scripted = ScriptedModel([ScriptedReply(content='[[LIST_DIR path="."/]]')] * 3)
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=FakeSandbox(), grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is False
assert result.error_kind == "timeout"
async def test_a_cancelled_run_stops_between_steps(_db, session_factory):
scripted = ScriptedModel(
[
ScriptedReply(content='[[LIST_DIR path="."/]]'),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=FakeSandbox(), grants=FakeGrants([_db.model_id])
)
request = await _request(_db, **DEVELOPER)
# The hook cancels after the first model call has been made.
calls = {"n": 0}
async def should_cancel() -> bool:
calls["n"] += 1
return calls["n"] > 1
request.should_cancel = should_cancel
result = await runner.run_turn(request)
await session.commit()
assert result.ok is False
assert result.error_kind == "cancelled"
# ---------------------------------------------------------------------------
# Errors: transient versus fatal
# ---------------------------------------------------------------------------
async def test_a_rate_limit_is_transient_and_reported(_db, session_factory):
scripted = ScriptedModel([ScriptedReply(raises=ProviderUnavailableError("429 rate limited"))])
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=FakeSandbox(), grants=FakeGrants([_db.model_id])
)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is False
assert result.error_kind == "rate_limit"
assert result.retryable is True
async def test_an_ungranted_model_is_refused_before_any_call(_db, session_factory):
scripted = ScriptedModel([ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]")])
grants = FakeGrants(["some-other-model"])
async with session_factory() as session:
runner = _runner(session, model=scripted, sandbox=FakeSandbox(), grants=grants)
result = await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
assert result.ok is False
assert result.error_kind == "auth"
assert result.retryable is False
# The model was never called: an ungranted employee spends nothing.
assert scripted.calls == 0
assert grants.checks == [_db.model_id]
def test_error_classification():
assert classify_error(ProviderUnavailableError("429 rate limit")) is ErrorKind.RATE_LIMIT
assert classify_error(ProviderUnavailableError("auth failed")) is ErrorKind.AUTH
assert classify_error(ProviderUnavailableError("connection refused")) is ErrorKind.NETWORK
assert classify_error(TimeoutError()) is ErrorKind.TIMEOUT
assert classify_error(ProviderUnavailableError("boom")) is ErrorKind.UNAVAILABLE
# ---------------------------------------------------------------------------
# Billing and secrets
# ---------------------------------------------------------------------------
async def test_every_model_call_writes_usage_events(_db, session_factory):
"""agybot forgot the manager chat; the engine does not."""
scripted = ScriptedModel(
[
ScriptedReply(content='[[LIST_DIR path="."/]]', input_tokens=11, output_tokens=7),
ScriptedReply(
content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]", input_tokens=13, output_tokens=9
),
]
)
from sqlalchemy import select
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=FakeSandbox(), grants=FakeGrants([_db.model_id])
)
await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
rows = (
(await session.execute(select(UsageEventRow).where(UsageEventRow.run_id == _db.run_id)))
.scalars()
.all()
)
assert len(rows) == 4 # two calls, input and output each
assert sum(r.quantity for r in rows if r.unit == "model_input_tokens") == 24
assert sum(r.quantity for r in rows if r.unit == "model_output_tokens") == 16
async def test_no_secret_from_settings_reaches_the_model(_db, session_factory, monkeypatch):
"""The prompt is built by the engine; the settings stay on the server."""
monkeypatch.setattr(settings, "jwt_secret", "super-secret-jwt-value")
monkeypatch.setattr(settings, "database_url", "postgresql+asyncpg://u:p@h/db")
scripted = ScriptedModel(
[
ScriptedReply(content='[[READ_FILE path="app/x.py"/]]'),
ScriptedReply(content="[[WORK_REPORT]]تمام[[/WORK_REPORT]]"),
]
)
sandbox = FakeSandbox(files={"/workspace/app/x.py": "print(1)"})
async with session_factory() as session:
runner = _runner(
session, model=scripted, sandbox=sandbox, grants=FakeGrants([_db.model_id])
)
await runner.run_turn(await _request(_db, **DEVELOPER))
await session.commit()
seen = scripted.all_text()
assert "super-secret-jwt-value" not in seen
assert "postgresql+asyncpg://u:p@h/db" not in seen
# The model still got the rules and the role it plays.
assert "قوانین پلتفرم" in seen
assert "WORK_REPORT" in seen
+184
View File
@@ -0,0 +1,184 @@
"""The toolkit's guards: path safety, the git refusal and bounded commands.
These are the checks that run *before* anything reaches the sandbox, so they
are tested as pure functions wherever possible — a path bug must be
reproducible without a container.
"""
from __future__ import annotations
import pytest
from app.application.engine.protocol import Tag, TagName, edit_blocks_of
from app.application.engine.tools import (
RUN_DEFAULT_TIMEOUT,
RUN_MAX_TIMEOUT,
ToolExecutor,
UnsafePathError,
safe_join,
tool_key_for,
)
from tests.engine_fakes import FakeSandbox
ROOT = "/workspace/.worktrees/task-1"
# ---------------------------------------------------------------------------
# Path safety
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"relative",
["app/main.py", "./app/main.py", "app/deep/nested/file.py", "a.txt"],
)
def test_a_normal_relative_path_is_allowed(relative):
assert safe_join(ROOT, relative) == f"{ROOT}/{relative.lstrip('./')}".replace(
"/./", "/"
)
@pytest.mark.parametrize(
"relative",
["/etc/passwd", "~/.ssh/id_rsa", "../../etc/passwd", "app/../../secrets", ".git/config"],
)
def test_an_escaping_path_is_refused(relative):
with pytest.raises(UnsafePathError):
safe_join(ROOT, relative)
def test_an_empty_path_is_refused():
with pytest.raises(UnsafePathError):
safe_join(ROOT, " ")
def test_the_git_directory_is_never_readable():
with pytest.raises(UnsafePathError):
safe_join(ROOT, ".git/objects/ab/cdef")
# ---------------------------------------------------------------------------
# The git refusal
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"command",
["git push origin dev", "git remote add origin x", "git config user.email a@b.c"],
)
async def test_git_commands_are_refused_before_the_sandbox(command):
sandbox = FakeSandbox()
executor = ToolExecutor(
sandbox=sandbox, session_external_id="sess-1", root=ROOT
)
outcome = await executor.execute(
Tag(name=TagName.RUN, attrs={"cmd": command}, body="")
)
assert outcome.ok is False
assert outcome.refused is True
assert "پلتفرم" in outcome.output
assert sandbox.executed == []
async def test_other_git_commands_are_allowed_through():
"""Only push/remote/config are refused: the platform owns the rest."""
sandbox = FakeSandbox(commands={"git": []})
executor = ToolExecutor(
sandbox=sandbox, session_external_id="sess-1", root=ROOT
)
outcome = await executor.execute(
Tag(name=TagName.RUN, attrs={"cmd": "git status --porcelain -uall"}, body="")
)
assert outcome.ok is True
assert len(sandbox.executed) == 1
# ---------------------------------------------------------------------------
# Bounds
# ---------------------------------------------------------------------------
async def test_run_clamps_the_timeout_to_the_ceiling():
sandbox = FakeSandbox()
executor = ToolExecutor(sandbox=sandbox, session_external_id="sess-1", root=ROOT)
await executor.execute(
Tag(name=TagName.RUN, attrs={"cmd": "sleep 1", "timeout": "99999"}, body="")
)
assert sandbox.executed[0]["timeout_seconds"] == RUN_MAX_TIMEOUT
async def test_run_defaults_to_a_sane_timeout():
sandbox = FakeSandbox()
executor = ToolExecutor(sandbox=sandbox, session_external_id="sess-1", root=ROOT)
await executor.execute(Tag(name=TagName.RUN, attrs={"cmd": "ls"}, body=""))
assert sandbox.executed[0]["timeout_seconds"] == RUN_DEFAULT_TIMEOUT
async def test_a_huge_output_is_truncated_to_its_tail():
sandbox = FakeSandbox(
commands={"noisy": [type("R", (), {"exit_code": 0, "stdout": "x" * 30000, "stderr": ""})()]}
)
executor = ToolExecutor(
sandbox=sandbox, session_external_id="sess-1", root=ROOT, max_output_chars=1000
)
outcome = await executor.execute(
Tag(name=TagName.RUN, attrs={"cmd": "noisy"}, body="")
)
assert outcome.ok is True
assert len(outcome.output) <= 1100
assert outcome.output.endswith("x" * 100)
# ---------------------------------------------------------------------------
# The tools themselves
# ---------------------------------------------------------------------------
async def test_write_then_read_round_trips_inside_the_root():
sandbox = FakeSandbox()
executor = ToolExecutor(sandbox=sandbox, session_external_id="sess-1", root=ROOT)
written = await executor.execute(
Tag(name=TagName.WRITE_FILE, attrs={"path": "app/a.py"}, body="print(1)")
)
read = await executor.execute(
Tag(name=TagName.READ_FILE, attrs={"path": "app/a.py"}, body="")
)
assert written.ok and read.ok
assert "print(1)" in read.output
assert sandbox.writes[0]["path"] == f"{ROOT}/app/a.py"
async def test_list_dir_and_search_run_quoted_commands():
sandbox = FakeSandbox()
executor = ToolExecutor(sandbox=sandbox, session_external_id="sess-1", root=ROOT)
await executor.execute(Tag(name=TagName.LIST_DIR, attrs={"path": "app"}, body=""))
await executor.execute(
Tag(name=TagName.SEARCH, attrs={"pattern": "TODO", "path": "app"}, body="")
)
assert sandbox.executed[0]["command"].startswith("ls -la -- ")
assert "grep -RIn -e TODO" in sandbox.executed[1]["command"]
async def test_a_missing_file_reports_not_found_without_raising():
sandbox = FakeSandbox()
executor = ToolExecutor(sandbox=sandbox, session_external_id="sess-1", root=ROOT)
outcome = await executor.execute(
Tag(name=TagName.READ_FILE, attrs={"path": "app/none.py"}, body="")
)
assert outcome.ok is False
assert "none.py" in outcome.output
def test_tool_keys_map_onto_the_existing_catalog():
assert tool_key_for(Tag(name=TagName.READ_FILE, attrs={}, body="")) == (
"sandbox.read_file"
)
assert tool_key_for(Tag(name=TagName.RUN, attrs={}, body="")) == (
"sandbox.run_command"
)
def test_edit_blocks_refuse_a_malformed_body():
tag = Tag(name=TagName.EDIT_FILE, attrs={"path": "a"}, body="just text")
with pytest.raises(ValueError):
edit_blocks_of(tag)