Coverage for src/local_deep_research/web/queue/manager.py: 100%

1"""Queue manager for handling research queue operations""" 

2 

3from loguru import logger 

4from sqlalchemy import func 

5from sqlalchemy.orm import Session 

6 

7from ...database.session_context import get_user_db_session 

8from ...database.models import QueuedResearch, ResearchHistory 

9from .processor_v2 import queue_processor 

10 

11 

12def _safe_commit(session: Session) -> None: 

13 """Commit the transaction, rolling back on failure. 

14 

15 Mirrors ``UserQueueService._safe_commit``. The ``get_user_db_session`` 

16 context manager only closes the session; it does not roll back a failed 

17 commit, which would leave the session in a dirty state and mask the 

18 original error on any subsequent operation. See issue #2055. 

19 

20 The rollback itself is guarded so that a rollback failure is logged at 

21 debug level rather than masking the original commit error — the same 

22 defensive shape as ``QueueProcessorV2._commit_with_safe_rollback``. 

23 """ 

24 try: 

25 session.commit() 

26 except Exception: 

27 logger.exception("Database commit failed, attempting rollback") 

28 try: 

29 session.rollback() 

30 except Exception: 

31 logger.debug( 

32 "Rollback after commit failure also failed", exc_info=True 

33 ) 

34 raise 

35 

36 

37class QueueManager: 

38 """Manages the research queue operations""" 

39 

40 @staticmethod 

41 def add_to_queue(username, research_id, query, mode, settings): 

42 """ 

43 Add a research to the queue 

44 

45 Args: 

46 username: User who owns the research 

47 research_id: UUID of the research 

48 query: Research query 

49 mode: Research mode 

50 settings: Research settings dictionary 

51 

52 Returns: 

53 int: Queue position 

54 """ 

55 with get_user_db_session(username) as session: 

56 # Get the next position in queue for this user 

57 max_position = ( 

58 session.query(func.max(QueuedResearch.position)) 

59 .filter_by(username=username) 

60 .scalar() 

61 or 0 

62 ) 

63 

64 queued_record = QueuedResearch( 

65 username=username, 

66 research_id=research_id, 

67 query=query, 

68 mode=mode, 

69 settings_snapshot=settings, 

70 position=max_position + 1, 

71 ) 

72 session.add(queued_record) 

73 _safe_commit(session) 

74 

75 logger.info( 

76 f"Added research {research_id} to queue at position {max_position + 1}" 

77 ) 

78 

79 # Send RESEARCH_QUEUED notification if enabled 

80 try: 

81 from ...settings import SettingsManager 

82 from ...notifications import send_queue_notification 

83 

84 settings_manager = SettingsManager(session) 

85 settings_snapshot = settings_manager.get_settings_snapshot() 

86 

87 send_queue_notification( 

88 username=username, 

89 research_id=research_id, 

90 query=query, 

91 settings_snapshot=settings_snapshot, 

92 position=max_position + 1, 

93 ) 

94 except Exception as e: 

95 logger.debug(f"Failed to send queued notification: {e}") 

96 

97 # Notify queue processor about the new queued research 

98 # Note: When using QueueManager, we don't have all parameters for direct execution 

99 # So it will fall back to queue mode 

100 try: 

101 queue_processor.notify_research_queued(username, research_id) 

102 except Exception: 

103 logger.warning("Failed to notify queue processor") 

104 

105 return max_position + 1 

106 

107 @staticmethod 

108 def get_queue_position(username, research_id): 

109 """ 

110 Get the current queue position for a research 

111 

112 Args: 

113 username: User who owns the research 

114 research_id: UUID of the research 

115 

116 Returns: 

117 int: Current queue position or None if not in queue 

118 """ 

119 with get_user_db_session(username) as session: 

120 queued = ( 

121 session.query(QueuedResearch) 

122 .filter_by(username=username, research_id=research_id) 

123 .first() 

124 ) 

125 

126 if not queued: 

127 return None 

128 

129 # Count how many are ahead in queue 

130 ahead_count = ( 

131 session.query(QueuedResearch) 

132 .filter( 

133 QueuedResearch.username == username, 

134 QueuedResearch.position < queued.position, 

135 ) 

136 .count() 

137 ) 

138 

139 return ahead_count + 1 

140 

141 @staticmethod 

142 def remove_from_queue(username, research_id): 

143 """ 

144 Remove a research from the queue 

145 

146 Args: 

147 username: User who owns the research 

148 research_id: UUID of the research 

149 

150 Returns: 

151 bool: True if removed, False if not found 

152 """ 

153 with get_user_db_session(username) as session: 

154 queued = ( 

155 session.query(QueuedResearch) 

156 .filter_by(username=username, research_id=research_id) 

157 .first() 

158 ) 

159 

160 if not queued: 

161 return False 

162 

163 position = queued.position 

164 session.delete(queued) 

165 

166 # Update positions of items behind in queue 

167 session.query(QueuedResearch).filter( 

168 QueuedResearch.username == username, 

169 QueuedResearch.position > position, 

170 ).update({QueuedResearch.position: QueuedResearch.position - 1}) 

171 

172 _safe_commit(session) 

173 logger.info(f"Removed research {research_id} from queue") 

174 return True 

175 

176 @staticmethod 

177 def get_user_queue(username): 

178 """ 

179 Get all queued researches for a user 

180 

181 Args: 

182 username: User to get queue for 

183 

184 Returns: 

185 list: List of queued research info 

186 """ 

187 with get_user_db_session(username) as session: 

188 queued_items = ( 

189 session.query(QueuedResearch) 

190 .filter_by(username=username) 

191 .order_by(QueuedResearch.position) 

192 .all() 

193 ) 

194 

195 result = [] 

196 for item in queued_items: 

197 # Get research info 

198 research = ( 

199 session.query(ResearchHistory) 

200 .filter_by(id=item.research_id) 

201 .first() 

202 ) 

203 

204 if research: 

205 result.append( 

206 { 

207 "research_id": item.research_id, 

208 "query": item.query, 

209 "mode": item.mode, 

210 "position": item.position, 

211 "created_at": item.created_at.isoformat() 

212 if item.created_at 

213 else None, 

214 "is_processing": item.is_processing, 

215 } 

216 ) 

217 

218 return result