feat(platform): add remote Local PC execution targets

autogpt_platform/backend/backend/api/features/chat/routes.py

@@ -4,14 +4,15 @@
 import logging
 from collections.abc import AsyncGenerator
 from datetime import datetime, timezone
-from typing import Annotated
+from typing import Annotated, Literal
 from uuid import uuid4
 
 from autogpt_libs import auth
 from fastapi import APIRouter, Depends, HTTPException, Query, Response, Security
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
+from backend.api.features.local_executor.gating import is_local_executor_enabled
 from backend.copilot import active_turns
 from backend.copilot import service as chat_service
 from backend.copilot import stream_registry, turn_queue
@@ -27,8 +28,10 @@
 from backend.copilot.model import (
     CHAT_STATUS_IDLE,
     CHAT_STATUS_RUNNING,
+    ChatSession,
     ChatSessionInfo,
-    ChatSessionMetadata,
+    LocalExecutionTargetMetadata,
+    PublicChatSessionMetadata,
     create_chat_session,
     delete_chat_session,
     get_chat_session_metadata,
@@ -71,6 +74,18 @@
 )
 from backend.copilot.service import strip_injected_context_for_display
 from backend.copilot.tools.e2b_sandbox import kill_sandbox
+from backend.copilot.tools.local_pc_machine import (
+    MachineConnectionStaleError,
+    MachineControlError,
+    MachineNotConnectedError,
+    MachineSessionBinding,
+    activate_machine_session,
+    attach_machine_session,
+    detach_machine_session,
+    get_machine_presence,
+)
+from backend.copilot.tools.local_pc_relay_protocol import RelayPresence
+from backend.copilot.tools.local_pc_shim import get_shim_manager
 from backend.copilot.tools.manage_presets import (
     PresetDeletedResponse,
     PresetListResponse,
@@ -258,6 +273,28 @@ class PeekPendingMessagesResponse(BaseModel):
     count: int
 
 
+class CloudExecutionTargetRequest(BaseModel):
+    kind: Literal["cloud"] = "cloud"
+
+
+class LocalExecutionTargetRequest(BaseModel):
+    kind: Literal["local"] = "local"
+    machine_id: str = Field(
+        min_length=1,
+        max_length=128,
+        pattern=r"^[A-Za-z0-9][A-Za-z0-9._:-]*$",
+    )
+    expected_connection_id: str = Field(min_length=1, max_length=128)
+    browse_id: str = Field(min_length=1, max_length=256)
+    directory_ref: str = Field(min_length=1, max_length=256)
+
+
+ExecutionTargetRequest = Annotated[
+    CloudExecutionTargetRequest | LocalExecutionTargetRequest,
+    Field(discriminator="kind"),
+]
+
+
 class CreateSessionRequest(BaseModel):
     """Request model for creating (or get-or-creating) a chat session.
 
@@ -284,6 +321,15 @@ class CreateSessionRequest(BaseModel):
 
     dry_run: bool = False
     builder_graph_id: str | None = Field(default=None, max_length=128)
+    execution_target: ExecutionTargetRequest = Field(
+        default_factory=CloudExecutionTargetRequest
+    )
+
+    @model_validator(mode="after")
+    def builder_sessions_use_cloud_execution(self) -> "CreateSessionRequest":
+        if self.builder_graph_id and self.execution_target.kind == "local":
+            raise ValueError("Builder-bound sessions do not support Local PC execution")
+        return self
 
 
 class CreateSessionResponse(BaseModel):
@@ -292,7 +338,7 @@ class CreateSessionResponse(BaseModel):
     id: str
     created_at: str
     user_id: str | None
-    metadata: ChatSessionMetadata = ChatSessionMetadata()
+    metadata: PublicChatSessionMetadata = PublicChatSessionMetadata()
 
 
 class ActiveStreamInfo(BaseModel):
@@ -321,7 +367,7 @@ class SessionDetailResponse(BaseModel):
     oldest_sequence: int | None = None
     total_prompt_tokens: int = 0
     total_completion_tokens: int = 0
-    metadata: ChatSessionMetadata = ChatSessionMetadata()
+    metadata: PublicChatSessionMetadata = PublicChatSessionMetadata()
 
 
 class SessionSummaryResponse(BaseModel):
@@ -371,6 +417,168 @@ class UpdateSessionPinnedRequest(BaseModel):
     is_pinned: bool
 
 
+def _bindings_match(
+    expected: MachineSessionBinding,
+    actual: MachineSessionBinding,
+) -> bool:
+    return expected == actual
+
+
+async def _best_effort_detach_local_session(
+    presence: RelayPresence,
+    session_id: str,
+) -> None:
+    try:
+        await detach_machine_session(presence, session_id)
+    except Exception:
+        logger.warning(
+            "[LocalPC] Failed to compensate machine attachment for session %s",
+            session_id[:12],
+            exc_info=True,
+        )
+
+
+async def _compensate_local_session_creation(
+    presence: RelayPresence,
+    session_id: str,
+    *,
+    user_id: str,
+    organization_id: str | None,
+) -> None:
+    await _best_effort_detach_local_session(presence, session_id)
+    try:
+        await delete_chat_session(
+            session_id,
+            user_id,
+            organization_id=organization_id,
+        )
+    except Exception:
+        logger.exception(
+            "[LocalPC] Failed to delete compensated session %s",
+            session_id[:12],
+        )
+
+
+async def _create_local_chat_session(
+    *,
+    user_id: str,
+    organization_id: str | None,
+    team_id: str | None,
+    dry_run: bool,
+    target: LocalExecutionTargetRequest,
+) -> ChatSession:
+    if not await is_local_executor_enabled(user_id):
+        raise HTTPException(status_code=404, detail="Local PC executor is not enabled")
+
+    config = ChatConfig()
+    if not config.use_local_pc_executor:
+        raise HTTPException(
+            status_code=409,
+            detail="Local PC execution is unavailable on this deployment",
+        )
+
+    try:
+        presence = await get_machine_presence(
+            user_id,
+            config.local_pc_executor_oauth_client_id,
+            target.machine_id,
+            expected_connection_id=target.expected_connection_id,
+        )
+    except (MachineNotConnectedError, MachineConnectionStaleError) as exc:
+        raise HTTPException(status_code=409, detail=str(exc)) from exc
+
+    session_id = str(uuid4())
+    try:
+        attached = await attach_machine_session(
+            presence,
+            session_id=session_id,
+            browse_id=target.browse_id,
+            directory_ref=target.directory_ref,
+        )
+        if attached.session_id != session_id or attached.revision != 1:
+            raise MachineControlError(
+                "INVALID_SESSION_BINDING",
+                "The Local PC executor returned an invalid session binding",
+            )
+
+        execution_target = LocalExecutionTargetMetadata(
+            machine_id=target.machine_id,
+            directory_ref=target.directory_ref,
+            allowed_root=attached.allowed_root,
+            root_fingerprint=attached.fingerprint,
+            root_grant=attached.root_grant,
+            revision=attached.revision,
+        )
+        session = await create_chat_session(
+            user_id,
+            dry_run=dry_run,
+            organization_id=organization_id,
+            team_id=team_id,
+            session_id=session_id,
+            execution_target=execution_target,
+        )
+        activated = await activate_machine_session(presence, attached)
+        if not _bindings_match(attached, activated):
+            raise MachineControlError(
+                "INVALID_SESSION_BINDING",
+                "The Local PC executor activated a different session binding",
+            )
+
+        try:
+            await get_shim_manager().wait_for(session_id, timeout=8.0)
+            hello = await get_shim_manager().get_hello_async(session_id)
+        except Exception as exc:
+            raise MachineControlError(
+                "SESSION_DATA_CHANNEL_UNAVAILABLE",
+                "The Local PC executor did not open the session data channel",
+            ) from exc
+        if (
+            hello is None
+            or hello.machine_id != target.machine_id
+            or hello.allowed_root != attached.allowed_root
+        ):
+            raise MachineControlError(
+                "SESSION_DATA_CHANNEL_MISMATCH",
+                "The Local PC executor opened a mismatched session data channel",
+            )
+        await detach_machine_session(presence, session_id)
+        return session
+    except BaseException as exc:
+        cleanup_task = asyncio.create_task(
+            _compensate_local_session_creation(
+                presence,
+                session_id,
+                user_id=user_id,
+                organization_id=organization_id,
+            )
+        )
+        try:
+            await asyncio.shield(cleanup_task)
+        except asyncio.CancelledError:
+            await cleanup_task
+        if isinstance(exc, asyncio.CancelledError):
+            raise
+        if isinstance(
+            exc,
+            (MachineNotConnectedError, MachineConnectionStaleError),
+        ):
+            raise HTTPException(status_code=409, detail=str(exc)) from exc
+        if isinstance(exc, MachineControlError):
+            status_code = (
+                409
+                if exc.code
+                in {
+                    "DIRECTORY_REFERENCE_INVALID",
+                    "DIRECTORY_UNAVAILABLE",
+                    "SESSION_ALREADY_ACTIVE",
+                    "SESSION_REVISION_MISMATCH",
+                }
+                else 503
+            )
+            raise HTTPException(status_code=status_code, detail=str(exc)) from exc
+        raise
+
+
 # ========== Routes ==========
 
 
@@ -477,8 +685,9 @@ async def create_session(
     Returns:
         CreateSessionResponse: Details of the resulting session.
     """
-    dry_run = request.dry_run if request else False
-    builder_graph_id = request.builder_graph_id if request else None
+    request = request or CreateSessionRequest()
+    dry_run = request.dry_run
+    builder_graph_id = request.builder_graph_id
 
     logger.info(
         f"Creating session with user_id: "
@@ -487,7 +696,15 @@ async def create_session(
         f"{f', builder_graph_id={builder_graph_id}' if builder_graph_id else ''}"
     )
 
-    if builder_graph_id:
+    if request.execution_target.kind == "local":
+        session = await _create_local_chat_session(
+            user_id=user_id,
+            organization_id=ctx.org_id,
+            team_id=ctx.team_id,
+            dry_run=dry_run,
+            target=request.execution_target,
+        )
+    elif builder_graph_id:
         session = await get_or_create_builder_session(
             user_id,
             builder_graph_id,
@@ -506,7 +723,7 @@ async def create_session(
         id=session.session_id,
         created_at=session.started_at.isoformat(),
         user_id=session.user_id,
-        metadata=session.metadata,
+        metadata=PublicChatSessionMetadata.from_internal(session.metadata),
     )
 
 
@@ -537,6 +754,43 @@ async def delete_session(
     Raises:
         HTTPException: 404 if session not found or not owned by user.
     """
+    session_metadata = await get_chat_session_metadata(session_id, user_id)
+    if session_metadata is None:
+        raise HTTPException(
+            status_code=404,
+            detail=f"Session {session_id} not found or access denied",
+        )
+
+    execution_target = session_metadata.metadata.execution_target
+    if execution_target.kind == "local":
+        try:
+            config = ChatConfig()
+            presence = await get_machine_presence(
+                user_id,
+                config.local_pc_executor_oauth_client_id,
+                execution_target.machine_id,
+            )
+            await detach_machine_session(presence, session_id)
+        except Exception:
+            logger.warning(
+                "[LocalPC] Failed to detach deleted session %s",
+                session_id[:12],
+                exc_info=True,
+            )
+        try:
+            shim = await get_shim_manager().get_or_create_shim_for_session(
+                session_id, timeout=1.0
+            )
+            await shim.kill()
+        except (ConnectionError, TimeoutError):
+            pass
+        except Exception:
+            logger.warning(
+                "[LocalPC] Failed to close deleted session data channel %s",
+                session_id[:12],
+                exc_info=True,
+            )
+
     deleted = await delete_chat_session(session_id, user_id, organization_id=ctx.org_id)
 
     if not deleted:
@@ -736,7 +990,7 @@ async def get_session(
         oldest_sequence=page.oldest_sequence,
         total_prompt_tokens=total_prompt,
         total_completion_tokens=total_completion,
-        metadata=page.session.metadata,
+        metadata=PublicChatSessionMetadata.from_internal(page.session.metadata),
     )