feat: schedule prompts in session service

This commit is contained in:
2026-06-11 08:18:44 -04:00
parent 60612065ff
commit 4169e7fc28
2 changed files with 187 additions and 9 deletions
+88 -3
View File
@@ -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
View File
@@ -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)