Coverage for core / src / sensorkit / backend / request.py: 90%
263 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2from __future__ import annotations
4import asyncio
5import contextlib
6import uuid
7from collections.abc import Awaitable, Coroutine, Generator
8from dataclasses import dataclass
9from datetime import UTC, datetime, timedelta
10from typing import Any, Callable, Literal, cast, final, overload, override
12from loguru import logger
13from pydantic import BaseModel
15from sensorkit.backend.base import StreamContext
16from sensorkit.backend.event import Event, EventMultiplexer
17from sensorkit.common.aio import cleanup_future, scoped_waiter
20class Call[R: BaseModel, V: BaseModel = R](Awaitable[V]):
21 """An invocation of a simple request."""
23 def __init__(self, coro: Coroutine[Any, Any, bytes], response_type: type[R]):
24 self._coro = coro
25 self._response_type = response_type
26 self._future: asyncio.Future[V] = asyncio.get_running_loop().create_future()
27 self._got_response = False
28 self.response: R | None = None
30 self._future.add_done_callback(cleanup_future)
32 @final
33 async def _invoke(self, timeout: float) -> R:
34 # Send request and await response.
35 async with asyncio.timeout(timeout):
36 recv = await self._coro
38 if self._response_type:
39 self.response = self._response_type.model_validate_json(recv)
41 self._got_response = True
43 return self.response
45 async def invoke(self, timeout: float = 5.0):
46 """Send the request and populate the future with the parsed response."""
47 try:
48 response = await self._invoke(timeout)
49 self._future.set_result(response)
50 return response
51 except BaseException as e:
52 self._future.set_exception(e)
53 raise
55 async def wait(self):
56 """Wait until the call's future is resolved."""
57 await self._future
59 def get_future(self) -> asyncio.Future[V]:
60 """Return the underlying future resolving to the call's final result."""
61 return self._future
63 def done(self):
64 """Return True if the call has completed (success or error)."""
65 return self._future.done()
67 def result(self) -> V:
68 """Return the call result, raising if the call has not completed or raised."""
69 return self._future.result()
71 async def run_to_completion(self) -> V:
72 """Invoke the call if not already sent, then await and return the final result."""
73 if not self._got_response:
74 await self.invoke()
76 await self._future
77 return self._future.result()
79 @override
80 def __await__(self) -> Generator[Any, None, V]:
81 return self.run_to_completion().__await__()
84type HandlerFunc[P: BaseModel | None, R: BaseModel | None] = Callable[[P], Coroutine[Any, Any, R]]
87class CallHandler[P: BaseModel | None, R: BaseModel | None]:
88 """Callable that handles incoming simple requests."""
90 def __init__(
91 self,
92 request: Request[P, R, R],
93 func: HandlerFunc[P, R],
94 ):
95 self._request = request
96 self._handler_func = func
98 def _run_handler(self, data: P):
99 return self._handler_func(data)
101 async def __call__(self, payload: bytes) -> bytes:
102 # Parse the input payload.
103 payload_type = self._request.payload
104 data: P = (
105 payload_type.model_validate_json(payload)
106 if self._request.payload is not None
107 else None
108 )
110 # Run the user handler func and get the response object.
111 response = await self._run_handler(data)
113 # Serialize the response payload.
114 response_payload = (
115 response.model_dump_json().encode()
116 if response is not None
117 else b""
118 )
120 return response_payload
123class ExtendedResponse(BaseModel):
124 """A long-running request response."""
125 call_id: uuid.UUID = None
126 call_state: Literal["running", "success", "failure"] = "running"
129class CallEvent(Event, ExtendedResponse):
130 """An event that occurs in the context of a long-running request."""
131 good_until: datetime | None
132 payload: Any = None
135class ExtendedCall[R: ExtendedResponse, V: BaseModel](Call[R, V]):
136 """A Call extension that supports long-running requests."""
138 def __init__(
139 self,
140 coro: Coroutine[Any, Any, bytes],
141 response_type: type[R],
142 result_type: type[V],
143 event_mux: EventMultiplexer,
144 ):
145 super().__init__(coro, response_type)
146 self._result_type = result_type
147 self._event_mux = event_mux
149 @override
150 async def invoke(self, timeout: float = 10.0):
151 # Consume the event stream.
152 context = contextlib.ExitStack()
153 queue = context.enter_context(self._event_mux.event_queue(CallEvent))
155 try:
156 await self._event_mux.wait_ready()
158 # Execute the initial request-response communication.
159 response = await self._invoke(timeout)
160 except (asyncio.CancelledError, Exception) as e:
161 context.close()
162 self._future.set_exception(e)
163 raise
165 if response.call_state == "failure":
166 # The handler rejected the call. The response is authoritative, so settle here instead
167 # of leaving the caller on the event stream's timeout.
168 try:
169 error = await self._rejection_error(queue, response.call_id)
170 except (asyncio.CancelledError, Exception) as e:
171 context.close()
172 self._future.set_exception(e)
173 raise
175 context.close()
176 self._future.set_exception(error)
177 raise error
179 # Start a background task to receive response progress and end events.
180 started = asyncio.Event()
182 self._task = asyncio.create_task(
183 self._await_response_events(
184 context,
185 queue,
186 response.call_id,
187 started,
188 timeout,
189 )
190 )
192 def extended_call_done(t: asyncio.Task):
193 # Note we must check the task exception even if the future is already done to avoid
194 # leaks and warnings.
195 if t.cancelled():
196 self._future.cancel()
197 elif err := t.exception():
198 if not self._future.done():
199 self._future.set_exception(err)
200 elif not self._future.done():
201 self._future.set_result(t.result())
203 self._task.add_done_callback(extended_call_done)
205 # Wait until the first event is received or the task ends for some reason, whichever
206 # comes first.
207 async with scoped_waiter(started.wait()) as wait:
208 await asyncio.wait(
209 [wait, self._task],
210 return_when=asyncio.FIRST_COMPLETED,
211 )
213 return response
215 async def _rejection_error(
216 self,
217 queue: asyncio.Queue[Event],
218 response_id: uuid.UUID,
219 timeout: float = 1.0,
220 ) -> CallError:
221 """Return the error for a rejected call, carrying the handler's reason where available.
223 The reason travels on the failure event rather than the response, and the event stream may
224 be a different transport than the request that just delivered the rejection. The event is
225 published ahead of that response, so it is normally waiting in the queue already; if it
226 does not arrive within timeout seconds, the rejection still stands and is reported without
227 a reason.
229 Args:
230 queue: The call's registered CallEvent queue.
231 response_id: The call id to match events against.
232 timeout: How long to wait for the failure event before giving up on the reason.
233 """
234 try:
235 async with asyncio.timeout(timeout):
236 while True:
237 event = await queue.get()
238 queue.task_done()
240 if not isinstance(event, CallEvent) or event.call_id != response_id:
241 continue
243 # Rejecting leaves the response done, so a handler is free to emit progress
244 # before its failure event. The reason rides on the latter.
245 if event.call_state == "failure":
246 return CallError(f"response failed: {event.payload}")
247 except TimeoutError:
248 return CallError("call rejected")
250 async def _await_response_events(
251 self,
252 context: contextlib.ExitStack,
253 queue: asyncio.Queue[Event],
254 response_id: uuid.UUID,
255 started: asyncio.Event,
256 timeout: float,
257 ) -> V:
258 with context:
259 while True:
260 async with asyncio.timeout(timeout):
261 # Consume the next matching ResponseEvent.
262 event = await queue.get()
263 queue.task_done()
265 if not isinstance(event, CallEvent) or event.call_id != response_id:
266 continue
268 match event.call_state:
269 case "success":
270 return cast(
271 V,
272 self._result_type.model_validate(event.payload)
273 if self._result_type is not None
274 else None
275 )
276 case "failure":
277 raise CallError(f"response failed: {event.payload}")
279 # Signal that at least one matching response event has been received.
280 started.set()
282 # Figure the next timeout. This relies on peer clock synchronization.
283 timeout = (event.good_until - datetime.now(UTC)).total_seconds()
285 # Unreachable.
286 raise RuntimeError
289type ExtendedHandlerFunc[P: BaseModel | None, R: BaseModel | None, V: BaseModel | None] = (
290 Callable[
291 [P, CallContext[R, V]],
292 Coroutine[Any, Any, None]
293 ]
294)
297class CallContext[R: ExtendedResponse | None, V: BaseModel | None]:
298 """Context for an incoming extended request currently being handled."""
300 def __init__(self, call_id: uuid.UUID, stream: StreamContext):
301 self.call_id = call_id
302 self._stream = stream
303 self.response = asyncio.get_running_loop().create_future()
304 self.finalized = False
306 def accept(self, *, response: R):
307 """Mark the call as accepted and set the initial response, allowing progress events to follow."""
308 if self.response.done():
309 raise CallHandlerResponseError(responded=True)
311 if response is None:
312 response = ExtendedResponse()
314 response.call_id = self.call_id
315 response.call_state = "running"
316 self.response.set_result(response)
318 def reject(self, *, response: R):
319 """Reject the call immediately, returning a failure response without further events."""
320 if self.response.done():
321 raise CallHandlerResponseError(responded=True)
323 if response is None:
324 response = ExtendedResponse()
326 response.call_id = self.call_id
327 response.call_state = "failure"
328 self.response.set_result(response)
330 async def progress(self, ttl: float, payload: Any = None):
331 """Emit a progress event, extending the caller's deadline by ttl seconds."""
332 if not self.response.done():
333 raise CallHandlerResponseError(responded=False)
335 await self._call_event(
336 CallEvent(
337 call_id=self.call_id,
338 call_state="running",
339 good_until=datetime.now(UTC) + timedelta(seconds=ttl),
340 payload=payload,
341 )
342 )
344 async def progress_from_task(
345 self,
346 task: asyncio.Task,
347 *,
348 cadence: float,
349 ttl: float,
350 payload_func: Callable[[], Any] | None = None,
351 ):
352 """Emit progress events at the given cadence until task completes, then await the task."""
353 done = None
354 fs = [task]
356 while not done:
357 await self.progress(ttl, payload_func() if payload_func else None)
358 done, _ = await asyncio.wait(fs, timeout=cadence)
360 # Raise if the task raised.
361 await task
363 async def succeed(self, *, result: V):
364 """Emit a success event carrying the final result, completing the extended call."""
365 if not self.response.done():
366 raise CallHandlerResponseError(responded=False)
368 if self.response.result().call_state == "failure":
369 raise CallHandlerError("Cannot send success result for rejected call")
371 await self._call_event(
372 CallEvent(
373 call_id=self.call_id,
374 call_state="success",
375 good_until=None,
376 payload=result,
377 ),
378 finalize=True,
379 )
381 async def fail(self, payload: Any = None):
382 """Emit a failure event, completing the extended call in a failed state."""
383 if not self.response.done():
384 raise CallHandlerResponseError(responded=False)
386 await self._call_event(
387 CallEvent(
388 call_id=self.call_id,
389 call_state="failure",
390 good_until=None,
391 payload=payload,
392 ),
393 finalize=True,
394 )
396 async def _call_event(self, event: CallEvent, finalize: bool = False):
397 if self.finalized:
398 raise CallHandlerError("Call has already been finalized")
400 self.finalized = finalize
402 await self._stream.publish_event(event.model_dump_json().encode())
405class ExtendedCallHandler[P: BaseModel | None, R: ExtendedResponse, V: BaseModel](
406 CallHandler[P, R]
407):
408 """Callable that handles incoming long-running requests."""
410 def __init__(
411 self,
412 request: Request[P, R, V],
413 func: ExtendedHandlerFunc[P, R, V],
414 stream: StreamContext,
415 ):
416 super().__init__(request, func) # noqa
417 self._handler_func = func
418 self._stream = stream
420 @override
421 async def _run_handler(self, data: P):
422 # Create the call context object and run the call handler func.
423 context = CallContext(uuid.uuid1(), self._stream)
424 handler_task = asyncio.create_task(self._handler_func(data, context))
426 # Handle the initial response and error paths.
427 await asyncio.wait(
428 [
429 handler_task,
430 context.response,
431 ],
432 return_when=asyncio.FIRST_COMPLETED,
433 )
435 if handler_task.done():
436 # The handler func exited already. If it raised, fail.
437 if e := handler_task.exception():
438 if not context.finalized:
439 await context.fail(str(e))
441 logger.warning(f"Exception in call handler: {e}")
442 raise e
444 if not context.response.done():
445 # The handler func did not call `accept` or `reject`.
446 await context.fail()
447 raise CallHandlerResponseError(responded=False)
449 # Get the initial response object.
450 response = context.response.result()
451 response.call_id = context.call_id
453 if response.call_state == "failure":
454 # The handler rejected the request. Make sure a fail event is sent.
455 if not context.finalized:
456 await context.fail("Request rejected")
458 # Ensure exceptions raised in the handler are propagated.
459 def call_handler_done(_):
460 if not context.finalized and not handler_task.cancelled():
461 payload = None
463 if e := handler_task.exception():
464 payload = str(e)
466 logger.debug(f"failing unfinalized call handler {e}")
467 _t = asyncio.create_task(context.fail(payload))
469 handler_task.add_done_callback(call_handler_done)
470 return response
473class CallError(Exception):
474 """Raised when a call method fails."""
477class CallHandlerError(CallError):
478 """Raised when a call method is called in an invalid state."""
481class CallHandlerResponseError(CallHandlerError):
482 def __init__(self, *, responded: bool):
483 self.responded = responded
484 problem = "has already called" if responded else "did not call"
485 super().__init__(f"Call handler {problem} `accept` or `reject`")
488@dataclass
489class Request[P: BaseModel | None, R: BaseModel | None, V: BaseModel | None]:
490 """A declaration of a callable request.
492 If the response type is `ExtendedResponse` or a subclass of it, the request is considered
493 "extended".
494 """
495 name: str
496 """The name of the request."""
497 payload: type[P]
498 """The type expected to be sent in the request call."""
499 response: type[R]
500 """The type expected to be received in the initial call response."""
501 result: type[V]
502 """The type of the payload of the success ResponseEvent in an extended request.
503 Always equal to the response type for simple requests."""
505 def is_extended(self):
506 """Return True if this is a long-running request using ExtendedResponse."""
507 return self.response and issubclass(self.response, ExtendedResponse)
509 def create_handler(
510 self,
511 func: HandlerFunc[P, R] | ExtendedHandlerFunc[P, R, V],
512 stream: StreamContext | None = None,
513 ) -> CallHandler[P, R]:
514 """Construct the appropriate CallHandler or ExtendedCallHandler for this request."""
515 if self.is_extended():
516 return ExtendedCallHandler(self, func, stream)
517 else:
518 return CallHandler(self, func)
520 @classmethod
521 @overload
522 def define[R: ExtendedResponse, P: BaseModel | None = None](
523 cls,
524 name: str,
525 *,
526 payload: type[P] | None = None,
527 response: type[R],
528 ) -> Request[P, R, None]:
529 """Define a long-running request that uses the entity event stream to track progress."""
531 @classmethod
532 @overload
533 def define[V: BaseModel, R: ExtendedResponse = ExtendedResponse, P: BaseModel | None = None](
534 cls,
535 name: str,
536 *,
537 payload: type[P] | None = None,
538 response: type[R] = ExtendedResponse,
539 result: type[V],
540 ) -> Request[P, R, V]:
541 """Define a long-running request that uses the entity event stream to track progress."""
543 @classmethod
544 @overload
545 def define[R: BaseModel | None = None, P: BaseModel | None = None](
546 cls,
547 name: str,
548 *,
549 payload: type[P] | None = None,
550 response: type[R] | None = None,
551 ) -> Request[P, R, R]:
552 """Define a simple request that requires an immediate response."""
554 @classmethod
555 def define(
556 cls,
557 name: str,
558 payload = None,
559 response = None,
560 result = None,
561 ):
562 if response and issubclass(response, ExtendedResponse):
563 # This is an extended request. No checks or defaults are needed in this case.
564 return cls(name, payload, response, result)
565 elif result:
566 # Here a result type is given but no response type. In this case, having a result means
567 # this is an extended request, so we default to the bare ExtendedResponse type.
568 if response is not None:
569 raise TypeError("simple request cannot have a result type")
571 return cls(name, payload, ExtendedResponse, result)
572 else:
573 # Simple requests always treat their response as the result.
574 return cls(name, payload, response, result or response)