Coverage for src/workstack/status/orchestrator.py: 69%

1"""Orchestrator for collecting and assembling status information.""" 

2 

3import logging 

4from concurrent.futures import ThreadPoolExecutor, TimeoutError, as_completed 

5from pathlib import Path 

6 

7from workstack.core.context import WorkstackContext 

8from workstack.status.collectors.base import StatusCollector 

9from workstack.status.models.status_data import StatusData, WorktreeInfo 

10 

11logger = logging.getLogger(__name__) 

12 

13 

14class StatusOrchestrator: 

15 """Coordinates all status collectors and assembles final data. 

16 

17 The orchestrator runs collectors in parallel with timeouts to ensure 

18 responsive output even if some collectors are slow or fail. 

19 """ 

20 

21 def __init__(self, collectors: list[StatusCollector], *, timeout_seconds: float = 2.0) -> None: 

22 """Create a status orchestrator. 

23 

24 Args: 

25 collectors: List of status collectors to run 

26 timeout_seconds: Maximum time to wait for each collector (default: 2.0) 

27 """ 

28 self.collectors = collectors 

29 self.timeout_seconds = timeout_seconds 

30 

31 def collect_status( 

32 self, ctx: WorkstackContext, worktree_path: Path, repo_root: Path 

33 ) -> StatusData: 

34 """Collect all status information in parallel. 

35 

36 Each collector runs in its own thread with a timeout. Failed or slow 

37 collectors will return None for their section. 

38 

39 Args: 

40 ctx: Workstack context with operations 

41 worktree_path: Path to the worktree 

42 repo_root: Path to repository root 

43 

44 Returns: 

45 StatusData with all collected information 

46 """ 

47 # Determine worktree info 

48 worktree_info = self._get_worktree_info(ctx, worktree_path, repo_root) 

49 

50 # Run collectors in parallel 

51 results: dict[str, object] = {} 

52 

53 with ThreadPoolExecutor(max_workers=5) as executor: 

54 # Submit all available collectors 

55 futures = {} 

56 for collector in self.collectors: 

57 if collector.is_available(ctx, worktree_path): 

58 future = executor.submit(collector.collect, ctx, worktree_path, repo_root) 

59 futures[future] = collector.name 

60 

61 # Collect results with timeout per collector 

62 # Use a separate timeout for as_completed (total time for all collectors) 

63 total_timeout = self.timeout_seconds * len(futures) if futures else 1.0 

64 

65 try: 

66 for future in as_completed(futures, timeout=total_timeout): 

67 collector_name = futures[future] 

68 try: 

69 result = future.result(timeout=0.1) # Should be immediate once complete 

70 results[collector_name] = result 

71 except TimeoutError: 

72 # Error boundary: Collector timeouts shouldn't fail entire command 

73 # Log for debugging but continue with other collectors 

74 logger.debug( 

75 f"Collector '{collector_name}' timed out after {self.timeout_seconds}s" 

76 ) 

77 results[collector_name] = None 

78 except Exception as e: 

79 # Error boundary: Individual collector failures shouldn't fail 

80 # entire command. This is an acceptable use of exception handling 

81 # at error boundaries per EXCEPTION_HANDLING.md - parallel 

82 # collectors should degrade gracefully 

83 logger.debug(f"Collector '{collector_name}' failed: {e}") 

84 results[collector_name] = None 

85 except TimeoutError: 

86 # Some collectors didn't complete in time 

87 # Mark incomplete collectors as None 

88 for future, collector_name in futures.items(): 

89 if future.running() or not future.done(): 

90 logger.debug(f"Collector '{collector_name}' did not complete in time") 

91 results[collector_name] = None 

92 

93 # Get related worktrees 

94 related_worktrees = self._get_related_worktrees(ctx, repo_root, worktree_path) 

95 

96 # Assemble StatusData - cast results to expected types 

97 # Results are either the correct type or None (from collector failures) 

98 from workstack.status.models.status_data import ( 

99 DependencyStatus, 

100 EnvironmentStatus, 

101 GitStatus, 

102 PlanStatus, 

103 PullRequestStatus, 

104 StackPosition, 

105 ) 

106 

107 git_result = results.get("git") 

108 stack_result = results.get("stack") 

109 pr_result = results.get("pr") 

110 env_result = results.get("environment") 

111 deps_result = results.get("dependencies") 

112 plan_result = results.get("plan") 

113 

114 return StatusData( 

115 worktree_info=worktree_info, 

116 git_status=git_result if isinstance(git_result, GitStatus) else None, 

117 stack_position=stack_result if isinstance(stack_result, StackPosition) else None, 

118 pr_status=pr_result if isinstance(pr_result, PullRequestStatus) else None, 

119 environment=env_result if isinstance(env_result, EnvironmentStatus) else None, 

120 dependencies=deps_result if isinstance(deps_result, DependencyStatus) else None, 

121 plan=plan_result if isinstance(plan_result, PlanStatus) else None, 

122 related_worktrees=related_worktrees, 

123 ) 

124 

125 def _get_worktree_info( 

126 self, ctx: WorkstackContext, worktree_path: Path, repo_root: Path 

127 ) -> WorktreeInfo: 

128 """Get basic worktree information. 

129 

130 Args: 

131 ctx: Workstack context 

132 worktree_path: Path to worktree 

133 repo_root: Path to repository root 

134 

135 Returns: 

136 WorktreeInfo with basic information 

137 """ 

138 # Check paths exist before resolution to avoid OSError 

139 is_root = False 

140 if worktree_path.exists() and repo_root.exists(): 

141 is_root = worktree_path.resolve() == repo_root.resolve() 

142 

143 name = "root" if is_root else worktree_path.name 

144 branch = ctx.git_ops.get_current_branch(worktree_path) 

145 

146 return WorktreeInfo(name=name, path=worktree_path, branch=branch, is_root=is_root) 

147 

148 def _get_related_worktrees( 

149 self, ctx: WorkstackContext, repo_root: Path, current_path: Path 

150 ) -> list[WorktreeInfo]: 

151 """Get list of other worktrees in the repository. 

152 

153 Args: 

154 ctx: Workstack context 

155 repo_root: Path to repository root 

156 current_path: Path to current worktree (excluded from results) 

157 

158 Returns: 

159 List of WorktreeInfo for other worktrees 

160 """ 

161 worktrees = ctx.git_ops.list_worktrees(repo_root) 

162 

163 # Check paths exist before resolution to avoid OSError 

164 if not current_path.exists(): 

165 return [] 

166 

167 current_resolved = current_path.resolve() 

168 

169 related = [] 

170 for wt in worktrees: 

171 # Skip if worktree path doesn't exist 

172 if not wt.path.exists(): 

173 continue 

174 

175 wt_resolved = wt.path.resolve() 

176 

177 # Skip current worktree 

178 if wt_resolved == current_resolved: 

179 continue 

180 

181 # Determine if this is the root worktree 

182 is_root = False 

183 if repo_root.exists(): 

184 is_root = wt_resolved == repo_root.resolve() 

185 

186 name = "root" if is_root else wt.path.name 

187 

188 related.append(WorktreeInfo(name=name, path=wt.path, branch=wt.branch, is_root=is_root)) 

189 

190 return related