agent-runtime-core 0.6.0__tar.gz → 0.7.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/PKG-INFO +1 -1
  2. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/__init__.py +10 -1
  3. agent_runtime_core-0.7.0/agent_runtime_core/contexts.py +348 -0
  4. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/pyproject.toml +1 -1
  5. agent_runtime_core-0.7.0/tests/test_contexts.py +263 -0
  6. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/.gitignore +0 -0
  7. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/LICENSE +0 -0
  8. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/README.md +0 -0
  9. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/config.py +0 -0
  10. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/events/__init__.py +0 -0
  11. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/events/base.py +0 -0
  12. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/events/memory.py +0 -0
  13. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/events/redis.py +0 -0
  14. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/events/sqlite.py +0 -0
  15. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/interfaces.py +0 -0
  16. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/llm/__init__.py +0 -0
  17. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/llm/anthropic.py +0 -0
  18. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/llm/litellm_client.py +0 -0
  19. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/llm/openai.py +0 -0
  20. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/persistence/__init__.py +0 -0
  21. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/persistence/base.py +0 -0
  22. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/persistence/file.py +0 -0
  23. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/persistence/manager.py +0 -0
  24. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/queue/__init__.py +0 -0
  25. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/queue/base.py +0 -0
  26. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/queue/memory.py +0 -0
  27. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/queue/redis.py +0 -0
  28. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/queue/sqlite.py +0 -0
  29. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/registry.py +0 -0
  30. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/runner.py +0 -0
  31. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/state/__init__.py +0 -0
  32. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/state/base.py +0 -0
  33. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/state/memory.py +0 -0
  34. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/state/redis.py +0 -0
  35. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/state/sqlite.py +0 -0
  36. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/steps.py +0 -0
  37. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/testing.py +0 -0
  38. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/tool_calling_agent.py +0 -0
  39. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/tracing/__init__.py +0 -0
  40. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/tracing/langfuse.py +0 -0
  41. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/agent_runtime_core/tracing/noop.py +0 -0
  42. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/__init__.py +0 -0
  43. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/test_events.py +0 -0
  44. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/test_imports.py +0 -0
  45. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/test_persistence.py +0 -0
  46. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/test_queue.py +0 -0
  47. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/test_state.py +0 -0
  48. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/test_steps.py +0 -0
  49. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/tests/test_testing.py +0 -0
  50. {agent_runtime_core-0.6.0 → agent_runtime_core-0.7.0}/uv.lock +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agent-runtime-core
3
- Version: 0.6.0
3
+ Version: 0.7.0
4
4
  Summary: Framework-agnostic Python library for executing AI agents with consistent patterns
5
5
  Project-URL: Homepage, https://github.com/makemore/agent-runtime-core
6
6
  Project-URL: Repository, https://github.com/makemore/agent-runtime-core
@@ -34,7 +34,7 @@ Example usage:
34
34
  return RunResult(final_output={"message": "Hello!"})
35
35
  """
36
36
 
37
- __version__ = "0.6.0"
37
+ __version__ = "0.7.0"
38
38
 
39
39
  # Core interfaces
40
40
  from agent_runtime_core.interfaces import (
@@ -91,6 +91,12 @@ from agent_runtime_core.steps import (
91
91
  StepCancelledError,
92
92
  )
93
93
 
94
+ # Concrete RunContext implementations for different use cases
95
+ from agent_runtime_core.contexts import (
96
+ InMemoryRunContext,
97
+ FileRunContext,
98
+ )
99
+
94
100
  # Testing utilities
95
101
  from agent_runtime_core.testing import (
96
102
  MockRunContext,
@@ -169,6 +175,9 @@ __all__ = [
169
175
  "ExecutionState",
170
176
  "StepExecutionError",
171
177
  "StepCancelledError",
178
+ # Concrete RunContext implementations
179
+ "InMemoryRunContext",
180
+ "FileRunContext",
172
181
  # Testing
173
182
  "MockRunContext",
174
183
  "MockLLMClient",
@@ -0,0 +1,348 @@
1
+ """
2
+ Concrete RunContext implementations for different use cases.
3
+
4
+ These implementations satisfy the RunContext protocol and can be used
5
+ directly with StepExecutor and agent runtimes.
6
+
7
+ Usage:
8
+ # For simple scripts (in-memory, no persistence)
9
+ ctx = InMemoryRunContext(run_id=uuid4())
10
+
11
+ # For scripts that need persistence across restarts
12
+ ctx = FileRunContext(run_id=uuid4(), checkpoint_dir="./checkpoints")
13
+
14
+ # Use with StepExecutor
15
+ from agent_runtime_core.steps import StepExecutor, Step
16
+ executor = StepExecutor(ctx)
17
+ results = await executor.run([
18
+ Step("fetch", fetch_data),
19
+ Step("process", process_data),
20
+ ])
21
+ """
22
+
23
+ import json
24
+ import os
25
+ from datetime import datetime
26
+ from pathlib import Path
27
+ from typing import Any, Callable, Optional
28
+ from uuid import UUID, uuid4
29
+
30
+ from agent_runtime_core.interfaces import EventType, Message, ToolRegistry
31
+
32
+
33
+ class InMemoryRunContext:
34
+ """
35
+ In-memory RunContext implementation.
36
+
37
+ Good for:
38
+ - Unit testing
39
+ - Simple scripts that don't need persistence
40
+ - Development and prototyping
41
+
42
+ State is lost when the process exits.
43
+
44
+ Example:
45
+ ctx = InMemoryRunContext(
46
+ run_id=uuid4(),
47
+ input_messages=[{"role": "user", "content": "Hello"}],
48
+ )
49
+
50
+ # Use with an agent
51
+ result = await my_agent.run(ctx)
52
+
53
+ # Or with StepExecutor
54
+ executor = StepExecutor(ctx)
55
+ results = await executor.run(steps)
56
+ """
57
+
58
+ def __init__(
59
+ self,
60
+ run_id: Optional[UUID] = None,
61
+ *,
62
+ conversation_id: Optional[UUID] = None,
63
+ input_messages: Optional[list[Message]] = None,
64
+ params: Optional[dict] = None,
65
+ metadata: Optional[dict] = None,
66
+ tool_registry: Optional[ToolRegistry] = None,
67
+ on_event: Optional[Callable[[str, dict], None]] = None,
68
+ ):
69
+ """
70
+ Initialize an in-memory run context.
71
+
72
+ Args:
73
+ run_id: Unique identifier for this run (auto-generated if not provided)
74
+ conversation_id: Associated conversation ID (optional)
75
+ input_messages: Input messages for this run
76
+ params: Additional parameters
77
+ metadata: Run metadata
78
+ tool_registry: Registry of available tools
79
+ on_event: Optional callback for events (for testing/debugging)
80
+ """
81
+ self._run_id = run_id or uuid4()
82
+ self._conversation_id = conversation_id
83
+ self._input_messages = input_messages or []
84
+ self._params = params or {}
85
+ self._metadata = metadata or {}
86
+ self._tool_registry = tool_registry or ToolRegistry()
87
+ self._cancelled = False
88
+ self._state: Optional[dict] = None
89
+ self._events: list[dict] = []
90
+ self._on_event = on_event
91
+
92
+ @property
93
+ def run_id(self) -> UUID:
94
+ """Unique identifier for this run."""
95
+ return self._run_id
96
+
97
+ @property
98
+ def conversation_id(self) -> Optional[UUID]:
99
+ """Conversation this run belongs to (if any)."""
100
+ return self._conversation_id
101
+
102
+ @property
103
+ def input_messages(self) -> list[Message]:
104
+ """Input messages for this run."""
105
+ return self._input_messages
106
+
107
+ @property
108
+ def params(self) -> dict:
109
+ """Additional parameters for this run."""
110
+ return self._params
111
+
112
+ @property
113
+ def metadata(self) -> dict:
114
+ """Metadata associated with this run."""
115
+ return self._metadata
116
+
117
+ @property
118
+ def tool_registry(self) -> ToolRegistry:
119
+ """Registry of available tools for this agent."""
120
+ return self._tool_registry
121
+
122
+ async def emit(self, event_type: EventType | str, payload: dict) -> None:
123
+ """Emit an event (stored in memory)."""
124
+ event_type_str = event_type.value if hasattr(event_type, 'value') else str(event_type)
125
+ event = {
126
+ "event_type": event_type_str,
127
+ "payload": payload,
128
+ "timestamp": datetime.utcnow().isoformat(),
129
+ }
130
+ self._events.append(event)
131
+ if self._on_event:
132
+ self._on_event(event_type_str, payload)
133
+
134
+ async def checkpoint(self, state: dict) -> None:
135
+ """Save a state checkpoint (in memory)."""
136
+ self._state = state
137
+
138
+ async def get_state(self) -> Optional[dict]:
139
+ """Get the last checkpointed state."""
140
+ return self._state
141
+
142
+ def cancelled(self) -> bool:
143
+ """Check if cancellation has been requested."""
144
+ return self._cancelled
145
+
146
+ def cancel(self) -> None:
147
+ """Request cancellation of this run."""
148
+ self._cancelled = True
149
+
150
+ @property
151
+ def events(self) -> list[dict]:
152
+ """Get all emitted events (for testing/debugging)."""
153
+ return self._events.copy()
154
+
155
+ def clear_events(self) -> None:
156
+ """Clear all events (for testing)."""
157
+ self._events.clear()
158
+
159
+
160
+ class FileRunContext:
161
+ """
162
+ File-based RunContext implementation with persistent checkpoints.
163
+
164
+ Good for:
165
+ - Scripts that need to resume after restart
166
+ - Long-running processes without a database
167
+ - Simple persistence without external dependencies
168
+
169
+ Checkpoints are saved as JSON files in the specified directory.
170
+
171
+ Example:
172
+ ctx = FileRunContext(
173
+ run_id=uuid4(),
174
+ checkpoint_dir="./checkpoints",
175
+ input_messages=[{"role": "user", "content": "Process this"}],
176
+ )
177
+
178
+ # Checkpoints are saved to ./checkpoints/{run_id}.json
179
+ executor = StepExecutor(ctx)
180
+ results = await executor.run(steps)
181
+
182
+ # To resume after restart, use the same run_id:
183
+ ctx = FileRunContext(run_id=previous_run_id, checkpoint_dir="./checkpoints")
184
+ results = await executor.run(steps, resume=True)
185
+ """
186
+
187
+ def __init__(
188
+ self,
189
+ run_id: Optional[UUID] = None,
190
+ *,
191
+ checkpoint_dir: str = "./checkpoints",
192
+ conversation_id: Optional[UUID] = None,
193
+ input_messages: Optional[list[Message]] = None,
194
+ params: Optional[dict] = None,
195
+ metadata: Optional[dict] = None,
196
+ tool_registry: Optional[ToolRegistry] = None,
197
+ on_event: Optional[Callable[[str, dict], None]] = None,
198
+ ):
199
+ """
200
+ Initialize a file-based run context.
201
+
202
+ Args:
203
+ run_id: Unique identifier for this run (auto-generated if not provided)
204
+ checkpoint_dir: Directory to store checkpoint files
205
+ conversation_id: Associated conversation ID (optional)
206
+ input_messages: Input messages for this run
207
+ params: Additional parameters
208
+ metadata: Run metadata
209
+ tool_registry: Registry of available tools
210
+ on_event: Optional callback for events
211
+ """
212
+ self._run_id = run_id or uuid4()
213
+ self._checkpoint_dir = Path(checkpoint_dir)
214
+ self._conversation_id = conversation_id
215
+ self._input_messages = input_messages or []
216
+ self._params = params or {}
217
+ self._metadata = metadata or {}
218
+ self._tool_registry = tool_registry or ToolRegistry()
219
+ self._cancelled = False
220
+ self._on_event = on_event
221
+ self._state_cache: Optional[dict] = None
222
+
223
+ # Ensure checkpoint directory exists
224
+ self._checkpoint_dir.mkdir(parents=True, exist_ok=True)
225
+
226
+ @property
227
+ def run_id(self) -> UUID:
228
+ """Unique identifier for this run."""
229
+ return self._run_id
230
+
231
+ @property
232
+ def conversation_id(self) -> Optional[UUID]:
233
+ """Conversation this run belongs to (if any)."""
234
+ return self._conversation_id
235
+
236
+ @property
237
+ def input_messages(self) -> list[Message]:
238
+ """Input messages for this run."""
239
+ return self._input_messages
240
+
241
+ @property
242
+ def params(self) -> dict:
243
+ """Additional parameters for this run."""
244
+ return self._params
245
+
246
+ @property
247
+ def metadata(self) -> dict:
248
+ """Metadata associated with this run."""
249
+ return self._metadata
250
+
251
+ @property
252
+ def tool_registry(self) -> ToolRegistry:
253
+ """Registry of available tools for this agent."""
254
+ return self._tool_registry
255
+
256
+ def _checkpoint_path(self) -> Path:
257
+ """Get the path to the checkpoint file for this run."""
258
+ return self._checkpoint_dir / f"{self._run_id}.json"
259
+
260
+ def _events_path(self) -> Path:
261
+ """Get the path to the events file for this run."""
262
+ return self._checkpoint_dir / f"{self._run_id}_events.jsonl"
263
+
264
+ async def emit(self, event_type: EventType | str, payload: dict) -> None:
265
+ """Emit an event (appended to events file)."""
266
+ event_type_str = event_type.value if hasattr(event_type, 'value') else str(event_type)
267
+ event = {
268
+ "event_type": event_type_str,
269
+ "payload": payload,
270
+ "timestamp": datetime.utcnow().isoformat(),
271
+ }
272
+
273
+ # Append to events file (JSONL format)
274
+ with open(self._events_path(), "a") as f:
275
+ f.write(json.dumps(event) + "\n")
276
+
277
+ if self._on_event:
278
+ self._on_event(event_type_str, payload)
279
+
280
+ async def checkpoint(self, state: dict) -> None:
281
+ """Save a state checkpoint to file."""
282
+ self._state_cache = state
283
+ checkpoint_data = {
284
+ "run_id": str(self._run_id),
285
+ "state": state,
286
+ "updated_at": datetime.utcnow().isoformat(),
287
+ }
288
+
289
+ # Write atomically using temp file
290
+ temp_path = self._checkpoint_path().with_suffix(".tmp")
291
+ with open(temp_path, "w") as f:
292
+ json.dump(checkpoint_data, f, indent=2)
293
+ temp_path.rename(self._checkpoint_path())
294
+
295
+ async def get_state(self) -> Optional[dict]:
296
+ """Get the last checkpointed state from file."""
297
+ if self._state_cache is not None:
298
+ return self._state_cache
299
+
300
+ checkpoint_path = self._checkpoint_path()
301
+ if not checkpoint_path.exists():
302
+ return None
303
+
304
+ try:
305
+ with open(checkpoint_path) as f:
306
+ data = json.load(f)
307
+ self._state_cache = data.get("state")
308
+ return self._state_cache
309
+ except (json.JSONDecodeError, IOError):
310
+ return None
311
+
312
+ def cancelled(self) -> bool:
313
+ """Check if cancellation has been requested."""
314
+ return self._cancelled
315
+
316
+ def cancel(self) -> None:
317
+ """Request cancellation of this run."""
318
+ self._cancelled = True
319
+
320
+ def get_events(self) -> list[dict]:
321
+ """Read all events from the events file."""
322
+ events_path = self._events_path()
323
+ if not events_path.exists():
324
+ return []
325
+
326
+ events = []
327
+ with open(events_path) as f:
328
+ for line in f:
329
+ line = line.strip()
330
+ if line:
331
+ try:
332
+ events.append(json.loads(line))
333
+ except json.JSONDecodeError:
334
+ pass
335
+ return events
336
+
337
+ def clear(self) -> None:
338
+ """Delete checkpoint and events files for this run."""
339
+ checkpoint_path = self._checkpoint_path()
340
+ events_path = self._events_path()
341
+
342
+ if checkpoint_path.exists():
343
+ checkpoint_path.unlink()
344
+ if events_path.exists():
345
+ events_path.unlink()
346
+
347
+ self._state_cache = None
348
+
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "agent-runtime-core"
7
- version = "0.6.0"
7
+ version = "0.7.0"
8
8
  description = "Framework-agnostic Python library for executing AI agents with consistent patterns"
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -0,0 +1,263 @@
1
+ """
2
+ Tests for concrete RunContext implementations.
3
+ """
4
+
5
+ import asyncio
6
+ import json
7
+ import tempfile
8
+ from pathlib import Path
9
+ from uuid import uuid4
10
+
11
+ import pytest
12
+
13
+ from agent_runtime_core.contexts import InMemoryRunContext, FileRunContext
14
+ from agent_runtime_core.interfaces import EventType, ToolRegistry
15
+ from agent_runtime_core.steps import Step, StepExecutor
16
+
17
+
18
+ class TestInMemoryRunContext:
19
+ """Tests for InMemoryRunContext."""
20
+
21
+ def test_init_defaults(self):
22
+ """Test initialization with defaults."""
23
+ ctx = InMemoryRunContext()
24
+
25
+ assert ctx.run_id is not None
26
+ assert ctx.conversation_id is None
27
+ assert ctx.input_messages == []
28
+ assert ctx.params == {}
29
+ assert ctx.metadata == {}
30
+ assert isinstance(ctx.tool_registry, ToolRegistry)
31
+ assert ctx.cancelled() is False
32
+
33
+ def test_init_with_values(self):
34
+ """Test initialization with custom values."""
35
+ run_id = uuid4()
36
+ conv_id = uuid4()
37
+ messages = [{"role": "user", "content": "Hello"}]
38
+ params = {"temperature": 0.7}
39
+ metadata = {"user_id": "123"}
40
+
41
+ ctx = InMemoryRunContext(
42
+ run_id=run_id,
43
+ conversation_id=conv_id,
44
+ input_messages=messages,
45
+ params=params,
46
+ metadata=metadata,
47
+ )
48
+
49
+ assert ctx.run_id == run_id
50
+ assert ctx.conversation_id == conv_id
51
+ assert ctx.input_messages == messages
52
+ assert ctx.params == params
53
+ assert ctx.metadata == metadata
54
+
55
+ @pytest.mark.asyncio
56
+ async def test_checkpoint_and_get_state(self):
57
+ """Test checkpointing and retrieving state."""
58
+ ctx = InMemoryRunContext()
59
+
60
+ # Initially no state
61
+ state = await ctx.get_state()
62
+ assert state is None
63
+
64
+ # Checkpoint some state
65
+ await ctx.checkpoint({"step": 1, "data": "test"})
66
+
67
+ # Retrieve state
68
+ state = await ctx.get_state()
69
+ assert state == {"step": 1, "data": "test"}
70
+
71
+ # Update state
72
+ await ctx.checkpoint({"step": 2, "data": "updated"})
73
+ state = await ctx.get_state()
74
+ assert state == {"step": 2, "data": "updated"}
75
+
76
+ @pytest.mark.asyncio
77
+ async def test_emit_events(self):
78
+ """Test event emission."""
79
+ events_received = []
80
+
81
+ def on_event(event_type, payload):
82
+ events_received.append((event_type, payload))
83
+
84
+ ctx = InMemoryRunContext(on_event=on_event)
85
+
86
+ await ctx.emit(EventType.RUN_STARTED, {"agent": "test"})
87
+ await ctx.emit("custom.event", {"data": "value"})
88
+
89
+ # Check callback was called
90
+ assert len(events_received) == 2
91
+ assert events_received[0] == ("run.started", {"agent": "test"})
92
+ assert events_received[1] == ("custom.event", {"data": "value"})
93
+
94
+ # Check events are stored
95
+ assert len(ctx.events) == 2
96
+ assert ctx.events[0]["event_type"] == "run.started"
97
+ assert ctx.events[1]["event_type"] == "custom.event"
98
+
99
+ def test_cancellation(self):
100
+ """Test cancellation."""
101
+ ctx = InMemoryRunContext()
102
+
103
+ assert ctx.cancelled() is False
104
+ ctx.cancel()
105
+ assert ctx.cancelled() is True
106
+
107
+
108
+ class TestFileRunContext:
109
+ """Tests for FileRunContext."""
110
+
111
+ @pytest.fixture
112
+ def temp_dir(self):
113
+ """Create a temporary directory for checkpoints."""
114
+ with tempfile.TemporaryDirectory() as tmpdir:
115
+ yield tmpdir
116
+
117
+ def test_init_creates_directory(self, temp_dir):
118
+ """Test that init creates the checkpoint directory."""
119
+ checkpoint_dir = Path(temp_dir) / "nested" / "checkpoints"
120
+ ctx = FileRunContext(checkpoint_dir=str(checkpoint_dir))
121
+
122
+ assert checkpoint_dir.exists()
123
+
124
+ @pytest.mark.asyncio
125
+ async def test_checkpoint_and_get_state(self, temp_dir):
126
+ """Test checkpointing and retrieving state from file."""
127
+ run_id = uuid4()
128
+ ctx = FileRunContext(run_id=run_id, checkpoint_dir=temp_dir)
129
+
130
+ # Initially no state
131
+ state = await ctx.get_state()
132
+ assert state is None
133
+
134
+ # Checkpoint some state
135
+ await ctx.checkpoint({"step": 1, "data": "test"})
136
+
137
+ # Verify file was created
138
+ checkpoint_path = Path(temp_dir) / f"{run_id}.json"
139
+ assert checkpoint_path.exists()
140
+
141
+ # Retrieve state
142
+ state = await ctx.get_state()
143
+ assert state == {"step": 1, "data": "test"}
144
+
145
+ # Create new context with same run_id - should load existing state
146
+ ctx2 = FileRunContext(run_id=run_id, checkpoint_dir=temp_dir)
147
+ state2 = await ctx2.get_state()
148
+ assert state2 == {"step": 1, "data": "test"}
149
+
150
+ @pytest.mark.asyncio
151
+ async def test_emit_events_to_file(self, temp_dir):
152
+ """Test event emission to file."""
153
+ run_id = uuid4()
154
+ ctx = FileRunContext(run_id=run_id, checkpoint_dir=temp_dir)
155
+
156
+ await ctx.emit(EventType.RUN_STARTED, {"agent": "test"})
157
+ await ctx.emit(EventType.STEP_COMPLETED, {"step": "fetch"})
158
+
159
+ # Verify events file was created
160
+ events_path = Path(temp_dir) / f"{run_id}_events.jsonl"
161
+ assert events_path.exists()
162
+
163
+ # Read events back
164
+ events = ctx.get_events()
165
+ assert len(events) == 2
166
+ assert events[0]["event_type"] == "run.started"
167
+ assert events[1]["event_type"] == "step.completed"
168
+
169
+ @pytest.mark.asyncio
170
+ async def test_clear(self, temp_dir):
171
+ """Test clearing checkpoint and events."""
172
+ run_id = uuid4()
173
+ ctx = FileRunContext(run_id=run_id, checkpoint_dir=temp_dir)
174
+
175
+ await ctx.checkpoint({"data": "test"})
176
+ await ctx.emit(EventType.RUN_STARTED, {})
177
+
178
+ checkpoint_path = Path(temp_dir) / f"{run_id}.json"
179
+ events_path = Path(temp_dir) / f"{run_id}_events.jsonl"
180
+
181
+ assert checkpoint_path.exists()
182
+ assert events_path.exists()
183
+
184
+ ctx.clear()
185
+
186
+ assert not checkpoint_path.exists()
187
+ assert not events_path.exists()
188
+ assert await ctx.get_state() is None
189
+
190
+
191
+ class TestContextsWithStepExecutor:
192
+ """Test that contexts work correctly with StepExecutor."""
193
+
194
+ @pytest.mark.asyncio
195
+ async def test_in_memory_context_with_executor(self):
196
+ """Test InMemoryRunContext with StepExecutor."""
197
+ ctx = InMemoryRunContext()
198
+
199
+ async def step1(ctx, state):
200
+ state["step1"] = "done"
201
+ return "step1_result"
202
+
203
+ async def step2(ctx, state):
204
+ assert state["step1"] == "done"
205
+ return "step2_result"
206
+
207
+ executor = StepExecutor(ctx)
208
+ results = await executor.run([
209
+ Step("step1", step1),
210
+ Step("step2", step2),
211
+ ])
212
+
213
+ assert results["step1"] == "step1_result"
214
+ assert results["step2"] == "step2_result"
215
+
216
+ @pytest.mark.asyncio
217
+ async def test_file_context_resume(self):
218
+ """Test FileRunContext resume capability."""
219
+ with tempfile.TemporaryDirectory() as temp_dir:
220
+ run_id = uuid4()
221
+ execution_count = {"step1": 0, "step2": 0}
222
+
223
+ async def step1(ctx, state):
224
+ execution_count["step1"] += 1
225
+ return "step1_result"
226
+
227
+ async def step2(ctx, state):
228
+ execution_count["step2"] += 1
229
+ # Simulate failure on first attempt
230
+ if execution_count["step2"] == 1:
231
+ raise RuntimeError("Simulated failure")
232
+ return "step2_result"
233
+
234
+ # First run - step1 succeeds, step2 fails
235
+ ctx1 = FileRunContext(run_id=run_id, checkpoint_dir=temp_dir)
236
+ executor1 = StepExecutor(ctx1)
237
+
238
+ with pytest.raises(Exception):
239
+ await executor1.run([
240
+ Step("step1", step1),
241
+ Step("step2", step2),
242
+ ])
243
+
244
+ assert execution_count["step1"] == 1
245
+ assert execution_count["step2"] == 1
246
+
247
+ # Second run - should resume from checkpoint, skip step1
248
+ ctx2 = FileRunContext(run_id=run_id, checkpoint_dir=temp_dir)
249
+ executor2 = StepExecutor(ctx2)
250
+
251
+ results = await executor2.run([
252
+ Step("step1", step1),
253
+ Step("step2", step2),
254
+ ], resume=True)
255
+
256
+ # step1 should NOT have been re-executed
257
+ assert execution_count["step1"] == 1
258
+ # step2 should have been retried
259
+ assert execution_count["step2"] == 2
260
+
261
+ assert results["step1"] == "step1_result"
262
+ assert results["step2"] == "step2_result"
263
+