abstra 3.23.3__py3-none-any.whl → 3.23.5__py3-none-any.whl

abstra_internals/controllers/execution/consumer.py

@@ -1,12 +1,10 @@
 from concurrent.futures import ThreadPoolExecutor
-from multiprocessing.connection import Client
 from uuid import uuid4
 
 from abstra_internals.controllers.execution.worker_process import process_main
 from abstra_internals.controllers.main import MainController
 from abstra_internals.environment import (
     EXECUTION_QUEUE_CONCURRENCY,
-    HOST,
     PROCESS_TIMEOUT_SECONDS,
 )
 from abstra_internals.logger import AbstraLogger
@@ -35,36 +33,38 @@ class ConsumerController:
             f"[ConsumerController] Starting loop with {self.concurrency} threads"
         )
 
-        with ThreadPoolExecutor(
+        thread_pool = ThreadPoolExecutor(
             max_workers=self.concurrency,
             thread_name_prefix="ExecutionConsumer",
-        ) as executor:
+        )
+
+        try:
             for msg in self.consumer.iter():
                 AbstraLogger.debug(
                     f"[ConsumerController] Submitting message [{msg.delivery_tag}] for processing"
                 )
-                executor.submit(
+                thread_pool.submit(
                     self.run_subprocess,
                     msg=msg,
                 )
+        finally:
+            AbstraLogger.warning(
+                "[ConsumerController] Consumer main loop exited, waiting for threads to finish"
+            )
+            thread_pool.shutdown(wait=True)
+            AbstraLogger.warning("[ConsumerController] All threads finished")
 
-        AbstraLogger.warning(
-            "[ConsumerController] Consumer main loop exited, waiting for threads to finish"
-        )
-        AbstraLogger.warning("[ConsumerController] All threads finished")
-
-        # If the loop exits and there are running executions, gracefull shutdown has failed
-        self.main_controller.fail_app_executions(
-            app_id=self.app_id,
-            reason="Failed to set status",
-        )
+            # If the loop exits and there are running executions, gracefull shutdown has failed
+            self.main_controller.fail_app_executions(
+                app_id=self.app_id,
+                reason="Failed to set status",
+            )
 
-    def run_subprocess(self, msg: QueueMessage) -> None:
-        connection = None
+    def run_subprocess(self, msg: QueueMessage):
         worker_id = str(uuid4())
-        try:
-            head_id = worker_id.split("-")[0]
+        head_id = worker_id.split("-")[0]
 
+        try:
             mp_context = self.main_controller.repositories.mp_context.get_context()
             stage = self.main_controller.get_stage(msg.preexecution.stage_id)
 
@@ -74,11 +74,13 @@ class ConsumerController:
                 )
                 self.consumer.threadsafe_ack(msg)
                 return
-            connection = Client(
-                (HOST, msg.preexecution.connection_port),
-                family="AF_INET",
-                authkey=msg.preexecution.connection_auth_key.encode(),
-            )
+
+            local_queue = None
+            if isinstance(
+                self.main_controller.repositories.producer, LocalProducerRepository
+            ):
+                local_queue = self.main_controller.repositories.producer.queue
+
             p = mp_context.Process(
                 target=process_main,
                 name=f"Worker-{head_id}",
@@ -89,13 +91,7 @@ class ConsumerController:
                     root_path=Settings.root_path,
                     server_port=Settings.server_port,
                     request=msg.preexecution.context,
-                    connection=connection,
-                    local_queue=self.main_controller.repositories.producer.queue
-                    if isinstance(
-                        self.main_controller.repositories.producer,
-                        LocalProducerRepository,
-                    )
-                    else None,
+                    local_queue=local_queue,
                 ),
             )
 
@@ -105,7 +101,6 @@ class ConsumerController:
             # If timeout is reached, the process is still alive
             if p.is_alive():
                 p.terminate()
-                p.join()  # Wait for termination to complete
                 raise NonCleanExit(
                     f"Run took too long to complete ({PROCESS_TIMEOUT_SECONDS} secs) and was terminated. Please make sure all your requests calls have a timeout set."
                 )
@@ -133,13 +128,6 @@ class ConsumerController:
             )
 
             self.consumer.threadsafe_nack(msg)
-        finally:
-            # Always close the connection when done
-            if connection is not None:
-                try:
-                    connection.close()
-                except Exception:
-                    pass  # Ignore errors when closing
 
             AbstraLogger.warning(
                 f"[ConsumerController] Message [{msg.delivery_tag}] has been negatively acknowledged"

abstra_internals/controllers/execution/execution.py

@@ -4,7 +4,10 @@ import traceback
 from pathlib import Path
 from typing import Literal, Optional, Tuple
 
-from abstra_internals.controllers.execution.execution_client import ExecutionClient
+from abstra_internals.controllers.execution.execution_client import (
+    ExecutionClient,
+    HeadlessClient,
+)
 from abstra_internals.controllers.execution.execution_client_form import ClientAbandoned
 from abstra_internals.controllers.sdk.sdk_context import SDKContext
 from abstra_internals.entities.execution import Execution
@@ -27,12 +30,12 @@ class ExecutionController:
         *,
         repositories: Repositories,
         stage: StageWithFile,
-        client: ExecutionClient,
-        context: ClientContext,
+        client: Optional[ExecutionClient] = None,
+        context: Optional[ClientContext] = None,
     ) -> None:
         self.repositories = repositories
         self.stage = stage
-        self.client = client
+        self.client = client or HeadlessClient()
         self.context = context
 
     def run(self):

abstra_internals/controllers/execution/execution_client.py

@@ -1,7 +1,4 @@
 import abc
-from multiprocessing.connection import Connection
-
-from abstra_internals.entities.execution_context import ClientContext
 
 
 class ExecutionClient(abc.ABC):
@@ -19,17 +16,11 @@ class ExecutionClient(abc.ABC):
 
 
 class HeadlessClient(ExecutionClient):
-    context: ClientContext
-
-    def __init__(self, context: ClientContext, conn: Connection):
-        self.context = context
-        self._conn = conn
-
     def handle_failure(self, e: Exception) -> None:
-        self._conn.close()
+        pass
 
     def handle_success(self) -> None:
-        self._conn.close()
+        pass
 
     def handle_start(self, execution_id: str):
         pass

abstra_internals/controllers/execution/execution_client_form.py

@@ -1,6 +1,7 @@
-from multiprocessing.connection import Connection
 from typing import Dict, Optional
 
+import flask_sock
+
 from abstra_internals.controllers.execution.execution_client import ExecutionClient
 from abstra_internals.entities.execution_context import FormContext
 from abstra_internals.entities.forms.form_entity import ButtonAction, RenderedForm
@@ -17,38 +18,38 @@ class FormClient(ExecutionClient):
 
     def __init__(
         self,
-        conn: Connection,
+        ws: flask_sock.Server,
         context: FormContext,
         production_mode: bool,
     ) -> None:
         self.context = context
         self._production_mode = production_mode
-        self._conn = conn
+        self._ws = ws
 
     # WEBSOCKET
 
     def _receive_message(self):
         try:
-            str_data = self._conn.recv()
+            str_data = self._ws.receive()
             deserialized = deserialize(str_data)
             if not isinstance(deserialized, dict):
                 raise ValueError("Invalid message format")
             return deserialized
-        except (BrokenPipeError, EOFError):
+        except flask_sock.ConnectionClosed:
             raise ClientAbandoned()
 
     def _send(self, msg: forms_contract.Message) -> None:
         str_data = serialize(msg.to_json())
         try:
-            self._conn.send(str_data)
-        except (BrokenPipeError, EOFError):
+            self._ws.send(str_data)
+        except flask_sock.ConnectionClosed:
             pass
 
     def _user_code_send(self, msg: forms_contract.Message) -> None:
         str_data = serialize(msg.to_json())
         try:
-            self._conn.send(str_data)
-        except (BrokenPipeError, EOFError):
+            self._ws.send(str_data)
+        except flask_sock.ConnectionClosed:
             raise ClientAbandoned()
 
     def _wait_for_message(
@@ -161,14 +162,12 @@ class FormClient(ExecutionClient):
         self._send(
             forms_contract.ExecutionEndedMessage(close_dto, self._production_mode)
         )
-        self._conn.close()
 
     def handle_success(self):
         close_dto = forms_contract.CloseDTO(exit_status="SUCCESS")
         self._send(
             forms_contract.ExecutionEndedMessage(close_dto, self._production_mode)
         )
-        self._conn.close()
 
     ## Testing
 

abstra_internals/controllers/execution/execution_client_hook.py

@@ -1,4 +1,3 @@
-from multiprocessing.connection import Connection
 from typing import Dict
 
 from abstra_internals.controllers.execution.execution_client import ExecutionClient
@@ -7,23 +6,18 @@ from abstra_internals.entities.execution_context import HookContext, Request, Re
 
 class HookClient(ExecutionClient):
     context: HookContext
-    conn: Connection
-    response: Response
 
-    def __init__(self, context: HookContext, conn: Connection) -> None:
+    def __init__(self, context: HookContext) -> None:
         self.context = context
-        self.conn = conn
-        self.response = Response(headers={}, status=200, body="")
 
     def handle_failure(self, e: Exception) -> None:
         self.set_response(500, "An exception occurred during execution.", {})
-        self.conn.send(self.response)
 
     def handle_success(self) -> None:
-        self.conn.send(self.response)
+        pass
 
     def set_response(self, status: int, body: str, headers: Dict[str, str]) -> None:
-        self.response = Response(
+        self.context.response = Response(
             status=status,
             body=body,
             headers=headers,

abstra_internals/controllers/execution/execution_stdio.py

@@ -13,10 +13,9 @@ from abstra_internals.env_masker import GLOBAL_MASKER
 from abstra_internals.environment import IS_PRODUCTION
 from abstra_internals.interface.sdk.user_exceptions import ExecutionNotFound
 from abstra_internals.logger import AbstraLogger
-from abstra_internals.utils import serialize
 
 
-class StdioController:
+class BroadcastController:
     listeners: List[flask_sock.Server] = []
     _lock = threading.Lock()
 
@@ -32,22 +31,12 @@ class StdioController:
                 cls.listeners.remove(listener)
 
     @classmethod
-    def broadcast(
-        cls,
-        *,
-        type: Literal["stdout", "stderr"],
-        execution_id: str,
-        stage_id: str,
-        log: str,
-    ):
-        msg = serialize(
-            dict(type=type, log=log, execution_id=execution_id, stage_id=stage_id)
-        )
+    def broadcast(cls, *, msg: str):
         for listener in cls.listeners:
             try:
                 listener.send(msg)
             except Exception:
-                StdioController.unregister(listener)
+                BroadcastController.unregister(listener)
 
     def __init__(
         self,

abstra_internals/controllers/execution/execution_test.py

@@ -1,4 +1,3 @@
-from multiprocessing import Pipe
 from pathlib import Path
 
 from abstra_internals.controllers.execution.execution import ExecutionController
@@ -25,7 +24,8 @@ class ExecutionControllerTest(BaseTest):
                 headers={"auth": "some_secret_token"},
                 query_params={"c": "3"},
                 method="GET",
-            )
+            ),
+            response=Response(headers={}, status=200, body=""),
         )
 
         self.project = self.repositories.project.load(include_disabled_stages=False)
@@ -39,8 +39,7 @@ class ExecutionControllerTest(BaseTest):
         self.project.add_stage(self.stage)
         self.repositories.project.save(self.project)
 
-        self.parent_conn, child_conn = Pipe()
-        self.hook_client = HookClient(self.context, conn=child_conn)
+        self.hook_client = HookClient(self.context)
 
     def test_run_initial_returns_dto(self):
         ExecutionController(
@@ -50,10 +49,7 @@ class ExecutionControllerTest(BaseTest):
             client=self.hook_client,
         ).run()
 
-        response = self.parent_conn.recv()
-        assert isinstance(response, Response)
-
-        if not response:
+        if not self.context.response:
             self.fail("Response was not set")
 
-        self.assertEqual(response.status, 200)
+        self.assertEqual(self.context.response.status, 200)

abstra_internals/controllers/execution/worker_process.py

@@ -1,20 +1,16 @@
 import threading
 import traceback
 from multiprocessing import Queue
-from multiprocessing.connection import Connection
 from typing import Optional
 
 from abstra_internals.controllers.execution.execution import ExecutionController
-from abstra_internals.controllers.execution.execution_client import (
-    ExecutionClient,
-    HeadlessClient,
-)
-from abstra_internals.controllers.execution.execution_client_form import FormClient
-from abstra_internals.controllers.execution.execution_client_hook import HookClient
 from abstra_internals.controllers.main import MainController
 from abstra_internals.entities.execution import ClientContext
-from abstra_internals.entities.execution_context import FormContext, HookContext
-from abstra_internals.environment import IS_PRODUCTION, set_SERVER_UUID, set_WORKER_UUID
+from abstra_internals.environment import (
+    IS_PRODUCTION,
+    set_SERVER_UUID,
+    set_WORKER_UUID,
+)
 from abstra_internals.logger import AbstraLogger
 from abstra_internals.repositories.factory import (
     build_editor_repositories,
@@ -25,24 +21,6 @@ from abstra_internals.settings import Settings
 from abstra_internals.stdio_patcher import StdioPatcher
 
 
-def make_client_from_context(
-    ctx: ClientContext, connection: Connection
-) -> ExecutionClient:
-    if isinstance(ctx, FormContext):
-        return FormClient(
-            context=ctx,
-            conn=connection,
-            production_mode=IS_PRODUCTION,
-        )
-    elif isinstance(ctx, HookContext):
-        return HookClient(
-            context=ctx,
-            conn=connection,
-        )
-    else:
-        return HeadlessClient(context=ctx, conn=connection)
-
-
 def process_main(
     *,
     root_path: str,
@@ -50,8 +28,7 @@ def process_main(
     worker_id: str,
     app_id: str,
     stage: StageWithFile,
-    connection: Connection,
-    request: ClientContext,
+    request: Optional[ClientContext] = None,
     local_queue: Optional[Queue] = None,
 ):
     thread = threading.current_thread()
@@ -66,7 +43,6 @@ def process_main(
             controller = MainController(repositories)
         else:
             AbstraLogger.init("local")
-            assert local_queue is not None, "Local queue is not initialized"
             repositories = build_editor_repositories(local_queue)
             controller = MainController(repositories)
 
@@ -75,15 +51,14 @@ def process_main(
 
         StdioPatcher.apply(controller)
 
-        execution_data = ExecutionController(
+        ExecutionController(
             repositories=controller.repositories,
             stage=stage,
-            client=make_client_from_context(request, connection),
+            client=None,
             context=request,
         ).run()
 
         AbstraLogger.debug("[ConsumerController Subprocess] Finished")
-        return execution_data
     except Exception as e:
         AbstraLogger.error(f"[ConsumerController Subprocess] Error: {e}")
         AbstraLogger.capture_exception(e)

abstra_internals/controllers/main.py

@@ -1,7 +1,6 @@
 import datetime
 import pkgutil
 import webbrowser
-from multiprocessing import Pipe
 from pathlib import Path
 from shutil import move
 from tempfile import mkdtemp, mktemp
@@ -12,7 +11,6 @@ import flask
 from abstra_internals.cloud_api import get_api_key_info, get_project_info
 from abstra_internals.consts.filepaths import TEST_DATA_FILEPATH
 from abstra_internals.controllers.execution.execution import ExecutionController
-from abstra_internals.controllers.execution.execution_client import HeadlessClient
 from abstra_internals.controllers.execution.execution_client_form import FormClient
 from abstra_internals.controllers.execution.execution_client_hook import HookClient
 from abstra_internals.credentials import (
@@ -69,7 +67,6 @@ from abstra_internals.utils.ai import AiWs
 from abstra_internals.utils.diff import compute_updated_code_from_replacements
 from abstra_internals.utils.file import module2path, path2module
 from abstra_internals.utils.validate import validate_json
-from abstra_internals.utils.websockets import bind_ws_with_connection
 
 
 class UnknownNodeTypeError(Exception):
@@ -1684,13 +1681,12 @@ class MainController:
         if not job:
             raise Exception(f"Job with id {id} not found")
 
-        _, child_conn = Pipe()
-        ctx = JobContext()
+        print(f"Running job {job.id} ({job.title})")
+
         return ExecutionController(
             repositories=self.repositories,
             stage=job,
-            context=ctx,
-            client=HeadlessClient(context=ctx, conn=child_conn),
+            context=JobContext(),
         ).run()
 
     def debug_run_hook(self, id: str, request: Request):
@@ -1717,9 +1713,14 @@ class MainController:
 
         context = HookContext(
             request=request,
+            response=Response(
+                body="",
+                headers={},
+                status=200,
+            ),
         )
-        parent_conn, child_conn = Pipe()
-        client = HookClient(context=context, conn=child_conn)
+
+        client = HookClient(context=context)
 
         run_data = ExecutionController(
             repositories=self.repositories,
@@ -1728,13 +1729,13 @@ class MainController:
             context=context,
         ).run()
 
-        response = parent_conn.recv()
-        assert isinstance(response, Response)
+        if context.response is None or client.context.response is None:
+            flask.abort(500)
 
         return {
-            "body": response.body,
-            "status": response.status,
-            "headers": response.headers,
+            "body": client.context.response.body,
+            "status": context.response.status,
+            "headers": context.response.headers,
             **run_data,
         }
 
@@ -1764,14 +1765,10 @@ class MainController:
         if not task_id:
             raise Exception("Task ID is required for script execution")
 
-        _, child_conn = Pipe()
-        ctx = ScriptContext(task_id=task_id)
-
         return ExecutionController(
             repositories=self.repositories,
             stage=script,
-            context=ctx,
-            client=HeadlessClient(context=ctx, conn=child_conn),
+            context=ScriptContext(task_id=task_id),
         ).run()
 
     def debug_run_form_with_ai(self, id, prompt: str, url_params: Dict[str, str] = {}):
@@ -1796,10 +1793,8 @@ class MainController:
         if not form:
             raise Exception(f"Form with id {id} not found")
 
-        parent_conn, child_conn = Pipe()
-        bind_ws_with_connection(ws, parent_conn, block=False)
         client = FormClient(
-            conn=child_conn,
+            ws=ws,  # type: ignore
             context=context,
             production_mode=False,
         )
@@ -1823,13 +1818,10 @@ class MainController:
         stage = self.create_job(title, str(tempfile), (0, 0))
         tempfile.write_text(code)
 
-        _, child_conn = Pipe()
-        ctx = JobContext()
         execution_result = ExecutionController(
             repositories=self.repositories,
             stage=stage,
-            context=ctx,
-            client=HeadlessClient(context=ctx, conn=child_conn),
+            context=JobContext(),
         ).run()
 
         self.delete_job(stage.id, remove_file=True)

abstra_internals/controllers/workflows.py

@@ -225,6 +225,39 @@ class WorkflowController:
         stage.workflow_transitions.append(transition)
         self.repos.project.save(project)
 
+    def delete_transition(self, transition_id: str):
+        """
+        Delete a transition from the workflow by its ID.
+        This method removes the specified transition from the workflow,
+        breaking the connection between two stages.
+
+        Args:
+            transition_id (str): ID of the transition to delete.
+
+        Raises:
+            ValueError: If the transition with the given ID does not exist.
+
+        Copywritings:
+            Delete a transition from the workflow
+            Deleting a transition from the workflow...
+        """
+        project = self.repos.project.load()
+        transition_found = False
+
+        for stage in project.workflow_stages:
+            for i, transition in enumerate(stage.workflow_transitions):
+                if transition.id == transition_id:
+                    stage.workflow_transitions.pop(i)
+                    transition_found = True
+                    break
+            if transition_found:
+                break
+
+        if not transition_found:
+            raise ValueError(f"Transition {transition_id} does not exist.")
+
+        self.repos.project.save(project)
+
     def update_workflow(self, workflow_state_dto: Dict):
         """
         Update the entire workflow configuration with new stages and transitions.

abstra_internals/entities/execution_context.py

@@ -43,6 +43,7 @@ class JobExecutionMock(ExecutionMock):
 
 class HookContext(Serializable):
     request: Request
+    response: Response
     sent_tasks: List[str] = Field(default_factory=list)
     legacy_thread_data: dict = Field(default_factory=dict)
     mock_execution: HookExecutionMock = Field(default_factory=HookExecutionMock)

abstra_internals/entities/execution_test.py

@@ -2,7 +2,7 @@ from datetime import datetime
 from unittest import TestCase
 
 from abstra_internals.entities.execution import Execution
-from abstra_internals.entities.execution_context import HookContext, Request
+from abstra_internals.entities.execution_context import HookContext, Request, Response
 from abstra_internals.repositories.project.project import HookStage
 from abstra_internals.utils import is_serializable
 from abstra_internals.utils.datetime import from_utc_iso_string
@@ -24,7 +24,9 @@ class ExecutionTest(TestCase):
             method="GET",
         )
 
-        context = HookContext(request=request)
+        context = HookContext(
+            request=request, response=Response(headers={}, status=200, body="")
+        )
 
         execution: Execution[HookContext] = Execution.create(
             stage_id=mock_stage.id, context=context
@@ -55,7 +57,8 @@ class ExecutionTest(TestCase):
                 headers={"auth": "secret_token"},
                 query_params={"c": "3"},
                 method="GET",
-            )
+            ),
+            response=Response(headers={}, status=200, body=""),
         )
 
         execution: Execution[HookContext] = Execution.create(

abstra_internals/entities/forms/widgets/file_upload.py

@@ -49,7 +49,7 @@ def download_to_path(url: str) -> pathlib.Path:
     execution = SDKContextStore.get_execution()
     save_dir = get_uploads_dir() / execution.id
     save_dir.mkdir(parents=True, exist_ok=True)
-    save_path = save_dir / url.split("/")[-1]
+    save_path = save_dir / pathlib.Path(url).name
 
     if url.startswith("http://") or url.startswith("https://"):
         with save_path.open("wb") as f, requests.get(url, stream=True) as r: