Eager initialization of action handlers: if IDE is active, all handlers, if not only handlers of actions which will be run

Author: Aksem

finecode_extension_runner/src/finecode_extension_runner/_services/run_action.py

@@ -456,6 +456,132 @@ async def resolve_func_args_with_di(
     return args
 
 
+def _get_handler_raw_config(
+    handler: domain.ActionHandlerDeclaration,
+    runner_context: context.RunnerContext,
+) -> dict[str, typing.Any]:
+    handler_global_config = runner_context.project.action_handler_configs.get(
+        handler.source, None
+    )
+    handler_raw_config = {}
+    if handler_global_config is not None:
+        handler_raw_config = handler_global_config
+    if handler_raw_config == {}:
+        # still empty, just assign
+        handler_raw_config = handler.config
+    else:
+        # not empty anymore, deep merge
+        handler_config_merger.merge(handler_raw_config, handler.config)
+    return handler_raw_config
+
+
+async def ensure_handler_instantiated(
+    handler: domain.ActionHandlerDeclaration,
+    handler_cache: domain.ActionHandlerCache,
+    action_exec_info: domain.ActionExecInfo,
+    runner_context: context.RunnerContext,
+) -> None:
+    """Ensure handler is instantiated and initialized, populating handler_cache.
+
+    If handler is already instantiated (handler_cache.instance is not None), this is
+    a no-op. Otherwise, imports the handler class, resolves DI, instantiates it,
+    calls on_initialize lifecycle hook if present, and caches the result.
+    """
+    if handler_cache.instance is not None:
+        return
+
+    handler_raw_config = _get_handler_raw_config(handler, runner_context)
+
+    logger.trace(f"Load action handler {handler.name}")
+    try:
+        action_handler = run_utils.import_module_member_by_source_str(
+            handler.source
+        )
+    except ModuleNotFoundError as error:
+        logger.error(
+            f"Source of action handler {handler.name} '{handler.source}'"
+            " could not be imported"
+        )
+        logger.error(error)
+        raise ActionFailedException(
+            f"Import of action handler '{handler.name}' failed: {handler.source}"
+        ) from error
+
+    def get_handler_config(param_type):
+        # validate config using pydantic
+        try:
+            config_type = pydantic_dataclass(param_type)
+        except pydantic.ValidationError as exception:
+            raise ActionFailedException(exception.errors()) from exception
+        return config_type(**handler_raw_config)
+
+    def get_process_executor(param_type):
+        return action_exec_info.process_executor
+
+    exec_info = domain.ActionHandlerExecInfo()
+    # save immediately in context to be able to shutdown it if the first execution
+    # is interrupted by stopping ER
+    handler_cache.exec_info = exec_info
+    if inspect.isclass(action_handler):
+        args = await resolve_func_args_with_di(
+            func=action_handler.__init__,
+            known_args={
+                "config": get_handler_config,
+                "process_executor": get_process_executor,
+            },
+            params_to_ignore=["self"],
+        )
+
+        if "lifecycle" in args:
+            exec_info.lifecycle = args["lifecycle"]
+
+        handler_instance = action_handler(**args)
+        handler_cache.instance = handler_instance
+
+        service_instances = [
+            instance
+            for instance in args.values()
+            if isinstance(instance, service.Service)
+        ]
+        handler_cache.used_services = service_instances
+        for service_instance in service_instances:
+            if service_instance not in runner_context.running_services:
+                runner_context.running_services[service_instance] = (
+                    domain.RunningServiceInfo(used_by=[])
+                )
+
+            runner_context.running_services[service_instance].used_by.append(
+                handler_instance
+            )
+
+    else:
+        # handler is a plain function, not a class — nothing to instantiate
+        handler_cache.exec_info = exec_info
+        exec_info.status = domain.ActionHandlerExecInfoStatus.INITIALIZED
+        return
+
+    if (
+        exec_info.lifecycle is not None
+        and exec_info.lifecycle.on_initialize_callable is not None
+    ):
+        logger.trace(f"Initialize {handler.name} action handler")
+        try:
+            initialize_callable_result = (
+                exec_info.lifecycle.on_initialize_callable()
+            )
+            if inspect.isawaitable(initialize_callable_result):
+                await initialize_callable_result
+        except Exception as e:
+            logger.error(
+                f"Failed to initialize action handler {handler.name}: {e}"
+            )
+            raise ActionFailedException(
+                f"Initialisation of action handler '{handler.name}' failed: {e}"
+            ) from e
+
+    exec_info.status = domain.ActionHandlerExecInfoStatus.INITIALIZED
+
+
 async def execute_action_handler(
     handler: domain.ActionHandlerDeclaration,
     payload: code_action.RunActionPayload | None,
@@ -475,19 +601,6 @@ async def execute_action_handler(
     start_time = time.time_ns()
     execution_result: code_action.RunActionResult | None = None
 
-    handler_global_config = runner_context.project.action_handler_configs.get(
-        handler.source, None
-    )
-    handler_raw_config = {}
-    if handler_global_config is not None:
-        handler_raw_config = handler_global_config
-    if handler_raw_config == {}:
-        # still empty, just assign
-        handler_raw_config = handler.config
-    else:
-        # not empty anymore, deep merge
-        handler_config_merger.merge(handler_raw_config, handler.config)
-
     if handler_cache.instance is not None:
         handler_instance = handler_cache.instance
         handler_run_func = handler_instance.run
@@ -497,92 +610,21 @@ async def execute_action_handler(
             f"R{run_id} | Instance of action handler {handler.name} found in cache"
         )
     else:
-        logger.trace(f"R{run_id} | Load action handler {handler.name}")
-        try:
+        await ensure_handler_instantiated(
+            handler=handler,
+            handler_cache=handler_cache,
+            action_exec_info=action_exec_info,
+            runner_context=runner_context,
+        )
+        if handler_cache.instance is not None:
+            handler_run_func = handler_cache.instance.run
+        else:
+            # handler is a plain function
             action_handler = run_utils.import_module_member_by_source_str(
                 handler.source
             )
-        except ModuleNotFoundError as error:
-            logger.error(
-                f"R{run_id} | Source of action handler {handler.name} '{handler.source}'"
-                " could not be imported"
-            )
-            logger.error(error)
-            raise ActionFailedException(
-                f"Import of action handler '{handler.name}' failed(Run {run_id}): {handler.source}"
-            ) from error
-
-        def get_handler_config(param_type):
-            # validate config using pydantic
-            try:
-                config_type = pydantic_dataclass(param_type)
-            except pydantic.ValidationError as exception:
-                raise ActionFailedException(exception.errors()) from exception
-            return config_type(**handler_raw_config)
-
-        def get_process_executor(param_type):
-            return action_exec_info.process_executor
-
-        exec_info = domain.ActionHandlerExecInfo()
-        # save immediately in context to be able to shutdown it if the first execution
-        # is interrupted by stopping ER
-        handler_cache.exec_info = exec_info
-        if inspect.isclass(action_handler):
-            args = await resolve_func_args_with_di(
-                func=action_handler.__init__,
-                known_args={
-                    "config": get_handler_config,
-                    "process_executor": get_process_executor,
-                },
-                params_to_ignore=["self"],
-            )
-
-            if "lifecycle" in args:
-                exec_info.lifecycle = args["lifecycle"]
-
-            handler_instance = action_handler(**args)
-            handler_cache.instance = handler_instance
-            handler_run_func = handler_instance.run
-
-            service_instances = [
-                instance
-                for instance in args.values()
-                if isinstance(instance, service.Service)
-            ]
-            handler_cache.used_services = service_instances
-            for service_instance in service_instances:
-                if service_instance not in runner_context.running_services:
-                    runner_context.running_services[service_instance] = (
-                        domain.RunningServiceInfo(used_by=[])
-                    )
-
-                runner_context.running_services[service_instance].used_by.append(
-                    handler_instance
-                )
-
-        else:
             handler_run_func = action_handler
-
-        if (
-            exec_info.lifecycle is not None
-            and exec_info.lifecycle.on_initialize_callable is not None
-        ):
-            logger.trace(f"R{run_id} | Initialize {handler.name} action handler")
-            try:
-                initialize_callable_result = (
-                    exec_info.lifecycle.on_initialize_callable()
-                )
-                if inspect.isawaitable(initialize_callable_result):
-                    await initialize_callable_result
-            except Exception as e:
-                logger.error(
-                    f"R{run_id} | Failed to initialize action handler {handler.name}: {e}"
-                )
-                raise ActionFailedException(
-                    f"Initialisation of action handler '{handler.name}' failed(Run {run_id}): {e}"
-                ) from e
-
-        exec_info.status = domain.ActionHandlerExecInfoStatus.INITIALIZED
+        exec_info = handler_cache.exec_info
 
     def get_run_payload(param_type):
         return payload

finecode_extension_runner/src/finecode_extension_runner/lsp_server.py

@@ -404,6 +404,7 @@ async def update_config(
                 )
                 for svc in config.get("services", [])
             ],
+            handlers_to_initialize=config.get("handlers_to_initialize"),
         )
         response = await services.update_config(
             request=request,

finecode_extension_runner/src/finecode_extension_runner/schemas.py

@@ -40,6 +40,10 @@ class UpdateConfigRequest(BaseSchema):
     actions: dict[str, Action]
     action_handler_configs: dict[str, dict[str, Any]]
     services: list[ServiceDeclaration] = field(default_factory=list)
+    # If provided, eagerly instantiate these handlers after config update.
+    # Keys are action names, values are lists of handler names within that action.
+    # None means no eager initialization (lazy, on first use).
+    handlers_to_initialize: dict[str, list[str]] | None = None
 
 
 @dataclass

finecode_extension_runner/src/finecode_extension_runner/services.py

@@ -15,6 +15,8 @@
     ActionFailedException,
     StopWithResponse,
     run_action_raw,
+    create_action_exec_info,
+    ensure_handler_instantiated,
 )
 from finecode_extension_runner.di import bootstrap as di_bootstrap
 
@@ -122,9 +124,78 @@ def current_env_name_getter() -> str:
         service_declarations=request.services,
     )
 
+    if request.handlers_to_initialize is not None:
+        await initialize_handlers(request.handlers_to_initialize)
+
     return schemas.UpdateConfigResponse()
 
 
+async def initialize_handlers(
+    handlers_by_action: dict[str, list[str]],
+) -> None:
+    """Eagerly instantiate and initialize handlers.
+
+    This is called after update_config to pre-instantiate handlers so that
+    services (like LSP services) are started early rather than on first use.
+
+    Args:
+        handlers_by_action: mapping of action name → list of handler names
+            to eagerly initialize.
+    """
+    if global_state.runner_context is None:
+        logger.warning("Cannot initialize handlers: runner context is not set")
+        return
+
+    runner_context = global_state.runner_context
+    project = runner_context.project
+
+    for action_name, handler_names in handlers_by_action.items():
+        action_def = project.actions.get(action_name)
+        if action_def is None:
+            logger.warning(
+                f"Action '{action_name}' not found, skipping handler initialization"
+            )
+            continue
+
+        if action_name in runner_context.action_cache_by_name:
+            action_cache = runner_context.action_cache_by_name[action_name]
+        else:
+            action_cache = domain.ActionCache()
+            runner_context.action_cache_by_name[action_name] = action_cache
+
+        if action_cache.exec_info is None:
+            action_cache.exec_info = create_action_exec_info(action_def)
+
+        handlers_to_init = [
+            h for h in action_def.handlers if h.name in handler_names
+        ]
+        for handler in handlers_to_init:
+            if handler.name in action_cache.handler_cache_by_name:
+                handler_cache = action_cache.handler_cache_by_name[handler.name]
+                if handler_cache.instance is not None:
+                    continue
+            else:
+                handler_cache = domain.ActionHandlerCache()
+                action_cache.handler_cache_by_name[handler.name] = handler_cache
+
+            try:
+                await ensure_handler_instantiated(
+                    handler=handler,
+                    handler_cache=handler_cache,
+                    action_exec_info=action_cache.exec_info,
+                    runner_context=runner_context,
+                )
+                logger.trace(
+                    f"Eagerly initialized handler '{handler.name}' "
+                    f"for action '{action_name}'"
+                )
+            except Exception as e:
+                logger.error(
+                    f"Failed to eagerly initialize handler '{handler.name}' "
+                    f"for action '{action_name}': {e}"
+                )
+
+
 def reload_action(action_name: str) -> None:
     if global_state.runner_context is None:
         # TODO: raise error
@@ -231,6 +302,7 @@ def shutdown_action_handler(
                 if isinstance(used_service, service.DisposableService):
                     try:
                         used_service.dispose()
+                        logger.trace(f"Disposed service: {used_service}")
                     except Exception as exception:
                         logger.error(f"Failed to dispose service: {used_service}")
                         logger.exception(exception)

src/finecode/cli_app/commands/run_cmd.py

@@ -156,7 +156,9 @@ async def run_actions(
 
         try:
             await run_service.start_required_environments(
-                actions_by_projects, ws_context, update_config_in_running_runners=True
+                actions_by_projects,
+                ws_context,
+                update_config_in_running_runners=True,
             )
         except run_service.StartingEnvironmentsFailed as exception:
             raise RunFailed(