"""Asynchronous commands sharing a VM until a policy requests a snapshot."""
from __future__ import annotations
import asyncio
from collections.abc import AsyncIterator, Sequence
from concurrent.futures import Future
from contextlib import asynccontextmanager, suppress
from functools import partial
from math import isfinite
from typing import Literal
from contree_client.models import DirectoryList, GrepResult, InstanceResult, OperationEvent
from contree_sdk.compat import Self
from contree_sdk.execution import AsyncExecutor, RunRequest
from contree_sdk.file_io import run_file_io
from contree_sdk.inspection import AsyncImageInspector
from contree_sdk.session.asyncio import ContreeAsyncSession
from contree_sdk.session.base import require_str
from contree_sdk.session.cleanup import await_cleanup
from contree_sdk.session.contracts import AsyncOperationContract, AsyncSubprocessContract
from contree_sdk.session.inspection import AsyncSessionInspector
from contree_sdk.session.lazy_state import LazyPhase, LazyState
from contree_sdk.session.live_files import prepare_live_sources
from contree_sdk.session.models import HistoryEntry
from contree_sdk.session.snapshot_policy import (
AbstractSnapshotPolicy,
CompositeSnapshotPolicy,
DurationSnapshotPolicy,
IdleSnapshotPolicy,
SnapshotEvent,
)
[docs]
class AsyncLazySession(AsyncExecutor): # noqa: PLR0904 - public lifecycle and inspection hooks
"""Share one VM across commands; save one history entry per snapshot.
The supplied session, its client and its store remain caller-owned. Do not
mutate that session outside this wrapper while a VM is active. Use this
object as a context manager or call close/abort to release its worker.
"""
def __init__(
self,
session: ContreeAsyncSession,
*,
snapshot_policy: AbstractSnapshotPolicy | None = None,
snapshot_timeout: float = 30,
vm_timeout: float = 3600,
) -> None:
if not isfinite(snapshot_timeout) or snapshot_timeout <= 0:
raise ValueError("snapshot_timeout must be positive")
if not isfinite(vm_timeout) or vm_timeout <= snapshot_timeout:
raise ValueError("vm_timeout must be finite and greater than snapshot_timeout")
self.session = session
self.session_id = session.session_id
self.snapshot_policy = snapshot_policy if snapshot_policy is not None else IdleSnapshotPolicy(10)
self.duration_policy = DurationSnapshotPolicy(vm_timeout - snapshot_timeout)
self.lifecycle_policy = CompositeSnapshotPolicy(self.snapshot_policy, self.duration_policy)
self.snapshot_timeout = snapshot_timeout
self.vm_timeout = vm_timeout
self.state = LazyState()
self.condition = asyncio.Condition()
self.operation: AsyncOperationContract | None = None
self.worker: asyncio.Task[None] | None = None
self.last_entry: HistoryEntry | None = None
self._snapshot_future: Future[None] | None = None
self._changed = asyncio.Event()
self._loop: asyncio.AbstractEventLoop | None = None
@property
def phase(self) -> LazyPhase:
return self.state.phase
@property
def error(self) -> BaseException | None:
return self.state.error
[docs]
def prepare_request(self, request: RunRequest) -> RunRequest:
"""Apply the wrapped session's policy and validate shared-VM semantics.
Returns:
The effective request.
Raises:
ValueError: The request needs a disposable VM or startup-only options.
"""
request = self.session.prepare_request(request)
if request.disposable:
raise ValueError("LazySession commands cannot be disposable; use a separate ContreeAsyncSession")
if (
request.preserve_env
or request.hostname is not None
or request.resources_limits is not None
or request.networking is not None
):
raise ValueError("preserve_env, hostname, resources_limits and networking require VM startup configuration")
return request
[docs]
def keepalive_request(self) -> RunRequest:
"""Supply a long-running main process available in the base image.
Returns:
The non-disposable request used to start each VM.
"""
return RunRequest(command="sleep", args=("infinity",), timeout=self.vm_timeout, disposable=False)
[docs]
async def observe(self, event: OperationEvent | None, error: Exception | None) -> None:
"""Observe the operation's sole reader without depending on handle.wait()."""
async with self.condition:
previous_phase = self.state.phase
if self.state.observe(event, error):
self.lifecycle_policy.notify(SnapshotEvent.COMMAND_FINISHED)
self._notify()
if previous_phase not in {"closed", "failed"} and isinstance(self.state.error, Exception):
raise self.state.error
def _notify(self) -> None:
self.condition.notify_all()
self._changed.set()
def _policy_ready(self, future: Future[None]) -> None:
if self._loop is not None and not self._loop.is_closed():
self._loop.call_soon_threadsafe(self._changed.set)
def _snapshot_requested(self) -> bool:
future = self._snapshot_future
if future is None or not future.done():
return False
future.result()
return True
async def _start(self) -> None:
self.state.begin()
if self.lifecycle_policy.policies[0] is not self.snapshot_policy:
self.lifecycle_policy.notify(SnapshotEvent.CLOSE)
self.lifecycle_policy = CompositeSnapshotPolicy(self.snapshot_policy, self.duration_policy)
self.lifecycle_policy.notify(SnapshotEvent.RESET)
self._snapshot_future = self.lifecycle_policy.should_snapshot()
self._loop = asyncio.get_running_loop()
self._snapshot_future.add_done_callback(self._policy_ready)
self.duration_policy.start_deferred()
self.lifecycle_policy.notify(SnapshotEvent.VM_STARTING)
self.operation = await self.session.spawn_request(self.keepalive_request())
context = self.operation.context
if context is None or context.request.disposable:
raise ValueError("the keepalive factory must create a bound, non-disposable operation")
effective_timeout = context.request.timeout_seconds
if effective_timeout is None or not isfinite(effective_timeout) or effective_timeout <= self.snapshot_timeout:
raise ValueError("keepalive timeout must be finite and greater than snapshot_timeout")
self.duration_policy.set_duration(effective_timeout - self.snapshot_timeout)
if self.duration_policy.should_snapshot().done():
raise TimeoutError("LazySession VM renewal deadline expired during startup")
self.operation.add_observer(self.observe)
await self.operation.__aenter__() # noqa: PLC2801 - context spans several commands
self.state.phase = "running"
if self.worker is None:
self.worker = asyncio.create_task(self._control(), name=f"lazy-{self.session_id}")
async def _admit_and_spawn(
self, request: RunRequest
) -> tuple[AsyncOperationContract, AsyncSubprocessContract, float | None]:
await self.session.ensure_ready()
request = self.prepare_request(request)
stdin = await self.session.file_transfer.read_stdin(request.stdin) if request.stdin is not None else None
files = None
if request.files:
prepared = await run_file_io(prepare_live_sources, request.files)
files = await self.session.file_transfer.prepare_files(prepared)
async with self.condition:
while True:
self.state.check()
starting = self.state.phase == "cold"
if starting:
try:
await self._start()
except BaseException as error:
self.state.fail(error)
self._notify()
raise
if self.state.phase == "running":
if not starting:
try:
self.state.wants_snapshot(self._snapshot_requested())
except BaseException as error:
self.state.fail(error)
self._notify()
raise
if self.state.phase == "running":
break
self._notify()
await self.condition.wait()
operation = self.operation
if operation is None:
raise RuntimeError("running LazySession has no operation")
if files:
await operation.apply_files(files)
self.state.reserve()
try:
self.lifecycle_policy.notify(SnapshotEvent.COMMAND_STARTED)
handle = await operation.run(
request.title,
shell=request.shell is not None,
args=request.args,
env=dict(request.env) if request.env is not None else None,
cwd=request.cwd,
uid=request.uid,
gid=request.gid,
stdin=stdin,
stdin_open=request.stdin_open,
truncate_output_at=request.truncate_output_at,
)
except BaseException as error:
self.state.creating -= 1
self.state.fail(error)
self._notify()
raise
self.state.created(handle.spid)
self._notify()
return operation, handle, request.timeout_seconds
async def _spawn(self, request: RunRequest) -> tuple[AsyncOperationContract, AsyncSubprocessContract, float | None]:
try:
return await self._admit_and_spawn(request)
except BaseException:
if self.state.error is not None:
self.lifecycle_policy.notify(SnapshotEvent.CLOSE)
if self.state.error is not None and self.operation is not None:
with suppress(Exception):
await self.operation.cancel()
with suppress(Exception):
await self.operation.shutdown()
raise
[docs]
async def spawn_request(self, request: RunRequest) -> AsyncSubprocessContract:
"""Start a subprocess; track its exit even if the caller never waits.
Returns:
The subprocess handle from the operation factory.
"""
return (await self._spawn(request))[1]
[docs]
async def spawn(self, command: str | None = None, **kwargs) -> AsyncSubprocessContract:
return await self.spawn_request(
RunRequest(command=command, disposable=kwargs.pop("disposable", False), **kwargs)
)
[docs]
async def execute(self, request: RunRequest) -> InstanceResult:
operation, handle, timeout = await self._spawn(request)
try:
result = await handle.wait(timeout=timeout)
except BaseException:
with suppress(Exception):
await asyncio.shield(operation.signal("SIGKILL", spid=handle.spid))
raise
async with self.condition:
if self.state.error is not None:
raise RuntimeError("LazySession failed while executing a command") from self.state.error
return result
[docs]
async def run(self, command: str | None = None, **kwargs) -> InstanceResult:
"""Run a shared-VM command. Options are RunRequest fields; disposable defaults to False.
Returns:
The subprocess result, including a nonzero exit code.
"""
return await self.execute(RunRequest(command=command, disposable=kwargs.pop("disposable", False), **kwargs))
async def _save(self, operation: AsyncOperationContract) -> None:
await operation.signal("SIGTERM")
await operation.wait(timeout=self.snapshot_timeout)
await operation.shutdown()
async with self.condition:
if self.state.error is not None:
raise RuntimeError("event stream failed during snapshot") from self.state.error
entry = await self.session.commit_result(operation, title="LazySession snapshot")
if entry is None:
raise RuntimeError("commit policy rejected the LazySession snapshot")
async with self.condition:
self.last_entry = entry
self.operation = None
self.lifecycle_policy.notify(SnapshotEvent.RESET)
self._snapshot_future = None
self.state.saved()
self._notify()
async def _control(self) -> None:
try: # noqa: PLW0717 - one failure boundary owns the worker lifecycle
while True:
async with self.condition:
self._changed.clear()
if self.state.phase in {"closed", "failed"}:
return
previous_phase = self.state.phase
if self.state.wants_snapshot(self._snapshot_requested()):
operation = self.operation
if operation is None:
raise RuntimeError("snapshot has no operation") # noqa: TRY301 - fail the worker
self.state.phase = "snapshotting"
else:
operation = None
if self.state.phase != previous_phase:
self.condition.notify_all()
if operation is None:
await self._changed.wait()
else:
await self._save(operation)
except BaseException as error:
async with self.condition:
self.state.fail(error)
self._notify()
if self.operation is not None:
with suppress(Exception):
await self.operation.cancel()
with suppress(Exception):
await self.operation.shutdown()
finally:
self.lifecycle_policy.notify(SnapshotEvent.CLOSE)
[docs]
async def snapshot(self, *, timeout: float | None = None) -> HistoryEntry | None:
"""Drain commands and inspections, save their common image, and reopen admission.
Snapshots are lightweight and can run after every command. They stop
the VM. An unchanged filesystem can return the same image UUID.
timeout limits this wait; an already requested snapshot continues.
Errors remain available as error and are raised on subsequent calls.
Saving failure or a closed wrapper raises RuntimeError. A wait that
expires raises TimeoutError; saving continues in the worker.
Returns:
The saved history entry, or None if no VM has run yet.
"""
async with self.condition:
return await self._snapshot_locked(timeout=timeout)
async def _snapshot_locked(self, *, timeout: float | None = None) -> HistoryEntry | None:
self.state.check()
if self.state.phase == "cold":
return self.last_entry
generation = self.state.snapshots
if self.state.phase == "running":
self.state.phase = "draining"
self._notify()
try:
await asyncio.wait_for(
self.condition.wait_for(
lambda: (
self.state.snapshots != generation
or self.state.error is not None
or self.state.phase == "closed"
)
),
timeout=timeout,
)
except asyncio.TimeoutError:
raise TimeoutError("LazySession snapshot is still pending") from None
if self.state.error is not None:
raise RuntimeError("LazySession snapshot failed") from self.state.error
if self.state.snapshots == generation:
raise RuntimeError("LazySession closed before snapshot completed")
return self.last_entry
[docs]
async def close(self, *, timeout: float | None = None) -> None:
"""Save pending work and close this wrapper, retaining caller-owned resources.
Raises:
RuntimeError: Saving failed; the cause is available as error.
TimeoutError: The caller's wait expired; closing continues in the worker.
"""
async with self.condition:
if self.state.phase == "closed":
return
self.state.closing = True
if self.state.phase == "cold":
self.state.phase = "closed"
self.lifecycle_policy.notify(SnapshotEvent.CLOSE)
elif self.state.phase == "running":
self.state.phase = "draining"
self._notify()
try:
await asyncio.wait_for(
self.condition.wait_for(lambda: self.state.phase in {"closed", "failed"}), timeout=timeout
)
except asyncio.TimeoutError:
raise TimeoutError("LazySession close is still pending") from None
error = self.state.error
if error is not None:
with suppress(Exception):
await self.abort()
raise RuntimeError("LazySession close could not save pending work") from error
if self.worker is not None:
await asyncio.shield(self.worker)
[docs]
async def abort(self) -> None:
"""Discard unsaved work and close. An already running snapshot finishes first."""
async with self.condition:
await self.condition.wait_for(lambda: self.state.phase != "snapshotting")
operation = self.operation
self.state.closing = True
self.state.phase = "closed"
self.lifecycle_policy.notify(SnapshotEvent.CLOSE)
self._notify()
try:
if operation is not None:
try:
await operation.cancel()
finally:
try:
await operation.shutdown()
except Exception as error:
if error is not self.state.error:
raise
finally:
if self.worker is not None:
await asyncio.shield(self.worker)
[docs]
async def inspect(self, *, timeout: float | None = None) -> AsyncSessionInspector:
"""Return a view that selects the current filesystem for each request.
Live requests do not stop the VM. Cold requests use the saved image;
pending staged files start a VM so the view includes those files.
timeout limits admission and VM startup waits, not HTTP I/O.
Use inspect_snapshot() for a stable image and its image_uuid.
Returns:
A session-bound inspector; creating it starts no VM.
"""
async with self.condition:
self.state.check()
return AsyncSessionInspector(
self.session.client,
partial(self.inspection_source, timeout=timeout),
partial(self.inspect_snapshot, timeout=timeout),
self.session.file_transfer.on_progress,
)
async def _inspection_source_locked(self, timeout: float | None) -> AsyncImageInspector | str:
await self.session.ensure_ready()
loop = asyncio.get_running_loop()
deadline = None if timeout is None else loop.time() + timeout
while True:
self.state.check()
started = False
if self.state.phase == "cold":
if not await self.session.pending_files():
return self.session.create_inspector(
require_str(self.session.image_uuid, "session has no resolved image")
)
try:
await self._start()
except BaseException as error:
self.state.fail(error)
self._notify()
raise
started = True
if self.state.phase == "running":
if not started:
try:
self.state.wants_snapshot(self._snapshot_requested())
except BaseException as error:
self.state.fail(error)
self._notify()
raise
if self.state.phase == "running":
if self.operation is None:
raise RuntimeError("running LazySession has no operation")
try:
self.lifecycle_policy.notify(SnapshotEvent.INSPECTION_STARTED)
except BaseException as error:
self.state.fail(error)
self._notify()
raise
self.state.inspecting += 1
return self.operation.uuid
self._notify()
remaining = None if deadline is None else deadline - loop.time()
try:
await asyncio.wait_for(self.condition.wait(), remaining)
except asyncio.TimeoutError:
raise TimeoutError("LazySession inspection admission is still pending") from None
async def _release_inspection(self) -> None:
async with self.condition:
self.state.inspecting -= 1
try:
if self.state.phase not in {"closed", "failed"}:
self.lifecycle_policy.notify(SnapshotEvent.INSPECTION_FINISHED)
except BaseException as error:
self.state.fail(error)
raise
finally:
self._notify()
[docs]
@asynccontextmanager
async def inspection_source(self, *, timeout: float | None = None) -> AsyncIterator[AsyncImageInspector | str]:
"""Reserve the current source until the request and its stream close.
Override this hook to customize source selection. Image inspectors are
pinned; operation UUIDs remain protected from automatic snapshot.
Yields:
A saved image inspector or an active operation UUID.
Raises:
RuntimeError: The session fails or closes before inspection starts.
TimeoutError: Admission or VM startup exceeds the timeout.
"""
try:
async with self.condition:
source = await self._inspection_source_locked(timeout)
except BaseException:
if self.state.error is not None:
self.lifecycle_policy.notify(SnapshotEvent.CLOSE)
if self.state.error is not None and self.operation is not None:
with suppress(Exception):
await self.operation.cancel()
with suppress(Exception):
await self.operation.shutdown()
raise
try:
if isinstance(source, str):
async with self.condition:
try:
await asyncio.wait_for(
self.condition.wait_for(
lambda: self.state.ready or self.state.error is not None or self.state.phase == "closed"
),
timeout=timeout,
)
except asyncio.TimeoutError:
raise TimeoutError("LazySession inspection is waiting for VM startup") from None
if self.state.error is not None:
raise RuntimeError("LazySession failed before inspection") from self.state.error
if self.state.phase == "closed":
raise RuntimeError("LazySession closed before inspection")
yield source
finally:
if isinstance(source, str):
await await_cleanup(self._release_inspection())
[docs]
async def inspect_snapshot(self, *, timeout: float | None = None) -> AsyncImageInspector:
"""Save active work and pin the resulting image for inspection.
Waits for admitted commands and inspections, then stops the VM.
A cold session starts no VM. Later commands cannot change the pinned image.
timeout limits the snapshot wait; saving continues if the wait ends early.
A failed snapshot raises instead of returning the previous image.
Staged attachments appear only after a command applies them.
Returns:
An inspector borrowing the underlying client and file transfer.
"""
async with self.condition:
self.state.check()
await self.session.ensure_ready()
cold = self.state.phase == "cold"
entry = await self._snapshot_locked(timeout=timeout)
image_uuid = (
entry.image_uuid
if not cold and entry is not None
else require_str(self.session.image_uuid, "session has no resolved image")
)
return self.session.create_inspector(image_uuid)
[docs]
async def list_dir(self, path: str = "/") -> DirectoryList:
"""List session directory entries relative to the filesystem root.
Returns:
Directory entries and their filesystem metadata.
"""
view = await self.inspect()
return await view.list(path)
[docs]
async def grep(
self,
pattern: str | Sequence[str],
*,
path: str | Sequence[str] | None = None,
glob: str | Sequence[str] | None = None,
max_count: int | None = None,
max_total: int | None = None,
case: Literal["sensitive", "insensitive", "smart"] | None = None,
before: int | None = None,
after: int | None = None,
) -> GrepResult:
"""Search session files using the current source's syntax and limits.
Returns:
Matching lines and the server's truncation indicator.
"""
view = await self.inspect()
return await view.grep(
pattern,
path=path,
glob=glob,
max_count=max_count,
max_total=max_total,
case=case,
before=before,
after=after,
)
[docs]
async def read_file(self, path: str) -> bytes:
"""Read current files without stopping an active VM.
Returns:
The complete file as bytes.
"""
view = await self.inspect()
return await view.read_file(path)
def _boundary(self) -> None:
self.state.check()
if self.state.phase != "cold":
raise RuntimeError("snapshot pending work before changing history")
[docs]
async def create_branch(self, name: str, *, from_branch: str | None = None) -> None:
async with self.condition:
self._boundary()
await self.session.create_branch(name, from_branch=from_branch)
[docs]
async def switch_branch(self, name: str) -> None:
async with self.condition:
self._boundary()
await self.session.switch_branch(name)
[docs]
async def rollback(self, steps: int = 1) -> None:
async with self.condition:
self._boundary()
await self.session.rollback(steps)
async def __aenter__(self) -> Self:
async with self.condition:
self.state.check()
return self
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
if exc_type is None:
await self.close()
else:
with suppress(Exception):
await self.abort()