Coverage for src/pipeline/context_pressure_handler.py: 100%

1"""ContextPressureHandler: Context pressure checkpoint and restart handling.

2

3 Extracted from AgentSessionRunner to separate checkpoint/restart logic from

4 session lifecycle management. This module handles:

5- Fetching checkpoints from agents via SDK client

6- Building continuation prompts with checkpoint context

7- Managing restart loop state (continuation count, prompts)

8

9 Design principles:

10- Protocol-based SDK client for testability

11- Explicit error type for pressure threshold detection

12- Timeout handling for checkpoint fetch operations

13"""

14

15 from __future__ import annotations

16

17 import asyncio

18 import logging

19 from dataclasses import dataclass

20 from typing import (

21 TYPE_CHECKING,

22 Any,

23 cast,

24)

25

26 from src.domain.prompts import (

27 build_continuation_prompt,

28 extract_checkpoint,

29)

30 from src.pipeline.message_stream_processor import ContextPressureError

31

32 if TYPE_CHECKING:

33 from src.core.protocols import (

34 SDKClientFactoryProtocol,

35 )

36

37

38 logger = logging.getLogger(__name__)

39

40# Default timeout for checkpoint fetch operations (30 seconds)

41 DEFAULT_CHECKPOINT_TIMEOUT_SECONDS = 30

42

43

44@dataclass

45 class ContextPressureConfig:

46 """Configuration for context pressure handling.

47

48 Attributes:

49 checkpoint_request_prompt: Prompt template to request checkpoint from agent.

50 continuation_template: Template for continuation prompt with checkpoint.

51 checkpoint_timeout_seconds: Timeout for checkpoint fetch operations.

52 """

53

54 checkpoint_request_prompt: str

55 continuation_template: str

56 checkpoint_timeout_seconds: float = DEFAULT_CHECKPOINT_TIMEOUT_SECONDS

57

58

59@dataclass

60 class CheckpointResult:

61 """Result of fetching a checkpoint from an agent.

62

63 Attributes:

64 checkpoint: Extracted checkpoint text (may be empty).

65 timed_out: Whether the fetch timed out.

66 """

67

68 checkpoint: str

69 timed_out: bool = False

70

71

72 class ContextPressureHandler:

73 """Handles context pressure detection and checkpoint/restart logic.

74

75 This handler encapsulates:

76 - Fetching checkpoints from agents before restart

77 - Building continuation prompts with checkpoint context

78 - Managing restart loop state

79

80 The handler is stateless per-call; restart state is managed by the caller.

81 """

82

83 def __init__(

84 self,

85 config: ContextPressureConfig,

86 sdk_client_factory: SDKClientFactoryProtocol,

87 ) -> None:

88 """Initialize the handler.

89

90 Args:

91 config: Context pressure configuration.

92 sdk_client_factory: Factory for creating SDK clients.

93 """

94 self.config = config

95 self.sdk_client_factory = sdk_client_factory

96

97 async def fetch_checkpoint(

98 self,

99 session_id: str,

100 issue_id: str,

101 options: object,

102 timeout_seconds: float | None = None,

103 ) -> CheckpointResult:

104 """Fetch checkpoint from agent before context restart.

105

106 Sends checkpoint_request_prompt to the current session and extracts

107 the checkpoint block from the response.

108

109 Args:

110 session_id: SDK session ID from ContextPressureError.

111 issue_id: Issue ID for logging.

112 options: SDK client options.

113 timeout_seconds: Optional override for checkpoint timeout.

114

115 Returns:

116 CheckpointResult with extracted checkpoint text.

117 """

118 effective_timeout = (

119 timeout_seconds

120 if timeout_seconds is not None

121 else self.config.checkpoint_timeout_seconds

122 )

123

124 logger.info(

125 "Session %s: requesting checkpoint from session %s...",

126 issue_id,

127 session_id[:8] if session_id else "unknown",

128 )

129

130 checkpoint_prompt = self.config.checkpoint_request_prompt

131 if not checkpoint_prompt:

132 logger.warning(

133 "Session %s: no checkpoint_request prompt configured, using empty checkpoint",

134 issue_id,

135 )

136 return CheckpointResult(checkpoint="", timed_out=False)

137

138 # Create client to query for checkpoint

139 client = self.sdk_client_factory.create(options)

140

141 response_text = ""

142 try:

143 async with asyncio.timeout(effective_timeout):

144 async with client:

145 await client.query(checkpoint_prompt, session_id=session_id)

146 async for message in client.receive_response():

147 # Extract text from AssistantMessage

148 content = getattr(message, "content", None)

149 if content is not None:

150 for block in cast("list[Any]", content):

151 text = getattr(block, "text", None)

152 if text is not None:

153 response_text += text

154 except TimeoutError:

155 logger.warning(

156 "Session %s: checkpoint fetch timed out after %.1fs, using empty checkpoint",

157 issue_id,

158 effective_timeout,

159 )

160 return CheckpointResult(checkpoint="", timed_out=True)

161 except Exception as e:

162 logger.warning(

163 "Session %s: checkpoint query failed: %s, using empty checkpoint",

164 issue_id,

165 e,

166 )

167 return CheckpointResult(checkpoint="", timed_out=False)

168

169 # Extract checkpoint from response

170 checkpoint = extract_checkpoint(response_text)

171 logger.debug(

172 "Session %s: extracted checkpoint (%d chars)",

173 issue_id,

174 len(checkpoint),

175 )

176 return CheckpointResult(checkpoint=checkpoint, timed_out=False)

177

178 def build_continuation_prompt(self, checkpoint: str) -> str:

179 """Build continuation prompt with checkpoint context.

180

181 Args:

182 checkpoint: Extracted checkpoint text from previous session.

183

184 Returns:

185 Formatted continuation prompt for restart.

186 """

187 continuation_template = self.config.continuation_template

188 if continuation_template:

189 return build_continuation_prompt(continuation_template, checkpoint)

190 # Fallback: just use checkpoint as prompt

191 return f"Continue from checkpoint:\n\n{checkpoint}"

192

193 async def handle_pressure_error(

194 self,

195 error: ContextPressureError,

196 issue_id: str,

197 options: object,

198 continuation_count: int,

199 remaining_time: float,

200 ) -> tuple[str, int]:

201 """Handle a ContextPressureError by fetching checkpoint and building continuation.

202

203 This is the main entry point for handling context pressure. It:

204 1. Calculates effective timeout based on remaining session time

205 2. Fetches checkpoint from the current session

206 3. Builds the continuation prompt

207 4. Logs restart information

208

209 Args:

210 error: The ContextPressureError that triggered the restart.

211 issue_id: Issue ID for logging.

212 options: SDK client options.

213 continuation_count: Current restart count (will be incremented).

214 remaining_time: Remaining session time in seconds.

215

216 Returns:

217 Tuple of (continuation_prompt, new_continuation_count).

218 """

219 # Calculate effective timeout bounded by remaining session time

220 effective_timeout = max(

221 0, min(remaining_time, self.config.checkpoint_timeout_seconds)

222 )

223

224 # Fetch checkpoint from current session

225 result = await self.fetch_checkpoint(

226 session_id=error.session_id,

227 issue_id=issue_id,

228 options=options,

229 timeout_seconds=effective_timeout,

230 )

231

232 new_count = continuation_count + 1

233 logger.info(

234 "Session %s: context restart #%d at %.1f%%",

235 issue_id,

236 new_count,

237 error.pressure_ratio * 100,

238 )

239

240 # Build continuation prompt

241 continuation_prompt = self.build_continuation_prompt(result.checkpoint)

242

243 return continuation_prompt, new_count

244

245

246# Re-export ContextPressureError for convenience

247 __all__ = [

248 "DEFAULT_CHECKPOINT_TIMEOUT_SECONDS",

249 "CheckpointResult",

250 "ContextPressureConfig",

251 "ContextPressureError",

252 "ContextPressureHandler",

253]