From 4169e7fc28e21fa086cd8bc81b476f46a776e771 Mon Sep 17 00:00:00 2001 From: Felix Martin Date: Thu, 11 Jun 2026 08:18:44 -0400 Subject: [PATCH] feat: schedule prompts in session service --- src/hqt/sessions/service.py | 91 +++++++++++++++++++++++++++++-- tests/test_sessions.py | 105 +++++++++++++++++++++++++++++++++--- 2 files changed, 187 insertions(+), 9 deletions(-) diff --git a/src/hqt/sessions/service.py b/src/hqt/sessions/service.py index 7d18a9c..84395fc 100644 --- a/src/hqt/sessions/service.py +++ b/src/hqt/sessions/service.py @@ -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) diff --git a/tests/test_sessions.py b/tests/test_sessions.py index 361b0f7..e3f5d04 100644 --- a/tests/test_sessions.py +++ b/tests/test_sessions.py @@ -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)