Coverage for src/local_deep_research/web/queue/lifecycle_cleanup.py: 94%

1from collections.abc import Collection 

2from dataclasses import dataclass 

3from datetime import UTC, datetime 

4 

5from sqlalchemy.orm import Session 

6 

7from ...database.models import QueuedResearch, QueueStatus, TaskMetadata 

8 

9 

10@dataclass(frozen=True, slots=True) 

11class QueuedResearchCleanupResult: 

12 cleaned_ids: frozenset[str] 

13 protected_ids: frozenset[str] 

14 

15 

16def reconcile_research_queue_status( 

17 db_session: Session, *, force: bool = False 

18) -> bool: 

19 """Recompute research capacity counters without committing the transaction.""" 

20 queued_tasks = ( 

21 db_session.query(TaskMetadata) 

22 .filter_by(task_type="research", status="queued") 

23 .count() 

24 ) 

25 active_tasks = ( 

26 db_session.query(TaskMetadata) 

27 .filter_by(task_type="research", status="processing") 

28 .count() 

29 ) 

30 queue_status = db_session.query(QueueStatus).first() 

31 if queue_status is None: 

32 db_session.add( 

33 QueueStatus( 

34 active_tasks=active_tasks, 

35 queued_tasks=queued_tasks, 

36 last_checked=datetime.now(UTC), 

37 ) 

38 ) 

39 return True 

40 

41 changed = ( 

42 queue_status.active_tasks != active_tasks 

43 or queue_status.queued_tasks != queued_tasks 

44 ) 

45 if changed or force: 

46 queue_status.active_tasks = active_tasks 

47 queue_status.queued_tasks = queued_tasks 

48 queue_status.last_checked = datetime.now(UTC) 

49 return changed or force 

50 

51 

52def cleanup_queued_research_state( 

53 db_session: Session, 

54 research_ids: Collection[str], 

55 include_claimed: bool = False, 

56 acting_user: str | None = None, 

57) -> QueuedResearchCleanupResult: 

58 """Remove queued research state while leaving transaction control to the caller. 

59 

60 With ``acting_user``, a queue row another user owns is left in place, and so 

61 is the task metadata of its research ID, which that row still needs. Such 

62 IDs are reported as protected. ``TaskMetadata`` has no owner column, so 

63 metadata that no queue row points at is removed either way. 

64 """ 

65 research_id_set = set(research_ids) 

66 if not research_id_set: 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true

67 return QueuedResearchCleanupResult(frozenset(), frozenset()) 

68 

69 claimed_ids: set[str] = set() 

70 if not include_claimed: 

71 claimed_ids = { 

72 research_id 

73 for (research_id,) in ( 

74 db_session.query(QueuedResearch.research_id) 

75 .filter( 

76 QueuedResearch.research_id.in_(research_id_set), 

77 QueuedResearch.is_processing.is_(True), 

78 ) 

79 .all() 

80 ) 

81 } 

82 cleanup_ids = research_id_set - claimed_ids 

83 if not cleanup_ids: 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true

84 return QueuedResearchCleanupResult(frozenset(), frozenset(claimed_ids)) 

85 

86 queued_rows = ( 

87 db_session.query(QueuedResearch) 

88 .filter(QueuedResearch.research_id.in_(cleanup_ids)) 

89 .all() 

90 ) 

91 foreign_ids: set[str] = set() 

92 if acting_user is not None: 

93 foreign_ids = { 

94 queued_row.research_id 

95 for queued_row in queued_rows 

96 if queued_row.username != acting_user 

97 } 

98 queued_rows = [ 

99 queued_row 

100 for queued_row in queued_rows 

101 if queued_row.username == acting_user 

102 ] 

103 usernames = {queued_row.username for queued_row in queued_rows} 

104 for queued_row in queued_rows: 

105 db_session.delete(queued_row) 

106 

107 metadata_ids = cleanup_ids - foreign_ids 

108 if metadata_ids: 

109 ( 

110 db_session.query(TaskMetadata) 

111 .filter( 

112 TaskMetadata.task_id.in_(metadata_ids), 

113 TaskMetadata.task_type == "research", 

114 ) 

115 .delete(synchronize_session=False) 

116 ) 

117 

118 for username in usernames: 

119 surviving_rows = ( 

120 db_session.query(QueuedResearch) 

121 .filter_by(username=username) 

122 .order_by(QueuedResearch.position, QueuedResearch.id) 

123 .all() 

124 ) 

125 for position, queued_row in enumerate(surviving_rows, start=1): 

126 queued_row.position = position 

127 

128 reconcile_research_queue_status(db_session, force=True) 

129 return QueuedResearchCleanupResult( 

130 frozenset(metadata_ids), frozenset(claimed_ids | foreign_ids) 

131 )