codeplain_REST_api.py

import socket
import time
from typing import Optional

import requests
from requests.adapters import HTTPAdapter
from requests.exceptions import ConnectionError, RequestException, Timeout
from urllib3.connection import HTTPConnection

import plain2code_exceptions
from plain2code_console import RETRY_COLOR
from plain2code_state import RunState

MAX_RETRIES = 4
RETRY_DELAY = 3

# A report the render does not consume must not be able to hold the process open
# at exit, and requests waits indefinitely by default.
REPORT_TIMEOUT_SECONDS = 5

# A render call can leave the connection with nothing on it for several minutes while the
# service works on the answer. Network equipment between the two, such as the address
# translation a CI runner sits behind, drops a connection that quiet and tells neither end.
# The answer then has nowhere to go and the client waits for it for as long as it is allowed
# to run. These probes put a packet on the connection often enough to keep it, and they end
# it when the other end has really gone.
KEEPALIVE_IDLE_SECONDS = 60
KEEPALIVE_INTERVAL_SECONDS = 30
KEEPALIVE_FAILED_PROBES = 4


def keepalive_socket_options():
    """The socket options that keep a long, quiet request alive, for the platforms that have them."""
    options = list(HTTPConnection.default_socket_options)
    options.append((socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1))
    # The option that sets the idle time is named differently per platform, and Windows has
    # none of the three. Whatever is missing is left out; the keepalive itself still applies.
    for idle_option in ("TCP_KEEPIDLE", "TCP_KEEPALIVE"):
        if hasattr(socket, idle_option):
            options.append((socket.IPPROTO_TCP, getattr(socket, idle_option), KEEPALIVE_IDLE_SECONDS))
            break
    if hasattr(socket, "TCP_KEEPINTVL"):
        options.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, KEEPALIVE_INTERVAL_SECONDS))
    if hasattr(socket, "TCP_KEEPCNT"):
        options.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, KEEPALIVE_FAILED_PROBES))
    return options


class KeepaliveAdapter(HTTPAdapter):
    """Opens every connection with keepalive on, whether it goes direct or through a proxy."""

    def init_poolmanager(self, *args, **kwargs):
        kwargs["socket_options"] = keepalive_socket_options()
        return super().init_poolmanager(*args, **kwargs)

    def proxy_manager_for(self, proxy, **proxy_kwargs):
        # A request that goes through a proxy is served from a separate pool, which the
        # options above never reach. A render behind a proxy is as quiet as any other and
        # needs the same keepalive.
        proxy_kwargs["socket_options"] = keepalive_socket_options()
        return super().proxy_manager_for(proxy, **proxy_kwargs)


def new_session():
    """A requests session whose connections carry keepalive."""
    session = requests.Session()
    session.mount("https://", KeepaliveAdapter())
    session.mount("http://", KeepaliveAdapter())
    return session


RETRY_ERROR_CODES = [
    "LLMInternalError",
]

# Mapping from API error codes to exception classes
ERROR_CODE_EXCEPTIONS = {
    "FunctionalRequirementTooComplex": plain2code_exceptions.FunctionalRequirementTooComplex,
    "ConflictingRequirements": plain2code_exceptions.ConflictingRequirements,
    "RenderingCreditBalanceTooLow": plain2code_exceptions.RenderingCreditBalanceTooLow,
    "LLMInternalError": plain2code_exceptions.LLMInternalError,
    "MissingResource": plain2code_exceptions.MissingResource,
    "PlainSyntaxError": plain2code_exceptions.PlainSyntaxError,
    "InternalServerError": plain2code_exceptions.InternalServerError,
}


class CodeplainAPI:

    def __init__(self, api_key, console):
        self.api_key = api_key
        self.console = console
        self.session = new_session()

    @property
    def api_url(self):
        return self._api_url

    @api_url.setter
    def api_url(self, value):
        self._api_url = value

    def _extend_payload_with_run_state(self, payload: dict, run_state: RunState):
        run_state.increment_call_count()
        payload["render_state"] = run_state.to_dict()

    def _handle_retry_logic(
        self, attempt: int, retry_delay: int, num_retries: int, error: Exception, silent: bool
    ) -> int:
        """
        Handles retry logic with exponential backoff.

        Args:
            attempt: Current attempt number
            retry_delay: Current retry delay in seconds
            error: The exception that occurred
            num_retries: Number of retries allowed
            error_type: Type of error for logging (e.g., "Network error", "Error")

        Returns:
            Updated retry_delay for next attempt

        Raises:
            The original exception if max retries exceeded
        """
        is_connection_error = isinstance(error, ConnectionError) or isinstance(error, Timeout)
        connection_error_type = "Network error" if is_connection_error else "Error"
        if attempt < num_retries:
            if not silent:
                self.console.debug(
                    f"↻ {connection_error_type} on attempt {attempt + 1}/{num_retries + 1}: {error}. "
                    f"Retrying in {retry_delay} seconds...",
                    color=RETRY_COLOR,
                )
            time.sleep(retry_delay)
            # Exponential backoff
            return retry_delay * 2
        else:
            if not silent:
                self.console.error(f"Max retries ({num_retries}) exceeded. Last error: {error}")
            if is_connection_error:
                raise plain2code_exceptions.NetworkConnectionError(
                    "Connection error: Failed to connect to API server.\n\nPlease check that your internet connection is working."
                )
            if isinstance(error, RequestException):
                raise Exception(f"Error rendering plain code: {error}\n") from error
            raise error

    def _raise_for_error_code(self, response_json):
        """Raise appropriate exception based on error code in response."""
        error_code = response_json.get("error_code")
        if error_code not in ERROR_CODE_EXCEPTIONS:
            return

        message = response_json.get("message", "")

        if error_code == "FunctionalRequirementTooComplex":
            raise plain2code_exceptions.FunctionalRequirementTooComplex(
                message, response_json.get("proposed_breakdown")
            )
        if error_code == "MissingResource":
            raise plain2code_exceptions.MissingResource(f"Missing resource: {message}\n")
        if error_code == "ConflictingRequirements":
            raise plain2code_exceptions.ConflictingRequirements(f"Conflicting requirements: {message}\n")
        if error_code == "RenderingCreditBalanceTooLow":
            raise plain2code_exceptions.RenderingCreditBalanceTooLow(message)
        if error_code == "LLMInternalError":
            raise plain2code_exceptions.LLMInternalError(f"LLM internal error: {message}\n")
        if error_code == "PlainSyntaxError":
            raise plain2code_exceptions.plain_syntax_error(message or "Unknown error")
        if error_code == "InternalServerError":
            raise plain2code_exceptions.InternalServerError(
                "Internal server error.\n\n"
                "Please report the error to support@codeplain.ai with the attached .log file."
            )
        exception_class = ERROR_CODE_EXCEPTIONS[error_code]
        raise exception_class(message)

    def post_request(
        self,
        endpoint_url,
        headers,
        payload,
        run_state: Optional[RunState],
        num_retries: int = MAX_RETRIES,
        silent: bool = False,
        timeout: Optional[int] = None,
    ):
        if run_state is not None:
            self._extend_payload_with_run_state(payload, run_state)

        retry_delay = RETRY_DELAY
        response_json = None

        for attempt in range(num_retries + 1):
            try:
                response = self.session.post(endpoint_url, headers=headers, json=payload, timeout=timeout)

                try:
                    response_json = response.json()
                except requests.exceptions.JSONDecodeError as e:
                    self.console.debug(f"Failed to decode JSON response: {e}. Response text: {response.text}")
                    raise Exception(f"Error rendering plain code: Failed to decode API response ({e}).\n") from e

                if response.status_code == requests.codes.bad_request and "error_code" in response_json:
                    self._raise_for_error_code(response_json)

                response.raise_for_status()
                return response_json

            except Exception as e:
                # For other errors, check if they should be retried
                if response_json is not None and "error_code" in response_json:
                    if response_json["error_code"] not in RETRY_ERROR_CODES:
                        raise e

                retry_delay = self._handle_retry_logic(attempt, retry_delay, num_retries, e, silent)

    def connection_check(self, client_version):
        endpoint_url = f"{self.api_url}/connection_check"
        headers = {"Content-Type": "application/json"}
        payload = {
            "api_key": self.api_key,
            "client_version": client_version,
        }
        return self.post_request(endpoint_url, headers, payload, None, num_retries=0, silent=True)

    def status(self):
        endpoint_url = f"{self.api_url}/status"
        headers = {"Content-Type": "application/json"}
        payload = {
            "api_key": self.api_key,
        }
        return self.post_request(endpoint_url, headers, payload, None, num_retries=0, silent=True)

    def render_finished(self, outcome: str, run_state: RunState):
        """Report the terminal state a render reached.

        Called once per render, whatever the outcome. The render state machine
        lives here, so the API cannot tell a finished render from one still
        going; it has to be told.

        No retries and no console output: a report has no result the render
        needs, and must not add minutes to a render that has already ended nor
        print retry noise over the exit summary.
        """
        endpoint_url = f"{self.api_url}/render_finished"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}
        payload = {
            "outcome": outcome,
        }

        return self.post_request(
            endpoint_url, headers, payload, run_state, num_retries=0, silent=True, timeout=REPORT_TIMEOUT_SECONDS
        )

    def render_functional_requirement(
        self,
        frid: str,
        plain_source_tree: dict,
        linked_resources: dict,
        existing_files_content: dict,
        memory_files_content: dict,
        module_name: str,
        required_modules: dict,
        include_unittests: bool,
        run_state: RunState,
    ) -> dict[str, str]:
        """
        Renders the content of a functionality based on the provided ID,
        corresponding sections from a Plain document, and existing files' content.

        Args:
            frid (str): The unique identifier for the functionality to be rendered.
            plain_source_tree (dict): A dictionary containing the plain source tree.
            linked_resources (dict): A dictionary where the keys represent filenames of linked
                                        resources and the values are dictionaries containing
                                        the content of the linked resources and the sections
                                        they are linked to.
            existing_files_content (dict): A dictionary where the keys represent filenames
                                        and the values are the content of those files.
            memory_files_content (dict): A dictionary where the keys represent memory filenames
                                        and the values are the content of those files.
            module_name (str): The name of the module to render the functionality for.
            required_modules (dict): A dictionary where the keys represent module names
                                     and the values are lists of functionalities implemented in those modules.
            run_state (RunState): The current state of the rendering process.
        Returns:
            dict[str, str]: A dictionary where the keys are filenames and the values
                            are the rendered code for those files.

        Raises:
            ValueError: If the frid is invalid or the necessary sections cannot be found.
        """
        endpoint_url = f"{self.api_url}/render_functional_requirement"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "plain_source_tree": plain_source_tree,
            "linked_resources": linked_resources,
            "existing_files_content": existing_files_content,
            "memory_files_content": memory_files_content,
            "module_name": module_name,
            "required_modules": required_modules,
            "include_unittests": include_unittests,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def distill_conformance_test_memory(
        self,
        frid,
        plain_source_tree,
        linked_resources,
        existing_files_content,
        memory_files_content,
        module_name,
        required_modules,
        fix_journals,
        run_state: RunState,
    ):
        endpoint_url = f"{self.api_url}/distill_conformance_test_memory"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "plain_source_tree": plain_source_tree,
            "linked_resources": linked_resources,
            "existing_files_content": existing_files_content,
            "memory_files_content": memory_files_content,
            "module_name": module_name,
            "required_modules": required_modules,
            "fix_journals": fix_journals,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def refactor_source_files_if_needed(
        self, frid, module_name: str, files_to_check, existing_files_content, run_state: RunState
    ):
        endpoint_url = f"{self.api_url}/refactor_source_files_if_needed"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "module_name": module_name,
            "files_to_check": list(files_to_check),
            "existing_files_content": existing_files_content,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def render_conformance_tests(
        self,
        frid,
        functional_requirement_id,
        plain_source_tree,
        linked_resources,
        existing_files_content,
        memory_files_content,
        module_name: str,
        required_modules,
        conformance_tests_folder_name,
        conformance_tests_json,
        all_acceptance_tests,
        run_state: RunState,
    ):
        endpoint_url = f"{self.api_url}/render_conformance_tests"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "functional_requirement_id": functional_requirement_id,
            "plain_source_tree": plain_source_tree,
            "linked_resources": linked_resources,
            "existing_files_content": existing_files_content,
            "memory_files_content": memory_files_content,
            "module_name": module_name,
            "required_modules": required_modules,
            "conformance_tests_folder_name": conformance_tests_folder_name,
            "conformance_tests_json": conformance_tests_json,
            "all_acceptance_tests": all_acceptance_tests,
        }

        response = self.post_request(endpoint_url, headers, payload, run_state)
        return response["patched_response_files"], response["conformance_tests_plan_summary_string"]

    def generate_folder_name_from_functional_requirement(
        self,
        frid,
        module_name: str,
        functional_requirement,
        existing_folder_names,
        run_state: RunState,
    ):
        endpoint_url = f"{self.api_url}/generate_folder_name_from_functional_requirement"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "module_name": module_name,
            "functional_requirement": functional_requirement,
            "existing_folder_names": existing_folder_names,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def fix_conformance_tests_issue(
        self,
        frid,
        functional_requirement_id,
        plain_source_tree,
        linked_resources,
        existing_files_content,
        memory_files_content,
        module_name: str,
        conformance_tests_module_name: str,
        required_modules,
        code_diff,
        conformance_tests_files,
        acceptance_tests,
        conformance_tests_issue,
        implementation_fix_count,
        conformance_tests_folder_name,
        current_testing_frid_high_level_implementation_plan: Optional[str],
        conflicting_requirements_count: int,
        fix_attempts_journal: list,
        run_state: RunState,
    ):
        endpoint_url = f"{self.api_url}/fix_conformance_tests_issue"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "functional_requirement_id": functional_requirement_id,
            "plain_source_tree": plain_source_tree,
            "linked_resources": linked_resources,
            "existing_files_content": existing_files_content,
            "memory_files_content": memory_files_content,
            "module_name": module_name,
            "conformance_tests_module_name": conformance_tests_module_name,
            "required_modules": required_modules,
            "code_diff": code_diff,
            "conformance_tests_files": conformance_tests_files,
            "conformance_tests_issue": conformance_tests_issue,
            "implementation_fix_count": implementation_fix_count,
            "conformance_tests_folder_name": conformance_tests_folder_name,
            "current_testing_frid_high_level_implementation_plan": current_testing_frid_high_level_implementation_plan,
            "conflicting_requirements_count": conflicting_requirements_count,
            "fix_attempts_journal": fix_attempts_journal,
        }

        if acceptance_tests is not None:
            payload["acceptance_tests"] = acceptance_tests

        return self.post_request(endpoint_url, headers, payload, run_state)

    def render_acceptance_tests(
        self,
        frid,
        plain_source_tree,
        linked_resources,
        existing_files_content,
        memory_files_content,
        conformance_tests_files,
        module_name: str,
        required_modules,
        acceptance_test,
        run_state: RunState,
    ):
        """
        Renders acceptance tests based on the provided parameters.

        Args:
            frid (str): The unique identifier for the functionality.
            plain_source_tree (dict): A dictionary containing the plain source tree.
            linked_resources (dict): A dictionary where the keys represent resource names
                                    and the values are the content of those resources.
            existing_files_content (dict): A dictionary where the keys represent code base
                                        filenames and the values are the content of those files.
            memory_files_content (dict): A dictionary where the keys represent memory filenames
                                        and the values are the content of those files.
            conformance_tests_files (dict): A dictionary containing conformance test files.
            required_modules (dict): A dictionary where the keys represent module names
                                     and the values are lists of functionalities implemented in those modules.
            acceptance_test (dict): A dictionary containing acceptance test information.

        Returns:
            dict: The rendered acceptance tests.

        Raises:
            Exception: If the request fails or returns an error.
        """
        endpoint_url = f"{self.api_url}/render_acceptance_tests"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "plain_source_tree": plain_source_tree,
            "linked_resources": linked_resources,
            "existing_files_content": existing_files_content,
            "memory_files_content": memory_files_content,
            "conformance_tests_files": conformance_tests_files,
            "module_name": module_name,
            "required_modules": required_modules,
            "acceptance_test": acceptance_test,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def analyze_rendering(
        self,
        frid,
        plain_source_tree,
        linked_resources,
        existing_files_content,
        module_name: str,
        required_modules,
        implementation_code_diff,
        fixed_implementation_code_diff,
        run_state: RunState,
    ):
        endpoint_url = f"{self.api_url}/analyze_rendering"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "plain_source_tree": plain_source_tree,
            "linked_resources": linked_resources,
            "existing_files_content": existing_files_content,
            "module_name": module_name,
            "required_modules": required_modules,
            "implementation_code_diff": implementation_code_diff,
            "fixed_implementation_code_diff": fixed_implementation_code_diff,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def finish_functional_requirement(self, frid, module_name: str, run_state: RunState):
        endpoint_url = f"{self.api_url}/finish_functional_requirement"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "module_name": module_name,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def fail_functional_requirement(self, frid, module_name: str, run_state: RunState):
        endpoint_url = f"{self.api_url}/fail_functional_requirement"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}

        payload = {
            "frid": frid,
            "module_name": module_name,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def summarize_finished_conformance_tests(
        self,
        frid,
        plain_source_tree,
        linked_resources,
        conformance_test_files_content,
        module_name: str,
        required_modules,
        run_state: RunState,
    ):
        endpoint_url = f"{self.api_url}/summarize_finished_conformance_tests"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}
        payload = {
            "frid": frid,
            "plain_source_tree": plain_source_tree,
            "linked_resources": linked_resources,
            "conformance_test_files_content": conformance_test_files_content,
            "module_name": module_name,
            "required_modules": required_modules,
        }

        return self.post_request(endpoint_url, headers, payload, run_state)

    def agent_start(self, task_type: str, task_params: dict, frid: str, module_name: str, run_state: RunState):
        """Start a server-side agent session; returns the first turn (tool calls or completion)."""
        endpoint_url = f"{self.api_url}/agent/start"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}
        payload = {
            "task_type": task_type,
            "task_params": task_params,
            "frid": frid,
            "module_name": module_name,
        }
        return self.post_request(endpoint_url, headers, payload, run_state)

    def agent_continue(
        self, session_id: str, tool_results: list[dict], frid: str, module_name: str, run_state: RunState
    ):
        """Feed tool results into an agent session and run its next turn."""
        endpoint_url = f"{self.api_url}/agent/continue"
        headers = {"X-API-Key": self.api_key, "Content-Type": "application/json"}
        payload = {
            "session_id": session_id,
            "tool_results": tool_results,
            "frid": frid,
            "module_name": module_name,
        }
        return self.post_request(endpoint_url, headers, payload, run_state)