1"""
2Sandbox Manager for Hosted Agent Infrastructure.
3 
4Use when: building background coding agents that need sandboxed execution
5environments with pre-built images, warm pools, and session snapshots.
6 
7This module provides composable building blocks for sandbox lifecycle
8management. Each class handles one concern (image building, warm pools,
9session coordination) and can be used independently or combined via
10SandboxManager.
11 
12Note: This is pseudocode demonstrating architectural patterns.
13Adapt for your specific infrastructure (Modal, Fly.io, etc.).
14"""
15 
16from dataclasses import dataclass, field
17from datetime import datetime, timedelta
18from typing import Optional, Callable, Any
19from enum import Enum
20import asyncio
21 
22__all__ = [
23    "SandboxState",
24    "UserIdentity",
25    "SandboxConfig",
26    "Sandbox",
27    "RepositoryImage",
28    "ImageBuilder",
29    "WarmSandbox",
30    "WarmPoolManager",
31    "SandboxManager",
32    "AgentSession",
33]
34 
35 
36class SandboxState(Enum):
37    """Sandbox lifecycle states."""
38    CREATING = "creating"
39    SYNCING = "syncing"
40    READY = "ready"
41    EXECUTING = "executing"
42    SNAPSHOTTING = "snapshotting"
43    TERMINATED = "terminated"
44 
45 
46@dataclass
47class UserIdentity:
48    """User identity for commit attribution.
49 
50    Use when: configuring sandbox git identity so commits are
51    attributed to the prompting user, not the app.
52    """
53    id: str
54    name: str
55    email: str
56    github_token: str
57 
58 
59@dataclass
60class SandboxConfig:
61    """Configuration for sandbox creation.
62 
63    Use when: defining resource limits and timeouts for a new sandbox
64    to prevent cost runaway and resource exhaustion.
65    """
66    repo_url: str
67    base_image: str
68    memory_mb: int = 4096
69    cpu_cores: int = 2
70    disk_gb: int = 10
71    timeout_hours: int = 4
72 
73 
74@dataclass
75class Sandbox:
76    """Represents a sandboxed execution environment.
77 
78    Use when: interacting with a running sandbox to execute commands,
79    read/write files, or take snapshots for session continuity.
80    """
81    id: str
82    config: SandboxConfig
83    state: SandboxState
84    created_at: datetime
85    snapshot_id: Optional[str] = None
86    current_user: Optional[UserIdentity] = None
87 
88    # Event handlers
89    on_state_change: Optional[Callable[[SandboxState], None]] = None
90 
91    async def execute_command(self, command: str) -> dict[str, Any]:
92        """Execute a command in the sandbox.
93 
94        Use when: running shell commands (git, build tools, tests)
95        inside the isolated environment.
96 
97        Returns:
98            dict with keys "stdout", "stderr", "exit_code".
99        """
100        # Implementation depends on infrastructure
101        pass
102 
103    async def read_file(self, path: str) -> str:
104        """Read a file from the sandbox filesystem.
105 
106        Use when: agent needs to inspect source code or config files.
107        Safe to call before git sync completes.
108        """
109        pass
110 
111    async def write_file(self, path: str, content: str) -> None:
112        """Write a file to the sandbox filesystem.
113 
114        Use when: agent needs to modify source code. Block this
115        until git sync completes to avoid write conflicts.
116        """
117        pass
118 
119    async def snapshot(self) -> str:
120        """Create a snapshot of current filesystem state.
121 
122        Use when: preserving session state before sandbox termination
123        so follow-up prompts can restore instantly.
124        """
125        self.state = SandboxState.SNAPSHOTTING
126        snapshot_id = await self._create_snapshot()
127        self.snapshot_id = snapshot_id
128        self.state = SandboxState.READY
129        return snapshot_id
130 
131    async def _create_snapshot(self) -> str:
132        """Create snapshot (infrastructure-specific)."""
133        pass
134 
135    async def restore(self, snapshot_id: str) -> None:
136        """Restore sandbox to a previous snapshot."""
137        pass
138 
139    async def terminate(self) -> None:
140        """Terminate the sandbox."""
141        self.state = SandboxState.TERMINATED
142 
143 
144@dataclass
145class RepositoryImage:
146    """Pre-built image for a repository.
147 
148    Use when: checking whether a cached environment image exists
149    and whether it is recent enough to use.
150    """
151    repo_url: str
152    image_id: str
153    commit_sha: str
154    built_at: datetime
155 
156    def is_stale(self, max_age: timedelta = timedelta(minutes=30)) -> bool:
157        """Check if image is older than max age."""
158        return datetime.utcnow() - self.built_at > max_age
159 
160 
161class ImageBuilder:
162    """Builds and manages repository images.
163 
164    Use when: setting up the periodic image build loop that
165    pre-bakes development environments for fast sandbox spin-up.
166    """
167 
168    def __init__(self, github_app_token_provider: Callable[[], str]) -> None:
169        self.token_provider = github_app_token_provider
170        self.images: dict[str, RepositoryImage] = {}
171 
172    async def build_image(self, repo_url: str) -> RepositoryImage:
173        """Build a new image for a repository.
174 
175        Use when: the current image is stale or no image exists yet.
176        Runs clone, dependency install, build, and cache warming.
177        """
178        print(f"Building image for {repo_url}...")
179 
180        # Get fresh token for clone
181        token = self.token_provider()
182 
183        # These operations run in build environment
184        build_steps: list[str] = [
185            # Clone repository
186            f"git clone https://x-access-token:{token}@github.com/{repo_url} /workspace",
187 
188            # Install dependencies
189            "cd /workspace && npm install",
190 
191            # Run build
192            "cd /workspace && npm run build",
193 
194            # Warm caches by running once
195            "cd /workspace && npm run dev &",
196            "sleep 5",  # Let dev server start
197            "cd /workspace && npm test -- --run || true",  # Run tests to warm cache
198        ]
199 
200        # Execute build steps (infrastructure-specific)
201        for step in build_steps:
202            await self._execute_build_step(step)
203 
204        # Get current commit
205        commit_sha: str = await self._get_commit_sha()
206 
207        # Create and store image
208        image = RepositoryImage(
209            repo_url=repo_url,
210            image_id=await self._finalize_image(),
211            commit_sha=commit_sha,
212            built_at=datetime.utcnow()
213        )
214 
215        self.images[repo_url] = image
216        return image
217 
218    def get_latest_image(self, repo_url: str) -> Optional[RepositoryImage]:
219        """Get the most recent image for a repository."""
220        return self.images.get(repo_url)
221 
222    async def _execute_build_step(self, command: str) -> None:
223        """Execute a build step (infrastructure-specific)."""
224        pass
225 
226    async def _get_commit_sha(self) -> str:
227        """Get current HEAD commit SHA."""
228        pass
229 
230    async def _finalize_image(self) -> str:
231        """Finalize and store the image, return image ID."""
232        pass
233 
234 
235@dataclass
236class WarmSandbox:
237    """A pre-warmed sandbox ready for use.
238 
239    Use when: tracking warm pool inventory and claiming a sandbox
240    for an incoming user session.
241    """
242    sandbox: Sandbox
243    repo_url: str
244    created_at: datetime
245    image_version: str
246    is_claimed: bool = False
247    sync_complete: bool = False
248 
249 
250class WarmPoolManager:
251    """Manages pools of pre-warmed sandboxes.
252 
253    Use when: reducing cold start latency by maintaining ready-to-use
254    sandboxes that are pre-synced to the latest code.
255    """
256 
257    def __init__(
258        self,
259        image_builder: ImageBuilder,
260        target_pool_size: int = 3,
261        max_age: timedelta = timedelta(minutes=25)
262    ) -> None:
263        self.image_builder = image_builder
264        self.target_size = target_pool_size
265        self.max_age = max_age
266        self.pools: dict[str, list[WarmSandbox]] = {}
267 
268    async def get_warm_sandbox(self, repo_url: str) -> Optional[WarmSandbox]:
269        """Get a pre-warmed sandbox if available.
270 
271        Use when: a user submits a prompt and needs a sandbox immediately.
272        Returns None if no valid warm sandbox is available.
273        """
274        if repo_url not in self.pools:
275            return None
276 
277        for warm in self.pools[repo_url]:
278            if not warm.is_claimed and self._is_valid(warm):
279                warm.is_claimed = True
280                return warm
281 
282        return None
283 
284    def _is_valid(self, warm: WarmSandbox) -> bool:
285        """Check if a warm sandbox is still valid."""
286        age: timedelta = datetime.utcnow() - warm.created_at
287        if age > self.max_age:
288            return False
289 
290        # Check if image is still current
291        current = self.image_builder.get_latest_image(warm.repo_url)
292        if not current or current.image_id != warm.image_version:
293            return False
294 
295        return True
296 
297    async def maintain_pool(self, repo_url: str) -> None:
298        """Ensure pool has target number of warm sandboxes.
299 
300        Use when: called periodically or after an image rebuild to
301        keep the warm pool populated.
302        """
303        if repo_url not in self.pools:
304            self.pools[repo_url] = []
305 
306        # Remove invalid sandboxes
307        valid: list[WarmSandbox] = [w for w in self.pools[repo_url] if self._is_valid(w)]
308        self.pools[repo_url] = valid
309 
310        # Count available (unclaimed) sandboxes
311        available: int = len([w for w in valid if not w.is_claimed])
312        needed: int = self.target_size - available
313 
314        # Create new warm sandboxes
315        for _ in range(max(0, needed)):
316            warm = await self._create_warm_sandbox(repo_url)
317            self.pools[repo_url].append(warm)
318 
319    async def _create_warm_sandbox(self, repo_url: str) -> WarmSandbox:
320        """Create a new warm sandbox."""
321        image: Optional[RepositoryImage] = self.image_builder.get_latest_image(repo_url)
322        if not image:
323            raise ValueError(f"No image available for {repo_url}")
324 
325        # Create sandbox from image
326        sandbox: Sandbox = await self._create_sandbox_from_image(image)
327 
328        warm = WarmSandbox(
329            sandbox=sandbox,
330            repo_url=repo_url,
331            created_at=datetime.utcnow(),
332            image_version=image.image_id,
333            sync_complete=False
334        )
335 
336        # Start syncing to latest in background
337        asyncio.create_task(self._sync_to_latest(warm))
338 
339        return warm
340 
341    async def _sync_to_latest(self, warm: WarmSandbox) -> None:
342        """Sync sandbox to latest commit on base branch."""
343        await warm.sandbox.execute_command("git fetch origin main")
344        await warm.sandbox.execute_command("git reset --hard origin/main")
345        warm.sync_complete = True
346 
347    async def _create_sandbox_from_image(self, image: RepositoryImage) -> Sandbox:
348        """Create a sandbox from an image (infrastructure-specific)."""
349        pass
350 
351 
352class SandboxManager:
353    """Main manager for sandbox lifecycle.
354 
355    Use when: orchestrating the full sandbox lifecycle including
356    image building, warm pools, and session management. This is the
357    top-level entry point that composes ImageBuilder and WarmPoolManager.
358    """
359 
360    def __init__(
361        self,
362        repositories: list[str],
363        github_app_token_provider: Callable[[], str],
364        build_interval: timedelta = timedelta(minutes=30)
365    ) -> None:
366        self.repositories = repositories
367        self.image_builder = ImageBuilder(github_app_token_provider)
368        self.warm_pool = WarmPoolManager(self.image_builder)
369        self.build_interval = build_interval
370        self.active_sessions: dict[str, Sandbox] = {}
371 
372    async def start_build_loop(self) -> None:
373        """Start the background image build loop.
374 
375        Use when: initializing the system. Runs indefinitely, rebuilding
376        images every build_interval to keep environments fresh.
377        """
378        while True:
379            for repo in self.repositories:
380                try:
381                    await self.image_builder.build_image(repo)
382                    await self.warm_pool.maintain_pool(repo)
383                except Exception as e:
384                    print(f"Failed to build {repo}: {e}")
385 
386            await asyncio.sleep(self.build_interval.total_seconds())
387 
388    async def start_session(
389        self,
390        repo_url: str,
391        user: UserIdentity,
392        snapshot_id: Optional[str] = None
393    ) -> Sandbox:
394        """Start a new session for a user.
395 
396        Use when: a user submits a prompt. Tries warm pool first,
397        then snapshot restore, then cold start as fallback.
398        """
399        # Try to get from warm pool first
400        warm: Optional[WarmSandbox] = await self.warm_pool.get_warm_sandbox(repo_url)
401 
402        if warm:
403            sandbox = warm.sandbox
404            # Wait for sync if not complete
405            if not warm.sync_complete:
406                await self._wait_for_sync(warm)
407        elif snapshot_id:
408            # Restore from previous session snapshot
409            sandbox = await self._restore_from_snapshot(snapshot_id)
410        else:
411            # Cold start from latest image
412            sandbox = await self._cold_start(repo_url)
413 
414        # Configure for user
415        await self._configure_for_user(sandbox, user)
416 
417        # Track session
418        session_id: str = f"{user.id}_{datetime.utcnow().isoformat()}"
419        self.active_sessions[session_id] = sandbox
420 
421        return sandbox
422 
423    async def on_user_typing(self, user: UserIdentity, repo_url: str) -> None:
424        """Called when user starts typing a prompt.
425 
426        Use when: implementing predictive warm-up. Starts preparing a
427        sandbox so it is ready by the time the user submits.
428        """
429        warm: Optional[WarmSandbox] = await self.warm_pool.get_warm_sandbox(repo_url)
430 
431        if not warm:
432            # Start warming one now
433            asyncio.create_task(self.warm_pool.maintain_pool(repo_url))
434 
435    async def end_session(self, session_id: str) -> Optional[str]:
436        """End a session and return snapshot ID for potential follow-up.
437 
438        Use when: a session completes. Always snapshots before termination
439        to prevent state loss.
440        """
441        if session_id not in self.active_sessions:
442            return None
443 
444        sandbox: Sandbox = self.active_sessions[session_id]
445 
446        # Create snapshot before terminating
447        snapshot_id: str = await sandbox.snapshot()
448 
449        # Terminate sandbox
450        await sandbox.terminate()
451 
452        del self.active_sessions[session_id]
453 
454        return snapshot_id
455 
456    async def _configure_for_user(
457        self,
458        sandbox: Sandbox,
459        user: UserIdentity
460    ) -> None:
461        """Configure sandbox for a specific user."""
462        sandbox.current_user = user
463 
464        # Set git identity
465        await sandbox.execute_command(
466            f'git config user.name "{user.name}"'
467        )
468        await sandbox.execute_command(
469            f'git config user.email "{user.email}"'
470        )
471 
472    async def _wait_for_sync(self, warm: WarmSandbox) -> None:
473        """Wait for sync to complete."""
474        while not warm.sync_complete:
475            await asyncio.sleep(0.1)
476 
477    async def _restore_from_snapshot(self, snapshot_id: str) -> Sandbox:
478        """Restore a sandbox from a snapshot."""
479        pass
480 
481    async def _cold_start(self, repo_url: str) -> Sandbox:
482        """Start a sandbox from cold (no warm pool available)."""
483        pass
484 
485 
486class AgentSession:
487    """Agent session with file read/write coordination.
488 
489    Use when: wrapping a Sandbox to enforce the pattern where reads
490    are allowed before sync completes but writes are blocked until
491    sync finishes, preventing write conflicts.
492    """
493 
494    def __init__(self, sandbox: Sandbox) -> None:
495        self.sandbox = sandbox
496        self.sync_complete: bool = False
497        self.pending_writes: list[tuple[str, str]] = []
498 
499    async def read_file(self, path: str) -> str:
500        """Read a file -- allowed even before sync completes.
501 
502        Use when: agent needs to research code immediately. Safe because
503        in large repos, files being worked on are unlikely to have
504        changed in the last 30 minutes since image build.
505        """
506        return await self.sandbox.read_file(path)
507 
508    async def write_file(self, path: str, content: str) -> None:
509        """Write a file -- blocks until sync is complete.
510 
511        Use when: agent needs to modify source code. Queues the write
512        and waits for git sync to finish to prevent conflicts.
513        """
514        if not self.sync_complete:
515            # Queue the write
516            self.pending_writes.append((path, content))
517            await self._wait_for_sync()
518 
519        await self.sandbox.write_file(path, content)
520 
521    def mark_sync_complete(self) -> None:
522        """Called when git sync is complete."""
523        self.sync_complete = True
524 
525    async def _wait_for_sync(self) -> None:
526        """Wait for sync to complete, then flush pending writes."""
527        while not self.sync_complete:
528            await asyncio.sleep(0.1)
529 
530        # Flush pending writes
531        for path, content in self.pending_writes:
532            await self.sandbox.write_file(path, content)
533        self.pending_writes.clear()
534 
535 
536if __name__ == "__main__":
537    async def _demo() -> None:
538        """Demonstrate sandbox manager usage end-to-end."""
539 
540        def get_github_token() -> str:
541            """Get GitHub App installation token."""
542            # Implementation: call GitHub API to get installation token
543            return "ghs_xxxx"
544 
545        # Initialize manager with target repositories
546        manager = SandboxManager(
547            repositories=[
548                "myorg/frontend",
549                "myorg/backend",
550                "myorg/shared-libs"
551            ],
552            github_app_token_provider=get_github_token
553        )
554 
555        # Start background build loop
556        asyncio.create_task(manager.start_build_loop())
557 
558        # Simulate user session
559        user = UserIdentity(
560            id="user123",
561            name="Alice Developer",
562            email="[email protected]",
563            github_token="gho_user_token"
564        )
565 
566        # User starts typing -- predictively warm a sandbox
567        await manager.on_user_typing(user, "myorg/frontend")
568 
569        # User submits prompt -- get sandbox
570        sandbox: Sandbox = await manager.start_session("myorg/frontend", user)
571 
572        # Create session wrapper for read/write coordination
573        session = AgentSession(sandbox)
574 
575        # Agent can read immediately (before sync completes)
576        readme: str = await session.read_file("/workspace/README.md")
577 
578        # Agent work happens here...
579 
580        # End session and get snapshot for follow-up
581        # Find the session_id that was generated during start_session
582        active_ids = list(manager.active_sessions.keys())
583        if active_ids:
584            session_id = active_ids[0]
585            snapshot_id: Optional[str] = await manager.end_session(session_id)
586            print(f"Session ended, snapshot: {snapshot_id}")
587        else:
588            print("No active session found")
589 
590    asyncio.run(_demo())
591