stacklink-sdk 0.2.0__py3-none-any.whl

stacklink/workbench.py

@@ -0,0 +1,1090 @@
+from __future__ import annotations
+
+import base64
+import time
+from typing import Any, Iterator, Literal, Mapping, TypedDict, cast
+from urllib.error import HTTPError
+from urllib.request import Request, urlopen
+
+from .errors import StacklinkError
+
+
+WorkbenchLanguage = Literal["python", "javascript", "typescript"]
+WorkbenchOutputMode = Literal["inline", "preview", "events", "artifact"]
+WorkbenchExecutionKind = Literal["workbench_code", "workbench_bash"]
+WorkbenchReplayEnvironmentMode = Literal["require_match", "restore_snapshot", "current"]
+WorkbenchExecutionStatus = Literal[
+    "queued",
+    "started",
+    "running",
+    "cancel_requested",
+    "completed",
+    "failed",
+    "cancelled",
+    "timed_out",
+]
+
+
+class _WorkbenchUnset:
+    pass
+
+
+_WORKBENCH_UNSET = _WorkbenchUnset()
+_REUSABLE_RUN_COMPARISON_WAIT_SECONDS = 120.0
+_REUSABLE_RUN_COMPARISON_POLL_SECONDS = 1.0
+
+
+class WorkbenchReusableRunRequest(TypedDict, total=False):
+    mode: Literal["auto", "off"]
+    seriesId: str
+    name: str
+    description: str
+    versionReason: str
+
+
+class WorkbenchSession(TypedDict, total=False):
+    id: str
+    platformSessionId: str
+    consumerAppId: str
+    tenantId: str
+    status: str
+    runtimeProvider: str | None
+    runtimeSessionId: str | None
+    sandboxSize: Literal["standard", "medium", "large"]
+    runtimeSandboxSize: Literal["standard", "medium", "large"] | None
+    runtimeGeneration: int
+    runtimeEnvironmentRestoredGeneration: int
+    activeExecutionCount: int
+    workspaceStorageKind: str
+    workspaceStorageRef: str | None
+    environmentManifest: dict[str, Any]
+    lastActiveAt: str
+    destroyedAt: str | None
+    createdAt: str
+    updatedAt: str
+
+
+class WorkbenchSessionUseResponse(TypedDict):
+    workbenchSession: WorkbenchSession
+    reused: bool
+
+
+class WorkbenchSessionDestroyResponse(TypedDict):
+    workbenchSession: WorkbenchSession
+    computeSuspended: bool
+    retainedFiles: bool
+
+
+class WorkbenchWorkspaceDeleteReceipt(TypedDict):
+    deleted: Literal[True]
+    workspaceKey: str
+    deletedWorkspaceObjects: int
+    deletedExecutionObjects: int
+    deletedObjects: int
+
+
+class WorkbenchDeletedDatabaseRows(TypedDict):
+    workbenchSessions: int
+    backgroundJobs: int
+    files: int
+    fileUploads: int
+    artifacts: int
+    executionLogs: int
+    executionEvents: int
+    executionEventStreams: int
+    runDefinitions: int
+    environmentSnapshots: int
+    usageEvents: int
+    costLedger: int
+    budgetPolicies: int
+    budgetReservations: int
+    runtimeHeartbeats: int
+    approvalRequests: int
+
+
+class WorkbenchSessionDeleteReceipt(TypedDict):
+    deleted: Literal[True]
+    workbenchSessionId: str
+    platformSessionId: str
+    computeDestroyed: bool
+    deletedArtifactObjects: int
+    deletedUploadObjects: int
+    deletedSnapshots: int
+    workspace: WorkbenchWorkspaceDeleteReceipt
+    deletedDatabaseRows: WorkbenchDeletedDatabaseRows
+    databaseDeleted: Literal[True]
+
+
+class WorkbenchRunInput(TypedDict, total=False):
+    language: WorkbenchLanguage
+    code: str
+    workbenchSessionId: str
+    timeoutMs: int
+    background: bool
+    notifyOnComplete: bool
+    outputMode: WorkbenchOutputMode
+    maxOutputBytes: int
+    reusableRun: WorkbenchReusableRunRequest
+    currentStep: str
+    currentStepMetric: str
+    thought: str
+    metadata: dict[str, Any]
+
+
+class WorkbenchBashInput(TypedDict, total=False):
+    command: str
+    workbenchSessionId: str
+    timeoutMs: int
+    background: bool
+    notifyOnComplete: bool
+    outputMode: WorkbenchOutputMode
+    maxOutputBytes: int
+    reusableRun: WorkbenchReusableRunRequest
+    currentStep: str
+    currentStepMetric: str
+    thought: str
+    metadata: dict[str, Any]
+
+
+class WorkbenchExecutionLog(TypedDict, total=False):
+    id: str
+    workbenchSessionId: str
+    operation: str
+    language: WorkbenchLanguage | None
+    status: WorkbenchExecutionStatus
+    durationMs: int | None
+    artifactId: str | None
+    timeoutMs: int | None
+    deadlineAt: str | None
+    cancelRequestedAt: str | None
+    cancelledAt: str | None
+    runtimeProvider: str | None
+    runtimeSessionId: str | None
+    runtimeProcessId: str | None
+    runtimeOutputCursor: str | None
+    approvalRequestId: str | None
+    idempotencyKey: str | None
+    runtimeStartedAt: str | None
+    startedAt: str
+    completedAt: str | None
+    metadata: dict[str, Any]
+    result: Any
+    reusableRun: "WorkbenchReusableRunState"
+
+
+class WorkbenchObservableResultTerminal(TypedDict):
+    status: Literal["completed", "failed", "cancelled", "timed_out"]
+    exitCode: int | None
+
+
+class WorkbenchObservableResultOutput(TypedDict):
+    orderedSha256: str
+    stdoutSha256: str
+    stderrSha256: str
+    structuredResultSha256: str
+    byteCount: int
+
+
+class WorkbenchObservableResultArtifact(TypedDict):
+    sha256: str
+    sizeBytes: int
+    contentType: str | None
+
+
+class WorkbenchObservableResultV1(TypedDict):
+    contractVersion: Literal["workbench.observable-result.v1"]
+    terminal: WorkbenchObservableResultTerminal
+    output: WorkbenchObservableResultOutput
+    artifact: WorkbenchObservableResultArtifact | None
+
+
+class WorkbenchReusableRunState(TypedDict, total=False):
+    mode: Literal["auto", "off"]
+    status: Literal["pending", "ready", "failed", "skipped"]
+    seriesId: str | None
+    name: str | None
+    description: str | None
+    versionReason: str | None
+    errorCode: str | None
+    runDefinitionId: str | None
+    version: int | None
+    replayOfRunDefinitionId: str | None
+    observableResultContractVersion: str | None
+    observableResult: WorkbenchObservableResultV1 | None
+    observableResultSha256: str | None
+    comparisonStatus: Literal["matched", "mismatched"] | None
+    comparison: "WorkbenchRunComparison | None"
+    comparedAt: str | None
+    nextAction: Literal["get_execution", "get_reusable_run"] | None
+
+
+class WorkbenchRuntimeDiagnostics(TypedDict, total=False):
+    code: str
+    message: str
+    runtimeStatus: str
+    exitCode: int
+    error: Any
+
+
+class WorkbenchRuntimeResult(TypedDict, total=False):
+    status: str
+    success: bool
+    exitCode: int | None
+    terminalReason: str
+    stdout: str
+    stderr: str
+    outputBytes: int
+    outputTruncated: bool
+    diagnostics: WorkbenchRuntimeDiagnostics
+
+
+class WorkbenchRuntimeIdentity(TypedDict, total=False):
+    provider: str
+    sessionId: str
+    workspaceKey: str
+
+
+class WorkbenchRuntimeEnvelope(TypedDict):
+    runtime: WorkbenchRuntimeIdentity
+    result: WorkbenchRuntimeResult
+
+
+class WorkbenchBackgroundExecutionResult(TypedDict):
+    executionKind: WorkbenchExecutionKind
+    executionId: str
+    status: Literal["queued", "running"]
+    runtimeProcessId: str | None
+    outputMode: Literal["events"]
+
+
+class WorkbenchOutputPreview(TypedDict):
+    outputMode: Literal["preview"]
+    outputBytes: int
+    preview: str
+
+
+class WorkbenchFile(TypedDict, total=False):
+    id: str
+    path: str
+    contentType: str
+    sizeBytes: int | str
+    storageKind: str
+    sha256: str | None
+    etag: str | None
+    contentBase64: str
+    contentEncoding: Literal["base64"]
+    createdAt: str
+    updatedAt: str
+
+
+class WorkbenchFileResponse(TypedDict):
+    file: WorkbenchFile
+
+
+class WorkbenchFileListResponse(TypedDict, total=False):
+    files: list[WorkbenchFile]
+    runtimeFiles: list[WorkbenchFile]
+
+
+class WorkbenchFileMutationResponse(TypedDict, total=False):
+    file: WorkbenchFile
+    uploadId: str
+    finalized: bool
+    deleted: bool
+    path: str
+
+
+class WorkbenchFileSignedUpload(TypedDict):
+    uploadId: str
+    path: str
+    contentType: str
+    sizeBytes: int
+    storageKind: Literal["r2_workspace"]
+    uploadUrl: str
+    expiresAt: str
+    requiredHeaders: dict[str, str]
+
+
+class WorkbenchFileUploadPrepareResponse(TypedDict):
+    upload: WorkbenchFileSignedUpload
+
+
+class WorkbenchFileSignedDownload(TypedDict):
+    file: WorkbenchFile
+    downloadUrl: str
+    expiresAt: str
+
+
+class WorkbenchArtifact(TypedDict, total=False):
+    id: str
+    artifactRef: str
+    path: str
+    label: str | None
+    contentType: str | None
+    sizeBytes: int | str | None
+    sha256: str | None
+    storageKind: str
+    status: str
+    downloadUrl: str
+    createdAt: str
+    updatedAt: str
+
+
+class WorkbenchArtifactResponse(TypedDict, total=False):
+    workbenchSession: WorkbenchSession
+    artifact: WorkbenchArtifact
+    deleted: bool
+
+
+class WorkbenchArtifactListResponse(TypedDict, total=False):
+    workbenchSession: WorkbenchSession
+    artifacts: list[WorkbenchArtifact]
+
+
+class WorkbenchArtifactDownloadResponse(TypedDict, total=False):
+    artifact: WorkbenchArtifact
+    downloadUrl: str
+    expiresAt: str
+    contentType: str
+
+
+class WorkbenchRunComparison(TypedDict, total=False):
+    contractVersion: Literal["workbench.observable-result.v1"]
+    status: Literal["pending", "pending_finalization", "matched", "mismatched"]
+    sourceExecutionId: str
+    replayExecutionId: str
+    terminalMatch: bool
+    outputMatch: bool
+    artifactMatch: bool
+    sourceObservableResultSha256: str
+    replayObservableResultSha256: str
+    nextAction: Literal["get_reusable_run"]
+
+
+class WorkbenchRunReplay(TypedDict):
+    executionId: str
+    status: WorkbenchExecutionStatus
+    comparison: WorkbenchRunComparison
+
+
+class WorkbenchRunDefinition(TypedDict, total=False):
+    id: str
+    seriesId: str
+    version: int
+    workbenchSessionId: str
+    platformSessionId: str
+    sourceExecutionId: str
+    environmentSnapshotId: str | None
+    name: str
+    description: str | None
+    kind: Literal["code", "bash"]
+    language: WorkbenchLanguage | None
+    source: str
+    publicationStatus: Literal["draft", "published", "retired"]
+    validationStatus: Literal["unvalidated", "matched", "mismatched"]
+    validatedAt: str | None
+    validatedByExecutionId: str | None
+    sourceObservableResult: dict[str, Any] | None
+    sourceObservableResultSha256: str | None
+    sourceCodeSha256: str | None
+    environmentManifestSha256: str | None
+    creationMode: Literal["auto", "migrated_manual"]
+    versionReason: str | None
+    environmentManifest: dict[str, Any]
+    metadata: dict[str, Any]
+    replayCount: int
+    replays: list[WorkbenchRunReplay]
+    publishedAt: str | None
+    retiredAt: str | None
+    createdAt: str
+
+
+class WorkbenchRunDefinitionResponse(TypedDict):
+    runDefinition: WorkbenchRunDefinition
+
+
+class WorkbenchRunDefinitionListResponse(TypedDict):
+    runDefinitions: list[WorkbenchRunDefinition]
+    limit: int
+    hasMore: bool
+    nextCursor: str | None
+
+
+class WorkbenchRunComparisonResponse(TypedDict):
+    runDefinition: WorkbenchRunDefinition
+    execution: "WorkbenchExecutionResponse"
+    environment: "WorkbenchRunEnvironmentResolution"
+    comparison: WorkbenchRunComparison
+
+
+class WorkbenchRunEnvironmentResolution(TypedDict):
+    mode: WorkbenchReplayEnvironmentMode
+    status: Literal["queued", "completed", "failed"]
+    matchedBeforeReplay: bool | None
+    restored: bool
+    snapshotId: str | None
+
+
+class WorkbenchRunActivationResponse(TypedDict, total=False):
+    action: Literal["publish", "rollback"]
+    runDefinition: WorkbenchRunDefinition
+    previousPublishedRunDefinitionId: str | None
+    previousPublishedVersion: int | None
+
+
+class WorkbenchExecutionEvent(TypedDict, total=False):
+    id: str
+    sequence: str | int
+    executionKind: WorkbenchExecutionKind
+    executionId: str
+    eventType: str
+    status: str | None
+    stream: Literal["stdout", "stderr"] | None
+    message: str | None
+    data: Any
+    metadata: dict[str, Any]
+    occurredAt: str
+
+
+class WorkbenchExecutionEventsPage(TypedDict, total=False):
+    events: list[WorkbenchExecutionEvent]
+    nextCursor: str | None
+
+
+class WorkbenchExecutionStateResponse(TypedDict, total=False):
+    execution: WorkbenchExecutionLog
+    events: list[WorkbenchExecutionEvent]
+    cancelled: bool
+    remoteCancellation: dict[str, Any]
+
+
+class WorkbenchExecutionResponse(TypedDict, total=False):
+    workbenchSession: WorkbenchSession
+    executionKind: WorkbenchExecutionKind
+    executionId: str
+    status: WorkbenchExecutionStatus
+    outputMode: WorkbenchOutputMode
+    execution: WorkbenchExecutionLog
+    executionLog: WorkbenchExecutionLog
+    result: WorkbenchRuntimeEnvelope | WorkbenchBackgroundExecutionResult | WorkbenchOutputPreview
+    artifact: WorkbenchArtifact
+
+
+class WorkbenchRun:
+    def __init__(self, client, platform_session_id: str, response: WorkbenchExecutionResponse):
+        self.client = client
+        self.platform_session_id = platform_session_id
+        self.response = response
+        self.execution_id = str(response["executionId"])
+        self.execution_kind = str(response["executionKind"])
+
+    def events(self, *, poll_interval_seconds: float = 0.5, limit: int = 200) -> Iterator[WorkbenchExecutionEvent]:
+        after_sequence: str | None = None
+        while True:
+            page = cast(WorkbenchExecutionEventsPage, self.client.executions.events(
+                self.platform_session_id,
+                self.execution_kind,
+                self.execution_id,
+                afterSequence=after_sequence,
+                limit=limit,
+            ))
+            for event in page.get("events", []):
+                yield event
+                sequence = event.get("sequence")
+                if isinstance(sequence, (str, int)):
+                    after_sequence = str(sequence)
+            if page.get("hasMore") is True:
+                next_cursor = page.get("nextCursor")
+                if not isinstance(next_cursor, str) or not next_cursor:
+                    raise StacklinkError("Execution event page reported hasMore without nextCursor.")
+                after_sequence = next_cursor
+                continue
+            current = cast(WorkbenchExecutionStateResponse, self.client.executions.get(
+                self.platform_session_id,
+                self.execution_kind,
+                self.execution_id,
+            ))
+            execution = current.get("execution", {})
+            if isinstance(execution, dict) and execution.get("status") in {
+                "completed",
+                "failed",
+                "cancelled",
+                "timed_out",
+            }:
+                return
+            time.sleep(poll_interval_seconds)
+
+    def wait(self, *, wait_timeout_ms: int | None = None, poll_interval_ms: int | None = None) -> WorkbenchExecutionStateResponse:
+        body = {
+            key: value
+            for key, value in {
+                "waitTimeoutMs": wait_timeout_ms,
+                "pollIntervalMs": poll_interval_ms,
+            }.items()
+            if value is not None
+        }
+        return cast(WorkbenchExecutionStateResponse, self.client.executions.wait(
+            self.platform_session_id,
+            self.execution_kind,
+            self.execution_id,
+            body,
+        ))
+
+    def cancel(self, reason: str | None = None) -> WorkbenchExecutionStateResponse:
+        return cast(WorkbenchExecutionStateResponse, self.client.executions.cancel(
+            self.platform_session_id,
+            self.execution_kind,
+            self.execution_id,
+            reason,
+        ))
+
+
+def _wait_for_workbench_response(
+    client,
+    request: Mapping[str, Any],
+    accepted: WorkbenchExecutionResponse,
+) -> WorkbenchExecutionResponse:
+    session_id = request.get("sessionId")
+    execution_id = accepted.get("executionId")
+    execution_kind = accepted.get("executionKind")
+    if (
+        not isinstance(session_id, str)
+        or not isinstance(execution_id, str)
+        or execution_kind not in {"workbench_code", "workbench_bash"}
+    ):
+        return accepted
+
+    while True:
+        waited = cast(WorkbenchExecutionStateResponse, client.executions.wait(
+            session_id,
+            execution_kind,
+            execution_id,
+            {"waitTimeoutMs": 10_000, "pollIntervalMs": 500},
+        ))
+        execution = waited.get("execution")
+        if not isinstance(execution, dict) or execution.get("status") not in {
+            "completed",
+            "failed",
+            "cancelled",
+            "timed_out",
+        }:
+            continue
+        response: WorkbenchExecutionResponse = {
+            **accepted,
+            "executionId": str(execution.get("executionId") or execution.get("id") or execution_id),
+            "executionKind": cast(WorkbenchExecutionKind, execution.get("executionKind") or execution_kind),
+            "status": cast(WorkbenchExecutionStatus, execution["status"]),
+            "execution": cast(WorkbenchExecutionLog, execution),
+            "executionLog": cast(WorkbenchExecutionLog, execution),
+        }
+        output_mode = execution.get("outputMode")
+        if output_mode in {"inline", "preview", "events", "artifact"}:
+            response["outputMode"] = cast(WorkbenchOutputMode, output_mode)
+        if "result" in execution:
+            response["result"] = cast(
+                WorkbenchRuntimeEnvelope | WorkbenchBackgroundExecutionResult | WorkbenchOutputPreview,
+                execution["result"],
+            )
+        artifact = execution.get("artifact")
+        if isinstance(artifact, dict):
+            response["artifact"] = cast(WorkbenchArtifact, artifact)
+        return response
+
+
+class WorkbenchFiles:
+    def __init__(self, sandbox: "SessionSandbox"):
+        self.sandbox = sandbox
+
+    def list(self) -> WorkbenchFileListResponse:
+        return cast(WorkbenchFileListResponse, self.sandbox.client.request(
+            "GET",
+            f"/v1/workbench/sessions/{self.sandbox.ensure_session_id()}/files",
+            query={"sessionId": self.sandbox.platform_session_id},
+        ))
+
+    def read(self, path: str) -> WorkbenchFileResponse:
+        return cast(WorkbenchFileResponse, self.sandbox.client.request(
+            "GET",
+            f"/v1/workbench/sessions/{self.sandbox.ensure_session_id()}/files",
+            query={"sessionId": self.sandbox.platform_session_id, "path": path},
+        ))
+
+    def upload(
+        self,
+        path: str,
+        content: str | bytes,
+        *,
+        content_type: str | None = None,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> WorkbenchFileResponse:
+        raw = content.encode("utf-8") if isinstance(content, str) else content
+        body = {
+            "sessionId": self.sandbox.platform_session_id,
+            "path": path,
+            "contentBase64": base64.b64encode(raw).decode("ascii"),
+            "metadata": dict(metadata or {}),
+        }
+        if content_type is not None:
+            body["contentType"] = content_type
+        return cast(WorkbenchFileResponse, self.sandbox.client.request(
+            "PUT",
+            f"/v1/workbench/sessions/{self.sandbox.ensure_session_id()}/files",
+            json_body=body,
+        ))
+
+    def prepare_upload(self, path: str, *, content_type: str, size_bytes: int) -> WorkbenchFileUploadPrepareResponse:
+        return cast(WorkbenchFileUploadPrepareResponse, self.sandbox.client.request(
+            "POST",
+            f"/v1/workbench/sessions/{self.sandbox.ensure_session_id()}/files/upload-url",
+            json_body={
+                "sessionId": self.sandbox.platform_session_id,
+                "path": path,
+                "contentType": content_type,
+                "sizeBytes": size_bytes,
+            },
+        ))
+
+    def finalize_upload(
+        self,
+        upload_id: str,
+        *,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> WorkbenchFileMutationResponse:
+        return cast(WorkbenchFileMutationResponse, self.sandbox.client.request(
+            "POST",
+            f"/v1/workbench/sessions/{self.sandbox.ensure_session_id()}/files/finalize-upload",
+            json_body={
+                "sessionId": self.sandbox.platform_session_id,
+                "uploadId": upload_id,
+                "metadata": dict(metadata or {}),
+            },
+        ))
+
+    def upload_large(
+        self,
+        path: str,
+        content: str | bytes,
+        *,
+        content_type: str | None = None,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> WorkbenchFileMutationResponse:
+        raw = content.encode("utf-8") if isinstance(content, str) else content
+        resolved_content_type = content_type or (
+            "text/plain; charset=utf-8" if isinstance(content, str) else "application/octet-stream"
+        )
+        prepared = self.prepare_upload(path, content_type=resolved_content_type, size_bytes=len(raw))
+        upload = prepared.get("upload") if isinstance(prepared, dict) else None
+        if (
+            not isinstance(upload, dict)
+            or not isinstance(upload.get("uploadUrl"), str)
+            or not isinstance(upload.get("uploadId"), str)
+        ):
+            raise StacklinkError("Stacklink did not return a signed Workbench upload URL.")
+        headers = upload.get("requiredHeaders") if isinstance(upload.get("requiredHeaders"), dict) else {}
+        request = Request(upload["uploadUrl"], data=raw, method="PUT", headers={str(key): str(value) for key, value in headers.items()})
+        try:
+            with urlopen(request, timeout=self.sandbox.client.timeout) as response:
+                if response.status >= 400:
+                    raise StacklinkError(f"Workbench file upload failed with {response.status}.")
+        except HTTPError as error:
+            raise StacklinkError(f"Workbench file upload failed with {error.code}.") from error
+        return self.finalize_upload(
+            upload["uploadId"],
+            metadata=metadata,
+        )
+
+    def download(self, path: str) -> bytes:
+        signed = self.sandbox.client.request(
+            "GET",
+            f"/v1/workbench/sessions/{self.sandbox.ensure_session_id()}/files/download-url",
+            query={"sessionId": self.sandbox.platform_session_id, "path": path},
+        )
+        download_url = signed.get("downloadUrl") if isinstance(signed, dict) else None
+        if not isinstance(download_url, str):
+            raise StacklinkError("Stacklink did not return a signed Workbench download URL.")
+        try:
+            with urlopen(Request(download_url, method="GET"), timeout=self.sandbox.client.timeout) as response:
+                return response.read()
+        except HTTPError as error:
+            raise StacklinkError(f"Workbench file download failed with {error.code}.") from error
+
+    def delete(self, path: str) -> WorkbenchFileMutationResponse:
+        return cast(WorkbenchFileMutationResponse, self.sandbox.client.request(
+            "DELETE",
+            f"/v1/workbench/sessions/{self.sandbox.ensure_session_id()}/files",
+            query={"sessionId": self.sandbox.platform_session_id, "path": path},
+        ))
+
+
+class WorkbenchRuns:
+    def __init__(self, sandbox: "SessionSandbox"):
+        self.sandbox = sandbox
+
+    def list(
+        self,
+        *,
+        series_id: str | None = None,
+        cursor: str | None = None,
+        limit: int | None = None,
+    ) -> WorkbenchRunDefinitionListResponse:
+        query = {
+            "sessionId": self.sandbox.platform_session_id,
+            "workbenchSessionId": self.sandbox.ensure_session_id(),
+        }
+        if series_id is not None:
+            query["seriesId"] = series_id
+        if cursor is not None:
+            query["cursor"] = cursor
+        if limit is not None:
+            query["limit"] = limit
+        return cast(
+            WorkbenchRunDefinitionListResponse,
+            self.sandbox.client.request("GET", "/v1/workbench/run-definitions", query=query),
+        )
+
+    def get(self, run_definition_id: str) -> WorkbenchRunDefinitionResponse:
+        return cast(
+            WorkbenchRunDefinitionResponse,
+            self.sandbox.client.request(
+                "GET",
+                f"/v1/workbench/run-definitions/{run_definition_id}",
+                query={"sessionId": self.sandbox.platform_session_id},
+            ),
+        )
+
+    def update(
+        self,
+        run_definition_id: str,
+        *,
+        name: str | _WorkbenchUnset = _WORKBENCH_UNSET,
+        description: str | None | _WorkbenchUnset = _WORKBENCH_UNSET,
+        version_reason: str | None | _WorkbenchUnset = _WORKBENCH_UNSET,
+    ) -> WorkbenchRunDefinitionResponse:
+        body: dict[str, Any] = {"sessionId": self.sandbox.platform_session_id}
+        if not isinstance(name, _WorkbenchUnset):
+            body["name"] = name
+        if not isinstance(description, _WorkbenchUnset):
+            body["description"] = description
+        if not isinstance(version_reason, _WorkbenchUnset):
+            body["versionReason"] = version_reason
+        return cast(
+            WorkbenchRunDefinitionResponse,
+            self.sandbox.client.request(
+                "PATCH",
+                f"/v1/workbench/run-definitions/{run_definition_id}",
+                json_body=body,
+            ),
+        )
+
+    def save(
+        self,
+        execution_id: str,
+        name: str,
+        *,
+        series_id: str | None = None,
+        description: str | None = None,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> WorkbenchRunDefinitionResponse:
+        body = {
+            "sessionId": self.sandbox.platform_session_id,
+            "workbenchSessionId": self.sandbox.ensure_session_id(),
+            "executionId": execution_id,
+            "name": name,
+            "metadata": dict(metadata or {}),
+        }
+        if series_id is not None:
+            body["seriesId"] = series_id
+        if description is not None:
+            body["description"] = description
+        return cast(
+            WorkbenchRunDefinitionResponse,
+            self.sandbox.client.request("POST", "/v1/workbench/run-definitions", json_body=body),
+        )
+
+    def replay(
+        self,
+        run_definition_id: str,
+        *,
+        background: bool = False,
+        timeout_ms: int | None = None,
+        environment_mode: WorkbenchReplayEnvironmentMode = "require_match",
+        idempotency_key: str,
+    ) -> WorkbenchRunComparisonResponse:
+        body = {
+            "sessionId": self.sandbox.platform_session_id,
+            "background": background,
+            "environmentMode": environment_mode,
+            "idempotencyKey": idempotency_key,
+        }
+        if timeout_ms is not None:
+            body["timeoutMs"] = timeout_ms
+        accepted = cast(WorkbenchRunComparisonResponse, self.sandbox.client.request(
+            "POST",
+            f"/v1/workbench/run-definitions/{run_definition_id}/replay",
+            json_body=body,
+        ))
+        if background:
+            return accepted
+        completed = _wait_for_workbench_response(
+            self.sandbox.client,
+            body,
+            accepted["execution"],
+        )
+        definition = accepted["runDefinition"]
+        execution_id = str(completed.get("executionId") or accepted["comparison"].get("replayExecutionId"))
+        comparison = accepted["comparison"]
+        comparison_deadline = time.monotonic() + _REUSABLE_RUN_COMPARISON_WAIT_SECONDS
+        while True:
+            current = self.get(run_definition_id).get("runDefinition", {})
+            replay = next(
+                (
+                    entry
+                    for entry in current.get("replays", [])
+                    if entry.get("executionId") == execution_id
+                ),
+                None,
+            )
+            if isinstance(replay, dict) and isinstance(replay.get("comparison"), dict):
+                comparison = cast(WorkbenchRunComparison, replay["comparison"])
+            if comparison.get("status") in {"matched", "mismatched"}:
+                break
+            if time.monotonic() >= comparison_deadline:
+                break
+            time.sleep(_REUSABLE_RUN_COMPARISON_POLL_SECONDS)
+        execution_log = completed.get("executionLog")
+        metadata = execution_log.get("metadata") if isinstance(execution_log, dict) else None
+        environment = metadata.get("replayEnvironment") if isinstance(metadata, dict) else None
+        return {
+            **accepted,
+            "execution": completed,
+            "environment": cast(
+                WorkbenchRunEnvironmentResolution,
+                environment if isinstance(environment, dict) else accepted["environment"],
+            ),
+            "comparison": cast(
+                WorkbenchRunComparison,
+                comparison,
+            ),
+        }
+
+    def publish(
+        self,
+        run_definition_id: str,
+        expected_published_version: int | None,
+    ) -> WorkbenchRunActivationResponse:
+        return cast(WorkbenchRunActivationResponse, self.sandbox.client.request(
+            "POST",
+            f"/v1/workbench/run-definitions/{run_definition_id}/publish",
+            json_body={
+                "sessionId": self.sandbox.platform_session_id,
+                "expectedPublishedVersion": expected_published_version,
+            },
+        ))
+
+    def rollback(
+        self,
+        run_definition_id: str,
+        expected_published_version: int | None,
+    ) -> WorkbenchRunActivationResponse:
+        return cast(WorkbenchRunActivationResponse, self.sandbox.client.request(
+            "POST",
+            f"/v1/workbench/run-definitions/{run_definition_id}/rollback",
+            json_body={
+                "sessionId": self.sandbox.platform_session_id,
+                "expectedPublishedVersion": expected_published_version,
+            },
+        ))
+
+
+class WorkbenchArtifacts:
+    def __init__(self, sandbox: "SessionSandbox"):
+        self.sandbox = sandbox
+
+    def list(self) -> WorkbenchArtifactListResponse:
+        return cast(WorkbenchArtifactListResponse, self.sandbox.client.request(
+            "GET",
+            "/v1/workbench/artifacts",
+            query={"workbenchSessionId": self.sandbox.ensure_session_id()},
+        ))
+
+    def promote(
+        self,
+        path: str,
+        *,
+        label: str | None = None,
+        content_type: str | None = None,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> WorkbenchArtifactResponse:
+        body: dict[str, Any] = {
+            "sessionId": self.sandbox.platform_session_id,
+            "workbenchSessionId": self.sandbox.ensure_session_id(),
+            "path": path,
+            "metadata": dict(metadata or {}),
+        }
+        if label is not None:
+            body["label"] = label
+        if content_type is not None:
+            body["contentType"] = content_type
+        return cast(WorkbenchArtifactResponse, self.sandbox.client.request("POST", "/v1/workbench/artifacts", json_body=body))
+
+    def download_url(self, artifact_id: str) -> WorkbenchArtifactDownloadResponse:
+        return cast(WorkbenchArtifactDownloadResponse, self.sandbox.client.workbench.artifact_download_url(artifact_id))
+
+    def download(self, artifact_id: str) -> bytes:
+        return self.sandbox.client.workbench.download_artifact(artifact_id)
+
+    def delete(self, artifact_id: str, *, path: str | None = None) -> WorkbenchArtifactResponse:
+        return cast(WorkbenchArtifactResponse, self.sandbox.client.request(
+            "DELETE",
+            f"/v1/workbench/artifacts/{artifact_id}",
+            query={
+                "sessionId": self.sandbox.platform_session_id,
+                "workbenchSessionId": self.sandbox.ensure_session_id(),
+                "path": path,
+            },
+        ))
+
+
+class SessionSandbox:
+    def __init__(self, client, platform_session_id: str):
+        self.client = client
+        self.platform_session_id = platform_session_id
+        self.workbench_session_id: str | None = None
+        self.files = WorkbenchFiles(self)
+        self.runs = WorkbenchRuns(self)
+        self.artifacts = WorkbenchArtifacts(self)
+
+    def use(
+        self,
+        workbench_session_id: str | None = None,
+        *,
+        create_new: bool = False,
+        metadata: Mapping[str, Any] | None = None,
+    ) -> WorkbenchSessionUseResponse:
+        response = cast(WorkbenchSessionUseResponse, self.client.request(
+            "POST",
+            "/v1/workbench/sessions",
+            json_body={
+                "sessionId": self.platform_session_id,
+                "workbenchSessionId": workbench_session_id,
+                "createNew": create_new,
+                "metadata": dict(metadata or {}),
+            },
+        ))
+        self.workbench_session_id = str(response["workbenchSession"]["id"])
+        return response
+
+    def ensure_session_id(self) -> str:
+        if self.workbench_session_id is None:
+            self.use()
+        assert self.workbench_session_id is not None
+        return self.workbench_session_id
+
+    def run(self, body: WorkbenchRunInput) -> WorkbenchRun:
+        request = {
+            **body,
+            "sessionId": self.platform_session_id,
+            "workbenchSessionId": body.get("workbenchSessionId") or self.ensure_session_id(),
+            "operation": "run_code",
+        }
+        response: WorkbenchExecutionResponse = self.client.request("POST", "/v1/workbench/run", json_body=request)
+        if body.get("background") is not True:
+            response = _wait_for_workbench_response(self.client, request, response)
+        self.workbench_session_id = str(response["workbenchSession"]["id"])
+        return WorkbenchRun(self.client, self.platform_session_id, response)
+
+    def bash(self, body: WorkbenchBashInput) -> WorkbenchRun:
+        request = {
+            **body,
+            "sessionId": self.platform_session_id,
+            "workbenchSessionId": body.get("workbenchSessionId") or self.ensure_session_id(),
+        }
+        response: WorkbenchExecutionResponse = self.client.request("POST", "/v1/workbench/bash", json_body=request)
+        if body.get("background") is not True:
+            response = _wait_for_workbench_response(self.client, request, response)
+        self.workbench_session_id = str(response["workbenchSession"]["id"])
+        return WorkbenchRun(self.client, self.platform_session_id, response)
+
+    def destroy(self) -> WorkbenchSessionDestroyResponse:
+        return cast(WorkbenchSessionDestroyResponse, self.client.workbench.destroy(
+            self.platform_session_id,
+            self.ensure_session_id(),
+        ))
+
+    def delete(self) -> WorkbenchSessionDeleteReceipt:
+        result = self.client.workbench.delete_session(
+            self.platform_session_id,
+            self.ensure_session_id(),
+        )
+        self.workbench_session_id = None
+        return result
+
+
+class WorkbenchResource:
+    def __init__(self, client):
+        self.client = client
+
+    def bind(self, platform_session_id: str) -> SessionSandbox:
+        return SessionSandbox(self.client, platform_session_id)
+
+    def models(self):
+        return self.client.request("GET", "/v1/workbench/models")
+
+    def run(self, body: dict[str, Any]):
+        response = cast(WorkbenchExecutionResponse, self.client.request("POST", "/v1/workbench/run", json_body=body))
+        if body.get("operation") == "run_code" and body.get("background") is not True:
+            return _wait_for_workbench_response(self.client, body, response)
+        return response
+
+    def bash(self, body: dict[str, Any]):
+        response = cast(WorkbenchExecutionResponse, self.client.request("POST", "/v1/workbench/bash", json_body=body))
+        if body.get("background") is not True:
+            return _wait_for_workbench_response(self.client, body, response)
+        return response
+
+    def destroy(self, session_id: str, workbench_session_id: str):
+        return self.run({
+            "sessionId": session_id,
+            "workbenchSessionId": workbench_session_id,
+            "operation": "destroy",
+        })
+
+    def delete_session(self, session_id: str, workbench_session_id: str) -> WorkbenchSessionDeleteReceipt:
+        return cast(WorkbenchSessionDeleteReceipt, self.client.request(
+            "DELETE",
+            f"/v1/workbench/sessions/{workbench_session_id}",
+            query={"sessionId": session_id},
+        ))
+
+    def sessions(self, session_id: str):
+        return self.client.request("GET", f"/v1/sessions/{session_id}/workbench")
+
+    def list_sessions(self, **query):
+        return self.client.request("GET", "/v1/workbench/sessions", query=query)
+
+    def get_session(self, workbench_session_id: str):
+        return self.client.request("GET", f"/v1/workbench/sessions/{workbench_session_id}")
+
+    def usage(self, workbench_session_id: str):
+        return self.client.request("GET", f"/v1/workbench/sessions/{workbench_session_id}/usage")
+
+    def list_artifacts(self, **query):
+        return self.client.request("GET", "/v1/workbench/artifacts", query=query)
+
+    def artifact_download_url(self, artifact_id: str):
+        return self.client.request("GET", f"/v1/workbench/artifacts/{artifact_id}/download-url")
+
+    def download_artifact(self, artifact_id: str) -> bytes:
+        return self.client.request_raw("GET", f"/v1/workbench/artifacts/{artifact_id}/download")
+
+    def list_budgets(self, **query):
+        return self.client.request("GET", "/v1/workbench/budgets", query=query)
+
+    def create_budget(self, body: dict[str, Any]):
+        return self.client.request("POST", "/v1/workbench/budgets", json_body=body)
+
+    def update_budget(self, budget_policy_id: str, body: dict[str, Any]):
+        return self.client.request("PATCH", f"/v1/workbench/budgets/{budget_policy_id}", json_body=body)
+
+    def delete_budget(self, budget_policy_id: str):
+        return self.client.request("DELETE", f"/v1/workbench/budgets/{budget_policy_id}")