writer 0.8.3rc21__py3-none-any.whl → 1.25.1rc1__py3-none-any.whl

writer/serve.py

@@ -10,29 +10,44 @@ import os
 import os.path
 import pathlib
 import socket
+import tempfile
 import textwrap
 import time
 import typing
-from contextlib import asynccontextmanager
+from contextlib import asynccontextmanager, suppress
 from importlib.machinery import ModuleSpec
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union, cast
+from typing import (
+    Any,
+    AsyncGenerator,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 from urllib.parse import urlsplit
 
+import orjson
 import uvicorn
-from fastapi import FastAPI, HTTPException, Request, Response
-from fastapi.responses import FileResponse, JSONResponse
+from fastapi import FastAPI, File, HTTPException, Request, Response, UploadFile
+from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from fastapi.routing import Mount
 from fastapi.staticfiles import StaticFiles
 from pydantic import ValidationError
 from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 
-from writer import VERSION, abstract, crypto
+from writer import VERSION, abstract
 from writer.ai import Graph
 from writer.app_runner import AppRunner
 from writer.ss_types import (
     AppProcessServerResponse,
     AutogenRequestBody,
     ComponentUpdateRequestPayload,
+    DeleteDataRequestBody,
     EventResponsePayload,
     HashRequestPayload,
     HashRequestResponsePayload,
@@ -41,6 +56,8 @@ from writer.ss_types import (
     InitResponseBodyRun,
     InitSessionRequestPayload,
     InitSessionResponsePayload,
+    RetrieveDataRequestBody,
+    RetrieveDataResponseBody,
     ServeMode,
     StateEnquiryResponsePayload,
     WriterEvent,
@@ -52,104 +69,13 @@ if typing.TYPE_CHECKING:
     from .auth import Auth, Unauthorized
 
 MAX_WEBSOCKET_MESSAGE_SIZE = 201 * 1024 * 1024
-logging.getLogger().setLevel(logging.INFO)
-
-
-class JobVault:
-    SCHEMES: List[str] = []
-    job_vault_implementations: List[Type["JobVault"]] = []
-
-    def __init__(self):
-        self.counter = 0
-        self.vault = {}
-
-    def generate_job_id(self):
-        self.counter += 1
-        return str(self.counter)
-
-    def set(self, job_id: str, value: Any):
-        self.vault[job_id] = value
-
-    def get(self, job_id: str):
-        return self.vault.get(job_id)
-
-    @classmethod
-    def register(cls, klass: Type["JobVault"]):
-        cls.job_vault_implementations.insert(0, klass)
-
-    @classmethod
-    def _get_matching_implementation(cls, connection_string):
-        for job_vault_implementation in cls.job_vault_implementations:
-            for scheme in job_vault_implementation.SCHEMES:
-                if connection_string.startswith(scheme):
-                    return job_vault_implementation
-
-    @classmethod
-    def create_vault(cls):
-        connection_string = os.getenv("WRITER_PERSISTENT_STORE")
-        if not connection_string:
-            return cls()
-
-        matching_implementation = cls._get_matching_implementation(connection_string)
-        if not matching_implementation:
-            supported_schemes = [
-                scheme
-                for implementation in JobVault.job_vault_implementations
-                for scheme in implementation.SCHEMES
-            ]
-            supported_schemes_msg = ", ".join(supported_schemes)
-            logging.error(f"No matching implementation found for { connection_string }. Falling back to in-memory JobVault. \
-                          Supported schemes: {supported_schemes_msg}.")
-            return cls()
-
-        try:
-            return matching_implementation()
-        except Exception as e:
-            logging.error(
-                f"There was an error connecting to { connection_string }. Falling back to in-memory JobVault. {repr(e)}"
-            )
-            return cls()
-
-
-class RedisJobVault(JobVault):
-    SCHEMES = ["redis://", "rediss://", "redis-socket://", "redis-sentinel://"]
-    DEFAULT_TTL = 86400
-
-    def __init__(self):
-        import redis  # type: ignore
-
-        super().__init__()
-        redis_connection_string = os.getenv("WRITER_PERSISTENT_STORE")
-        self.redis_client = redis.from_url(
-            redis_connection_string, decode_responses=True, socket_timeout=30
-        )
-        self.counter_key = "job_counter"
-        if not self.redis_client.exists(self.counter_key):
-            self.redis_client.set(self.counter_key, 0)
-
-    def generate_job_id(self):
-        job_id = self.redis_client.incr(self.counter_key)
-        return str(job_id)
-
-    def set(self, job_id: str, value: Any):
-        ttl = RedisJobVault.DEFAULT_TTL
-        env_ttl = os.getenv("WRITER_PERSISTENT_STORE_TTL")
-        if env_ttl is not None:
-            ttl = int(env_ttl)
-        json_str = json.dumps(value)
-        self.redis_client.set(f"job:{job_id}", json_str, ex=ttl)
-
-    def get(self, job_id: str):
-        json_str = self.redis_client.get(f"job:{job_id}")
-        if not json_str:
-            return None
-        return json.loads(json_str)
+BLUEPRINT_API_EXECUTION_TIMEOUT_SECONDS = int(os.getenv("AGENT_BUILDER_BLUEPRINT_API_EXECUTION_TIMEOUT", "600"))
+BLUEPRINT_API_RETRY_TIMEOUT = int(os.getenv("AGENT_BUILDER_BLUEPRINT_API_RETRY_TIMEOUT", "10000"))
 
 
 class WriterState(typing.Protocol):
     app_runner: AppRunner
     writer_app: bool
-    job_vault: JobVault
     is_server_static_mounted: bool
     meta: Union[Dict[str, Any], Callable[[], Dict[str, Any]]]  # meta tags for SEO
     opengraph_tags: Union[
@@ -175,8 +101,7 @@ def get_asgi_app(
     enable_remote_edit: bool = False,
     enable_server_setup: bool = True,
     on_load: Optional[Callable] = None,
-    on_shutdown: Optional[Callable] = None,
-    enable_jobs_api: bool = False,
+    on_shutdown: Optional[Callable] = None
 ) -> WriterFastAPI:
     """
     Builds an ASGI server that can be injected into another ASGI application
@@ -199,6 +124,7 @@ def get_asgi_app(
 
     _fix_mimetype()
     app_runner = AppRunner(user_app_path, serve_mode)
+    pending_tasks: Set[asyncio.Task] = set()
 
     @asynccontextmanager
     async def lifespan(asgi_app: FastAPI):
@@ -219,6 +145,13 @@ def get_asgi_app(
         except asyncio.CancelledError:
             pass
 
+        for pending_task in pending_tasks.copy():
+            pending_task.cancel()
+            try:
+                await pending_task
+            except asyncio.CancelledError:
+                pass
+
         app_runner.shut_down()
         if on_shutdown is not None:
             on_shutdown()
@@ -291,8 +224,65 @@ def get_asgi_app(
 
     @app.get("/api/health")
     async def health():
+        app_runner = app.state.app_runner
+        
+        # Check user app process
+        if app_runner.app_process is None or not app_runner.app_process.is_alive():
+            return JSONResponse(
+                status_code=503,
+                content={"status": "error", "message": "User app process is not running"}
+            )
+        
+        # Check project saver process (only in edit mode)
+        if app_runner.mode == "edit":
+            project_saver = app_runner.wf_project_context.write_files_async_process
+            if project_saver is None or not project_saver.is_alive():
+                return JSONResponse(
+                    status_code=503,
+                    content={"status": "error", "message": "Project saver process is not running"}
+                )
+        
         return {"status": "ok"}
 
+    @app.get("/api/export")
+    async def export_zip():
+        if serve_mode != "edit":
+            raise HTTPException(status_code=403, detail="Invalid mode.")
+        exported_zip_stream = app_runner.export_zip()
+        return StreamingResponse(
+            exported_zip_stream,
+            media_type="application/x-zip-compressed",
+            headers={
+                "Content-Disposition": "attachment; filename=exported_agent.zip"
+            }
+        )
+
+    @app.post("/api/import")
+    async def import_zip(file: UploadFile = File(...)):
+        if serve_mode != "edit":
+            raise HTTPException(status_code=403, detail="Invalid mode.")
+        if not file.filename or not file.filename.endswith(".zip"):
+            raise HTTPException(status_code=400, detail="Only .zip files are supported.")
+
+        MAX_FILE_SIZE = 200 * 1024 * 1024
+
+        try:
+            with tempfile.NamedTemporaryFile(delete=False) as tmp:
+                # Stream file to disk to avoid memory issues
+                size = 0
+                while chunk := await file.read(8192):
+                    size += len(chunk)
+                    if size > MAX_FILE_SIZE:
+                        tmp.close()
+                        os.unlink(tmp.name)
+                        raise HTTPException(status_code=413, detail=f"File too large. Max file size: {MAX_FILE_SIZE}")
+                    tmp.write(chunk)
+                tmp_path = tmp.name
+            await app_runner.import_zip(tmp_path)
+            os.remove(tmp_path)
+        except ValueError:
+            raise HTTPException(status_code=400, detail="Invalid upload.")
+
     @app.post("/api/autogen")
     async def autogen(requestBody: AutogenRequestBody, request: Request):
         import writer.autogen
@@ -303,6 +293,38 @@ def get_asgi_app(
             agent_token_header
             )
 
+    @app.post("/api/data/retrieve")
+    async def retrieve_data(requestBody: RetrieveDataRequestBody) -> RetrieveDataResponseBody:
+        from writer.keyvalue_storage import writer_kv_storage
+
+        all_keys = writer_kv_storage.get_data_keys()
+
+        keys_to_fetch = []
+        for key in all_keys:
+            if key in requestBody.skip_keys:
+                continue
+            if requestBody.key_contains and requestBody.key_contains not in key:
+                continue
+            keys_to_fetch.append(key)
+        
+        async def fetch_value(key: str):
+            return key, await asyncio.to_thread(writer_kv_storage.get, key, "data")
+
+        kv_pairs = await asyncio.gather(*(fetch_value(key) for key in keys_to_fetch))
+
+        return RetrieveDataResponseBody(result={k: v["data"] for k, v in kv_pairs})
+
+    @app.post("/api/data/delete")
+    async def delete_data(requestBody: DeleteDataRequestBody) -> None:
+        from writer.keyvalue_storage import writer_kv_storage
+
+        async def delete_key(key: str):
+            return key, await asyncio.to_thread(writer_kv_storage.delete, key)
+
+        await asyncio.gather(*(delete_key(key) for key in requestBody.keys))
+
+        return None
+
     @app.post("/api/init")
     async def init(
         initBody: InitRequestBody, request: Request, response: Response
@@ -362,99 +384,284 @@ def get_asgi_app(
         except json.JSONDecodeError:
             raise HTTPException(status_code=400, detail="Cannot parse the payload.")
         return payload
+    
+    def has_api_trigger(app_runner: AppRunner, blueprint_id: str) -> bool:
+        # Check if blueprint has at least one API trigger component
+        if not app_runner.bmc_components:
+            return False
+        return any(
+            comp["type"] == "blueprints_apitrigger" and comp.get("parentId") == blueprint_id
+            for comp in app_runner.bmc_components.values()
+        )
+
+    @app.get("/private/api/blueprints")
+    async def get_blueprints(request: Request):
+        """
+        Returns a list of blueprints available in the agent.
+        """
+        if not app_runner.bmc_components:
+            return JSONResponse(content=[])
+
+        blueprints = [
+            {
+                "id": comp["id"],
+                "key": comp.get("content", {}).get("key")
+            }
+            for comp in app_runner.bmc_components.values()
+            if comp["type"] == "blueprints_blueprint"
+            and has_api_trigger(app_runner, comp["id"])
+        ]
 
-    @app.post("/api/job/blueprint/{blueprint_key}")
-    async def create_blueprint_job(blueprint_key: str, request: Request, response: Response):
-        if not enable_jobs_api:
-            raise HTTPException(status_code=404)
+        return JSONResponse(content=blueprints)
 
-        crypto.verify_message_authorization_signature(f"create_job_{blueprint_key}", request)
+    @app.get("/private/api/cron-triggers")
+    async def get_cron_triggers(request: Request):
+        """
+        Returns a list of Cron Trigger blocks.
+        """
+        if not app_runner.bmc_components:
+            return JSONResponse(content=[], status_code=200)
+        
+        definition = abstract.templates["blueprints_crontrigger"].writer
+
+        cron_triggers = [
+            {
+                "id": comp.get("id"),
+                "blueprint_id": comp.get("parentId"),
+                "name": comp.get("content", {}).get("alias") or definition["name"],
+                "cron_expression": comp.get("content", {}).get("cronExpression", ""),
+                "timezone": comp.get("content", {}).get("timezone", "UTC"),
+            }
+            for comp in app_runner.bmc_components.values()
+            if comp.get("type") == "blueprints_crontrigger"
+        ]
 
-        def serialize_result(data):
+        return JSONResponse(content=cron_triggers, status_code=200)
+
+    @app.post("/private/api/blueprint/{blueprint_id}")
+    async def create_blueprint_job(blueprint_id: str, request: Request, response: Response, branch_id: Optional[str] = None):
+        # Keep-alive interval for SSE streaming
+        KEEPALIVE_INTERVAL = 15
+        payload = await _get_payload_as_json(request)
+
+        # --- Session initialization ---
+
+        async def init_session_and_validate(
+            app_runner: AppRunner,
+            cookies: Dict[str, Any],
+            headers: Dict[str, Any],
+        ) -> str:
+            # Initialize session with passed cookies/headers
+            sess_resp = await app_runner.init_session(InitSessionRequestPayload(
+                cookies=cookies, headers=headers, proposedSessionId=None
+            ))
+            if not sess_resp or not sess_resp.payload:
+                raise RuntimeError("Cannot initialize session.")
+            sid = sess_resp.payload.sessionId
+            if not await app_runner.check_session(sid):
+                raise RuntimeError("Cannot initialize session.")
+            return sid
+
+        # --- Blueprint discovery logic ---
+
+        def check_blueprint(app_runner: AppRunner, blueprint_id: str) -> bool:
+            # Locate blueprint component by its key
+            if not app_runner.bmc_components:
+                return False
+            return blueprint_id in app_runner.bmc_components
+
+        # --- Result serialization (recursive) ---
+
+        def serialize_result(data: Any) -> Any:
+            # Convert blueprint output into JSON-serializable structure
+            if isinstance(data, (str, int, float, bool, type(None))):
+                return data
             if isinstance(data, list):
                 return [serialize_result(item) for item in data]
             if isinstance(data, dict):
                 return {k: serialize_result(v) for k, v in data.items()}
-            if isinstance(data, (str, int, float, bool, type(None))):
-                return data
             try:
                 return json.loads(json.dumps(data))
             except (TypeError, OverflowError):
-                return f"Can't be displayed. Value of type: {str(type(data))}."
+                return f"Can't be displayed. Value of type: {type(data)}."
+
+        # --- SSE formatting utilities ---
+
+        async def format_event(event_type: str, data: Dict[str, Any]) -> str:
+            # Format a proper Server-Sent Event chunk
+            return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
+
+        async def format_keepalive() -> str:
+            # Send a SSE comment line as keep-alive (spec compliant)
+            return ": keep-alive\n\n"
 
-        def update_job(job_id: str, job_info: dict):
-            current_job_info = app.state.job_vault.get(job_id)
-            if not current_job_info:
-                raise RuntimeError("Job not found.")
-            merged_info = current_job_info | {"finished_at": int(time.time())} | job_info
-            app.state.job_vault.set(job_id, merged_info)
+        # --- The main worker logic that produces events ---
 
-        def job_done_callback(task: asyncio.Task, job_id: str):
+        async def event_logic(queue: asyncio.Queue):
             try:
-                apsr: Optional[AppProcessServerResponse] = task.result()
-                if apsr is None or apsr.status != "ok":
-                    update_job(job_id, {"status": "error"})
+                await queue.put(await format_event("status", {"status": "in progress", "created_at": int(time.time())}))
+                await queue.put(await format_event("status", {"status": "initializing", "msg": "Initializing session..."}))
+
+                # Validate session & credentials
+                session_id = await init_session_and_validate(
+                    app_runner, dict(request.cookies), dict(request.headers)
+                )
+
+                await queue.put(await format_event("status", {"status": "validating", "msg": "Validating blueprint..."}))
+
+                if not app_runner.bmc_components:
+                    raise RuntimeError("No blueprints defined in the agent.")
+
+                blueprint_exists = check_blueprint(app_runner, blueprint_id)
+                if not blueprint_exists:
+                    await queue.put(await format_event("error", {
+                        "msg": f"Blueprint '{blueprint_id}' was not found.",
+                        "finished_at": int(time.time())
+                    }))
                     return
-                result = None
-                if apsr.payload and apsr.payload.result:
-                    result = apsr.payload.result.get("result")
-                update_job(job_id, {"status": "complete", "result": serialize_result(result)})
-            except Exception as e:
-                update_job(job_id, {"status": "error"})
-                raise e
 
-        app_response = await app_runner.init_session(
-            InitSessionRequestPayload(
-                cookies=dict(request.cookies), headers=dict(request.headers), proposedSessionId=None
-            )
-        )
+                if not branch_id and not has_api_trigger(app_runner, blueprint_id):
+                    await queue.put(await format_event("error", {
+                        "msg": f"Blueprint '{blueprint_id}' lacks an API trigger.",
+                        "finished_at": int(time.time())
+                    }))
+                    return
 
-        if not app_response or not app_response.payload:
-            raise HTTPException(status_code=500, detail="Cannot initialize session.")
-        session_id = app_response.payload.sessionId
-        is_session_ok = await app_runner.check_session(session_id)
-        if not is_session_ok:
-            raise HTTPException(status_code=500, detail="Cannot initialize session.")
+                if branch_id:
+                    block = app_runner.bmc_components.get(branch_id)
+                    if not block:
+                        await queue.put(await format_event("error", {
+                            "msg": f"Block '{branch_id}' was not found.",
+                            "finished_at": int(time.time())
+                        }))
+                        return
+
+                await queue.put(await format_event("status", {"status": "executing", "msg": (f"Executing branch: {branch_id}..." if branch_id else f"Executing blueprint: {blueprint_id}...")}))
+
+                if branch_id:
+                    task = asyncio.create_task(
+                        app_runner.handle_event(
+                            session_id,
+                            WriterEvent(
+                                type="wf-run-blueprint-via-api",
+                                isSafe=True,
+                                handler="run_blueprint_via_api",
+                                payload={
+                                    "blueprint_id": blueprint_id,
+                                    "trigger_type": "Cron",
+                                    "branch_id": branch_id,
+                                    **(payload or {})
+                                },
+                            )
+                        )
+                    )
+                else:
+                    task = asyncio.create_task(
+                        app_runner.handle_event(
+                            session_id,
+                            WriterEvent(
+                                type="wf-run-blueprint-via-api",
+                                isSafe=True,
+                                handler="run_blueprint_via_api",
+                                payload={
+                                    "blueprint_id": blueprint_id,
+                                    "trigger_type": "API",
+                                    **(payload or {})
+                                },
+                            )
+                        )
+                    )
 
-        loop = asyncio.get_running_loop()
-        task = loop.create_task(
-            app_runner.handle_event(
-                session_id,
-                WriterEvent(
-                    type="wf-builtin-run",
-                    isSafe=True,
-                    handler=f"$runBlueprint_{blueprint_key}",
-                    payload=await _get_payload_as_json(request),
-                ),
-            )
-        )
+                await queue.put(await format_event("status", {"status": "running", "msg": ("Branch is running. Awaiting output..." if branch_id else "Blueprint is running. Awaiting output...")}))
 
-        job_id = app.state.job_vault.generate_job_id()
-        app.state.job_vault.set(
-            job_id, {"id": job_id, "status": "in progress", "created_at": int(time.time())}
-        )
-        task.add_done_callback(lambda t: job_done_callback(t, job_id))
-        return {"id": job_id, "token": crypto.get_hash(f"get_job_{job_id}")}
+                # Await blueprint execution with timeout protection
+                apsr = await asyncio.wait_for(task, timeout=BLUEPRINT_API_EXECUTION_TIMEOUT_SECONDS)
 
-    @app.get("/api/job/{job_id}")
-    async def get_blueprint_job(job_id: str, request: Request, response: Response):
-        if not enable_jobs_api:
-            raise HTTPException(status_code=404)
+                await queue.put(await format_event("status", {"status": "processing", "msg": "Processing blueprint result..."}))
 
-        crypto.verify_message_authorization_signature(f"get_job_{job_id}", request)
-        job = app.state.job_vault.get(job_id)
+                if not apsr or apsr.status != "ok":
+                    raise RuntimeError("Blueprint execution failed.")
 
-        if not job:
-            return JSONResponse(status_code=404, content={"id": job_id, "status": "not found"})
+                if apsr.payload and apsr.payload.result:
+                    task_status = apsr.payload.result.get("ok", False)
+                    result = serialize_result(
+                        apsr.payload.result.get("result")
+                    )
+                else:
+                    task_status = False
+                    result = "No result returned from blueprint execution."
+
+                if not task_status:
+                    await queue.put(await format_event("error", {
+                        "msg": result,
+                        "finished_at": int(time.time())
+                    }))
+                else:
+                    await queue.put(await format_event("artifact", {
+                        "artifact": result,
+                        "finished_at": int(time.time())
+                    }))
 
-        status_code = 200
-        if job.get("status") == "error":
-            status_code = 400
+            except Exception as e:
+                # Bubble up any unexpected error as 'error' SSE event
+                await queue.put(await format_event("error", {
+                    "msg": f"Agent Builder internal error: {str(e)}",
+                    "finished_at": int(time.time())
+                }))
+            finally:
+                # Always mark stream completion for consumer
+                await queue.put("data: [DONE]\n\n")
 
-        return JSONResponse(status_code=status_code, content=job)
+        # --- The streaming loop that multiplexes events and keep-alives ---
+
+        async def merged_stream() -> AsyncGenerator[str, None]:
+            # Type annotation required by mypy
+            queue: asyncio.Queue = asyncio.Queue()
+            producer_task = asyncio.create_task(event_logic(queue))
+
+            yield f"retry: {BLUEPRINT_API_RETRY_TIMEOUT}\n\n"
+
+            try:
+                while True:
+                    try:
+                        result = await asyncio.wait_for(
+                            queue.get(),
+                            timeout=KEEPALIVE_INTERVAL
+                            )
+                        if result == "data: [DONE]\n\n":
+                            return
+                        yield result
+                    except asyncio.TimeoutError:
+                        yield await format_keepalive()
+                    except asyncio.CancelledError:
+                        # Client disconnected, break streaming loop
+                        break
+            finally:
+                # Always cancel producer to prevent orphaned task
+                producer_task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await producer_task
+
+        return StreamingResponse(
+            merged_stream(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "Access-Control-Allow-Origin": "*",
+                "Access-Control-Allow-Headers": "Cache-Control",
+            },
+        )
 
     # Streaming
 
+    async def _send_json_or_queue(session_id: str, data: Any, websocket: WebSocket):
+        try:
+            binary_data = orjson.dumps(data)
+            await websocket.send_bytes(binary_data)
+        except (RuntimeError, WebSocketDisconnect):
+            await app_runner.queue_message(session_id, data)
+
     async def _stream_session_init(websocket: WebSocket):
         """
         Waits for the client to provide a session id to initialise the stream.
@@ -480,8 +687,6 @@ def get_asgi_app(
         Handles incoming requests from client.
         """
 
-        pending_tasks: Set[asyncio.Task] = set()
-
         try:
             while True:
                 req_message_raw = await websocket.receive_json()
@@ -523,18 +728,9 @@ def get_asgi_app(
                     pending_tasks.add(new_task)
                     new_task.add_done_callback(pending_tasks.discard)
         except WebSocketDisconnect:
-            pass
+            return
         except asyncio.CancelledError:
             raise
-        finally:
-            # Cancel pending tasks
-
-            for pending_task in pending_tasks.copy():
-                pending_task.cancel()
-                try:
-                    await pending_task
-                except asyncio.CancelledError:
-                    pass
 
     async def _handle_incoming_event(
         websocket: WebSocket, session_id: str, req_message: WriterWebsocketIncoming
@@ -564,7 +760,7 @@ def get_asgi_app(
             res_payload = typing.cast(EventResponsePayload, apsr.payload).model_dump()
         if res_payload is not None:
             response.payload = res_payload
-        await websocket.send_json(response.model_dump())
+        await _send_json_or_queue(session_id, response.model_dump(), websocket)
 
     async def _handle_incoming_edit_message(
         websocket: WebSocket, session_id: str, req_message: WriterWebsocketIncoming
@@ -582,11 +778,10 @@ def get_asgi_app(
             await app_runner.queue_announcement_async(
                 "componentUpdate", req_message.payload["components"], session_id
             )
-        elif req_message.type == "collaborationUpdate":
-            pass
-            # await app_runner.queue_announcement_async(
-            #     "collaborationAnnouncement", req_message.payload, exclude_session_id=session_id
-            # )
+        elif req_message.type == "collaborationPing":
+            await app_runner.queue_announcement_async(
+                "collaborationUpdate", req_message.payload, exclude_session_id=session_id
+            )
         elif req_message.type == "codeSaveRequest":
             app_runner.save_code(
                 session_id, req_message.payload["code"], req_message.payload["path"]
@@ -631,8 +826,10 @@ def get_asgi_app(
                 response.payload = {"sourceFiles": app_runner.source_files}
             except Exception as error:
                 response.payload = {"error": str(error)}
+        elif req_message.type == "writerVaultUpdate":
+            await app_runner.writer_vault_refresh(session_id)
 
-        await websocket.send_json(response.model_dump())
+        await _send_json_or_queue(session_id, response.model_dump(), websocket)
 
     async def _handle_keep_alive_message(
         websocket: WebSocket, session_id: str, req_message: WriterWebsocketIncoming
@@ -640,7 +837,7 @@ def get_asgi_app(
         response = WriterWebsocketOutgoing(
             messageType="keepAliveResponse", trackingId=req_message.trackingId, payload=None
         )
-        await websocket.send_json(response.model_dump())
+        await _send_json_or_queue(session_id, response.model_dump(), websocket)
 
     async def _handle_state_enquiry_message(
         websocket: WebSocket, session_id: str, req_message: WriterWebsocketIncoming
@@ -657,7 +854,7 @@ def get_asgi_app(
             res_payload = typing.cast(StateEnquiryResponsePayload, apsr.payload).model_dump()
         if res_payload is not None:
             response.payload = res_payload
-        await websocket.send_json(response.model_dump())
+        await _send_json_or_queue(session_id, response.model_dump(), websocket)
 
     async def _handle_hash_request(
         websocket: WebSocket, session_id: str, req_message: WriterWebsocketIncoming
@@ -673,13 +870,14 @@ def get_asgi_app(
         )
         if apsr is not None and apsr.payload is not None:
             response.payload = typing.cast(HashRequestResponsePayload, apsr.payload).model_dump()
-        await websocket.send_json(response.model_dump())
+        await _send_json_or_queue(session_id, response.model_dump(), websocket)
 
     async def _stream_outgoing_announcements(websocket: WebSocket, session_id: str):
         """
         Handles outgoing communications to the client (announcements).
         """
 
+        WEBSOCKET_CODE_UPDATE_CODE = 4001
         session_queue: asyncio.Queue = asyncio.Queue()
         app_runner.announcement_queues[session_id] = session_queue
 
@@ -689,11 +887,15 @@ def get_asgi_app(
                 announcement = WriterWebsocketOutgoing(
                     messageType="announcement", trackingId=-1, payload=announcement_data
                 )
-                await websocket.send_json(announcement.dict())
+                if websocket.application_state == WebSocketState.CONNECTED:
+                    await websocket.send_json(announcement.dict())
                 if announcement_data.get("type") == "codeUpdate":
+                    await websocket.close(WEBSOCKET_CODE_UPDATE_CODE, "Code update.")
                     return
         except WebSocketDisconnect:
             pass
+        except asyncio.CancelledError:
+            raise
         finally:
             if app_runner.announcement_queues.get(session_id) is None:
                 return
@@ -719,6 +921,14 @@ def get_asgi_app(
         if not is_session_ok:
             await websocket.close(code=1008)  # Invalid permissions
             return
+        
+        try:
+            queued_messages = await app_runner.retrieve_messages(session_id)
+            for message in queued_messages:
+                await websocket.send_json(message)
+            await app_runner.clear_messages(session_id)
+        except (WebSocketDisconnect, RuntimeError):
+            return
 
         task1 = asyncio.create_task(_stream_incoming_requests(websocket, session_id))
         task2 = asyncio.create_task(_stream_outgoing_announcements(websocket, session_id))
@@ -765,14 +975,10 @@ def get_asgi_app(
             )
         )
 
-    JobVault.register(RedisJobVault)
-
     # Return
     if enable_server_setup is True:
         _execute_server_setup_hook(user_app_path)
 
-    app.state.job_vault = JobVault.create_vault()
-
     return app
 
 
@@ -820,8 +1026,7 @@ def serve(
     port: Optional[int],
     host,
     enable_remote_edit=False,
-    enable_server_setup=False,
-    enable_jobs_api=False,
+    enable_server_setup=False
 ):
     """Initialises the web server."""
 
@@ -846,8 +1051,7 @@ def serve(
         mode,
         enable_remote_edit,
         on_load=on_load,
-        enable_server_setup=enable_server_setup,
-        enable_jobs_api=enable_jobs_api,
+        enable_server_setup=enable_server_setup
     )
     log_level = "warning"
     uvicorn.run(