修改
This commit is contained in:
hupenglong1
@@ -0,0 +1,211 @@
|
||||
"""Persist background bash task records.
|
||||
|
||||
Sibling to run_state_store.py: a SQLite table tracking bg shell tasks the
|
||||
agent spawns via bash(run_in_background=true). The store survives backend
|
||||
restarts so polling can be resumed; the actual remote process is owned by
|
||||
nohup on the Jupyter terminal and is independent of this store.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from threading import RLock
|
||||
from typing import Any
|
||||
|
||||
|
||||
ACTIVE_BG_STATUSES = {'running'}
|
||||
TERMINAL_BG_STATUSES = {'completed', 'failed', 'cancelled'}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BgTaskSpec:
|
||||
"""Payload an agent-side handler hands to the backend on registration."""
|
||||
|
||||
task_id: str
|
||||
account_key: str
|
||||
session_id: str
|
||||
run_id: str
|
||||
pid: int
|
||||
task_dir: str
|
||||
output_path: str
|
||||
pid_path: str
|
||||
exit_code_path: str
|
||||
command: str
|
||||
started_at: float
|
||||
wait_for_completion: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BgTaskStatus:
|
||||
"""Snapshot the agent reads back via bash_status."""
|
||||
|
||||
task_id: str
|
||||
status: str
|
||||
pid: int
|
||||
started_at: float
|
||||
finished_at: float | None
|
||||
exit_code: int | None
|
||||
output_path: str
|
||||
output_preview: str
|
||||
auto_resumed: bool
|
||||
|
||||
|
||||
class BashBgStore:
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
self.db_path = db_path
|
||||
self._lock = RLock()
|
||||
|
||||
def record_start(self, spec: BgTaskSpec) -> None:
|
||||
now = time.time()
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
insert into bash_bg_tasks (
|
||||
task_id, account_key, session_id, run_id, pid,
|
||||
task_dir, output_path, pid_path, exit_code_path,
|
||||
command, status, started_at, updated_at,
|
||||
wait_for_completion, auto_resumed
|
||||
)
|
||||
values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'running', ?, ?, ?, 0)
|
||||
on conflict(task_id) do nothing
|
||||
""",
|
||||
(
|
||||
spec.task_id,
|
||||
spec.account_key,
|
||||
spec.session_id,
|
||||
spec.run_id,
|
||||
spec.pid,
|
||||
spec.task_dir,
|
||||
spec.output_path,
|
||||
spec.pid_path,
|
||||
spec.exit_code_path,
|
||||
spec.command,
|
||||
spec.started_at,
|
||||
now,
|
||||
1 if spec.wait_for_completion else 0,
|
||||
),
|
||||
)
|
||||
|
||||
def mark_completed(
|
||||
self,
|
||||
task_id: str,
|
||||
*,
|
||||
exit_code: int,
|
||||
finished_at: float | None = None,
|
||||
) -> None:
|
||||
finished = finished_at if finished_at is not None else time.time()
|
||||
status = 'completed' if exit_code == 0 else 'failed'
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
update bash_bg_tasks
|
||||
set status = ?, exit_code = ?, finished_at = ?, updated_at = ?
|
||||
where task_id = ? and status = 'running'
|
||||
""",
|
||||
(status, exit_code, finished, finished, task_id),
|
||||
)
|
||||
|
||||
def mark_killed(self, task_id: str) -> None:
|
||||
now = time.time()
|
||||
with self._connect() as conn:
|
||||
conn.execute(
|
||||
"""
|
||||
update bash_bg_tasks
|
||||
set status = 'cancelled', finished_at = ?, updated_at = ?
|
||||
where task_id = ? and status = 'running'
|
||||
""",
|
||||
(now, now, task_id),
|
||||
)
|
||||
|
||||
def mark_auto_resumed(self, task_id: str) -> bool:
|
||||
"""Set auto_resumed=1 atomically; return True if we won the race."""
|
||||
with self._connect() as conn:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
update bash_bg_tasks
|
||||
set auto_resumed = 1, updated_at = ?
|
||||
where task_id = ? and auto_resumed = 0
|
||||
""",
|
||||
(time.time(), task_id),
|
||||
)
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def list_active(self) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"select * from bash_bg_tasks where status = 'running'"
|
||||
).fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def list_for_session(
|
||||
self,
|
||||
account_key: str,
|
||||
session_id: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
with self._connect() as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
select * from bash_bg_tasks
|
||||
where account_key = ? and session_id = ?
|
||||
order by started_at desc
|
||||
""",
|
||||
(account_key, session_id),
|
||||
).fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def get(self, task_id: str) -> dict[str, Any] | None:
|
||||
with self._connect() as conn:
|
||||
row = conn.execute(
|
||||
"select * from bash_bg_tasks where task_id = ?",
|
||||
(task_id,),
|
||||
).fetchone()
|
||||
return dict(row) if row is not None else None
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
with self._lock:
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
conn = sqlite3.connect(self.db_path, timeout=10)
|
||||
conn.row_factory = sqlite3.Row
|
||||
self._init_schema(conn)
|
||||
return conn
|
||||
|
||||
def _init_schema(self, conn: sqlite3.Connection) -> None:
|
||||
conn.execute('pragma journal_mode=wal')
|
||||
conn.execute(
|
||||
"""
|
||||
create table if not exists bash_bg_tasks (
|
||||
task_id text primary key,
|
||||
account_key text not null,
|
||||
session_id text not null,
|
||||
run_id text not null default '',
|
||||
pid integer not null,
|
||||
task_dir text not null,
|
||||
output_path text not null,
|
||||
pid_path text not null,
|
||||
exit_code_path text not null,
|
||||
command text not null,
|
||||
status text not null,
|
||||
started_at real not null,
|
||||
finished_at real,
|
||||
exit_code integer,
|
||||
updated_at real not null,
|
||||
wait_for_completion integer not null default 1,
|
||||
auto_resumed integer not null default 0
|
||||
)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
create index if not exists idx_bash_bg_session
|
||||
on bash_bg_tasks(account_key, session_id, started_at)
|
||||
"""
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
create index if not exists idx_bash_bg_active
|
||||
on bash_bg_tasks(status)
|
||||
"""
|
||||
)
|
||||
Reference in New Issue
Block a user