wehub-resource-sync
407 lines
15 KiB
Python
407 lines
15 KiB
Python
# 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
|