can1357
a9b7df8821
- Added `_normalize_host_tool_event` to `RpcClient` to remap transport tool events to executed host tools. - Updated worker tool-end handling to verify `is_error` is false before marking terminal actions complete. - Added test coverage in `test_client.py` and `test_worker.py` for host tool event normalization and errored review submission behavior.
833 lines
31 KiB
Python
833 lines
31 KiB
Python
"""Per-task RpcClient driver.
|
|
|
|
The orchestrator calls `run_task(...)` from within an asyncio loop. The
|
|
function spins up `RpcClient` on a worker thread, drives the kickoff/follow-up
|
|
prompt, and returns when the agent emits `agent_end`.
|
|
|
|
Host tools call back into the orchestrator's GitHub client and DB. Because the
|
|
RpcClient runs in its own subprocess and the host-tool callbacks are dispatched
|
|
on the RpcClient's stdout-reader thread, the callbacks block until coroutines
|
|
scheduled onto the parent loop complete (`asyncio.run_coroutine_threadsafe`).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import grp
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from omp_rpc import (
|
|
MessageUpdateEvent,
|
|
RpcClient,
|
|
RpcError,
|
|
RpcProcessExitError,
|
|
ToolExecutionEndEvent,
|
|
)
|
|
|
|
from robomp import host_tools, persona, pragmas
|
|
from robomp.cancellation import register_cancel_hook, unregister_cancel_hook
|
|
from robomp.config import Settings
|
|
from robomp.db import Database, issue_key
|
|
from robomp.git_ops import DirtyState, inspect_dirty_state
|
|
from robomp.github_backend import GitHubBackend
|
|
from robomp.github_client import CommentInfo, IssueInfo, PullRequestInfo, RepoInfo
|
|
from robomp.host_tools import AbortController, ToolBindings, _git_identity_env
|
|
from robomp.natives_cache import NativesCache
|
|
from robomp.natives_cache import compute_key as natives_compute_key
|
|
from robomp.sandbox import GitTransport, Workspace, _prepare_slot_runtime_env, _safe_directory_env
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class TaskInputs:
|
|
"""Common context shared by every task type."""
|
|
|
|
settings: Settings
|
|
db: Database
|
|
github: GitHubBackend
|
|
git_transport: GitTransport
|
|
repo: RepoInfo
|
|
issue: IssueInfo
|
|
workspace: Workspace
|
|
delivery_id: str
|
|
attempts: int = 0
|
|
slot_uid: int | None = None
|
|
natives_cache: NativesCache | None = None
|
|
|
|
|
|
@dataclass(slots=True, frozen=True)
|
|
class ThreadMessage:
|
|
"""One entry in the conversation a directive carries to the agent."""
|
|
|
|
kind: str # issue_body | pr_body | comment | review_comment | review
|
|
author: str
|
|
body: str
|
|
created_at: str
|
|
path: str | None = None # review_comment only
|
|
line: int | None = None # review_comment only
|
|
state: str | None = None # review only (APPROVED / CHANGES_REQUESTED / COMMENTED)
|
|
|
|
|
|
@dataclass(slots=True, frozen=True)
|
|
class DirectiveInfo:
|
|
"""A maintainer's `@bot` mention captured as an authoritative instruction.
|
|
|
|
`thread` is the full conversation context (issue/PR body + every prior
|
|
comment + every review) up to the moment the directive fired.
|
|
"""
|
|
|
|
body: str
|
|
author: str
|
|
thread: tuple[ThreadMessage, ...] = ()
|
|
pragmas: tuple[tuple[str, str], ...] = ()
|
|
authorizes_impl: bool = False
|
|
|
|
|
|
def _resolve_pragma_overrides(
|
|
directive: DirectiveInfo | None,
|
|
settings: Settings,
|
|
) -> tuple[str | None, pragmas.ThinkingLevel | None]:
|
|
"""Return `(model_override, thinking_override)` for the current directive.
|
|
|
|
`None` for either means "no override, use the settings default". Aliases
|
|
that don't match anything in the pool / level set are dropped (caller logs
|
|
the discard at the callsite that has access to issue_key).
|
|
"""
|
|
if directive is None or not directive.pragmas:
|
|
return None, None
|
|
model_value = pragmas.pragma_value(directive.pragmas, "model")
|
|
thinking_value = pragmas.pragma_value(directive.pragmas, "thinking")
|
|
model_override = pragmas.resolve_model_alias(model_value, settings.model_pool) if model_value else None
|
|
thinking_override = pragmas.resolve_thinking_level(thinking_value) if thinking_value else None
|
|
return model_override, thinking_override
|
|
|
|
|
|
_SCRUBBED_ENV_KEYS: tuple[str, ...] = (
|
|
# Secrets that MUST NOT reach the agent subprocess; an agent with the
|
|
# `bash` tool could otherwise `printenv` them out of roboomp's env.
|
|
"GITHUB_TOKEN",
|
|
"GITHUB_WEBHOOK_SECRET",
|
|
"ROBOMP_REPLAY_TOKEN",
|
|
"ROBOMP_GH_PROXY_HMAC_KEY",
|
|
)
|
|
|
|
_AGENT_HOME = Path("/srv/agent-home")
|
|
_AGENT_HOME_STAGE = Path("/srv/agent-home-stage")
|
|
|
|
|
|
def _stage_agent_home() -> None:
|
|
"""Copy late-appearing staged agent config into the runtime HOME."""
|
|
if not _AGENT_HOME_STAGE.exists():
|
|
return
|
|
|
|
for rel in (Path(".agent"), Path(".omp/agent")):
|
|
src = _AGENT_HOME_STAGE / rel
|
|
if not src.exists():
|
|
continue
|
|
|
|
dst = _AGENT_HOME / rel
|
|
try:
|
|
if os.path.lexists(dst):
|
|
if dst.is_dir() and not dst.is_symlink():
|
|
shutil.rmtree(dst)
|
|
else:
|
|
dst.unlink()
|
|
dst.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copytree(src, dst, dirs_exist_ok=True)
|
|
except OSError as exc:
|
|
log.warning("Failed to stage agent home path %s: %s", rel, exc)
|
|
|
|
if not _AGENT_HOME.exists():
|
|
return
|
|
|
|
chown_to_root = os.geteuid() == 0
|
|
for root, dirs, files in os.walk(_AGENT_HOME):
|
|
root_path = Path(root)
|
|
if root_path == _AGENT_HOME / ".omp":
|
|
# ~/.omp/run is slot-writable daemon presence state, not template
|
|
# config; keep it out of the read-only normalization below.
|
|
dirs[:] = [d for d in dirs if d != "run"]
|
|
try:
|
|
root_path.chmod(0o755)
|
|
if chown_to_root:
|
|
os.chown(root_path, 0, 0)
|
|
except OSError as exc:
|
|
log.warning("Failed to normalize agent home directory %s: %s", root_path, exc)
|
|
|
|
for name in dirs:
|
|
path = root_path / name
|
|
try:
|
|
path.chmod(0o755)
|
|
if chown_to_root:
|
|
os.chown(path, 0, 0)
|
|
except OSError as exc:
|
|
log.warning("Failed to normalize agent home directory %s: %s", path, exc)
|
|
|
|
for name in files:
|
|
path = root_path / name
|
|
try:
|
|
path.chmod(0o644)
|
|
if chown_to_root:
|
|
os.chown(path, 0, 0)
|
|
except OSError as exc:
|
|
log.warning("Failed to normalize agent home file %s: %s", path, exc)
|
|
|
|
|
|
def _ensure_agent_run_dir() -> None:
|
|
"""Keep ``~/.omp/run`` writable by every sandbox slot.
|
|
|
|
omp registers daemon project presence under ``~/.omp/run`` at startup,
|
|
nesting per-project dirs (``daemons/<hash>/clients``) that any slot user
|
|
must be able to create or enter regardless of which slot made them first.
|
|
The tree stays group ``omp``, setgid, group-writable; slot subprocesses
|
|
spawn with umask 0002 so their entries inherit group write.
|
|
"""
|
|
if os.geteuid() != 0:
|
|
return
|
|
run_dir = _AGENT_HOME / ".omp" / "run"
|
|
try:
|
|
gid = grp.getgrnam("omp").gr_gid
|
|
except KeyError:
|
|
return
|
|
try:
|
|
run_dir.mkdir(parents=True, exist_ok=True)
|
|
for root, _dirs, files in os.walk(run_dir):
|
|
root_path = Path(root)
|
|
os.chown(root_path, -1, gid)
|
|
root_path.chmod(0o2770)
|
|
for name in files:
|
|
file_path = root_path / name
|
|
os.chown(file_path, -1, gid)
|
|
file_path.chmod(0o660)
|
|
except OSError as exc:
|
|
log.warning("Failed to prepare agent run dir %s: %s", run_dir, exc)
|
|
|
|
|
|
def _build_extra_env(settings: Settings) -> dict[str, str]:
|
|
"""Build the env overlay passed to the omp subprocess.
|
|
|
|
`omp_rpc` merges this dict on top of `os.environ`, so overlaying empty
|
|
strings for the sensitive keys is what actually masks them in the
|
|
child — `del` on the parent's env would not help us here.
|
|
"""
|
|
del settings # kept for future hooks (model-specific env, etc.)
|
|
_stage_agent_home()
|
|
_ensure_agent_run_dir()
|
|
env = dict.fromkeys(_SCRUBBED_ENV_KEYS, "")
|
|
if _AGENT_HOME.is_dir():
|
|
env["HOME"] = str(_AGENT_HOME)
|
|
return env
|
|
|
|
|
|
_TERMINAL_TRIAGE_TOOLS: frozenset[str] = frozenset({"gh_open_pr", "mark_unable_to_reproduce", "abort_task"})
|
|
_TERMINAL_REVIEW_TOOLS: frozenset[str] = frozenset({"submit_pr_review", "abort_task"})
|
|
_PR_REQUIRING_CLASSIFICATIONS: frozenset[str] = frozenset({"bug", "documentation"})
|
|
|
|
|
|
def _needs_completion_reminder(
|
|
*,
|
|
task_kind: str,
|
|
inputs: TaskInputs,
|
|
bindings: ToolBindings,
|
|
tools_called: set[str],
|
|
) -> bool:
|
|
"""True iff a task turn ended before reaching its terminal tool."""
|
|
if bindings.abort is not None and bindings.abort.triggered:
|
|
return False
|
|
if task_kind == "review_pr":
|
|
return not (tools_called & _TERMINAL_REVIEW_TOOLS)
|
|
if task_kind != "triage_issue":
|
|
return False
|
|
row = inputs.db.get_issue(bindings.issue_key)
|
|
if row is None or row.classification not in _PR_REQUIRING_CLASSIFICATIONS:
|
|
return False
|
|
return not (tools_called & _TERMINAL_TRIAGE_TOOLS)
|
|
|
|
|
|
def _probe_workspace_dirty(workspace: Workspace, slot_uid: int | None) -> DirtyState:
|
|
"""Return the workspace's dirty state, swallowing inspection errors.
|
|
|
|
`inspect_dirty_state` already swallows individual git failures, but the
|
|
function itself can still raise if e.g. the workspace was wiped between
|
|
`_drive_turn` and the post-turn check. The reminder loop treats any
|
|
exception as "clean enough" so a corrupted workspace doesn't pin the
|
|
agent in a reminder loop.
|
|
"""
|
|
try:
|
|
return inspect_dirty_state(
|
|
workspace.repo_dir,
|
|
slot_uid=slot_uid,
|
|
safe_directory=workspace.repo_dir,
|
|
)
|
|
except Exception as exc: # noqa: BLE001 — best-effort post-turn probe
|
|
log.debug("workspace dirty probe failed", extra={"error": str(exc)})
|
|
return DirtyState(uncommitted=0, unpushed=0, summary="")
|
|
|
|
|
|
def _drive_turn(
|
|
client: RpcClient,
|
|
initial_prompt: str,
|
|
*,
|
|
task_kind: str,
|
|
inputs: TaskInputs,
|
|
bindings: ToolBindings,
|
|
tools_called: set[str],
|
|
) -> Any:
|
|
"""Run the initial prompt and, if the agent stopped early, send reminders.
|
|
|
|
Returns the final `Turn` (last `prompt_and_wait` result), or `None` when
|
|
the agent intentionally pulled the plug via `abort_task`.
|
|
"""
|
|
settings = inputs.settings
|
|
max_reminders = settings.task_completion_max_reminders
|
|
|
|
def _run(prompt: str) -> Any:
|
|
try:
|
|
return client.prompt_and_wait(prompt, timeout=settings.task_timeout_seconds)
|
|
except (RpcError, RpcProcessExitError):
|
|
# Did the agent intentionally pull the plug via `abort_task`?
|
|
# If so, swallow — the abort path is a clean exit, not a
|
|
# failure that should surface in the dashboard or trigger
|
|
# a comment to the reporter. Anything else propagates.
|
|
if bindings.abort is not None and bindings.abort.triggered:
|
|
log.info(
|
|
"rpc_aborted_by_tool",
|
|
extra={"issue": bindings.issue_key, "task": task_kind, "reason": bindings.abort.reason},
|
|
)
|
|
return None
|
|
raise
|
|
|
|
turn = _run(initial_prompt)
|
|
if turn is None:
|
|
return None
|
|
|
|
reminders_used = 0
|
|
while reminders_used < max_reminders:
|
|
needs_completion = _needs_completion_reminder(
|
|
task_kind=task_kind, inputs=inputs, bindings=bindings, tools_called=tools_called
|
|
)
|
|
if not needs_completion:
|
|
if task_kind == "review_pr":
|
|
break
|
|
dirty = _probe_workspace_dirty(inputs.workspace, inputs.slot_uid)
|
|
if not dirty.is_dirty:
|
|
break
|
|
reminders_used += 1
|
|
if needs_completion:
|
|
log.warning(
|
|
"rpc_completion_reminder",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"task": task_kind,
|
|
"attempt": reminders_used,
|
|
"max": max_reminders,
|
|
},
|
|
)
|
|
reminder = (
|
|
persona.review_completion_reminder(repo=inputs.repo, issue=inputs.issue, workspace=inputs.workspace)
|
|
if task_kind == "review_pr"
|
|
else persona.completion_reminder(repo=inputs.repo, issue=inputs.issue, workspace=inputs.workspace)
|
|
)
|
|
else:
|
|
assert dirty is not None
|
|
log.warning(
|
|
"rpc_dirty_state_reminder",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"task": task_kind,
|
|
"attempt": reminders_used,
|
|
"max": max_reminders,
|
|
"uncommitted": dirty.uncommitted,
|
|
"unpushed": dirty.unpushed,
|
|
},
|
|
)
|
|
reminder = persona.dirty_state_reminder(
|
|
repo=inputs.repo,
|
|
issue=inputs.issue,
|
|
workspace=inputs.workspace,
|
|
dirty=dirty,
|
|
)
|
|
next_turn = _run(reminder)
|
|
if next_turn is None:
|
|
return None
|
|
turn = next_turn
|
|
|
|
if reminders_used and _needs_completion_reminder(
|
|
task_kind=task_kind, inputs=inputs, bindings=bindings, tools_called=tools_called
|
|
):
|
|
log.warning(
|
|
"rpc_completion_unfinished",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"task": task_kind,
|
|
"reminders": reminders_used,
|
|
"tools_called": sorted(tools_called),
|
|
},
|
|
)
|
|
if reminders_used and task_kind != "review_pr":
|
|
final_dirty = _probe_workspace_dirty(inputs.workspace, inputs.slot_uid)
|
|
if final_dirty.is_dirty:
|
|
log.warning(
|
|
"rpc_dirty_state_unfinished",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"task": task_kind,
|
|
"reminders": reminders_used,
|
|
"uncommitted": final_dirty.uncommitted,
|
|
"unpushed": final_dirty.unpushed,
|
|
},
|
|
)
|
|
return turn
|
|
|
|
|
|
def _has_prior_session(session_dir: Path) -> bool:
|
|
"""Return True iff `session_dir` already contains an omp JSONL transcript.
|
|
|
|
pi's `coding-agent` writes one `*.jsonl` per session into `--session-dir`.
|
|
The presence of any such file is the signal that `--continue` will pick
|
|
up the most recent transcript (`SessionManager.continueRecent`) rather
|
|
than starting fresh.
|
|
"""
|
|
try:
|
|
return any(session_dir.glob("*.jsonl"))
|
|
except OSError:
|
|
return False
|
|
|
|
|
|
def _build_prompt(
|
|
task_kind: str,
|
|
inputs: TaskInputs,
|
|
*,
|
|
comment: CommentInfo | None,
|
|
pr_number: int | None,
|
|
review_payload: dict[str, Any] | None,
|
|
pr: PullRequestInfo | None = None,
|
|
directive: DirectiveInfo | None = None,
|
|
thread: tuple[ThreadMessage, ...] = (),
|
|
resuming: bool = False,
|
|
) -> str:
|
|
if task_kind == "triage_issue":
|
|
if resuming:
|
|
return persona.resume_triage(repo=inputs.repo, issue=inputs.issue, workspace=inputs.workspace)
|
|
if directive is not None:
|
|
return persona.kickoff_directive(
|
|
repo=inputs.repo,
|
|
issue=inputs.issue,
|
|
workspace=inputs.workspace,
|
|
directive=directive,
|
|
)
|
|
return persona.kickoff(repo=inputs.repo, issue=inputs.issue, workspace=inputs.workspace)
|
|
if task_kind == "review_pr":
|
|
assert pr is not None
|
|
return persona.kickoff_pr_review(repo=inputs.repo, pr=pr, workspace=inputs.workspace)
|
|
if task_kind == "handle_comment":
|
|
assert comment is not None
|
|
issue_row = inputs.db.get_issue(issue_key(inputs.repo.full_name, inputs.issue.number))
|
|
if issue_row is None:
|
|
pr_status = "no PR opened yet"
|
|
elif issue_row.pr_number is None:
|
|
pr_status = "no PR opened yet"
|
|
elif issue_row.state == "merged":
|
|
pr_status = f"PR #{issue_row.pr_number} was merged"
|
|
elif issue_row.state in ("closed", "abandoned"):
|
|
pr_status = f"PR #{issue_row.pr_number} was closed without merge"
|
|
else:
|
|
pr_status = f"PR #{issue_row.pr_number} is open"
|
|
if directive is not None:
|
|
return persona.directive(
|
|
repo=inputs.repo,
|
|
issue=inputs.issue,
|
|
workspace=inputs.workspace,
|
|
comment=comment,
|
|
directive=directive,
|
|
pr_status=pr_status,
|
|
pr_number=pr_number,
|
|
)
|
|
return persona.followup_comment(
|
|
repo=inputs.repo,
|
|
issue=inputs.issue,
|
|
workspace=inputs.workspace,
|
|
comment=comment,
|
|
pr_status=pr_status,
|
|
pr_number=pr_number,
|
|
thread=thread,
|
|
)
|
|
if task_kind == "handle_review":
|
|
assert review_payload is not None
|
|
path = str(review_payload.get("path") or "")
|
|
start = review_payload.get("start_line") or review_payload.get("line")
|
|
end = review_payload.get("line") or review_payload.get("original_line")
|
|
if isinstance(start, int) and isinstance(end, int) and start != end:
|
|
line_range = f":L{start}-L{end}"
|
|
elif isinstance(end, int):
|
|
line_range = f":L{end}"
|
|
else:
|
|
line_range = ""
|
|
body = str(review_payload.get("body") or "")
|
|
author = str(review_payload.get("author") or "")
|
|
return persona.followup_review(
|
|
repo=inputs.repo,
|
|
workspace=inputs.workspace,
|
|
pr_number=int(pr_number or 0),
|
|
comment_author=author,
|
|
comment_body=body,
|
|
comment_path=path,
|
|
comment_line_range=line_range,
|
|
)
|
|
raise ValueError(f"unknown task kind: {task_kind!r}")
|
|
|
|
|
|
def _run_rpc_blocking(
|
|
inputs: TaskInputs,
|
|
*,
|
|
task_kind: str,
|
|
prompt: str,
|
|
loop: asyncio.AbstractEventLoop,
|
|
bindings: ToolBindings,
|
|
directive: DirectiveInfo | None = None,
|
|
) -> str | None:
|
|
"""Run a full RPC turn synchronously. Returns final assistant text (or None)."""
|
|
settings = inputs.settings
|
|
|
|
tools_called: set[str] = set()
|
|
|
|
def _on_tool_end(event: ToolExecutionEndEvent) -> None:
|
|
tool_name = event.tool_name
|
|
# `tool_name` is transport-normalized by omp_rpc: an xd:// device
|
|
# dispatch (`write xd://submit_pr_review`) reports the host tool that
|
|
# ran, so terminal-action detection can match on host-tool names. A
|
|
# failed execution (`is_error`) does not count as reaching the
|
|
# terminal action — a rejected submit must still trigger the
|
|
# completion reminder.
|
|
ok = event.result is not None and not event.is_error
|
|
if ok:
|
|
tools_called.add(tool_name)
|
|
log.info(
|
|
"tool_end",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"tool": tool_name,
|
|
"ok": ok,
|
|
},
|
|
)
|
|
|
|
def _on_msg(event: MessageUpdateEvent) -> None:
|
|
ev = event.assistant_message_event
|
|
if isinstance(ev, dict) and ev.get("type") == "text_delta":
|
|
log.debug("delta", extra={"issue": bindings.issue_key, "delta": str(ev.get("delta", ""))[:200]})
|
|
|
|
rpc_env = _build_extra_env(settings)
|
|
rpc_env.update(_prepare_slot_runtime_env(inputs.workspace, inputs.slot_uid))
|
|
rpc_env.update(_safe_directory_env(bindings.workspace.repo_dir))
|
|
rpc_env.update(_git_identity_env(inputs.settings.resolved_author_name, inputs.settings.git_author_email))
|
|
# Bare worktrees have no node_modules; install (idempotently) so the agent
|
|
# can resolve workspace packages (@oh-my-pi/pi-*) and actually run tests.
|
|
host_tools.ensure_workspace_dependencies(bindings)
|
|
resuming = _has_prior_session(bindings.workspace.session_dir)
|
|
extra_args: tuple[str, ...] = ("--continue",) if resuming else ()
|
|
log.info(
|
|
"rpc_resume",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"task": task_kind,
|
|
"resuming": resuming,
|
|
"session_dir": str(bindings.workspace.session_dir),
|
|
"attempts": inputs.attempts,
|
|
},
|
|
)
|
|
model_override, thinking_override = _resolve_pragma_overrides(directive, settings)
|
|
chosen_model = model_override or settings.pick_model()
|
|
chosen_thinking = thinking_override or settings.thinking_level
|
|
log.info(
|
|
"rpc_model_pick",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"model": chosen_model,
|
|
"pool": list(settings.model_pool),
|
|
"thinking": chosen_thinking,
|
|
"pragma_model": model_override,
|
|
"pragma_thinking": thinking_override,
|
|
},
|
|
)
|
|
inputs.db.set_event_model(inputs.delivery_id, chosen_model)
|
|
append_system_prompt = (
|
|
persona.system_append_pr_review(
|
|
repo=inputs.repo,
|
|
issue=inputs.issue,
|
|
workspace=inputs.workspace,
|
|
bot_login=inputs.settings.bot_login,
|
|
)
|
|
if task_kind == "review_pr"
|
|
else persona.system_append(
|
|
repo=inputs.repo,
|
|
issue=inputs.issue,
|
|
workspace=inputs.workspace,
|
|
bot_login=inputs.settings.bot_login,
|
|
)
|
|
)
|
|
|
|
with RpcClient(
|
|
executable=settings.omp_command,
|
|
cwd=bindings.workspace.repo_dir,
|
|
session_dir=bindings.workspace.session_dir,
|
|
env=rpc_env,
|
|
no_session=False,
|
|
no_title=True,
|
|
model=chosen_model,
|
|
provider=settings.provider,
|
|
thinking=chosen_thinking if chosen_thinking != "off" else None,
|
|
append_system_prompt=append_system_prompt,
|
|
custom_tools=host_tools.build(bindings),
|
|
request_timeout=settings.request_timeout_seconds,
|
|
startup_timeout=60.0,
|
|
max_event_history=50_000,
|
|
extra_args=extra_args,
|
|
user=inputs.slot_uid,
|
|
group=inputs.slot_uid if inputs.slot_uid is not None else None,
|
|
extra_groups=["omp"] if inputs.slot_uid is not None else None,
|
|
) as client:
|
|
# Arm cancellation: from this point the API can kill the omp subprocess
|
|
# out from under us, which makes `prompt_and_wait` raise an `RpcError`
|
|
# we'll let propagate. The `with` exit calls `client.stop()` again, but
|
|
# it's idempotent.
|
|
#
|
|
# NOTE: omp_rpc.RpcClient.stop() has a bug where it sets `_stopping=True`
|
|
# before the stdout reader loop notices the closed pipe, so the reader's
|
|
# `if not self._stopping` guard skips `_mark_closed()` entirely.
|
|
# `_wait_for_agent_end` then blocks on `_event_condition` until the hard
|
|
# timeout because `_closed_error` is never set. We work around it here
|
|
# by calling `_mark_closed()` ourselves after stop returns — this is
|
|
# idempotent (it no-ops when `_closed_error` is already set).
|
|
def _cancel_hook() -> None:
|
|
try:
|
|
client.stop()
|
|
finally:
|
|
# Private API, but the only way to unblock `_wait_for_agent_end`
|
|
# without waiting for the request timeout. Idempotent.
|
|
client._mark_closed( # noqa: SLF001
|
|
RpcProcessExitError("cancelled by operator")
|
|
)
|
|
|
|
if bindings.abort is not None:
|
|
bindings.abort.stop = _cancel_hook
|
|
register_cancel_hook(_cancel_hook)
|
|
try:
|
|
client.install_headless_ui()
|
|
client.on_tool_execution_end(_on_tool_end)
|
|
client.on_message_update(_on_msg)
|
|
|
|
phases = persona.seed_phases(task_kind)
|
|
if phases:
|
|
try:
|
|
if task_kind in ("triage_issue", "review_pr") and not resuming:
|
|
# Fresh kickoff tasks seed the full plan.
|
|
client.set_todos(phases)
|
|
elif task_kind in ("triage_issue", "review_pr"):
|
|
# Resumed kickoff tasks keep prior todo state from the
|
|
# JSONL transcript; re-seeding would clobber progress.
|
|
log.info(
|
|
"set_todos skipped (resume)",
|
|
extra={"issue": bindings.issue_key, "task": task_kind},
|
|
)
|
|
else:
|
|
# Follow-up: keep prior phases (e.g. Reproduce / Fix / PR)
|
|
# so the agent still sees the context, but append the
|
|
# follow-up phase at the end.
|
|
existing = list(client.get_todos())
|
|
merged = [
|
|
{
|
|
"id": p.id,
|
|
"name": p.name,
|
|
"tasks": [
|
|
{
|
|
"id": t.id,
|
|
"content": t.content,
|
|
"status": t.status,
|
|
"notes": t.notes,
|
|
"details": t.details,
|
|
"blocker": t.blocker,
|
|
}
|
|
for t in p.tasks
|
|
],
|
|
}
|
|
for p in existing
|
|
] + phases
|
|
client.set_todos(merged)
|
|
except RpcError as exc:
|
|
log.warning("set_todos failed", extra={"err": str(exc)})
|
|
|
|
log.info(
|
|
"rpc_start",
|
|
extra={"issue": bindings.issue_key, "task": task_kind, "branch": bindings.workspace.branch},
|
|
)
|
|
hard_timeout_seconds = settings.task_timeout_seconds + settings.task_timeout_hard_grace_seconds
|
|
hard_timeout_fired = threading.Event()
|
|
|
|
def _hard_stop() -> None:
|
|
hard_timeout_fired.set()
|
|
log.warning(
|
|
"rpc_hard_timeout",
|
|
extra={"issue": bindings.issue_key, "task": task_kind, "timeout": hard_timeout_seconds},
|
|
)
|
|
try:
|
|
_cancel_hook()
|
|
except Exception:
|
|
log.exception(
|
|
"rpc hard timeout stop failed", extra={"issue": bindings.issue_key, "task": task_kind}
|
|
)
|
|
|
|
hard_timer = threading.Timer(hard_timeout_seconds, _hard_stop)
|
|
hard_timer.daemon = True
|
|
hard_timer.start()
|
|
try:
|
|
turn = _drive_turn(
|
|
client,
|
|
prompt,
|
|
task_kind=task_kind,
|
|
inputs=inputs,
|
|
bindings=bindings,
|
|
tools_called=tools_called,
|
|
)
|
|
if turn is None:
|
|
return None
|
|
finally:
|
|
hard_timer.cancel()
|
|
if hard_timeout_fired.is_set():
|
|
raise TimeoutError("omp task exceeded hard timeout")
|
|
if turn is not None and turn.assistant_message is not None:
|
|
stop_reason = turn.assistant_message.get("stopReason")
|
|
if stop_reason == "error":
|
|
error_msg = turn.assistant_message.get("errorMessage") or "model returned error"
|
|
raise RuntimeError(f"omp agent error (stopReason=error): {error_msg}")
|
|
log.info(
|
|
"rpc_done",
|
|
extra={
|
|
"issue": bindings.issue_key,
|
|
"task": task_kind,
|
|
"messages": len(turn.messages),
|
|
"events": len(turn.events),
|
|
},
|
|
)
|
|
return turn.assistant_text
|
|
finally:
|
|
unregister_cancel_hook()
|
|
|
|
|
|
async def run_task(
|
|
*,
|
|
task_kind: str,
|
|
inputs: TaskInputs,
|
|
comment: CommentInfo | None = None,
|
|
pr_number: int | None = None,
|
|
review_payload: dict[str, Any] | None = None,
|
|
pr: PullRequestInfo | None = None,
|
|
directive: DirectiveInfo | None = None,
|
|
thread: tuple[ThreadMessage, ...] = (),
|
|
) -> str | None:
|
|
"""Async wrapper that runs the synchronous RPC driver on a worker thread."""
|
|
review_mode = task_kind == "review_pr" or inputs.workspace.branch.startswith("review/pr-")
|
|
loop = asyncio.get_running_loop()
|
|
bindings = ToolBindings(
|
|
db=inputs.db,
|
|
github=inputs.github,
|
|
git_transport=inputs.git_transport,
|
|
repo=inputs.repo,
|
|
issue=inputs.issue,
|
|
workspace=inputs.workspace,
|
|
loop=loop,
|
|
settings=inputs.settings,
|
|
author_name=inputs.settings.resolved_author_name,
|
|
author_email=inputs.settings.git_author_email,
|
|
inbound_thread_number=pr_number,
|
|
inbound_is_pr=pr_number is not None,
|
|
review_mode=review_mode,
|
|
impl_authorized=bool(directive is not None and directive.authorizes_impl),
|
|
slot_uid=inputs.slot_uid,
|
|
abort=AbortController(),
|
|
)
|
|
resuming = _has_prior_session(inputs.workspace.session_dir)
|
|
prompt = _build_prompt(
|
|
task_kind,
|
|
inputs,
|
|
comment=comment,
|
|
pr_number=pr_number,
|
|
review_payload=review_payload,
|
|
pr=pr,
|
|
directive=directive,
|
|
thread=thread,
|
|
resuming=resuming,
|
|
)
|
|
try:
|
|
result = await asyncio.to_thread(
|
|
_run_rpc_blocking,
|
|
inputs,
|
|
task_kind=task_kind,
|
|
prompt=prompt,
|
|
loop=loop,
|
|
bindings=bindings,
|
|
directive=directive,
|
|
)
|
|
except BaseException:
|
|
# Failed/aborted task: NEVER capture, the artifacts may be inconsistent
|
|
# with the source state and would poison the cache.
|
|
raise
|
|
else:
|
|
await asyncio.to_thread(_capture_natives_cache, inputs)
|
|
return result
|
|
|
|
|
|
def _capture_natives_cache(inputs: TaskInputs) -> None:
|
|
"""Best-effort: store the workspace's fresh natives under its current key.
|
|
|
|
Runs after a successful task on a worker thread. ANY failure is logged
|
|
and swallowed — cache errors NEVER fail a task.
|
|
"""
|
|
cache = inputs.natives_cache
|
|
if cache is None:
|
|
return
|
|
workspace = inputs.workspace
|
|
native_dir = workspace.repo_dir / "packages" / "natives" / "native"
|
|
if not native_dir.exists():
|
|
return
|
|
try:
|
|
key = natives_compute_key(workspace.repo_dir)
|
|
except Exception as exc:
|
|
log.debug(
|
|
"natives_cache capture key compute failed",
|
|
extra={"workspace": workspace.workspace_key, "err": str(exc)},
|
|
)
|
|
return
|
|
try:
|
|
stored = cache.capture(
|
|
workspace.repo_full_name,
|
|
key,
|
|
native_dir,
|
|
source_workspace=workspace.workspace_key,
|
|
)
|
|
except Exception as exc:
|
|
log.warning(
|
|
"natives_cache capture failed",
|
|
extra={"workspace": workspace.workspace_key, "key": key, "err": str(exc)},
|
|
)
|
|
return
|
|
log.info(
|
|
"natives_cache",
|
|
extra={
|
|
"action": "stored" if stored is not None else "skip",
|
|
"workspace": workspace.workspace_key,
|
|
"repo": workspace.repo_full_name,
|
|
"key": key,
|
|
"cache_dir": str(stored) if stored else None,
|
|
},
|
|
)
|
|
|
|
|
|
__all__ = ["DirectiveInfo", "TaskInputs", "ThreadMessage", "run_task"]
|