chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,406 @@
|
||||
# mypy: ignore-errors
|
||||
# standalone example with sys.path sibling imports that mypy cannot follow
|
||||
"""Temporal session manager workflow.
|
||||
|
||||
A long-lived singleton workflow that acts as the sole orchestrator for agent
|
||||
session lifecycles. It starts and stops agent workflows, and maintains a
|
||||
registry of active sessions so that TUI clients can list, resume, rename,
|
||||
and destroy sessions without any filesystem persistence.
|
||||
|
||||
The manager is started once (well-known workflow ID ``session-manager``) and
|
||||
lives forever. All lifecycle operations — create, destroy, rename, fork — go
|
||||
through the manager so the registry is always consistent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
from temporalio import activity, workflow
|
||||
from temporalio.exceptions import ApplicationError
|
||||
from temporalio.workflow import ParentClosePolicy
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from pydantic import BaseModel, field_validator, model_serializer
|
||||
from temporal_sandbox_agent import ( # type: ignore[import-not-found]
|
||||
TASK_QUEUE,
|
||||
AgentRequest,
|
||||
AgentWorkflow,
|
||||
SwitchBackendSignal,
|
||||
SwitchToLocalBackend,
|
||||
WorkflowSnapshot,
|
||||
)
|
||||
from temporalio.client import Client
|
||||
from temporalio.contrib.openai_agents import OpenAIAgentsPlugin
|
||||
from temporalio.contrib.pydantic import pydantic_data_converter
|
||||
|
||||
from agents import trace
|
||||
from agents.sandbox import Manifest
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
MANAGER_WORKFLOW_ID = "session-manager"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data types
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DaytonaBackendConfig(BaseModel):
|
||||
type: Literal["daytona"] = "daytona"
|
||||
|
||||
@model_serializer(mode="wrap")
|
||||
def _serialize_always_include_type(self, handler: Any) -> dict[str, Any]:
|
||||
data: dict[str, Any] = handler(self)
|
||||
data["type"] = self.type
|
||||
return data
|
||||
|
||||
|
||||
class DockerBackendConfig(BaseModel):
|
||||
type: Literal["docker"] = "docker"
|
||||
|
||||
@model_serializer(mode="wrap")
|
||||
def _serialize_always_include_type(self, handler: Any) -> dict[str, Any]:
|
||||
data: dict[str, Any] = handler(self)
|
||||
data["type"] = self.type
|
||||
return data
|
||||
|
||||
|
||||
class E2BBackendConfig(BaseModel):
|
||||
type: Literal["e2b"] = "e2b"
|
||||
|
||||
@model_serializer(mode="wrap")
|
||||
def _serialize_always_include_type(self, handler: Any) -> dict[str, Any]:
|
||||
data: dict[str, Any] = handler(self)
|
||||
data["type"] = self.type
|
||||
return data
|
||||
|
||||
|
||||
class LocalBackendConfig(BaseModel):
|
||||
type: Literal["local"] = "local"
|
||||
workspace_root: Path | None = None
|
||||
|
||||
@model_serializer(mode="wrap")
|
||||
def _serialize_always_include_type(self, handler: Any) -> dict[str, Any]:
|
||||
data: dict[str, Any] = handler(self)
|
||||
data["type"] = self.type
|
||||
return data
|
||||
|
||||
@field_validator("workspace_root")
|
||||
@classmethod
|
||||
def _must_be_absolute(cls, v: Path | None) -> Path | None:
|
||||
if v is not None and not v.is_absolute():
|
||||
raise ValueError("workspace_root must be an absolute path")
|
||||
return v
|
||||
|
||||
|
||||
BackendConfig = DaytonaBackendConfig | DockerBackendConfig | E2BBackendConfig | LocalBackendConfig
|
||||
|
||||
|
||||
class SessionInfo(BaseModel):
|
||||
workflow_id: str
|
||||
title: str
|
||||
created_at: datetime
|
||||
cwd: str = ""
|
||||
backend: BackendConfig = DaytonaBackendConfig()
|
||||
parent_workflow_id: str | None = None
|
||||
fork_count: int = 0
|
||||
snapshot_id: str | None = None
|
||||
|
||||
|
||||
class CreateSessionRequest(BaseModel):
|
||||
cwd: str
|
||||
manifest: Manifest | None = None
|
||||
backend: BackendConfig = DaytonaBackendConfig()
|
||||
|
||||
|
||||
class RenameRequest(BaseModel):
|
||||
workflow_id: str
|
||||
title: str
|
||||
|
||||
|
||||
class ForkSessionRequest(BaseModel):
|
||||
source_workflow_id: str
|
||||
title: str | None = None # defaults to "{original title} (fork #N)"
|
||||
target_backend: BackendConfig | None = None
|
||||
|
||||
|
||||
class SwitchBackendRequest(BaseModel):
|
||||
source_workflow_id: str
|
||||
target_backend: BackendConfig
|
||||
|
||||
|
||||
class _SwitchWorkflowBackendArgs(BaseModel):
|
||||
"""Activity args for switch_workflow_backend."""
|
||||
|
||||
workflow_id: str
|
||||
signal: SwitchBackendSignal
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _default_manifest(
|
||||
backend: BackendConfig,
|
||||
) -> Manifest:
|
||||
"""Return the default workspace manifest for the given backend config."""
|
||||
if isinstance(backend, DaytonaBackendConfig):
|
||||
return Manifest(root="/home/daytona/workspace")
|
||||
if isinstance(backend, DockerBackendConfig):
|
||||
return Manifest(root="/workspace")
|
||||
if isinstance(backend, E2BBackendConfig):
|
||||
return Manifest() # E2B resolves workspace root relative to the sandbox home
|
||||
root = str(backend.workspace_root) if backend.workspace_root else "/workspace"
|
||||
return Manifest(root=root)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Activities
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@activity.defn
|
||||
async def pause_workflow(workflow_id: str) -> None:
|
||||
"""Pause the agent workflow and wait for its session to fully stop."""
|
||||
client = await Client.connect("localhost:7233", data_converter=pydantic_data_converter)
|
||||
handle = client.get_workflow_handle(workflow_id)
|
||||
await handle.execute_update(AgentWorkflow.pause)
|
||||
|
||||
|
||||
@activity.defn
|
||||
async def switch_workflow_backend(args: _SwitchWorkflowBackendArgs) -> None:
|
||||
"""Switch the agent workflow's backend and wait for it to take effect."""
|
||||
client = await Client.connect("localhost:7233", data_converter=pydantic_data_converter)
|
||||
handle = client.get_workflow_handle(args.workflow_id)
|
||||
await handle.execute_update(AgentWorkflow.switch_backend, args.signal)
|
||||
|
||||
|
||||
@activity.defn
|
||||
async def query_workflow_snapshot(workflow_id: str) -> WorkflowSnapshot:
|
||||
"""Query the target workflow for its run state and conversation history."""
|
||||
client = await Client.connect("localhost:7233", data_converter=pydantic_data_converter)
|
||||
handle = client.get_workflow_handle(workflow_id)
|
||||
return await handle.query(AgentWorkflow.get_snapshot)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Workflow
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@workflow.defn
|
||||
class SessionManagerWorkflow:
|
||||
"""Registry and orchestrator for agent sessions.
|
||||
|
||||
* ``create_session`` — starts a new agent child workflow and registers it.
|
||||
* ``destroy_session`` — signals the agent workflow to terminate and
|
||||
removes it from the registry.
|
||||
* ``list_sessions`` — query returning all active sessions.
|
||||
* ``rename_session`` — signal to update a session title.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._sessions: dict[str, SessionInfo] = {}
|
||||
self._shutdown = False
|
||||
|
||||
# -- Main loop (lives forever) -----------------------------------------
|
||||
|
||||
@workflow.run
|
||||
async def run(self) -> None:
|
||||
await workflow.wait_condition(lambda: self._shutdown)
|
||||
|
||||
# -- Lifecycle: create & destroy (updates for request-response) ---------
|
||||
|
||||
@workflow.update
|
||||
async def create_session(self, request: CreateSessionRequest) -> str:
|
||||
"""Start a new agent workflow and register it. Returns the workflow ID."""
|
||||
workflow_id = f"sandbox-agent-{workflow.uuid4()}"
|
||||
|
||||
manifest = request.manifest
|
||||
if manifest is None:
|
||||
manifest = _default_manifest(request.backend)
|
||||
|
||||
with OpenAIAgentsPlugin().tracing_context():
|
||||
with trace("Temporal Sandbox Sandbox Agent"):
|
||||
await workflow.start_child_workflow(
|
||||
AgentWorkflow.run,
|
||||
AgentRequest(
|
||||
messages=[],
|
||||
cwd=request.cwd,
|
||||
backend=request.backend.type,
|
||||
history=[],
|
||||
manifest=manifest,
|
||||
),
|
||||
id=workflow_id,
|
||||
task_queue=TASK_QUEUE,
|
||||
parent_close_policy=ParentClosePolicy.ABANDON,
|
||||
)
|
||||
self._sessions[workflow_id] = SessionInfo(
|
||||
workflow_id=workflow_id,
|
||||
title=f"Session {workflow_id[-8:]}",
|
||||
created_at=workflow.now(),
|
||||
cwd=request.cwd,
|
||||
backend=request.backend,
|
||||
)
|
||||
return workflow_id
|
||||
|
||||
@workflow.update
|
||||
async def fork_session(self, request: ForkSessionRequest) -> str:
|
||||
"""Fork an existing session into a new workflow with identical state.
|
||||
|
||||
Pauses the source workflow, queries its RunState and conversation
|
||||
history, then starts a new child workflow seeded with that state.
|
||||
When ``target_backend`` differs from the source, the sandbox session
|
||||
state is not carried over (it is backend-specific), but the portable
|
||||
snapshot is extracted so the new backend can create a fresh session
|
||||
from the same workspace filesystem state.
|
||||
"""
|
||||
source = self._sessions.get(request.source_workflow_id)
|
||||
if source is None:
|
||||
raise ApplicationError(f"Source session {request.source_workflow_id} not found")
|
||||
|
||||
# Pause the source workflow so its session stops naturally
|
||||
await workflow.execute_activity(
|
||||
pause_workflow,
|
||||
request.source_workflow_id,
|
||||
start_to_close_timeout=timedelta(minutes=11),
|
||||
)
|
||||
|
||||
# Fetch the source workflow's state via activity
|
||||
workflow_snapshot: WorkflowSnapshot = await workflow.execute_activity(
|
||||
query_workflow_snapshot,
|
||||
request.source_workflow_id,
|
||||
start_to_close_timeout=timedelta(seconds=30),
|
||||
)
|
||||
|
||||
target_config = (
|
||||
request.target_backend if request.target_backend is not None else source.backend
|
||||
)
|
||||
cross_backend = target_config.type != source.backend.type
|
||||
|
||||
# Determine fork title
|
||||
source.fork_count += 1
|
||||
if cross_backend:
|
||||
title = request.title or f"{source.title} [{target_config.type}]"
|
||||
else:
|
||||
title = request.title or f"{source.title} (fork #{source.fork_count})"
|
||||
|
||||
# Always pass the portable snapshot so the forked session can seed
|
||||
# its workspace. Never carry session_state — a fork creates an
|
||||
# independent session seeded from the snapshot, not a resume of the
|
||||
# source session.
|
||||
snapshot = workflow_snapshot.snapshot
|
||||
|
||||
manifest = _default_manifest(target_config)
|
||||
|
||||
# Start the forked workflow with the source's run state and history
|
||||
workflow_id = f"sandbox-agent-{workflow.uuid4()}"
|
||||
await workflow.start_child_workflow(
|
||||
AgentWorkflow.run,
|
||||
AgentRequest(
|
||||
messages=[],
|
||||
cwd=source.cwd,
|
||||
backend=target_config.type,
|
||||
sandbox_session_state=None,
|
||||
snapshot=snapshot,
|
||||
previous_response_id=workflow_snapshot.previous_response_id,
|
||||
history=workflow_snapshot.history,
|
||||
manifest=manifest,
|
||||
),
|
||||
id=workflow_id,
|
||||
task_queue=TASK_QUEUE,
|
||||
parent_close_policy=ParentClosePolicy.ABANDON,
|
||||
)
|
||||
|
||||
self._sessions[workflow_id] = SessionInfo(
|
||||
workflow_id=workflow_id,
|
||||
title=title,
|
||||
created_at=workflow.now(),
|
||||
cwd=source.cwd,
|
||||
backend=target_config,
|
||||
parent_workflow_id=request.source_workflow_id,
|
||||
snapshot_id=workflow_snapshot.sandbox_session_state.snapshot.id
|
||||
if workflow_snapshot.sandbox_session_state
|
||||
else None,
|
||||
)
|
||||
return workflow_id
|
||||
|
||||
@workflow.update
|
||||
async def switch_backend(self, request: SwitchBackendRequest) -> str:
|
||||
"""Switch a session to a different sandbox backend in-place.
|
||||
|
||||
Signals the agent workflow to change its backend for subsequent turns.
|
||||
The workflow stays the same — no fork, no new child workflow. The
|
||||
portable snapshot is preserved so the workspace can be carried over;
|
||||
the backend-specific session state is cleared by the agent workflow.
|
||||
"""
|
||||
source = self._sessions.get(request.source_workflow_id)
|
||||
if source is None:
|
||||
raise ApplicationError(f"Session {request.source_workflow_id} not found")
|
||||
|
||||
if isinstance(request.target_backend, LocalBackendConfig):
|
||||
target: Literal["daytona", "docker", "e2b"] | SwitchToLocalBackend = (
|
||||
SwitchToLocalBackend(
|
||||
workspace_root=str(request.target_backend.workspace_root)
|
||||
if request.target_backend.workspace_root
|
||||
else "/workspace",
|
||||
)
|
||||
)
|
||||
else:
|
||||
target = request.target_backend.type
|
||||
await workflow.execute_activity(
|
||||
switch_workflow_backend,
|
||||
_SwitchWorkflowBackendArgs(
|
||||
workflow_id=request.source_workflow_id,
|
||||
signal=SwitchBackendSignal(target=target),
|
||||
),
|
||||
start_to_close_timeout=timedelta(seconds=30),
|
||||
)
|
||||
|
||||
source.backend = request.target_backend
|
||||
return request.source_workflow_id
|
||||
|
||||
@workflow.update
|
||||
async def destroy_session(self, workflow_id: str) -> None:
|
||||
"""Signal the agent workflow to destroy and remove it from the registry."""
|
||||
handle = workflow.get_external_workflow_handle(workflow_id)
|
||||
await handle.signal(AgentWorkflow.destroy)
|
||||
self._sessions.pop(workflow_id, None)
|
||||
|
||||
# -- Metadata: queries and signals --------------------------------------
|
||||
|
||||
@workflow.query
|
||||
def list_sessions(self) -> list[SessionInfo]:
|
||||
"""Return all active sessions, newest first."""
|
||||
return sorted(
|
||||
self._sessions.values(),
|
||||
key=lambda s: s.created_at,
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
@workflow.signal
|
||||
async def rename_session(self, request: RenameRequest) -> None:
|
||||
"""Update the title of an existing session."""
|
||||
if request.workflow_id in self._sessions:
|
||||
self._sessions[request.workflow_id].title = request.title
|
||||
|
||||
@workflow.signal
|
||||
async def update_snapshot_id(self, request: RenameRequest) -> None:
|
||||
"""Update the cached snapshot_id for a session.
|
||||
|
||||
Reuses RenameRequest where ``title`` carries the snapshot ID.
|
||||
"""
|
||||
if request.workflow_id in self._sessions:
|
||||
self._sessions[request.workflow_id].snapshot_id = request.title
|
||||
|
||||
@workflow.signal
|
||||
async def shutdown(self) -> None:
|
||||
"""Terminate the manager workflow (rarely needed)."""
|
||||
self._shutdown = True
|
||||
Reference in New Issue
Block a user