feat: schedule prompts in session service
This commit is contained in:
@@ -10,7 +10,7 @@ from pathlib import Path
|
||||
from sqlalchemy.orm import Session as DBSession
|
||||
from sqlalchemy.orm import selectinload, sessionmaker
|
||||
|
||||
from hqt.db.models import Harness, Project, Session
|
||||
from hqt.db.models import Harness, Project, ScheduledPrompt, Session
|
||||
from hqt.errors import ServiceError
|
||||
from hqt.git import worktree
|
||||
from hqt.git.worktree import WorktreeState
|
||||
@@ -24,6 +24,10 @@ log = logging.getLogger(__name__)
|
||||
CAPTURE_MAX_ATTEMPTS = 6
|
||||
CAPTURE_RETRY_INTERVAL = 0.5
|
||||
ACTIVITY_RECENT_SECS = 5
|
||||
PROMPT_PENDING = "pending"
|
||||
PROMPT_SENDING = "sending"
|
||||
PROMPT_SENT = "sent"
|
||||
PROMPT_FAILED = "failed"
|
||||
|
||||
|
||||
async def _capture_sleep(t: float) -> None:
|
||||
@@ -35,6 +39,7 @@ class SessionInfo:
|
||||
session: Session
|
||||
alive: bool
|
||||
status: str = "dead"
|
||||
scheduled_prompt: ScheduledPrompt | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -414,6 +419,57 @@ class SessionService:
|
||||
sess.nickname = (nickname or "").strip() or None
|
||||
db.commit()
|
||||
|
||||
def schedule_prompt(
|
||||
self, session_id: int, prompt: str, due_at: datetime
|
||||
) -> ScheduledPrompt:
|
||||
"""Create or replace the pending scheduled prompt for a session."""
|
||||
if not prompt.strip():
|
||||
raise ServiceError("Prompt cannot be empty")
|
||||
with self.factory() as db:
|
||||
sess = db.get(Session, session_id)
|
||||
if sess is None:
|
||||
raise ServiceError("Session not found")
|
||||
scheduled = (
|
||||
db.query(ScheduledPrompt).filter_by(session_id=session_id).one_or_none()
|
||||
)
|
||||
if scheduled is None:
|
||||
scheduled = ScheduledPrompt(session_id=session_id, prompt=prompt)
|
||||
db.add(scheduled)
|
||||
else:
|
||||
scheduled.prompt = prompt
|
||||
scheduled.due_at = due_at
|
||||
scheduled.status = PROMPT_PENDING
|
||||
scheduled.error = None
|
||||
scheduled.sent_at = None
|
||||
db.commit()
|
||||
db.refresh(scheduled)
|
||||
db.expunge(scheduled)
|
||||
return scheduled
|
||||
|
||||
def cancel_scheduled_prompt(self, session_id: int) -> None:
|
||||
"""Delete any scheduled prompt for a session."""
|
||||
with self.factory() as db:
|
||||
scheduled = (
|
||||
db.query(ScheduledPrompt).filter_by(session_id=session_id).one_or_none()
|
||||
)
|
||||
if scheduled is None:
|
||||
return
|
||||
db.delete(scheduled)
|
||||
db.commit()
|
||||
|
||||
def get_scheduled_prompt(self, session_id: int) -> ScheduledPrompt | None:
|
||||
"""Return a session's pending scheduled prompt, if one exists."""
|
||||
with self.factory() as db:
|
||||
scheduled = (
|
||||
db.query(ScheduledPrompt)
|
||||
.filter_by(session_id=session_id, status=PROMPT_PENDING)
|
||||
.one_or_none()
|
||||
)
|
||||
if scheduled is None:
|
||||
return None
|
||||
db.expunge(scheduled)
|
||||
return scheduled
|
||||
|
||||
async def delete_session(
|
||||
self, session_id: int, remove_worktree: bool = False, force: bool = False
|
||||
) -> None:
|
||||
@@ -490,6 +546,21 @@ class SessionService:
|
||||
age = now - wi.last_activity
|
||||
return "active" if age <= ACTIVITY_RECENT_SECS else "idle"
|
||||
|
||||
def _pending_prompt_map(
|
||||
self, db: DBSession, session_ids: list[int]
|
||||
) -> dict[int, ScheduledPrompt]:
|
||||
if not session_ids:
|
||||
return {}
|
||||
prompts = (
|
||||
db.query(ScheduledPrompt)
|
||||
.filter(
|
||||
ScheduledPrompt.session_id.in_(session_ids),
|
||||
ScheduledPrompt.status == PROMPT_PENDING,
|
||||
)
|
||||
.all()
|
||||
)
|
||||
return {prompt.session_id: prompt for prompt in prompts}
|
||||
|
||||
async def list_sessions(
|
||||
self, project_id: int, now: float | None = None
|
||||
) -> list[SessionInfo]:
|
||||
@@ -504,13 +575,19 @@ class SessionService:
|
||||
)
|
||||
names = [s.tmux_session_name for s in sessions]
|
||||
info_map = await self.tmux.poll_info(names)
|
||||
prompt_map = self._pending_prompt_map(db, [s.id for s in sessions])
|
||||
result: list[SessionInfo] = []
|
||||
for s in sessions:
|
||||
self._maybe_capture_missing_session_id(db, s)
|
||||
wi = info_map.get(s.tmux_session_name)
|
||||
status = await self._status_for(s, wi, now)
|
||||
result.append(
|
||||
SessionInfo(session=s, alive=bool(wi and wi.alive), status=status)
|
||||
SessionInfo(
|
||||
session=s,
|
||||
alive=bool(wi and wi.alive),
|
||||
status=status,
|
||||
scheduled_prompt=prompt_map.get(s.id),
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
@@ -534,13 +611,21 @@ class SessionService:
|
||||
)
|
||||
names = [s.tmux_session_name for s in sessions]
|
||||
info_map = await self.tmux.poll_info(names)
|
||||
prompt_map = self._pending_prompt_map(db, [s.id for s in sessions])
|
||||
result: list[SessionInfo] = []
|
||||
for s in sessions:
|
||||
self._maybe_capture_missing_session_id(db, s)
|
||||
wi = info_map.get(s.tmux_session_name)
|
||||
alive = bool(wi and wi.alive)
|
||||
status = await self._status_for(s, wi, now)
|
||||
result.append(SessionInfo(session=s, alive=alive, status=status))
|
||||
result.append(
|
||||
SessionInfo(
|
||||
session=s,
|
||||
alive=alive,
|
||||
status=status,
|
||||
scheduled_prompt=prompt_map.get(s.id),
|
||||
)
|
||||
)
|
||||
name = s.nickname or s.tmux_session_name
|
||||
await self.tmux.set_window_label(
|
||||
s.tmux_session_name, window_label(status, name, alive=alive)
|
||||
|
||||
+99
-6
@@ -1,4 +1,5 @@
|
||||
import pytest
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
@@ -6,7 +7,7 @@ from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy.pool import StaticPool
|
||||
|
||||
from hqt.db.models import Base, Harness, Project
|
||||
from hqt.db.models import Base, Harness, Project, ScheduledPrompt, Session
|
||||
from hqt.errors import ServiceError
|
||||
from hqt.sessions.service import SessionInfo, SessionService
|
||||
from hqt.tmux.manager import SpawnResult, TmuxManager
|
||||
@@ -43,6 +44,8 @@ def tmux():
|
||||
m.poll_info = AsyncMock(return_value={})
|
||||
m.capture_pane = AsyncMock(return_value="")
|
||||
m.attach = AsyncMock()
|
||||
m.send_text = AsyncMock(return_value=True)
|
||||
m.send_enter = AsyncMock(return_value=True)
|
||||
m.window_exists = AsyncMock(return_value=True)
|
||||
m.respawn_verified = AsyncMock(
|
||||
return_value=SpawnResult(ok=True, window_id=None, error="")
|
||||
@@ -93,6 +96,97 @@ async def test_list_sessions(service, db, tmux):
|
||||
assert result[0].alive is True
|
||||
|
||||
|
||||
def test_schedule_prompt_creates_pending_row(service, db):
|
||||
sess = _seed_session(db)
|
||||
due_at = datetime(2026, 6, 11, 12)
|
||||
|
||||
prompt = service.schedule_prompt(sess.id, "run the tests", due_at)
|
||||
|
||||
assert prompt.session_id == sess.id
|
||||
assert prompt.prompt == "run the tests"
|
||||
assert prompt.due_at == due_at
|
||||
assert prompt.status == "pending"
|
||||
assert prompt.error is None
|
||||
rows = db.query(ScheduledPrompt).all()
|
||||
assert len(rows) == 1
|
||||
|
||||
|
||||
def test_schedule_prompt_replaces_existing_row_for_session(service, db):
|
||||
sess = _seed_session(db)
|
||||
first_due = datetime(2026, 6, 11, 12)
|
||||
second_due = datetime(2026, 6, 11, 13)
|
||||
first = service.schedule_prompt(sess.id, "old prompt", first_due)
|
||||
|
||||
second = service.schedule_prompt(sess.id, "new prompt", second_due)
|
||||
|
||||
rows = db.query(ScheduledPrompt).all()
|
||||
assert len(rows) == 1
|
||||
assert second.id == first.id
|
||||
assert rows[0].id == first.id
|
||||
assert rows[0].prompt == "new prompt"
|
||||
assert rows[0].due_at == second_due
|
||||
assert rows[0].status == "pending"
|
||||
assert rows[0].error is None
|
||||
|
||||
|
||||
def test_schedule_prompt_rejects_empty_prompt(service):
|
||||
due_at = datetime(2026, 6, 11, 12)
|
||||
|
||||
with pytest.raises(ServiceError, match="Prompt cannot be empty"):
|
||||
service.schedule_prompt(1, " ", due_at)
|
||||
|
||||
|
||||
def test_schedule_prompt_rejects_missing_session(service):
|
||||
due_at = datetime(2026, 6, 11, 12)
|
||||
|
||||
with pytest.raises(ServiceError, match="Session not found"):
|
||||
service.schedule_prompt(999, "hello", due_at)
|
||||
|
||||
|
||||
def test_cancel_scheduled_prompt_deletes_row(service, db):
|
||||
sess = _seed_session(db)
|
||||
due_at = datetime(2026, 6, 11, 12)
|
||||
service.schedule_prompt(sess.id, "run the tests", due_at)
|
||||
|
||||
service.cancel_scheduled_prompt(sess.id)
|
||||
|
||||
assert db.query(ScheduledPrompt).all() == []
|
||||
|
||||
|
||||
def test_get_scheduled_prompt_returns_pending_only(service, db):
|
||||
sess = _seed_session(db)
|
||||
due_at = datetime(2026, 6, 11, 12)
|
||||
service.schedule_prompt(sess.id, "run the tests", due_at)
|
||||
|
||||
prompt = service.get_scheduled_prompt(sess.id)
|
||||
assert prompt is not None
|
||||
assert prompt.status == "pending"
|
||||
|
||||
row = db.query(ScheduledPrompt).filter_by(session_id=sess.id).one()
|
||||
row.status = "failed"
|
||||
db.commit()
|
||||
|
||||
assert service.get_scheduled_prompt(sess.id) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_sessions_includes_pending_scheduled_prompt(service, db, tmux):
|
||||
from hqt.tmux.runner import WindowInfo
|
||||
|
||||
sess = _seed_session(db)
|
||||
due_at = datetime(2026, 6, 11, 12)
|
||||
service.schedule_prompt(sess.id, "run the tests", due_at)
|
||||
tmux.poll_info.return_value = {
|
||||
sess.tmux_session_name: WindowInfo(alive=True, last_activity=1000)
|
||||
}
|
||||
|
||||
result = await service.list_sessions(project_id=1)
|
||||
|
||||
assert result[0].scheduled_prompt is not None
|
||||
assert result[0].scheduled_prompt.prompt == "run the tests"
|
||||
assert result[0].scheduled_prompt.status == "pending"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_session(service, tmux):
|
||||
await service.create_session(project_id=1, harness_name="claude-code")
|
||||
@@ -985,16 +1079,15 @@ async def test_create_session_result_propagates_spawn_failure(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _seed_session(db, nickname=None):
|
||||
"""Create a session for project 1 / harness 'claude-code' (seeded by the db fixture)."""
|
||||
from hqt.db.models import Harness, Session
|
||||
|
||||
def _seed_session(db, project_id: int = 1, nickname=None) -> Session:
|
||||
"""Create a session for the seeded 'claude-code' harness."""
|
||||
harness = db.query(Harness).filter_by(name="claude-code").first()
|
||||
sess = Session(
|
||||
project_id=1,
|
||||
project_id=project_id,
|
||||
harness_id=harness.id,
|
||||
nickname=nickname,
|
||||
tmux_session_name="hqt-rename",
|
||||
harness_session_id="sess-seeded",
|
||||
archived=False,
|
||||
)
|
||||
db.add(sess)
|
||||
|
||||
Reference in New Issue
Block a user