Coverage for src/infra/agent_runtime.py: 100%

1"""Agent runtime configuration builder.

2

3 This module provides AgentRuntimeBuilder for centralized agent runtime setup.

4 It consolidates duplicated configuration logic from AgentSessionRunner and

5 RunCoordinator into a single, testable builder.

6

7 Design principles:

8- Builder pattern with fluent API for configuration

9- All SDK imports local to build() method for lazy-import guarantees

10- Returns AgentRuntime dataclass with all components needed for session

11"""

12

13 from __future__ import annotations

14

15 import logging

16 import os

17 from dataclasses import dataclass, field

18 from typing import TYPE_CHECKING

19

20 from src.infra.tools.env import SCRIPTS_DIR, get_lock_dir

21

22 logger = logging.getLogger(__name__)

23

24 if TYPE_CHECKING:

25 from pathlib import Path

26

27 from src.core.protocols import DeadlockMonitorProtocol, SDKClientFactoryProtocol

28 from src.infra.hooks import FileReadCache, LintCache

29

30

31@dataclass

32 class AgentRuntime:

33 """Runtime configuration bundle for agent sessions.

34

35 Contains all components needed to run an agent session:

36 - SDK options for client creation

37 - Caches for file read and lint deduplication

38 - Environment variables for agent execution

39 - Hook lists for PreToolUse, PostToolUse, and Stop events

40

41 Attributes:

42 options: SDK ClaudeAgentOptions for session creation.

43 file_read_cache: Cache for blocking redundant file reads.

44 lint_cache: Cache for blocking redundant lint commands.

45 env: Environment variables for agent execution.

46 pre_tool_hooks: List of PreToolUse hook callables.

47 post_tool_hooks: List of PostToolUse hook callables.

48 stop_hooks: List of Stop hook callables.

49 """

50

51 options: object

52 file_read_cache: FileReadCache

53 lint_cache: LintCache

54 env: dict[str, str]

55 pre_tool_hooks: list[object] = field(default_factory=list)

56 post_tool_hooks: list[object] = field(default_factory=list)

57 stop_hooks: list[object] = field(default_factory=list)

58

59

60 class AgentRuntimeBuilder:

61 """Builder for agent runtime configuration.

62

63 Provides a fluent API for configuring agent sessions. Each .with_*()

64 method returns self for chaining. Call .build() to create the

65 AgentRuntime with all configured components.

66

67 Example:

68 runtime = (

69 AgentRuntimeBuilder(repo_path, agent_id, factory)

70 .with_hooks(deadlock_monitor=monitor)

71 .with_env()

72 .with_mcp()

73 .with_disallowed_tools()

74 .build()

75 )

76 # Use runtime.options to create SDK client

77

78 Attributes:

79 repo_path: Path to the repository root.

80 agent_id: Unique agent identifier for lock management.

81 sdk_client_factory: Factory for creating SDK options and matchers.

82 """

83

84 def __init__(

85 self,

86 repo_path: Path,

87 agent_id: str,

88 sdk_client_factory: SDKClientFactoryProtocol,

89 ) -> None:

90 """Initialize the builder.

91

92 Args:

93 repo_path: Path to the repository root.

94 agent_id: Unique agent identifier for lock management.

95 sdk_client_factory: Factory for creating SDK options and matchers.

96 """

97 self._repo_path = repo_path

98 self._agent_id = agent_id

99 self._sdk_client_factory = sdk_client_factory

100

101 # Lint tools configuration

102 self._lint_tools: set[str] | frozenset[str] | None = None

103

104 # Hook configuration

105 self._deadlock_monitor: DeadlockMonitorProtocol | None = None

106 self._include_stop_hook: bool = True

107 self._include_mala_disallowed_tools_hook: bool = True

108

109 # Environment and options

110 self._env: dict[str, str] | None = None

111 self._mcp_servers: object | None = None

112 self._disallowed_tools: list[str] | None = None

113

114 def with_hooks(

115 self,

116 *,

117 deadlock_monitor: DeadlockMonitorProtocol | None = None,

118 include_stop_hook: bool = True,

119 include_mala_disallowed_tools_hook: bool = True,

120 ) -> AgentRuntimeBuilder:

121 """Configure hook behavior.

122

123 Args:

124 deadlock_monitor: Optional DeadlockMonitor for lock event hooks.

125 include_stop_hook: Whether to include stop hook (default True).

126 include_mala_disallowed_tools_hook: Whether to include the

127 block_mala_disallowed_tools hook (default True). Set False

128 for fixer agents which don't need this restriction.

129

130 Returns:

131 Self for chaining.

132 """

133 self._deadlock_monitor = deadlock_monitor

134 self._include_stop_hook = include_stop_hook

135 self._include_mala_disallowed_tools_hook = include_mala_disallowed_tools_hook

136 return self

137

138 def with_env(self, extra: dict[str, str] | None = None) -> AgentRuntimeBuilder:

139 """Configure environment variables.

140

141 Builds standard environment with PATH, LOCK_DIR, AGENT_ID, etc.

142 Merges with os.environ and any extra variables provided.

143

144 Args:

145 extra: Additional environment variables to include.

146

147 Returns:

148 Self for chaining.

149 """

150 self._env = {

151 **os.environ,

152 "PATH": f"{SCRIPTS_DIR}:{os.environ.get('PATH', '')}",

153 "LOCK_DIR": str(get_lock_dir()),

154 "AGENT_ID": self._agent_id,

155 "REPO_NAMESPACE": str(self._repo_path),

156 "MCP_TIMEOUT": "300000",

157 }

158 if extra:

159 self._env.update(extra)

160 return self

161

162 def with_mcp(self, servers: object | None = None) -> AgentRuntimeBuilder:

163 """Configure MCP servers.

164

165 Args:

166 servers: MCP server configuration. If None, uses get_mcp_servers().

167

168 Returns:

169 Self for chaining.

170 """

171 if servers is not None:

172 self._mcp_servers = servers

173 else:

174 from src.infra.mcp import get_mcp_servers

175

176 self._mcp_servers = get_mcp_servers(self._repo_path)

177 return self

178

179 def with_disallowed_tools(

180 self, tools: list[str] | None = None

181 ) -> AgentRuntimeBuilder:

182 """Configure disallowed tools.

183

184 Args:

185 tools: List of disallowed tool names. If None, uses get_disallowed_tools().

186

187 Returns:

188 Self for chaining.

189 """

190 if tools is not None:

191 self._disallowed_tools = tools

192 else:

193 from src.infra.mcp import get_disallowed_tools

194

195 self._disallowed_tools = get_disallowed_tools()

196 return self

197

198 def with_lint_tools(

199 self, lint_tools: set[str] | frozenset[str] | None = None

200 ) -> AgentRuntimeBuilder:

201 """Configure lint tools for cache.

202

203 Args:

204 lint_tools: Set of lint tool names. If None, uses defaults.

205

206 Returns:

207 Self for chaining.

208 """

209 self._lint_tools = lint_tools

210 return self

211

212 def build(self) -> AgentRuntime:

213 """Build the agent runtime configuration.

214

215 Creates all caches, hooks, environment, and SDK options. All SDK

216 imports happen here to preserve lazy-import guarantees.

217

218 Returns:

219 AgentRuntime with all configured components.

220

221 Raises:

222 RuntimeError: If required configuration is missing.

223 """

224 # Import hooks locally for lazy-import guarantees

225 from src.infra.hooks import (

226 FileReadCache,

227 LintCache,

228 block_dangerous_commands,

229 block_mala_disallowed_tools,

230 make_file_read_cache_hook,

231 make_lint_cache_hook,

232 make_lock_enforcement_hook,

233 make_lock_event_hook,

234 make_lock_wait_hook,

235 make_stop_hook,

236 )

237

238 # Create caches

239 file_read_cache = FileReadCache()

240 lint_cache = LintCache(

241 repo_path=self._repo_path,

242 lint_tools=self._lint_tools,

243 )

244

245 # Build pre-tool hooks (order matters)

246 pre_tool_hooks: list[object] = [

247 block_dangerous_commands,

248 make_lock_enforcement_hook(self._agent_id, str(self._repo_path)),

249 make_file_read_cache_hook(file_read_cache),

250 make_lint_cache_hook(lint_cache),

251 ]

252

253 # Conditionally add mala disallowed tools hook (not needed for fixer agents)

254 if self._include_mala_disallowed_tools_hook:

255 pre_tool_hooks.insert(1, block_mala_disallowed_tools)

256

257 post_tool_hooks: list[object] = []

258 stop_hooks: list[object] = []

259

260 if self._include_stop_hook:

261 stop_hooks.append(make_stop_hook(self._agent_id))

262

263 # Add deadlock monitor hooks if configured

264 if self._deadlock_monitor is not None:

265 logger.info("Wiring deadlock monitor hooks: agent_id=%s", self._agent_id)

266 # Import LockEvent types here to inject into hooks

267 from src.core.models import LockEvent, LockEventType

268

269 monitor = self._deadlock_monitor

270 # PreToolUse hook for real-time WAITING detection on lock-wait.sh

271 pre_tool_hooks.append(

272 make_lock_wait_hook(

273 agent_id=self._agent_id,

274 emit_event=monitor.handle_event,

275 repo_namespace=str(self._repo_path),

276 lock_event_class=LockEvent,

277 lock_event_type_enum=LockEventType,

278 )

279 )

280 # PostToolUse hook for ACQUIRED/RELEASED events

281 post_tool_hooks.append(

282 make_lock_event_hook(

283 agent_id=self._agent_id,

284 emit_event=monitor.handle_event,

285 repo_namespace=str(self._repo_path),

286 lock_event_class=LockEvent,

287 lock_event_type_enum=LockEventType,

288 )

289 )

290 else:

291 logger.info("No deadlock monitor configured; skipping lock event hooks")

292

293 # Build environment if not explicitly set

294 if self._env is None:

295 self.with_env()

296 env = self._env or {}

297

298 # Build hooks dict using factory

299 make_matcher = self._sdk_client_factory.create_hook_matcher

300 hooks_dict: dict[str, list[object]] = {

301 "PreToolUse": [make_matcher(None, pre_tool_hooks)],

302 }

303 if stop_hooks:

304 hooks_dict["Stop"] = [make_matcher(None, stop_hooks)]

305 if post_tool_hooks:

306 hooks_dict["PostToolUse"] = [make_matcher(None, post_tool_hooks)]

307

308 logger.debug(

309 "Built hooks: PreToolUse=%d PostToolUse=%d Stop=%d",

310 len(pre_tool_hooks),

311 len(post_tool_hooks),

312 len(stop_hooks),

313 )

314

315 # Build SDK options

316 options = self._sdk_client_factory.create_options(

317 cwd=str(self._repo_path),

318 mcp_servers=self._mcp_servers,

319 disallowed_tools=self._disallowed_tools,

320 env=env,

321 hooks=hooks_dict,

322 )

323

324 return AgentRuntime(

325 options=options,

326 file_read_cache=file_read_cache,

327 lint_cache=lint_cache,

328 env=env,

329 pre_tool_hooks=pre_tool_hooks,

330 post_tool_hooks=post_tool_hooks,

331 stop_hooks=stop_hooks,

332 )