import logging
import os
import sys
import time
from collections import defaultdict, deque
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import ExitStack
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Any,
    DefaultDict,
    Dict,
    Iterator,
    List,
    Optional,
    Set,
    Tuple,
    Union,
    cast,
)

import dagster._check as check
from dagster import DagsterInstance
from dagster._core.launcher.base import LaunchRunContext
from dagster._core.remote_representation import CodeLocationOrigin
from dagster._core.remote_representation.origin import RegisteredCodeLocationOrigin
from dagster._core.utils import FuturesAwareThreadPoolExecutor
from dagster._grpc.client import DagsterGrpcClient
from dagster._grpc.types import CancelExecutionRequest
from dagster._serdes import deserialize_value, serialize_value
from dagster._time import get_current_datetime, get_current_timestamp
from dagster._utils.cached_method import cached_method
from dagster._utils.container import retrieve_containerized_utilization_metrics
from dagster._utils.error import SerializableErrorInfo, serializable_error_info_from_exc_info
from dagster._utils.interrupts import raise_interrupts_as
from dagster._utils.merger import merge_dicts
from dagster._utils.typed_dict import init_optional_typeddict
from dagster_cloud_cli.core.errors import DagsterCloudHTTPError, raise_http_error
from dagster_cloud_cli.core.workspace import CodeLocationDeployData, get_instance_ref_for_user_code

from dagster_cloud.api.dagster_cloud_api import (
    AgentHeartbeat,
    AgentUtilizationMetrics,
    BatchDagsterCloudUploadApiResponse,
    DagsterCloudApi,
    DagsterCloudApiErrorResponse,
    DagsterCloudApiGrpcResponse,
    DagsterCloudApiRequest,
    DagsterCloudApiResponse,
    DagsterCloudApiSuccess,
    DagsterCloudApiThreadTelemetry,
    DagsterCloudApiUnknownCommandResponse,
    DagsterCloudUploadApiResponse,
    TimestampedError,
)
from dagster_cloud.batching import Batcher
from dagster_cloud.instance import DagsterCloudAgentInstance
from dagster_cloud.util.errors import truncate_serialized_error
from dagster_cloud.workspace.user_code_launcher import (
    DagsterCloudUserCodeLauncher,
    UserCodeLauncherEntry,
)

from ..opentelemetry.observers.execution_observer import observe_execution
from ..util import SERVER_HANDLE_TAG, compressed_namedtuple_upload_file, is_isolated_run
from ..version import __version__
from .instrumentation.constants import DAGSTER_CLOUD_AGENT_METRIC_PREFIX
from .instrumentation.run_launch import extract_run_attributes
from .instrumentation.schedule import inspect_schedule_result
from .instrumentation.sensor import inspect_sensor_result
from .queries import (
    ADD_AGENT_HEARTBEATS_MUTATION,
    DEPLOYMENTS_QUERY,
    GET_USER_CLOUD_REQUESTS_QUERY,
    WORKSPACE_ENTRIES_QUERY,
)

if TYPE_CHECKING:
    import datetime

CHECK_WORKSPACE_INTERVAL_SECONDS = 5

# Interval at which Agent heartbeats are sent to the host cloud
AGENT_HEARTBEAT_INTERVAL_SECONDS = 30

# Interval at which to record utilization metrics
AGENT_UTILIZATION_METRICS_INTERVAL_SECONDS = 60

AGENT_HEARTBEAT_ERROR_LIMIT = 25  # Send at most 25 errors

DEFAULT_PENDING_REQUESTS_LIMIT = 100

SLEEP_INTERVAL_SECONDS = float(os.getenv("DAGSTER_CLOUD_AGENT_SLEEP_INTERVAL_SECONDS", "0.5"))


def UPLOAD_API_RESPONSE_BATCHING_ENABLED():
    return os.getenv("DAGSTER_CLOUD_AGENT_UPLOAD_API_RESPONSE_BATCHING_ENABLED") == "true"


DEPLOYMENT_INFO_QUERY = """
    query DeploymentInfo {
         deploymentInfo {
             deploymentName
         }
     }
"""

AGENT_VERSION_LABEL = "version"
AGENT_DEFAULT_MAX_THREADS_PER_CORE = 10
AGENT_MAX_THREADPOOL_WORKERS = int(
    os.getenv(
        "DAGSTER_CLOUD_AGENT_MAX_THREADPOOL_WORKERS",
        str((os.cpu_count() or 1) * AGENT_DEFAULT_MAX_THREADS_PER_CORE),
    )
)

LIVENESS_CHECK_INTERVAL_SECONDS = float(
    os.getenv("DAGSTER_CLOUD_AGENT_LIVENESS_CHECK_INTERVAL", "15.0")
)


class DagsterCloudAgent:
    def __init__(
        self,
        instance: DagsterCloudAgentInstance,
        pending_requests_limit: int = DEFAULT_PENDING_REQUESTS_LIMIT,
    ):
        self._logger = logging.getLogger("dagster_cloud.agent")
        self._instance: DagsterCloudAgentInstance = instance

        self._batcher: DefaultDict[
            str, Batcher[Tuple[str, DagsterCloudUploadApiResponse], None]
        ] = defaultdict(self._batcher_factory)

        if self._logger.isEnabledFor(logging.DEBUG):
            self._logger.info("Starting Dagster Cloud agent with debug logging...")
        else:
            self._logger.info("Starting Dagster Cloud agent...")

        self._exit_stack = ExitStack()
        self._iteration = 0

        self._executor = self._exit_stack.enter_context(
            FuturesAwareThreadPoolExecutor(
                max_workers=AGENT_MAX_THREADPOOL_WORKERS,
                thread_name_prefix="dagster_cloud_agent_worker",
            )
        )
        self._request_ids_to_futures: Dict[str, Future] = {}
        self._utilization_metrics = init_optional_typeddict(AgentUtilizationMetrics)

        self._last_heartbeat_time: Optional[datetime.datetime] = None

        self._last_workspace_check_time = None

        self._errors: deque = deque(
            maxlen=AGENT_HEARTBEAT_ERROR_LIMIT
        )  # (SerializableErrorInfo, timestamp) tuples

        self._pending_requests: List[Dict[str, Any]] = []
        self._locations_with_pending_requests: Set[Tuple[str, str, bool]] = set()
        self._ready_requests: List[Dict[str, Any]] = []

        self._location_query_times: Dict[Tuple[str, str, bool], float] = {}
        self._pending_requests_limit = check.int_param(
            pending_requests_limit, "pending_requests_limit"
        )
        self._active_deployments: Set[Tuple[str, bool]] = (  # deployment_name, is_branch_deployment
            set()
        )

        self._last_liveness_check_time = None

    def __enter__(self):
        return self

    def __exit__(self, _exception_type, _exception_value, _traceback):
        self._exit_stack.close()

    def _batcher_factory(
        self,
    ) -> Batcher[Tuple[str, DagsterCloudUploadApiResponse], None]:
        return Batcher(
            "upload_api_response",
            self._batch_upload_api_response,
            max_wait_ms=50,
            max_batch_size=32,
        )

    @property
    def _active_deployment_names(self):
        return [deployment[0] for deployment in self._active_deployments]

    @property
    def _active_full_deployment_names(self):
        return [deployment[0] for deployment in self._active_deployments if not deployment[1]]

    def _check_initial_deployment_names(self):
        if self._instance.deployment_names:
            result = self._instance.organization_scoped_graphql_client().execute(
                DEPLOYMENTS_QUERY,
                variable_values={"deploymentNames": self._instance.deployment_names},
            )
            deployments = result["data"]["deployments"]
            existing_deployment_names = {deployment["deploymentName"] for deployment in deployments}
            requested_deployment_names = set(self._instance.deployment_names)
            missing_deployment_names = requested_deployment_names.difference(
                existing_deployment_names
            )

            if missing_deployment_names:
                deployment_str = f"deployment{'s' if len(missing_deployment_names) > 1 else ''} {', '.join(missing_deployment_names)}"
                raise Exception(
                    f"Agent is configured to serve an invalid {deployment_str}. Check your"
                    " agent configuration to make sure it is serving the correct deployment.",
                )

    def _update_agent_resource_limits(
        self, user_code_launcher: DagsterCloudUserCodeLauncher
    ) -> None:
        # The agent should have environment variables defining its resource requests and limits.
        # However, the agent may be running in a container with resource limits that are different
        # For example, on k8s there are ways to effect change on the cpu limit, like mutating admission webhooks.
        # Since the effective cgroup limits precede actual hosts resources when it comes to actual behavior for
        # throttling and oom kills, we attempt to obtain these and fallback on the environment variables.
        container_utilization_metrics = self._utilization_metrics.get("container_utilization", {})
        memory_limit = container_utilization_metrics.get("memory_limit")
        if not memory_limit:
            memory_limit = os.getenv("DAGSTER_CLOUD_AGENT_MEMORY_LIMIT")
            self._logger.info(
                "Cannot obtain cgroup memory limit, using environment value: "
                f"DAGSTER_CLOUD_AGENT_MEMORY_LIMIT={memory_limit}"
            )

        cpu_cfs_period_us = container_utilization_metrics.get("cpu_cfs_period_us")
        cpu_cfs_quota_us = container_utilization_metrics.get("cpu_cfs_quota_us")

        cpu_limit = None
        if cpu_cfs_quota_us and cpu_cfs_period_us:
            cpu_limit = (
                1000.0 * cpu_cfs_quota_us
            ) / cpu_cfs_period_us  # cpu_limit expressed in milliseconds of cpu

        if not cpu_limit:
            cpu_limit = os.getenv("DAGSTER_CLOUD_AGENT_CPU_LIMIT")
            self._logger.info(
                "Cannot obtain CPU CFS values, using environment value: "
                f"DAGSTER_CLOUD_AGENT_CPU_LIMIT={cpu_limit}"
            )

        if not user_code_launcher.user_code_deployment_type.supports_utilization_metrics:
            self._logger.info(
                (
                    f"Cannot interpret resource limits for agent type {user_code_launcher.user_code_deployment_type.value}."
                    "Skipping utilization metrics retrieval."
                )
            )
            return

        limits = {
            "cpu_limit": cpu_limit,
            "memory_limit": memory_limit,
        }

        cpu_request = os.getenv("DAGSTER_CLOUD_AGENT_CPU_REQUEST")
        memory_request = os.getenv("DAGSTER_CLOUD_AGENT_MEMORY_REQUEST")
        if cpu_request:
            limits["cpu_request"] = cpu_request
        if memory_request:
            limits["memory_request"] = memory_request

        self._utilization_metrics["resource_limits"][
            user_code_launcher.user_code_deployment_type.value
        ] = limits  # type: ignore

    def _update_utilization_metrics(self, user_code_launcher: DagsterCloudUserCodeLauncher):
        container_utilization_metrics = retrieve_containerized_utilization_metrics(
            logger=self._logger,
            previous_measurement_timestamp=self._utilization_metrics["container_utilization"][
                "measurement_timestamp"
            ],
            previous_cpu_usage=self._utilization_metrics["container_utilization"]["cpu_usage"],
        )
        self._utilization_metrics["request_utilization"].update(
            self._executor.get_current_utilization_metrics()
        )
        self._utilization_metrics["container_utilization"].update(container_utilization_metrics)

        self._update_agent_resource_limits(user_code_launcher)

        self._logger.info(f"Current utilization metrics: {self._utilization_metrics}")

    def run_loop(
        self,
        user_code_launcher,
        agent_uuid,
    ):
        heartbeat_interval_seconds = AGENT_HEARTBEAT_INTERVAL_SECONDS

        if (
            not self._instance.includes_branch_deployments
            and not self._instance.deployment_names
            and not self._instance.include_all_serverless_deployments
        ):
            self._logger.info(
                "Deployment name was not set - checking to see if it can be fetched from the"
                " server..."
            )
            # Fetch the deployment name from the server if it isn't set (only true
            # for old agents, and only will work if there's a single deployment in the org)
            result = self._instance.graphql_client.execute(DEPLOYMENT_INFO_QUERY)
            deployment_name = result["data"]["deploymentInfo"]["deploymentName"]
            self._instance = self._exit_stack.enter_context(
                DagsterInstance.from_ref(self._instance.ref_for_deployment(deployment_name))  # type: ignore  # (instance subclass)
            )

        self._check_initial_deployment_names()

        self._check_update_workspace(
            user_code_launcher,
            upload_all=user_code_launcher.upload_snapshots_on_startup,
        )

        self._logger.info(
            f"Will start polling for requests from {self._instance.dagster_cloud_url} once user code has"
            " been loaded."
        )

        heartbeat_error_size_limit = int(os.getenv("DAGSTER_CLOUD_AGENT_ERROR_SIZE_LIMIT", "10000"))

        while True:
            try:
                for error in self.run_iteration(user_code_launcher):
                    if error:
                        self._logger.error(str(error))
                        self._errors.appendleft(
                            (
                                truncate_serialized_error(
                                    error, heartbeat_error_size_limit, max_depth=3
                                ),
                                get_current_datetime(),
                            )
                        )
            except Exception:
                error_info = serializable_error_info_from_exc_info(sys.exc_info())
                self._logger.error(f"Caught error:\n{error_info}")
                self._errors.appendleft(
                    (
                        truncate_serialized_error(
                            error_info, heartbeat_error_size_limit, max_depth=3
                        ),
                        get_current_datetime(),
                    )
                )

            # Check for any received interrupts
            with raise_interrupts_as(KeyboardInterrupt):
                pass

            if user_code_launcher.ready_to_serve_requests:
                try:
                    self._check_add_heartbeat(agent_uuid, heartbeat_interval_seconds)
                except Exception:
                    self._logger.exception("Failed to add heartbeat")

            # Check for any received interrupts
            with raise_interrupts_as(KeyboardInterrupt):
                pass

            try:
                self._check_update_workspace(user_code_launcher, upload_all=False)

            except Exception:
                self._logger.error(
                    f"Failed to check for workspace updates: \n{serializable_error_info_from_exc_info(sys.exc_info())}"
                )

            self._write_liveness_sentinel_if_overdue()

            # Check for any received interrupts
            with raise_interrupts_as(KeyboardInterrupt):
                time.sleep(SLEEP_INTERVAL_SECONDS)

    def _write_liveness_sentinel_if_overdue(self):
        if self._last_liveness_check_time is False:
            return

        now = time.time()
        if self._last_liveness_check_time is None:
            self._logger.info("Starting liveness sentinel")
        elif self._last_liveness_check_time + LIVENESS_CHECK_INTERVAL_SECONDS > now:
            return

        # Write to a sentinel file to indicate that we've finished our initial
        # reconciliation - this is used to indicate that we're ready to
        # serve requests
        try:
            if not os.access("/opt", os.W_OK):
                self._logger.warning("Disabling liveness sentinel - /opt is not writable")
                self._last_liveness_check_time = False
                return
            Path("/opt/liveness_sentinel.txt").touch(exist_ok=True)
            self._last_liveness_check_time = now
        except Exception as e:
            self._logger.error(f"Failed to write liveness sentinel and disabling it: {e}")
            self._last_liveness_check_time = False

    def _check_update_workspace(self, user_code_launcher, upload_all):
        curr_time = get_current_datetime()

        if (
            self._last_workspace_check_time
            and (curr_time - self._last_workspace_check_time).total_seconds()
            < CHECK_WORKSPACE_INTERVAL_SECONDS
        ):
            return

        self._last_workspace_check_time = curr_time
        self._query_for_workspace_updates(user_code_launcher, upload_all=upload_all)

    def _check_add_heartbeat(
        self,
        agent_uuid,
        heartbeat_interval_seconds,
    ):
        curr_time = get_current_datetime()

        if (
            self._last_heartbeat_time
            and (curr_time - self._last_heartbeat_time).total_seconds() < heartbeat_interval_seconds
        ):
            return

        errors = [
            TimestampedError(
                timestamp=timestamp.timestamp(),
                error=error,
            )
            for (error, timestamp) in self._errors
            if timestamp.timestamp() > curr_time.timestamp() - 60 * 60 * 24
        ]

        run_worker_statuses_dict = self._instance.user_code_launcher.get_cloud_run_worker_statuses(
            self._active_deployment_names
        )

        code_server_heartbeats_dict = self._instance.user_code_launcher.get_grpc_server_heartbeats()

        agent_image_tag = os.getenv("DAGSTER_CLOUD_AGENT_IMAGE_TAG")
        if self._instance.user_code_launcher.agent_metrics_enabled:
            num_running_requests = self._utilization_metrics["request_utilization"][
                "num_running_requests"
            ]
            max_concurrent_requests = self._utilization_metrics["request_utilization"][
                "max_concurrent_requests"
            ]
            self._logger.info(
                f"Current agent threadpool utilization: {num_running_requests}/{max_concurrent_requests} threads"
            )

        self._last_heartbeat_time = curr_time

        heartbeats = {
            deployment_name: AgentHeartbeat(
                timestamp=curr_time.timestamp(),
                agent_id=agent_uuid,
                agent_label=self._instance.dagster_cloud_api_agent_label,
                agent_type=(
                    type(self._instance.user_code_launcher).__name__
                    if self._instance.user_code_launcher
                    else None
                ),
                metadata=merge_dicts(
                    {AGENT_VERSION_LABEL: __version__},
                    {"image_tag": agent_image_tag} if agent_image_tag else {},
                    {
                        "utilization_metrics": self._utilization_metrics
                        if self._instance.user_code_launcher.agent_metrics_enabled
                        else {}
                    },
                ),
                run_worker_statuses=run_worker_statuses_dict[deployment_name],
                code_server_heartbeats=code_server_heartbeats_dict.get(deployment_name, []),
                agent_queues_config=self._instance.agent_queues_config,
            )
            for deployment_name in self._active_deployment_names
        }

        serialized_agent_heartbeats = [
            {
                "deploymentName": deployment_name,
                "serializedAgentHeartbeat": serialize_value(heartbeat),
            }
            for deployment_name, heartbeat in heartbeats.items()
        ]

        serialized_errors = [serialize_value(error) for error in errors]
        try:
            self._instance.organization_scoped_graphql_client().execute(
                ADD_AGENT_HEARTBEATS_MUTATION,
                variable_values={
                    "serializedAgentHeartbeats": serialized_agent_heartbeats,
                    "serializedErrors": serialized_errors,
                },
                idempotent_mutation=True,
            )
        except DagsterCloudHTTPError as e:
            if e.response.status_code == 413:
                heartbeats_size = sum([len(heartbeat) for heartbeat in serialized_agent_heartbeats])
                errors_size = sum([len(error) for error in serialized_errors])

                error_message = f"Request with {heartbeats_size} heartbeat bytes and {errors_size} error bytes was too large, submitting a smaller heartbeat without messages and stack traces."
                self._logger.exception(error_message)

                serialized_agent_heartbeats = [
                    {
                        "deploymentName": deployment_name,
                        "serializedAgentHeartbeat": serialize_value(
                            heartbeat.without_messages_and_errors()
                        ),
                    }
                    for deployment_name, heartbeat in heartbeats.items()
                ]

                self._instance.organization_scoped_graphql_client().execute(
                    ADD_AGENT_HEARTBEATS_MUTATION,
                    variable_values={
                        "serializedAgentHeartbeats": serialized_agent_heartbeats,
                        "serializedErrors": [
                            serialize_value(
                                TimestampedError(
                                    timestamp=curr_time.timestamp(),
                                    error=SerializableErrorInfo(
                                        error_message,
                                        stack=[],
                                        cls_name=None,
                                    ),
                                )
                            )
                        ],
                    },
                    idempotent_mutation=True,
                )
            else:
                raise

    @property
    def executor(self) -> ThreadPoolExecutor:
        return self._executor

    @property
    def request_ids_to_futures(self) -> Dict[str, Future]:
        return self._request_ids_to_futures

    def _upload_outdated_workspace_entries(
        self,
        deployment_name: str,
        is_branch_deployment: bool,
        user_code_launcher: DagsterCloudUserCodeLauncher,
    ):
        result = self._instance.graphql_client_for_deployment(deployment_name).execute(
            WORKSPACE_ENTRIES_QUERY,
            variable_values={
                "deploymentNames": [deployment_name],
                "includeAllServerlessDeployments": False,
                "agentQueues": self._instance.agent_queues_config.queues,
            },
        )
        entries = result["data"]["deployments"][0]["workspaceEntries"]
        now = time.time()

        upload_metadata = {}

        for entry in entries:
            location_name = entry["locationName"]
            code_location_deploy_data = deserialize_value(
                entry["serializedDeploymentMetadata"], CodeLocationDeployData
            )
            if entry["hasOutdatedData"]:
                # Spin up a server for this location and upload its metadata to Cloud
                # (Bump the TTL counter as well to leave the server up)
                self._location_query_times[
                    (deployment_name, location_name, is_branch_deployment)
                ] = now
                upload_metadata[(deployment_name, location_name)] = UserCodeLauncherEntry(
                    code_location_deploy_data=code_location_deploy_data,
                    update_timestamp=float(entry["metadataTimestamp"]),
                )

        user_code_launcher.add_upload_metadata(upload_metadata)

    def _has_ttl(self, user_code_launcher, is_branch_deployment):
        # branch deployments always have TTLs, other deployments only if you asked for it specifically
        return is_branch_deployment or user_code_launcher.server_ttl_enabled_for_full_deployments

    def _get_ttl_seconds(self, is_branch_deployment):
        return (
            self._instance.user_code_launcher.branch_deployment_ttl_seconds
            if is_branch_deployment
            else self._instance.user_code_launcher.full_deployment_ttl_seconds
        )

    def _get_locations_with_ttl_to_query(self, user_code_launcher) -> List[Tuple[str, str]]:
        now = time.time()

        # For the deployments with TTLs, decide which locations to consider
        # Include the location if:
        # - a) There's a pending request in the queue for it
        # - b) It's TTL hasn't expired since the last time somebody asked for it
        # Always include locations in a), and add locations from b) until you hit a limit
        location_candidates: Set[Tuple[str, str, float]] = {
            (deployment, location, -1.0)  # Score below 0 so that they're at the front of the list
            for deployment, location, is_branch_deployment in self._locations_with_pending_requests
            if self._has_ttl(user_code_launcher, is_branch_deployment)
        }

        num_locations_to_query = self._instance.user_code_launcher.server_ttl_max_servers

        if len(location_candidates) > num_locations_to_query:
            self._logger.warning(
                f"Temporarily keeping {len(location_candidates)} servers with pending requests "
                f"running, which is more than the configured {num_locations_to_query} servers."
            )
            return [(deployment, location) for deployment, location, _score in location_candidates]

        for location_entry, query_time in self._location_query_times.items():
            if location_entry in self._locations_with_pending_requests:
                # Skip locations that are already in location_candidates due to having
                # pending requests
                continue

            deployment_name, location, is_branch_deployment = location_entry

            if not self._has_ttl(user_code_launcher, is_branch_deployment):
                continue

            time_since_last_query = now - query_time

            if time_since_last_query >= self._get_ttl_seconds(is_branch_deployment):
                continue

            location_candidates.add((deployment_name, location, time_since_last_query))

        sorted_results = sorted(location_candidates, key=lambda x: x[2])

        # sort by time since last query asending and return the first N
        filtered_results = sorted_results[:num_locations_to_query]

        num_left_out = len(sorted_results) - num_locations_to_query

        if num_left_out > 0:
            self._logger.warning(
                f"{len(sorted_results)} locations have been queried within TTL, but "
                f"filtering out {num_left_out} to stay within {num_locations_to_query}"
            )

        return [(deployment, location) for deployment, location, _score in filtered_results]

    def _query_for_workspace_updates(
        self,
        user_code_launcher: DagsterCloudUserCodeLauncher,
        upload_all: bool,
    ):
        locations_with_ttl_to_query = self._get_locations_with_ttl_to_query(user_code_launcher)

        deployments_to_query = {key[0] for key in locations_with_ttl_to_query}

        if locations_with_ttl_to_query:
            locations_str = ", ".join(
                f"{deployment}:{location}" for deployment, location in locations_with_ttl_to_query
            )
            self._logger.debug(f"Querying for the following locations with TTL: {locations_str}")

        # If you have specified a non-branch deployment and no TTL, always consider it
        if self._instance.deployment_names:
            deployments_to_query = deployments_to_query.union(set(self._instance.deployment_names))

        # Create mapping of
        # - location name => deployment metadata
        deployment_map: Dict[Tuple[str, str], UserCodeLauncherEntry] = {}
        all_locations: Set[Tuple[str, str]] = set()

        self._active_deployments = set()

        if deployments_to_query or self._instance.include_all_serverless_deployments:
            result = self._instance.organization_scoped_graphql_client().execute(
                WORKSPACE_ENTRIES_QUERY,
                variable_values={
                    "deploymentNames": list(deployments_to_query),
                    "includeAllServerlessDeployments": self._instance.include_all_serverless_deployments,
                    "agentQueues": self._instance.agent_queues_config.queues,
                },
            )

            for deployment_result in result["data"]["deployments"]:
                deployment_name = deployment_result["deploymentName"]
                is_branch_deployment = deployment_result["isBranchDeployment"]

                self._active_deployments.add((deployment_name, is_branch_deployment))

                entries = deployment_result["workspaceEntries"]

                for entry in entries:
                    location_name = entry["locationName"]

                    location_key = (deployment_name, location_name)

                    all_locations.add(location_key)
                    code_location_deploy_data = deserialize_value(
                        entry["serializedDeploymentMetadata"], CodeLocationDeployData
                    )

                    # The GraphQL can return a mix of deployments with TTLs (for example, all
                    # branch deployments, and also full deployments if you have that configured)
                    # and locations without TTLs (for example, full deployments). Deployments
                    # without TTLs always include all of their locations. Deployments with TTLs
                    # only include the locations within locations_with_ttl_to_query.
                    if not self._has_ttl(
                        user_code_launcher, is_branch_deployment
                    ) or location_key in cast(Set[Tuple[str, str]], locations_with_ttl_to_query):
                        deployment_map[location_key] = UserCodeLauncherEntry(
                            code_location_deploy_data=code_location_deploy_data,
                            update_timestamp=float(entry["metadataTimestamp"]),
                        )

        if len(deployment_map):
            update_str = ", ".join(
                f"{deployment}:{location}" for deployment, location in deployment_map.keys()
            )
            self._logger.debug(f"Reconciling with the following locations: {update_str}")
        else:
            self._logger.debug("Reconciling with no locations")

        if upload_all:
            user_code_launcher.add_upload_metadata(deployment_map)
        else:
            user_code_launcher.update_grpc_metadata(deployment_map)

        # Tell run worker monitoring which deployments it should care about
        user_code_launcher.update_run_worker_monitoring_deployments(self._active_deployment_names)

        # In the rare event that there are pending requests that are no longer in the workspace at
        # all (if, say, a location is removed while requests are enqueued), they should be forcibly
        # moved to ready so that they don't stay pending forever - callsites will get an error
        # about the location not existing, but that's preferable to slowly timing out
        pending_requests_copy = self._pending_requests.copy()
        self._pending_requests = []
        for json_request in pending_requests_copy:
            location_name = self._get_location_from_request(json_request)
            deployment_name = json_request["deploymentName"]
            if (deployment_name, location_name) not in all_locations:
                self._ready_requests.append(json_request)
            else:
                self._pending_requests.append(json_request)

    def _get_grpc_client(
        self,
        user_code_launcher: DagsterCloudUserCodeLauncher,
        deployment_name: str,
        location_name: str,
    ) -> DagsterGrpcClient:
        endpoint = user_code_launcher.get_grpc_endpoint(deployment_name, location_name)
        return endpoint.create_client()

    def _get_location_origin_from_request(
        self,
        request: DagsterCloudApiRequest,
    ) -> Optional[CodeLocationOrigin]:
        """Derive the location from the specific argument passed in to a dagster_cloud_api call."""
        api_name = request.request_api
        if api_name in {
            DagsterCloudApi.GET_EXTERNAL_EXECUTION_PLAN,
            DagsterCloudApi.GET_SUBSET_EXTERNAL_PIPELINE_RESULT,
        }:
            external_pipeline_origin = request.request_args.job_origin
            return external_pipeline_origin.repository_origin.code_location_origin
        elif api_name in {
            DagsterCloudApi.GET_EXTERNAL_PARTITION_CONFIG,
            DagsterCloudApi.GET_EXTERNAL_PARTITION_TAGS,
            DagsterCloudApi.GET_EXTERNAL_PARTITION_NAMES,
            DagsterCloudApi.GET_EXTERNAL_PARTITION_SET_EXECUTION_PARAM_DATA,
            DagsterCloudApi.GET_EXTERNAL_SCHEDULE_EXECUTION_DATA,
            DagsterCloudApi.GET_EXTERNAL_SENSOR_EXECUTION_DATA,
        }:
            return request.request_args.repository_origin.code_location_origin
        elif api_name == DagsterCloudApi.GET_EXTERNAL_NOTEBOOK_DATA:
            return request.request_args.code_location_origin
        elif api_name == DagsterCloudApi.PING_LOCATION:
            return RegisteredCodeLocationOrigin(request.request_args.location_name)
        else:
            return None

    @cached_method
    def _get_user_code_instance_ref(self, deployment_name: str):
        return get_instance_ref_for_user_code(self._instance.ref_for_deployment(deployment_name))

    def _handle_api_request(
        self,
        request: DagsterCloudApiRequest,
        deployment_name: str,
        is_branch_deployment: bool,
        user_code_launcher: DagsterCloudUserCodeLauncher,
    ) -> Union[DagsterCloudApiSuccess, DagsterCloudApiGrpcResponse]:
        api_name = request.request_api

        code_location_origin = self._get_location_origin_from_request(request)
        location_name = code_location_origin.location_name if code_location_origin else None

        if api_name == DagsterCloudApi.PING_LOCATION:
            # Do nothing - this request only exists to bump TTL for the location
            return DagsterCloudApiSuccess()
        elif api_name == DagsterCloudApi.CHECK_FOR_WORKSPACE_UPDATES:
            # Dagster Cloud has requested that we upload new metadata for any out of date locations in
            # the workspace
            self._upload_outdated_workspace_entries(
                deployment_name, is_branch_deployment, user_code_launcher
            )
            return DagsterCloudApiSuccess()
        elif api_name == DagsterCloudApi.GET_EXTERNAL_EXECUTION_PLAN:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )
            serialized_snapshot_or_error = client.execution_plan_snapshot(
                execution_plan_snapshot_args=request.request_args._replace(
                    instance_ref=self._get_user_code_instance_ref(deployment_name)
                )
            )
            return DagsterCloudApiGrpcResponse(
                serialized_response_or_error=serialized_snapshot_or_error
            )

        elif api_name == DagsterCloudApi.GET_SUBSET_EXTERNAL_PIPELINE_RESULT:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )

            serialized_subset_result_or_error = client.external_pipeline_subset(
                pipeline_subset_snapshot_args=request.request_args
            )

            return DagsterCloudApiGrpcResponse(
                serialized_response_or_error=serialized_subset_result_or_error
            )
        elif api_name == DagsterCloudApi.GET_EXTERNAL_PARTITION_CONFIG:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )
            serialized_partition_config_or_error = client.external_partition_config(
                partition_args=request.request_args,
            )
            return DagsterCloudApiGrpcResponse(
                serialized_response_or_error=serialized_partition_config_or_error
            )
        elif api_name == DagsterCloudApi.GET_EXTERNAL_PARTITION_TAGS:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )
            serialized_partition_tags_or_error = client.external_partition_tags(
                partition_args=request.request_args,
            )
            return DagsterCloudApiGrpcResponse(
                serialized_response_or_error=serialized_partition_tags_or_error
            )
        elif api_name == DagsterCloudApi.GET_EXTERNAL_PARTITION_NAMES:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )
            serialized_partition_names_or_error = client.external_partition_names(
                partition_names_args=request.request_args,
            )
            return DagsterCloudApiGrpcResponse(
                serialized_response_or_error=serialized_partition_names_or_error
            )
        elif api_name == DagsterCloudApi.GET_EXTERNAL_PARTITION_SET_EXECUTION_PARAM_DATA:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )
            serialized_partition_execution_params_or_error = (
                client.external_partition_set_execution_params(
                    partition_set_execution_param_args=request.request_args
                )
            )
            return DagsterCloudApiGrpcResponse(
                serialized_response_or_error=serialized_partition_execution_params_or_error
            )
        elif api_name == DagsterCloudApi.GET_EXTERNAL_SCHEDULE_EXECUTION_DATA:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )

            args = request.request_args._replace(
                instance_ref=self._get_user_code_instance_ref(deployment_name)
            )

            schedule_attributes = {
                "schedule": args.schedule_name,
                "repository": args.repository_origin.repository_name,
                "location": args.repository_origin.code_location_origin.location_name,
                "deployment": deployment_name,
            }

            with observe_execution(
                opentelemetry=self._instance.opentelemetry,
                event_key=f"{DAGSTER_CLOUD_AGENT_METRIC_PREFIX}.schedule.evaluation",
                short_description="schedule evaluation requests",
                attributes=schedule_attributes,
                result_evaluator_callback=inspect_schedule_result,
            ) as observer:
                serialized_schedule_data_or_error = client.external_schedule_execution(
                    external_schedule_execution_args=args,
                )
                observer.evaluate_result(
                    serialized_data_or_error=serialized_schedule_data_or_error,
                )
                return DagsterCloudApiGrpcResponse(
                    serialized_response_or_error=serialized_schedule_data_or_error
                )

        elif api_name == DagsterCloudApi.GET_EXTERNAL_SENSOR_EXECUTION_DATA:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )

            args = request.request_args._replace(
                instance_ref=self._get_user_code_instance_ref(deployment_name)
            )

            sensor_attributes = {
                "sensor": args.sensor_name,
                "repository": args.repository_origin.repository_name,
                "location": args.repository_origin.code_location_origin.location_name,
                "deployment": deployment_name,
            }

            with observe_execution(
                opentelemetry=self._instance.opentelemetry,
                event_key=f"{DAGSTER_CLOUD_AGENT_METRIC_PREFIX}.sensor.evaluation",
                short_description="sensor evaluation requests",
                attributes=sensor_attributes,
                result_evaluator_callback=inspect_sensor_result,
            ) as observer:
                serialized_sensor_data_or_error = client.external_sensor_execution(
                    sensor_execution_args=args,
                )
                observer.evaluate_result(serialized_sensor_data_or_error)
                return DagsterCloudApiGrpcResponse(
                    serialized_response_or_error=serialized_sensor_data_or_error
                )
        elif api_name == DagsterCloudApi.GET_EXTERNAL_NOTEBOOK_DATA:
            client = self._get_grpc_client(
                user_code_launcher, deployment_name, cast(str, location_name)
            )
            response = client.external_notebook_data(request.request_args.notebook_path)
            return DagsterCloudApiGrpcResponse(serialized_response_or_error=response.decode())
        elif api_name == DagsterCloudApi.LAUNCH_RUN:
            run = request.request_args.dagster_run

            with DagsterInstance.from_ref(
                self._get_user_code_instance_ref(deployment_name)
            ) as scoped_instance:
                scoped_instance.report_engine_event(
                    f"{self._instance.agent_display_name} is launching run {run.run_id}",
                    run,
                    cls=self.__class__,
                )

                scoped_instance.add_run_tags(
                    run.run_id,
                    merge_dicts(
                        (
                            {"dagster/agent_label": self._instance.dagster_cloud_api_agent_label}
                            if self._instance.dagster_cloud_api_agent_label
                            else {}
                        ),
                        {"dagster/agent_id": self._instance.instance_uuid},
                    ),
                )

                run_attributes = extract_run_attributes(deployment_name, run)
                with observe_execution(
                    opentelemetry=self._instance.opentelemetry,
                    event_key=f"{DAGSTER_CLOUD_AGENT_METRIC_PREFIX}.run.launches",
                    short_description="run execution requests",
                    attributes=run_attributes,
                ) as observer:
                    launcher = scoped_instance.get_run_launcher_for_run(run)  # type: ignore  # (instance subclass)

                    if is_isolated_run(run):
                        launcher.launch_run(LaunchRunContext(dagster_run=run, workspace=None))
                    else:
                        scoped_instance.report_engine_event(
                            f"Launching {run.run_id} without an isolated run environment.",
                            run,
                            cls=self.__class__,
                        )

                        run_location_name = cast(
                            str,
                            run.remote_job_origin.repository_origin.code_location_origin.location_name,
                        )

                        server = user_code_launcher.get_grpc_server(
                            deployment_name, run_location_name
                        )

                        # Record the server handle that we launched it on to for run monitoring
                        scoped_instance.add_run_tags(
                            run.run_id, new_tags={SERVER_HANDLE_TAG: str(server.server_handle)}
                        )

                        launcher.launch_run_from_grpc_client(
                            scoped_instance, run, server.server_endpoint.create_client()
                        )
                    observer.evaluate_result(run=run)

                return DagsterCloudApiSuccess()
        elif api_name == DagsterCloudApi.TERMINATE_RUN:
            # With isolated agents enabled:
            # Run workers now poll for run status. We don't use the run launcher to terminate.
            # Once min agent version is bumped, we can deprecate this command.
            # For backcompat, we use the run launcher to terminate unless the user opts in.
            run = request.request_args.dagster_run

            with DagsterInstance.from_ref(
                self._instance.ref_for_deployment(deployment_name)
            ) as scoped_instance:
                if self._instance.is_using_isolated_agents:
                    scoped_instance.report_engine_event(
                        f"{self._instance.agent_display_name} received request to mark run as canceling",
                        run,
                        cls=self.__class__,
                    )
                    scoped_instance.report_run_canceling(run)
                else:
                    scoped_instance.report_engine_event(
                        f"{self._instance.agent_display_name} received request to terminate run",
                        run,
                        cls=self.__class__,
                    )
                    if is_isolated_run(run):
                        launcher = scoped_instance.get_run_launcher_for_run(run)  # type: ignore  # (instance subclass)
                        launcher.terminate(run.run_id)
                    else:
                        run_location_name = cast(
                            str,
                            run.remote_job_origin.repository_origin.code_location_origin.location_name,
                        )

                        server = user_code_launcher.get_grpc_server(
                            deployment_name, run_location_name
                        )
                        client = server.server_endpoint.create_client()

                        scoped_instance.report_run_canceling(run)
                        client.cancel_execution(CancelExecutionRequest(run_id=run.run_id))

            return DagsterCloudApiSuccess()

        else:
            check.assert_never(api_name)
            raise Exception(f"Unexpected dagster cloud api call {api_name}")

    def _process_api_request(
        self,
        json_request: Dict,
        user_code_launcher: DagsterCloudUserCodeLauncher,
        submitted_to_executor_timestamp: float,
    ) -> Optional[SerializableErrorInfo]:
        thread_start_run_timestamp = get_current_timestamp()
        api_result: Optional[DagsterCloudApiResponse] = None
        error_info: Optional[SerializableErrorInfo] = None

        request_id = json_request["requestId"]
        request_api = json_request["requestApi"]
        request_body = json_request["requestBody"]
        deployment_name = json_request["deploymentName"]
        is_branch_deployment = json_request["isBranchDeployment"]

        request: Union[str, DagsterCloudApiRequest] = DagsterCloudApiRequest.format_request(
            request_id, request_api
        )

        if request_api not in DagsterCloudApi.__members__:
            api_result = DagsterCloudApiUnknownCommandResponse(request_api=request_api)
            self._logger.warning(
                f"Ignoring request {json_request}: Unknown command. This is likely due to running an "
                "older version of the agent."
            )
        else:
            try:
                request = deserialize_value(request_body, DagsterCloudApiRequest)
                self._logger.info(f"Received request {request}.")
                api_result = self._handle_api_request(
                    request, deployment_name, is_branch_deployment, user_code_launcher
                )
            except Exception:
                error_info = serializable_error_info_from_exc_info(sys.exc_info())
                api_result = DagsterCloudApiErrorResponse(error_infos=[error_info])

                self._logger.error(f"Error serving request {json_request}: {error_info}")

        self._logger.info(f"Finished processing request {request}.")

        thread_finished_request_time = get_current_timestamp()
        thread_telemetry = DagsterCloudApiThreadTelemetry(
            submitted_to_executor_timestamp=submitted_to_executor_timestamp,
            thread_start_run_timestamp=thread_start_run_timestamp,
            thread_end_handle_api_request_timestamp=thread_finished_request_time,
        )

        api_result = api_result.with_thread_telemetry(thread_telemetry)

        assert api_result

        upload_response = DagsterCloudUploadApiResponse(
            request_id=request_id,
            request_api=request_api,
            response=api_result,
        )

        self._logger.info(f"Uploading response for request {request}.")

        if UPLOAD_API_RESPONSE_BATCHING_ENABLED():
            self._batcher[deployment_name].submit((deployment_name, upload_response))
        else:
            upload_api_response(self._instance, deployment_name, upload_response)

        self._logger.info(f"Finished uploading response for request {request}.")

        return error_info

    def _batch_upload_api_response(
        self, upload_response_batch: List[Tuple[str, DagsterCloudUploadApiResponse]]
    ) -> List[None]:
        deployment_names = set(deployment_name for deployment_name, _ in upload_response_batch)
        assert len(deployment_names) == 1
        batch_upload_api_response(
            self._instance,
            next(iter(deployment_names)),
            [resp for _, resp in upload_response_batch],
        )
        return [None for _ in upload_response_batch]

    def _get_location_from_request(self, json_request: Dict[str, Any]) -> Optional[str]:
        request_api = json_request["requestApi"]
        request_body = json_request["requestBody"]
        if request_api not in DagsterCloudApi.__members__:
            return None

        request = deserialize_value(request_body, DagsterCloudApiRequest)
        location_origin = self._get_location_origin_from_request(request)
        if not location_origin:
            return None

        return location_origin.location_name

    def run_iteration(
        self, user_code_launcher: DagsterCloudUserCodeLauncher
    ) -> Iterator[Optional[SerializableErrorInfo]]:
        if not user_code_launcher.ready_to_serve_requests:
            return

        num_pending_requests = len(self._pending_requests)

        if num_pending_requests < self._pending_requests_limit:
            # limit (implicit default 10) applied separately for requests to branch deployments, and for each full deployment
            result = self._instance.organization_scoped_graphql_client().execute(
                GET_USER_CLOUD_REQUESTS_QUERY,
                {
                    "forBranchDeployments": self._instance.includes_branch_deployments,
                    "forFullDeployments": self._active_full_deployment_names,
                    "agentQueues": self._instance.agent_queues_config.queues,
                },
            )
            json_requests = result["data"]["userCloudAgent"]["popUserCloudAgentRequests"]

            self._logger.debug(
                f"Iteration #{self._iteration}: Adding {len(json_requests)} requests to be"
                f" processed. Currently {num_pending_requests} waiting for server to be ready"
            )
            self._pending_requests.extend(json_requests)

        else:
            self._logger.warning(
                f"Iteration #{self._iteration}: Waiting to pull requests from the queue since there are"
                f" already {len(self._pending_requests)} in the queue"
            )

        invalid_requests = []
        self._locations_with_pending_requests = set()

        # Determine which pending requests are now ready (their locations have been loaded, or the
        # request does not correspond to a particular location)
        for json_request in self._pending_requests:
            deployment_name = json_request["deploymentName"]
            is_branch_deployment = json_request["isBranchDeployment"]
            location_name = self._get_location_from_request(json_request)
            if location_name:
                self._location_query_times[
                    (deployment_name, location_name, is_branch_deployment)
                ] = time.time()

                if not user_code_launcher.has_grpc_endpoint(deployment_name, location_name):
                    # Next completed periodic workspace update will make the location up to date
                    # - keep this in the queue until then
                    invalid_requests.append(json_request)
                    self._locations_with_pending_requests.add(
                        (deployment_name, location_name, is_branch_deployment)
                    )
                    continue

            self._ready_requests.append(json_request)

        # Any invalid requests go back in the pending queue - the next workspace update will
        # ensure that the usercodelauncher spins up locations for those requests
        self._pending_requests = invalid_requests

        # send all ready requests to the threadpool
        for json_request in self._ready_requests:
            request_id = json_request["requestId"]
            submitted_to_executor_timestamp = get_current_timestamp()
            future = self._executor.submit(
                self._process_api_request,
                json_request,
                user_code_launcher,
                submitted_to_executor_timestamp,
            )

            self._request_ids_to_futures[request_id] = future

        self._ready_requests = []

        # Process futures that are done
        # Create a shallow copy of the futures dict to modify it while iterating
        for request_id, future in self._request_ids_to_futures.copy().items():
            if future.done():
                response: Optional[SerializableErrorInfo] = None

                try:
                    response = future.result(timeout=0)
                except:
                    response = serializable_error_info_from_exc_info(sys.exc_info())

                # Do not process a request again once we have its result
                del self._request_ids_to_futures[request_id]

                # Yield the error information from the future
                if response:
                    yield response

        if self._instance.user_code_launcher.agent_metrics_enabled and (
            self._utilization_metrics["container_utilization"]["measurement_timestamp"] is None
            or (
                get_current_timestamp()
                - self._utilization_metrics["container_utilization"]["measurement_timestamp"]
                > AGENT_UTILIZATION_METRICS_INTERVAL_SECONDS
            )
        ):
            self._update_utilization_metrics(user_code_launcher)
        self._iteration += 1

        yield None


def batch_upload_api_response(
    instance: DagsterCloudAgentInstance,
    deployment_name: str,
    batch: List[DagsterCloudUploadApiResponse],
):
    with compressed_namedtuple_upload_file(BatchDagsterCloudUploadApiResponse(batch=batch)) as f:
        resp = instance.requests_managed_retries_session.put(
            instance.dagster_cloud_upload_api_response_url,
            headers=instance.headers_for_deployment(deployment_name),
            files={"api_response_batch.tmp": f},
            timeout=instance.dagster_cloud_api_timeout,
            proxies=instance.dagster_cloud_api_proxies,
        )
        raise_http_error(resp)


def upload_api_response(
    instance: DagsterCloudAgentInstance,
    deployment_name: str,
    upload_response: DagsterCloudUploadApiResponse,
):
    with compressed_namedtuple_upload_file(upload_response) as f:
        resp = instance.requests_managed_retries_session.put(
            instance.dagster_cloud_upload_api_response_url,
            headers=instance.headers_for_deployment(deployment_name),
            files={"api_response.tmp": f},
            timeout=instance.dagster_cloud_api_timeout,
            proxies=instance.dagster_cloud_api_proxies,
        )
        raise_http_error(resp)
