licos-agent-runtime 0.3.4__tar.gz → 0.3.6__tar.gz

{licos_agent_runtime-0.3.4 → licos_agent_runtime-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: licos-agent-runtime
-Version: 0.3.4
+Version: 0.3.6
 Summary: LICOS Python agent runtime for LangGraph/LangChain projects
 Requires-Python: >=3.10
 Requires-Dist: fastapi<1,>=0.121
@@ -15,7 +15,7 @@ Requires-Dist: langgraph-sdk<1,>=0.3
 Requires-Dist: langgraph<2,>=1.1
 Requires-Dist: langsmith<1,>=0.4
 Requires-Dist: licos-dev-sdk>=0.4.9
-Requires-Dist: licos-platform-sdk>=0.5.7
+Requires-Dist: licos-platform-sdk>=0.5.8
 Requires-Dist: pydantic<3,>=2.12
 Requires-Dist: pyyaml<7,>=6
 Requires-Dist: uvicorn<1,>=0.38

{licos_agent_runtime-0.3.4 → licos_agent_runtime-0.3.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "licos-agent-runtime"
-version = "0.3.4"
+version = "0.3.6"
 description = "LICOS Python agent runtime for LangGraph/LangChain projects"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -21,7 +21,7 @@ dependencies = [
     "langgraph-checkpoint-postgres>=3.0,<4",
     "langsmith>=0.4,<1",
     "httpx>=0.27,<1",
-    "licos-platform-sdk>=0.5.7",
+    "licos-platform-sdk>=0.5.8",
     "licos-dev-sdk>=0.4.9",
     "filelock>=3.16,<4",
     "PyYAML>=6,<7",

{licos_agent_runtime-0.3.4 → licos_agent_runtime-0.3.6}/src/licos_agent_runtime/app.py

@@ -26,6 +26,8 @@ from .service import GraphService
 from .session_files import session_user_id
 from .stream_runner import RunOpt, agent_stream_handler, workflow_stream_handler
 from .trace_records import clear_run_records, get_run_record, list_run_records
+from .workflow_durability import WorkflowInstanceConflict, WorkflowInstanceStore, load_workflow_durability_config
+from .workflow_protocol import workflow_end_event, workflow_log_id, workflow_start_event
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +54,28 @@ def create_app(service: GraphService | None = None) -> FastAPI:
     subscriptions = EventSubscriptions(runtime_service)
     interactive_runs: dict[str, tuple[asyncio.Task, str, float]] = {}
 
+    def durable_store() -> WorkflowInstanceStore:
+        config = load_workflow_durability_config()
+        if config is None or is_agent_project():
+            raise HTTPException(status_code=404, detail="Durable WORKFLOW is not enabled")
+        return WorkflowInstanceStore(config)
+
+    async def prepare_durable(payload: dict[str, Any], ctx: Any) -> Any:
+        config = load_workflow_durability_config() if not is_agent_project() else None
+        if config is None:
+            return None
+        store = WorkflowInstanceStore(config)
+        try:
+            cached = await store.begin(payload, ctx)
+        except WorkflowInstanceConflict as exc:
+            raise HTTPException(status_code=409, detail=str(exc)) from exc
+        except PermissionError as exc:
+            raise HTTPException(status_code=403, detail=str(exc)) from exc
+        except ValueError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+        ctx.metadata["workflow_store"] = store
+        return cached
+
     @asynccontextmanager
     async def lifespan(app: FastAPI):
         try:
@@ -172,6 +196,9 @@ def create_app(service: GraphService | None = None) -> FastAPI:
                     watch_task.cancel()
                     runtime_service.approvals.unwatch(ctx.run_id)
             else:
+                cached = await prepare_durable(payload, ctx)
+                if cached is not None:
+                    return cached
                 result = await runtime_service.run_with_timeout(payload, ctx)
             if isinstance(result, dict) and result.get("status") == "timeout":
                 status_code = 504
@@ -199,6 +226,8 @@ def create_app(service: GraphService | None = None) -> FastAPI:
             error_code = str(exc.code)
             error_message = str(exc)
             raise
+        except HTTPException:
+            raise
         except Exception as exc:
             status_code = 500
             error_response = runtime_service.error_classifier.get_error_response(
@@ -264,6 +293,16 @@ def create_app(service: GraphService | None = None) -> FastAPI:
             workspace_id=ctx.identity.workspace_id,
         )
 
+        cached = await prepare_durable(payload, ctx)
+        if cached is not None:
+            async def cached_events():
+                log_id = workflow_log_id()
+                yield runtime_service._sse_event(workflow_start_event(run_id=ctx.run_id, log_id=log_id))
+                yield runtime_service._sse_event(workflow_end_event(
+                    run_id=ctx.run_id, log_id=log_id, output=cached, time_cost_ms=0,
+                ))
+            return StreamingResponse(cached_events(), media_type="text/event-stream")
+
         if is_agent_project():
             stream_generator = agent_stream_handler(
                 payload=payload,
@@ -300,6 +339,72 @@ def create_app(service: GraphService | None = None) -> FastAPI:
         except PermissionError as exc:
             raise HTTPException(status_code=403, detail=str(exc)) from exc
 
+    @app.get("/workflow/instances/{instance_id}")
+    async def workflow_instance_status(instance_id: str, request: Request) -> dict[str, Any]:
+        ctx = new_context(method="workflow_instance_status", headers=request.headers)
+        request_context.set(ctx)
+        try:
+            row = await durable_store().get(instance_id, ctx)
+            return {
+                "instance_id": instance_id,
+                "status": row["status"],
+                "graph_version": row["graph_version"],
+                "last_run_id": row["last_run_id"],
+                "output": json.loads(row["output_data"]) if row["output_data"] else None,
+                "error": row["error_message"],
+            }
+        except LookupError as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        except PermissionError as exc:
+            raise HTTPException(status_code=403, detail=str(exc)) from exc
+        except ValueError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    @app.post("/workflow/instances/{instance_id}/resume")
+    async def workflow_instance_resume(instance_id: str, request: Request) -> dict[str, Any]:
+        ctx = new_context(method="workflow_resume", headers=request.headers)
+        request_context.set(ctx)
+        store = durable_store()
+        try:
+            await ensure_credit_available(
+                "GRAPH_RUN",
+                request_id=ctx.run_id,
+                current_user_id=ctx.identity.current_user_id,
+                current_user_org_id=ctx.identity.current_user_org_id,
+                owner_user_id=ctx.identity.user_id,
+                owner_org_id=ctx.identity.owner_org_id,
+                workspace_id=ctx.identity.workspace_id,
+            )
+            await store.claim_resume(instance_id, ctx)
+            ctx.metadata["workflow_store"] = store
+            return await runtime_service.run_with_timeout(None, ctx)
+        except WorkflowInstanceConflict as exc:
+            raise HTTPException(status_code=409, detail=str(exc)) from exc
+        except LookupError as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        except PermissionError as exc:
+            raise HTTPException(status_code=403, detail=str(exc)) from exc
+        except ValueError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+
+    @app.post("/workflow/instances/{instance_id}/cancel")
+    async def workflow_instance_cancel(instance_id: str, request: Request) -> dict[str, Any]:
+        ctx = new_context(method="workflow_instance_cancel", headers=request.headers)
+        request_context.set(ctx)
+        try:
+            result = await durable_store().cancel(instance_id, ctx)
+            if result["status"] == "cancelled" and result["run_id"] in runtime_service.running_tasks:
+                runtime_service.cancel_run(result["run_id"], ctx)
+            return result
+        except WorkflowInstanceConflict as exc:
+            raise HTTPException(status_code=409, detail=str(exc)) from exc
+        except LookupError as exc:
+            raise HTTPException(status_code=404, detail=str(exc)) from exc
+        except PermissionError as exc:
+            raise HTTPException(status_code=403, detail=str(exc)) from exc
+        except ValueError as exc:
+            raise HTTPException(status_code=400, detail=str(exc)) from exc
+
     @app.get("/agent/approvals/{run_id}")
     async def agent_approvals(run_id: str, request: Request) -> dict[str, Any]:
         ctx = new_context(method="approval_list", headers=request.headers)

{licos_agent_runtime-0.3.4 → licos_agent_runtime-0.3.6}/src/licos_agent_runtime/run_config.py

@@ -9,7 +9,7 @@ def init_run_config(graph: Any, ctx: Context) -> dict[str, Any]:
     del graph
     return {
         "configurable": {
-            "thread_id": ctx.run_id,
+            "thread_id": ctx.metadata.get("workflow_instance_id", ctx.run_id),
             "run_id": ctx.run_id,
         },
         "metadata": {

{licos_agent_runtime-0.3.4 → licos_agent_runtime-0.3.6}/src/licos_agent_runtime/service.py

@@ -66,6 +66,7 @@ from .workflow_protocol import (
     workflow_start_event,
 )
 from .workflow_canvas import canvas_submit_response, workflow_canvas_response, workflow_node_code_locations
+from .workflow_durability import WorkflowInstanceStore, load_workflow_durability_config
 
 logger = logging.getLogger(__name__)
 
@@ -200,7 +201,8 @@ class GraphService:
         self.agent_module = agent_module
         self.graph_module = graph_module
         self.running_tasks: dict[str, asyncio.Task] = {}
-        self._run_owners: dict[str, str] = {}
+        self._run_owners: dict[str, tuple[str, str | None, str | None]] = {}
+        self._run_contexts: dict[str, Context] = {}
         self.error_classifier = ErrorClassifier()
         self._agent_stream_runner = AgentStreamRunner()
         self._workflow_stream_runner = WorkflowStreamRunner()
@@ -246,6 +248,8 @@ class GraphService:
 
             if self._graph is None:
                 self._graph = get_graph_instance(self.graph_module, ctx)
+            if load_workflow_durability_config() and getattr(self._graph, "checkpointer", None) is None:
+                raise ValueError("Durable WORKFLOW must compile its graph with workflow_checkpointer()")
             return self._graph
 
     @staticmethod
@@ -270,15 +274,20 @@ class GraphService:
         stream_runner = self._get_stream_runner()
         yield from stream_runner.stream(payload, graph, run_config, ctx)
 
-    async def run(self, payload: dict[str, Any], ctx: Context | None = None) -> dict[str, Any]:
+    async def run(self, payload: dict[str, Any] | None, ctx: Context | None = None) -> dict[str, Any]:
         ctx = ctx or new_context("run")
         run_id = ctx.run_id
         agent_project = is_agent_project()
+        workflow_store: WorkflowInstanceStore | None = ctx.metadata.get("workflow_store")
+        heartbeat_task: asyncio.Task[Any] | None = None
         recorder = RunTraceRecorder(ctx, project_kind="agent" if agent_project else "workflow")
         logger.info("Starting run: run_id=%s", run_id)
 
         context_token = request_context.set(ctx)
         try:
+            if workflow_store is not None:
+                await workflow_store.assert_running(ctx)
+                heartbeat_task = asyncio.create_task(workflow_store.heartbeat(ctx, asyncio.current_task()))
             if agent_project:
                 request = payload
                 snapshot = self._agent_snapshot(ctx)
@@ -302,8 +311,9 @@ class GraphService:
                     ) if snapshot and snapshot.harness else await invocation
                     turn.commit(extract_answer(result))
             else:
-                payload = materialize_input_files(payload)
-                recorder.set_root_input(payload)
+                if payload is not None:
+                    payload = materialize_input_files(payload)
+                recorder.set_root_input(payload if payload is not None else {"resume": True})
                 graph = self._get_graph(ctx)
                 self._attach_trace_code_locations(recorder, graph)
                 run_config = append_trace_callback(init_run_config(graph, ctx), recorder)
@@ -314,9 +324,13 @@ class GraphService:
                 result = await enrich_workflow_output_file_urls(result, ctx)
                 if isinstance(result, dict):
                     result.setdefault("run_id", run_id)
+                    if workflow_store is not None:
+                        await workflow_store.update_status(ctx, "completed", output=result)
                     _save_record_safely(recorder, output=result)
                     return result
                 response = {"run_id": run_id, "result": result}
+                if workflow_store is not None:
+                    await workflow_store.update_status(ctx, "completed", output=response)
                 _save_record_safely(recorder, output=response)
                 return response
             if result is None:
@@ -329,11 +343,17 @@ class GraphService:
             _save_record_safely(recorder, output=response)
             return response
         except asyncio.CancelledError:
+            if workflow_store is not None:
+                status = ("cancelled" if ctx.metadata.get("workflow_cancel_requested") else
+                          "failed" if ctx.metadata.get("workflow_timeout") else "interrupted")
+                await workflow_store.update_status(ctx, status, error="Execution timeout" if status == "failed" else None)
             logger.info("Run cancelled: run_id=%s", run_id)
             response = {"status": "cancelled", "run_id": run_id, "message": "Execution was cancelled"}
             _save_record_safely(recorder, status="error", output=response, error="cancelled")
             raise
         except Exception as exc:
+            if workflow_store is not None:
+                await workflow_store.update_status(ctx, "failed", error=str(exc))
             _save_record_safely(recorder, status="error", error=exc)
             err = self.error_classifier.classify(exc, {"node_name": "run", "run_id": run_id})
             logger.error(
@@ -344,8 +364,14 @@ class GraphService:
             )
             raise
         finally:
-            request_context.reset(context_token)
-            self.approvals.forget_run(run_id)
+            try:
+                if heartbeat_task is not None:
+                    heartbeat_task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await heartbeat_task
+            finally:
+                request_context.reset(context_token)
+                self.approvals.forget_run(run_id)
             self.running_tasks.pop(run_id, None)
             telemetry.flush()
 
@@ -358,7 +384,7 @@ class GraphService:
         except Exception as exc:
             logger.warning("Failed to resolve workflow code locations: trace_id=%s error=%s", recorder.trace_id, exc)
 
-    async def run_with_timeout(self, payload: dict[str, Any], ctx: Context) -> dict[str, Any]:
+    async def run_with_timeout(self, payload: dict[str, Any] | None, ctx: Context) -> dict[str, Any]:
         active = self.running_tasks.get(ctx.run_id)
         if active is not None and not active.done():
             raise ValueError(f"run_id is already active: {ctx.run_id}")
@@ -367,25 +393,30 @@ class GraphService:
             snapshot = self._agent_snapshot(ctx)
             if snapshot and snapshot.harness:
                 timeout = min(timeout, snapshot.harness.execution.max_duration_seconds)
-            if ctx.identity.user_token:
-                self._run_owners[ctx.run_id] = session_user_id(ctx)
+        if ctx.identity.user_token:
+            self._run_owners[ctx.run_id] = self._run_owner(ctx)
+        self._run_contexts[ctx.run_id] = ctx
         task = asyncio.create_task(self.run(payload, ctx))
         self.running_tasks[ctx.run_id] = task
         try:
             completed, _ = await asyncio.wait({task}, timeout=timeout)
             if completed:
                 return await task
+            ctx.metadata["workflow_timeout"] = True
             return {
                 "status": "timeout",
                 "run_id": ctx.run_id,
                 "message": f"Execution timeout: exceeded {timeout:g} seconds",
             }
         finally:
-            if not task.done():
-                task.cancel()
-                with suppress(asyncio.CancelledError):
-                    await task
-            self._run_owners.pop(ctx.run_id, None)
+            try:
+                if not task.done():
+                    task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await task
+            finally:
+                self._run_owners.pop(ctx.run_id, None)
+                self._run_contexts.pop(ctx.run_id, None)
 
     async def stream_sse(
         self,
@@ -446,7 +477,7 @@ class GraphService:
         recorder = RunTraceRecorder(ctx, project_kind="agent")
 
         if ctx.identity.user_token:
-            self._run_owners[run_id] = session_user_id(ctx)
+            self._run_owners[run_id] = self._run_owner(ctx)
 
         yield self._sse_event(message_start_event(local_message_id, server_message_id))
 
@@ -598,35 +629,45 @@ class GraphService:
         run_opt: RunOpt,
     ) -> AsyncIterable[str]:
         run_id = ctx.run_id
+        workflow_store: WorkflowInstanceStore | None = ctx.metadata.get("workflow_store")
+        heartbeat_task: asyncio.Task[Any] | None = None
         log_id = workflow_log_id()
         start_perf = time.perf_counter()
         event_index = 0
-        payload = materialize_input_files(payload)
         recorder = RunTraceRecorder(ctx, project_kind="workflow")
-        recorder.set_root_input(payload)
         final_output: Any = {}
         stream_error: BaseException | str | None = None
         status_code = 200
         error_code: str | None = None
-
-        yield self._sse_event(workflow_start_event(run_id=run_id, log_id=log_id), event_index)
-        event_index += 1
+        saw_end = False
+        if ctx.identity.user_token:
+            self._run_owners[run_id] = self._run_owner(ctx)
+        self._run_contexts[run_id] = ctx
+        if workflow_store is not None:
+            heartbeat_task = asyncio.create_task(workflow_store.heartbeat(ctx, asyncio.current_task()))
 
         queue: asyncio.Queue[dict[str, Any] | BaseException | None] = asyncio.Queue()
 
         async def produce() -> None:
             try:
-                async for event in self._workflow_platform_events(payload, ctx, run_opt, log_id, start_perf, recorder):
+                if workflow_store is not None:
+                    await workflow_store.assert_running(ctx)
+                prepared_payload = materialize_input_files(payload)
+                recorder.set_root_input(prepared_payload)
+                async for event in self._workflow_platform_events(prepared_payload, ctx, run_opt, log_id, start_perf, recorder):
                     await queue.put(event)
             except BaseException as exc:
                 await queue.put(exc)
             finally:
                 await queue.put(None)
 
+        context_token = request_context.set(ctx)
         producer = asyncio.create_task(produce())
         ping_interval = run_opt.workflow_ping_interval_seconds
 
         try:
+            yield self._sse_event(workflow_start_event(run_id=run_id, log_id=log_id), event_index)
+            event_index += 1
             while True:
                 try:
                     if ping_interval and ping_interval > 0:
@@ -646,6 +687,8 @@ class GraphService:
                     code = stream_error_code(item) or WORKFLOW_DEFAULT_ERROR_CODE
                     message = str(item) or item.__class__.__name__
                     stream_error = item
+                    if workflow_store is not None:
+                        await workflow_store.update_status(ctx, "failed", error=message)
                     status_code = 500
                     error_code = str(code)
                     yield self._sse_event(
@@ -657,38 +700,66 @@ class GraphService:
 
                 if item.get("type") == "workflow_end":
                     final_output = item.get("output", {})
+                    if workflow_store is not None:
+                        await workflow_store.update_status(ctx, "completed", output=final_output)
+                    saw_end = True
                 if item.get("type") == "error":
                     stream_error = item.get("message") or "workflow error"
+                    if workflow_store is not None:
+                        await workflow_store.update_status(ctx, "failed", error=str(stream_error))
                     status_code = 500
                     error_code = str(item.get("code") or WORKFLOW_DEFAULT_ERROR_CODE)
                 yield self._sse_event(item, event_index)
                 event_index += 1
+                if item.get("type") in {"error", "workflow_end"}:
+                    break
         except asyncio.CancelledError:
             producer.cancel()
-            stream_error = "cancelled"
+            cancelled = workflow_store is None or bool(ctx.metadata.get("workflow_cancel_requested"))
+            stream_error = "cancelled" if cancelled else "interrupted"
+            if workflow_store is not None:
+                await workflow_store.update_status(ctx, "cancelled" if cancelled else "interrupted")
             status_code = 499
-            error_code = "CANCELLED"
+            error_code = "CANCELLED" if cancelled else "INTERRUPTED"
             yield self._sse_event(
-                workflow_error_event(run_id=run_id, log_id=log_id, code=1, message="cancelled"),
+                workflow_error_event(run_id=run_id, log_id=log_id, code=1, message=str(stream_error)),
                 event_index,
             )
         finally:
-            with suppress(asyncio.CancelledError):
-                await producer
-            if stream_error is not None:
-                _save_record_safely(recorder, status="error", output=final_output, error=stream_error)
-            else:
-                _save_record_safely(recorder, output=final_output)
-            observability.record_request_end(
-                ctx,
-                route="/stream_run",
-                status_code=status_code,
-                latency_ms=_elapsed_perf_ms(start_perf),
-                error_code=error_code,
-                error_message=stream_error,
-            )
-            self.running_tasks.pop(run_id, None)
-            telemetry.flush()
+            try:
+                if not producer.done():
+                    producer.cancel()
+                with suppress(asyncio.CancelledError):
+                    await producer
+                if workflow_store is not None and not saw_end and stream_error is None:
+                    stream_error = "workflow stream ended without a terminal event"
+                    await workflow_store.update_status(ctx, "failed", error=str(stream_error))
+            finally:
+                try:
+                    if heartbeat_task is not None:
+                        heartbeat_task.cancel()
+                        with suppress(asyncio.CancelledError):
+                            await heartbeat_task
+                finally:
+                    try:
+                        if stream_error is not None:
+                            _save_record_safely(recorder, status="error", output=final_output, error=stream_error)
+                        else:
+                            _save_record_safely(recorder, output=final_output)
+                        observability.record_request_end(
+                            ctx,
+                            route="/stream_run",
+                            status_code=status_code,
+                            latency_ms=_elapsed_perf_ms(start_perf),
+                            error_code=error_code,
+                            error_message=stream_error,
+                        )
+                    finally:
+                        self.running_tasks.pop(run_id, None)
+                        self._run_owners.pop(run_id, None)
+                        self._run_contexts.pop(run_id, None)
+                        request_context.reset(context_token)
+                        telemetry.flush()
 
     async def _workflow_platform_events(
         self,
@@ -754,6 +825,7 @@ class GraphService:
                         code=WORKFLOW_DEFAULT_ERROR_CODE,
                         message=str(error),
                     )
+                    return
 
             final_output = await enrich_workflow_output_file_urls(final_output, ctx)
             yield workflow_end_event(
@@ -806,6 +878,10 @@ class GraphService:
                 return value.strip()
         return node_id or "workflow_node"
 
+    @staticmethod
+    def _run_owner(ctx: Context) -> tuple[str, str | None, str | None]:
+        return (session_user_id(ctx), ctx.identity.workspace_id, ctx.identity.project_id)
+
     def cancel_run(self, run_id: str, ctx: Context | None = None) -> dict[str, Any]:
         logger.info("Attempting to cancel run_id=%s", run_id)
         task = self.running_tasks.get(run_id)
@@ -821,12 +897,15 @@ class GraphService:
                 "run_id": run_id,
                 "message": "Task has already completed.",
             }
-        if ctx is not None and is_agent_project():
+        if ctx is not None:
             if not ctx.identity.user_token:
-                raise ValueError("AGENT cancellation requires an authenticated Bearer token")
+                raise ValueError("Cancellation requires an authenticated Bearer token")
             owner = self._run_owners.get(run_id)
-            if owner is None or session_user_id(ctx) != owner:
-                raise PermissionError("Only the requesting user may cancel this AGENT run")
+            if owner is None or self._run_owner(ctx) != owner:
+                raise PermissionError("Only the requesting user may cancel this run")
+        request_ctx = self._run_contexts.get(run_id)
+        if request_ctx is not None:
+            request_ctx.metadata["workflow_cancel_requested"] = True
         task.cancel()
         return {
             "status": "success",
@@ -897,6 +976,12 @@ class GraphService:
             parser = LangGraphParser(graph)
             result["nodes"] = parser.get_nodes_metadata()
             result["agent_config"] = read_agent_config()
+        else:
+            durability = load_workflow_durability_config()
+            result["durability"] = {
+                "enabled": durability is not None,
+                "graph_version": durability.graph_version if durability else None,
+            }
         return result
 
     def agent_config_info(self) -> dict[str, Any]:

{licos_agent_runtime-0.3.4 → licos_agent_runtime-0.3.6}/src/licos_agent_runtime/stream_runner.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+from contextlib import suppress
 from dataclasses import dataclass
 from typing import Any, AsyncIterable, Callable, Iterable
 
@@ -119,6 +120,8 @@ async def _registered_stream(
     run_opt: RunOpt | None = None,
 ) -> AsyncIterable[str]:
     queue: asyncio.Queue[str | None] = asyncio.Queue()
+    detached = False
+    durable = ctx.metadata.get("workflow_store") is not None
 
     async def produce() -> None:
         try:
@@ -126,28 +129,39 @@ async def _registered_stream(
             if run_opt is not None:
                 kwargs["run_opt"] = run_opt
             async for item in stream_sse_func(payload, **kwargs):
-                await queue.put(item)
+                if not detached:
+                    await queue.put(item)
         except asyncio.CancelledError:
-            await queue.put(sse_event_func({"status": "cancelled", "run_id": run_id}))
+            if not detached:
+                await queue.put(sse_event_func({"status": "cancelled", "run_id": run_id}))
             raise
         except Exception as exc:
-            await queue.put(sse_event_func(error_classifier.get_error_response(exc, {"run_id": run_id})))
+            if not detached:
+                await queue.put(sse_event_func(error_classifier.get_error_response(exc, {"run_id": run_id})))
         finally:
-            await queue.put(None)
+            if not detached:
+                await queue.put(None)
 
     task = asyncio.create_task(produce())
     register_task_func(run_id, task)
 
-    while True:
-        item = await queue.get()
-        if item is None:
-            break
-        yield item
-
-    if task.done():
-        try:
-            exception = task.exception()
-        except asyncio.CancelledError:
-            exception = None
-        if exception is not None:
-            raise exception
+    try:
+        while True:
+            item = await queue.get()
+            if item is None:
+                break
+            yield item
+
+        if task.done():
+            try:
+                exception = task.exception()
+            except asyncio.CancelledError:
+                exception = None
+            if exception is not None:
+                raise exception
+    finally:
+        detached = True
+        if not durable and not task.done():
+            task.cancel()
+            with suppress(asyncio.CancelledError):
+                await task