IActionRunner.run_action_iter and its implementation. Sending partial results using yield in run action and using explicit partial result sender

Author: Aksem

finecode_builtin_handlers/src/finecode_builtin_handlers/format.py

@@ -29,7 +29,7 @@ async def run(
         self,
         payload: format_action.FormatRunPayload,
         run_context: format_action.FormatRunContext,
-    ) -> format_action.FormatRunResult:
+    ):
         run_meta = run_context.meta
         file_uris: list[ResourceUri]
 
@@ -66,14 +66,14 @@ async def run(
         format_files_action_instance = self.action_runner.get_action_by_source(
             format_files_action.FormatFilesAction
         )
-        format_result = await self.action_runner.run_action(
+        async for partial in self.action_runner.run_action_iter(
             action=format_files_action_instance,
             payload=format_files_action.FormatFilesRunPayload(
                 file_paths=file_uris,
                 save=payload.save,
             ),
             meta=run_meta,
-        )
-        return format_action.FormatRunResult(
-            result_by_file_path=format_result.result_by_file_path
-        )
+        ):
+            yield format_action.FormatRunResult(
+                result_by_file_path=partial.result_by_file_path
+            )

finecode_builtin_handlers/src/finecode_builtin_handlers/lint.py

@@ -31,7 +31,7 @@ async def run(
         self,
         payload: lint_action.LintRunPayload,
         run_context: lint_action.LintRunContext,
-    ) -> lint_action.LintRunResult:
+    ):
         run_meta = run_context.meta
         file_uris: list[ResourceUri]
 
@@ -68,9 +68,9 @@ async def run(
         lint_files_action_instance = self.action_runner.get_action_by_source(
             lint_files_action.LintFilesAction
         )
-        lint_result = await self.action_runner.run_action(
+        async for partial in self.action_runner.run_action_iter(
             action=lint_files_action_instance,
             payload=lint_files_action.LintFilesRunPayload(file_paths=file_uris),
             meta=run_meta,
-        )
-        return lint_action.LintRunResult(messages=lint_result.messages)
+        ):
+            yield lint_action.LintRunResult(messages=partial.messages)

finecode_extension_api/src/finecode_extension_api/code_action.py

@@ -88,6 +88,20 @@ def update(self, result: RunActionResult) -> None:
             self._current_result.update(result)
 
 
+class PartialResultSender(typing.Protocol):
+    """Handler-facing interface for sending partial results to the client."""
+
+    async def send(self, result: RunActionResult) -> None: ...
+
+
+class _NoOpPartialResultSender:
+    async def send(self, result: RunActionResult) -> None:
+        pass
+
+
+_NOOP_SENDER = _NoOpPartialResultSender()
+
+
 class RunActionContext(typing.Generic[RunPayloadType]):
     # data object to save data between action steps(only during one run, after run data
     # is removed). Keep it simple, without business logic, just data storage, but you
@@ -101,12 +115,14 @@ def __init__(
         initial_payload: RunPayloadType,
         meta: RunActionMeta,
         info_provider: RunContextInfoProvider,
+        partial_result_sender: PartialResultSender = _NOOP_SENDER,
     ) -> None:
         self.run_id = run_id
         self.initial_payload = initial_payload
         self.meta = meta
         self.exit_stack = contextlib.AsyncExitStack()
         self._info_provider = info_provider
+        self.partial_result_sender = partial_result_sender
 
     @property
     def current_result(self) -> RunActionResult | None:
@@ -148,12 +164,14 @@ def __init__(
         initial_payload: RunPayloadType,
         meta: RunActionMeta,
         info_provider: RunContextInfoProvider,
+        partial_result_sender: PartialResultSender = _NOOP_SENDER,
     ) -> None:
         super().__init__(
             run_id=run_id,
             initial_payload=initial_payload,
             meta=meta,
             info_provider=info_provider,
+            partial_result_sender=partial_result_sender,
         )
         self.partial_result_scheduler = partialresultscheduler.PartialResultScheduler()
 

finecode_extension_api/src/finecode_extension_api/interfaces/iactionrunner.py

@@ -1,3 +1,4 @@
+import collections.abc
 import typing
 
 from finecode_extension_api import code_action, service
@@ -36,6 +37,13 @@ async def run_action(
         meta: code_action.RunActionMeta,
     ) -> ResultT: ...
 
+    def run_action_iter(
+        self,
+        action: ActionDeclaration[code_action.Action[PayloadT, typing.Any, ResultT]],
+        payload: PayloadT,
+        meta: code_action.RunActionMeta,
+    ) -> collections.abc.AsyncIterator[ResultT]: ...
+
     def get_actions_names(self) -> list[str]: ...
 
 

finecode_extension_runner/src/finecode_extension_runner/_services/run_action.py

@@ -40,6 +40,25 @@ def __init__(self, response: schemas.RunActionResponse) -> None:
         self.response = response
 
 
+class _TrackingPartialResultSender:
+    """Wraps partial_result_sender.schedule_sending with state tracking."""
+
+    def __init__(
+        self,
+        token: int | str,
+        send_func: collections.abc.Callable[
+            [int | str, code_action.RunActionResult], collections.abc.Awaitable[None]
+        ],
+    ) -> None:
+        self._token = token
+        self._send_func = send_func
+        self.has_sent = False
+
+    async def send(self, result: code_action.RunActionResult) -> None:
+        self.has_sent = True
+        await self._send_func(self._token, result)
+
+
 def set_partial_result_sender(send_func: typing.Callable) -> None:
     global partial_result_sender
     partial_result_sender = partial_result_sender_module.PartialResultSender(
@@ -60,6 +79,7 @@ async def run_action(
     meta: code_action.RunActionMeta,
     partial_result_token: int | str | None = None,
     run_id: int | None = None,
+    partial_result_queue: asyncio.Queue | None = None,
 ) -> code_action.RunActionResult | None:
     # design decisions:
     # - keep payload unchanged between all subaction runs.
@@ -105,14 +125,22 @@ async def run_action(
 
     run_context: code_action.RunActionContext | AsyncPlaceholderContext
     run_context_info = code_action.RunContextInfoProvider(is_concurrent_execution=execute_handlers_concurrently)
+    if partial_result_token is not None:
+        tracking_sender = _TrackingPartialResultSender(
+            token=partial_result_token,
+            send_func=partial_result_sender.schedule_sending,
+        )
+    else:
+        tracking_sender = None
     if action_exec_info.run_context_type is not None:
         constructor_args = await resolve_func_args_with_di(
             action_exec_info.run_context_type.__init__,
             known_args={
                 "run_id": lambda _: run_id,
                 "initial_payload": lambda _: payload,
                 "meta": lambda _: meta,
-                "info_provider": lambda _: run_context_info
+                "info_provider": lambda _: run_context_info,
+                "partial_result_sender": lambda _: tracking_sender or code_action._NOOP_SENDER,
             },
             params_to_ignore=["self"],
         )
@@ -163,6 +191,8 @@ async def run_action(
                         action_cache=action_cache,
                         action_exec_info=action_exec_info,
                         runner_context=runner_context,
+                        partial_result_token=partial_result_token,
+                        tracking_sender=tracking_sender,
                     )
 
                 parts = [part async for part in payload]
@@ -181,6 +211,12 @@ async def run_action(
                 try:
                     async with asyncio.TaskGroup() as tg:
                         for part in parts:
+                            if part not in run_context.partial_result_scheduler.coroutines_by_key:
+                                logger.warning(
+                                    f"R{run_id} | No coroutines scheduled for part {part} "
+                                    f"of action '{action_def.name}', skipping"
+                                )
+                                continue
                             part_coros = (
                                 run_context.partial_result_scheduler.coroutines_by_key[part]
                             )
@@ -193,6 +229,7 @@ async def run_action(
                                     partial_result_sender,
                                     action_def.name,
                                     run_id,
+                                    partial_result_queue=partial_result_queue,
                                 )
                             else:
                                 coro = run_subresult_coros_sequentially(
@@ -202,6 +239,7 @@ async def run_action(
                                     partial_result_sender,
                                     action_def.name,
                                     run_id,
+                                    partial_result_queue=partial_result_queue,
                                 )
                             subresult_task = tg.create_task(coro)
                             subresults_tasks.append(subresult_task)
@@ -246,6 +284,8 @@ async def run_action(
                                         action_cache=action_cache,
                                         action_exec_info=action_exec_info,
                                         runner_context=runner_context,
+                                        partial_result_token=partial_result_token,
+                                        tracking_sender=tracking_sender,
                                     )
                                 )
                                 handlers_tasks.append(handler_task)
@@ -276,6 +316,8 @@ async def run_action(
                                 action_cache=action_cache,
                                 action_exec_info=action_exec_info,
                                 runner_context=runner_context,
+                                partial_result_token=partial_result_token,
+                                tracking_sender=tracking_sender,
                             )
                         except ActionFailedException as exception:
                             raise exception
@@ -314,6 +356,10 @@ async def run_action(
             f"Unexpected result type: {type(action_result).__name__}"
         )
 
+    if partial_result_queue is not None and action_result is not None:
+        await partial_result_queue.put(action_result)
+        return None
+
     return action_result
 
 
@@ -592,6 +638,8 @@ async def execute_action_handler(
     action_exec_info: domain.ActionExecInfo,
     action_cache: domain.ActionCache,
     runner_context: context.RunnerContext,
+    partial_result_token: int | str | None = None,
+    tracking_sender: _TrackingPartialResultSender | None = None,
 ) -> code_action.RunActionResult:
     logger.trace(f"R{run_id} | Run {handler.name} on {str(payload)[:100]}...")
     if handler.name in action_cache.handler_cache_by_name:
@@ -647,8 +695,30 @@ def get_run_context(param_type):
         # there is also `inspect.iscoroutinefunction` but it cannot recognize coroutine
         # functions which are class methods. Use `isawaitable` on result instead.
         call_result = handler_run_func(**args)
-        if inspect.isawaitable(call_result):
-            execution_result = await call_result
+        if inspect.isasyncgen(call_result):
+            execution_result = None
+            async for partial_result in call_result:
+                if partial_result_token is not None:
+                    await partial_result_sender.schedule_sending(
+                        partial_result_token, partial_result
+                    )
+                if execution_result is None:
+                    result_type_pydantic = pydantic_dataclass(type(partial_result))
+                    execution_result = result_type_pydantic(
+                        **dataclasses.asdict(partial_result)
+                    )
+                else:
+                    execution_result.update(partial_result)
+            if partial_result_token is not None:
+                await partial_result_sender.send_all_immediately()
+                execution_result = None  # partials already sent
+        elif inspect.isawaitable(call_result):
+            handler_result = await call_result
+            if tracking_sender is not None and tracking_sender.has_sent:
+                await partial_result_sender.send_all_immediately()
+                execution_result = None
+            else:
+                execution_result = handler_result
         else:
             execution_result = call_result
     except Exception as exception:
@@ -684,6 +754,7 @@ async def run_subresult_coros_concurrently(
     partial_result_sender: partial_result_sender_module.PartialResultSender,
     action_name: str,
     run_id: int,
+    partial_result_queue: asyncio.Queue | None = None,
 ) -> code_action.RunActionResult | None:
     coros_tasks: list[asyncio.Task] = []
     try:
@@ -725,7 +796,10 @@ async def run_subresult_coros_concurrently(
             else:
                 action_subresult.update(coro_result)
 
-    if send_partial_results:
+    if partial_result_queue is not None:
+        await partial_result_queue.put(action_subresult)
+        return None
+    elif send_partial_results:
         await partial_result_sender.schedule_sending(
             partial_result_token, action_subresult
         )
@@ -741,6 +815,7 @@ async def run_subresult_coros_sequentially(
     partial_result_sender: partial_result_sender_module.PartialResultSender,
     action_name: str,
     run_id: int,
+    partial_result_queue: asyncio.Queue | None = None,
 ) -> code_action.RunActionResult | None:
     action_subresult: code_action.RunActionResult | None = None
     for coro in coros:
@@ -761,7 +836,10 @@ async def run_subresult_coros_sequentially(
             else:
                 action_subresult.update(coro_result)
 
-    if send_partial_results:
+    if partial_result_queue is not None:
+        await partial_result_queue.put(action_subresult)
+        return None
+    elif send_partial_results:
         await partial_result_sender.schedule_sending(
             partial_result_token, action_subresult
         )

finecode_extension_runner/src/finecode_extension_runner/impls/action_runner.py

@@ -1,10 +1,13 @@
+import asyncio
 import collections.abc
 import typing
 from finecode_extension_api import code_action
 from finecode_extension_api.interfaces import iactionrunner
 
 from finecode_extension_runner import domain, run_utils
 
+_SENTINEL = object()
+
 
 class ActionRunner(iactionrunner.IActionRunner):
     def __init__(self, run_action_func: typing.Callable[[domain.ActionDeclaration, code_action.RunActionPayload, code_action.RunActionMeta], collections.abc.Coroutine[None, None, code_action.RunActionResult]],
@@ -30,6 +33,36 @@ async def run_action(
         except Exception as exception:
             raise iactionrunner.ActionRunFailed(str(exception)) from exception
 
+    @typing.override
+    async def run_action_iter(
+        self,
+        action: iactionrunner.ActionDeclaration[iactionrunner.ActionT],
+        payload: code_action.RunActionPayload,
+        meta: code_action.RunActionMeta,
+    ) -> collections.abc.AsyncIterator[code_action.RunActionResult]:
+        queue: asyncio.Queue = asyncio.Queue()
+
+        async def producer():
+            try:
+                await self._run_action_func(action, payload, meta, partial_result_queue=queue)
+            finally:
+                await queue.put(_SENTINEL)
+
+        task = asyncio.create_task(producer())
+        try:
+            while True:
+                item = await queue.get()
+                if item is _SENTINEL:
+                    break
+                yield item
+        finally:
+            if not task.done():
+                task.cancel()
+                try:
+                    await task
+                except (asyncio.CancelledError, Exception):
+                    pass
+
     @typing.override
     def get_actions_names(self) -> list[str]:
         return list(self._actions_getter().keys())