remove timeout from on_session_end

examples/frontdesk/frontdesk_agent.py

@@ -202,6 +202,8 @@ async def on_session_end(ctx: JobContext) -> None:
     else:
         ctx.tagger.fail(reason="Appointment was not booked")
 
+    logger.info("session tags: %s", ctx.tagger.tags)
+
 
 @server.rtc_session(on_session_end=on_session_end)
 async def frontdesk_agent(ctx: JobContext):

livekit-agents/livekit/agents/ipc/job_proc_lazy_main.py

@@ -40,6 +40,7 @@
     InitializeRequest,
     ShutdownRequest,
     StartJobRequest,
+    UserEntrypointDone,
 )
 
 
@@ -327,6 +328,8 @@ def log_exception(t: asyncio.Task[Any]) -> None:
         if session := self._job_ctx._primary_agent_session:
             await session.aclose()
 
+        await self._client.send(UserEntrypointDone())
+
         await self._job_ctx._on_session_end()
 
         if self._session_end_fnc:

livekit-agents/livekit/agents/ipc/job_thread_executor.py

@@ -69,6 +69,7 @@ def __init__(
         self._inference_executor = inference_executor
         self._inference_tasks: list[asyncio.Task[None]] = []
         self._id = utils.shortuuid("THEXEC_")
+        self._user_entrypoint_done_fut: asyncio.Future[None] | None = None
 
     @property
     def id(self) -> str:
@@ -188,14 +189,16 @@ async def aclose(self) -> None:
             return
 
         self._closing = True
+        self._user_entrypoint_done_fut = asyncio.Future[None]()
+
         with contextlib.suppress(utils.aio.duplex_unix.DuplexClosed):
             await channel.asend_message(self._pch, proto.ShutdownRequest())
 
         try:
-            if self._main_atask:
-                await asyncio.wait_for(
-                    asyncio.shield(self._main_atask), timeout=self._opts.close_timeout
-                )
+            await asyncio.wait_for(
+                asyncio.shield(self._user_entrypoint_done_fut),
+                timeout=self._opts.close_timeout,
+            )
         except asyncio.TimeoutError:
             logger.error("job shutdown is taking too much time..", extra=self.logging_extra())
 
@@ -281,6 +284,10 @@ async def _monitor_task(self) -> None:
             if isinstance(msg, proto.Exiting):
                 logger.debug("job exiting", extra={"reason": msg.reason, **self.logging_extra()})
 
+            if isinstance(msg, proto.UserEntrypointDone):
+                if self._user_entrypoint_done_fut and not self._user_entrypoint_done_fut.done():
+                    self._user_entrypoint_done_fut.set_result(None)
+
             if isinstance(msg, proto.InferenceRequest):
                 self._inference_tasks.append(asyncio.create_task(self._do_inference_task(msg)))
 

livekit-agents/livekit/agents/ipc/proto.py

@@ -205,6 +205,21 @@ def read(self, b: io.BytesIO) -> None:
         pass
 
 
+@dataclass
+class UserEntrypointDone:
+    """sent by the subprocess when the user's entrypoint has finished and shutdown begins.
+    This signals that the shutdown timeout should stop - the remaining work (on_session_end)
+    should not be subject to the timeout."""
+
+    MSG_ID: ClassVar[int] = 10
+
+    def write(self, b: io.BytesIO) -> None:
+        pass
+
+    def read(self, b: io.BytesIO) -> None:
+        pass
+
+
 IPC_MESSAGES = {
     InitializeRequest.MSG_ID: InitializeRequest,
     InitializeResponse.MSG_ID: InitializeResponse,
@@ -216,4 +231,5 @@ def read(self, b: io.BytesIO) -> None:
     InferenceRequest.MSG_ID: InferenceRequest,
     InferenceResponse.MSG_ID: InferenceResponse,
     DumpStackTraceRequest.MSG_ID: DumpStackTraceRequest,
+    UserEntrypointDone.MSG_ID: UserEntrypointDone,
 }

livekit-agents/livekit/agents/ipc/supervised_proc.py

@@ -102,6 +102,7 @@ def __init__(
         self._kill_sent = False
         self._initialize_fut = asyncio.Future[None]()
         self._lock = asyncio.Lock()
+        self._user_entrypoint_done_fut: asyncio.Future[None] | None = None
 
     @abstractmethod
     def _create_process(self, cch: socket.socket, log_cch: socket.socket) -> mp.Process: ...
@@ -244,14 +245,16 @@ async def aclose(self) -> None:
             return
 
         self._closing = True
+        self._user_entrypoint_done_fut = asyncio.Future[None]()
+
         with contextlib.suppress(duplex_unix.DuplexClosed):
             await channel.asend_message(self._pch, proto.ShutdownRequest())
 
         try:
-            if self._supervise_atask:
-                await asyncio.wait_for(
-                    asyncio.shield(self._supervise_atask), timeout=self._opts.close_timeout
-                )
+            await asyncio.wait_for(
+                asyncio.shield(self._user_entrypoint_done_fut),
+                timeout=self._opts.close_timeout,
+            )
         except asyncio.TimeoutError:
             logger.error(
                 "process did not exit in time, killing process",
@@ -393,6 +396,10 @@ async def _read_ipc_task(
                     extra={"reason": msg.reason, **self.logging_extra()},
                 )
 
+            if isinstance(msg, proto.UserEntrypointDone):
+                if self._user_entrypoint_done_fut and not self._user_entrypoint_done_fut.done():
+                    self._user_entrypoint_done_fut.set_result(None)
+
             ipc_ch.send_nowait(msg)
 
     @log_exceptions(logger=logger)