WAL in ER

Author: Aksem

finecode_extension_api/src/finecode_extension_api/code_action.py

@@ -55,6 +55,7 @@ class DevEnv(enum.StrEnum):
 class RunActionMeta:
     trigger: RunActionTrigger
     dev_env: DevEnv
+    wal_run_id: str = ""
 
 
 class RunReturnCode(enum.IntEnum):

finecode_extension_runner/src/finecode_extension_runner/_services/run_action.py

@@ -7,7 +7,7 @@
 from loguru import logger
 
 import finecode_extension_runner.start as runner_start
-from finecode_extension_runner import global_state, logs
+from finecode_extension_runner import er_wal, global_state, logs
 
 
 @click.group()
@@ -19,6 +19,7 @@ def main():
 @main.command()
 @click.option("--trace", "trace", is_flag=True, default=False)
 @click.option("--debug", "debug", is_flag=True, default=False)
+@click.option("--wal", "wal", is_flag=True, default=False)
 @click.option(
     "--project-path",
     "project_path",
@@ -29,6 +30,7 @@ def main():
 def start(
     trace: bool,
     debug: bool,
+    wal: bool,
     project_path: Path,
     env_name: str | None,
 ):
@@ -55,6 +57,8 @@ def start(
     global_state.log_level = "INFO" if trace is False else "TRACE"
     global_state.project_dir_path = project_path
     global_state.env_name = env_name
+    if wal:
+        global_state.wal_writer = er_wal.ErWalWriter()
 
     log_file_path = (project_path
         / ".venvs"

finecode_extension_runner/src/finecode_extension_runner/cli.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import dataclasses
+import enum
+import pathlib
+import sys
+import typing
+
+from finecode_extension_runner import wal as shared_wal
+
+
+class ErWalEventType(enum.StrEnum):
+    RUN_ACCEPTED = "run.accepted"
+    RUN_DISPATCHED = "run.dispatched"
+    RUN_COMPLETED = "run.completed"
+    RUN_FAILED = "run.failed"
+    HANDLER_STARTED = "handler.started"
+    HANDLER_COMPLETED = "handler.completed"
+    HANDLER_FAILED = "handler.failed"
+    PARTIAL_RESULT_FIRST_SENT = "partial_result.first_sent"
+    PARTIAL_RESULT_FINAL_SENT = "partial_result.final_sent"
+
+
+def default_wal_dir_path() -> pathlib.Path:
+    venv_dir_path = pathlib.Path(sys.executable).parent.parent
+    return venv_dir_path / "state" / "finecode" / "wal" / "er"
+
+
+def _serialize_payload(payload: typing.Any | None) -> dict[str, typing.Any]:
+    if payload is None:
+        return {}
+    if dataclasses.is_dataclass(payload):
+        return dataclasses.asdict(payload)
+    if isinstance(payload, dict):
+        return payload
+    return {"value": str(payload)}
+
+
+@dataclasses.dataclass
+class ErWalConfig:
+    dir_path: pathlib.Path | None = None
+    max_segment_bytes: int = shared_wal.DEFAULT_MAX_SEGMENT_BYTES
+    max_segments: int = shared_wal.DEFAULT_MAX_SEGMENTS
+
+
+class ErWalWriter:
+    def __init__(self, config: ErWalConfig | None = None) -> None:
+        effective_config = config or ErWalConfig()
+        self._dir_path = effective_config.dir_path or default_wal_dir_path()
+        self._writer = shared_wal.WalWriter(
+            shared_wal.WalConfig(
+                enabled=True,
+                dir_path=self._dir_path,
+                max_segment_bytes=effective_config.max_segment_bytes,
+                max_segments=effective_config.max_segments,
+                writer_id_prefix="er",
+            )
+        )
+
+    def append(
+        self,
+        *,
+        event_type: ErWalEventType | str,
+        wal_run_id: str,
+        action_name: str,
+        project_path: str,
+        trigger: str,
+        dev_env: str,
+        payload: typing.Any | None = None,
+    ) -> None:
+        serialized_event_type = (
+            event_type.value if isinstance(event_type, ErWalEventType) else event_type
+        )
+        self._writer.append(
+            event_type=serialized_event_type,
+            wal_run_id=wal_run_id,
+            action_name=action_name,
+            project_path=project_path,
+            trigger=trigger,
+            dev_env=dev_env,
+            payload=_serialize_payload(payload),
+        )
+
+    def close(self) -> None:
+        self._writer.close()
+
+
+def emit_run_event(
+    wal_writer: ErWalWriter | None,
+    *,
+    event_type: ErWalEventType,
+    wal_run_id: str,
+    action_name: str,
+    project_path: pathlib.Path | str,
+    trigger: str,
+    dev_env: str,
+    payload: typing.Any | None = None,
+) -> None:
+    if wal_writer is None:
+        return
+    wal_writer.append(
+        event_type=event_type,
+        wal_run_id=wal_run_id,
+        action_name=action_name,
+        project_path=str(project_path),
+        trigger=trigger,
+        dev_env=dev_env,
+        payload=payload,
+    )

finecode_extension_runner/src/finecode_extension_runner/er_wal.py

@@ -2,6 +2,7 @@
 from typing import Literal
 
 import finecode_extension_runner.context as context
+from finecode_extension_runner import er_wal
 
 runner_context: context.RunnerContext | None = None
 # it's the same as `runner_context.project.dir_path`, but it's available from the start of
@@ -10,3 +11,4 @@
 log_level: Literal["TRACE", "INFO"] = "INFO"
 env_name: str = ""
 log_file_path: Path | None = None
+wal_writer: er_wal.ErWalWriter | None = None

finecode_extension_runner/src/finecode_extension_runner/global_state.py

@@ -23,7 +23,7 @@
 from finecode_extension_api.interfaces import ifileeditor
 from pydantic.dataclasses import dataclass as pydantic_dataclass
 
-from finecode_extension_runner import schemas, services
+from finecode_extension_runner import er_wal, global_state, schemas, services
 from finecode_extension_runner._services import run_action as run_action_service
 from finecode_extension_runner._services import merge_results as merge_results_service
 from finecode_extension_runner.di import resolver
@@ -318,6 +318,8 @@ def create_lsp_server() -> lsp_server.LanguageServer:
 
     def on_process_exit():
         logger.info("Exit extension runner")
+        if global_state.wal_writer is not None:
+            global_state.wal_writer.close()
         services.shutdown_all_action_handlers()
         services.exit_all_action_handlers()
 
@@ -353,6 +355,8 @@ def _on_initialized(ls: CustomLanguageServer, params: types.InitializedParams):
 
 def _on_shutdown(ls: CustomLanguageServer, params):
     logger.info("Shutdown extension runner")
+    if global_state.wal_writer is not None:
+        global_state.wal_writer.close()
     services.shutdown_all_action_handlers()
 
     logger.debug("Stop Finecode async tasks")
@@ -367,6 +371,8 @@ def _on_shutdown(ls: CustomLanguageServer, params):
 
 def _on_exit(ls: lsp_server.LanguageServer, params):
     logger.info("Exit extension runner")
+    if global_state.wal_writer is not None:
+        global_state.wal_writer.close()
 
 
 def uri_to_path(uri: str) -> pathlib.Path:
@@ -537,6 +543,28 @@ async def run_action(
     options: dict[str, typing.Any] | None,
 ):
     logger.trace(f"Run action: {action_name}")
+    wal_run_id = (options or {}).get("wal_run_id")
+    if not isinstance(wal_run_id, str) or wal_run_id.strip() == "":
+        return {"error": "Missing required wal_run_id in run options"}
+
+    meta = (options or {}).get("meta") or {}
+    trigger = meta.get("trigger", "unknown")
+    dev_env = meta.get("dev_env", "unknown")
+    project_path = global_state.project_dir_path or pathlib.Path(".")
+    er_wal.emit_run_event(
+        global_state.wal_writer,
+        event_type=er_wal.ErWalEventType.RUN_ACCEPTED,
+        wal_run_id=wal_run_id,
+        action_name=action_name,
+        project_path=project_path,
+        trigger=trigger,
+        dev_env=dev_env,
+        payload={
+            "partial_result_token": (options or {}).get("partial_result_token"),
+            "progress_token": (options or {}).get("progress_token"),
+        },
+    )
+
     request = schemas.RunActionRequest(action_name=action_name, params=params)
 
     # use pydantic dataclass to convert dict to dataclass instance recursively
@@ -552,6 +580,16 @@ async def run_action(
         if isinstance(exception, services.StopWithResponse):
             status = "stopped"
             response = exception.response
+            er_wal.emit_run_event(
+                global_state.wal_writer,
+                event_type=er_wal.ErWalEventType.RUN_FAILED,
+                wal_run_id=wal_run_id,
+                action_name=action_name,
+                project_path=project_path,
+                trigger=trigger,
+                dev_env=dev_env,
+                payload={"error": "stopped"},
+            )
         else:
             error_msg = ""
             if isinstance(exception, services.ActionFailedException):
@@ -561,6 +599,16 @@ async def run_action(
                 logger.error("Unhandled exception in action run:")
                 logger.exception(exception)
                 error_msg = f"{type(exception)}: {str(exception)}"
+            er_wal.emit_run_event(
+                global_state.wal_writer,
+                event_type=er_wal.ErWalEventType.RUN_FAILED,
+                wal_run_id=wal_run_id,
+                action_name=action_name,
+                project_path=project_path,
+                trigger=trigger,
+                dev_env=dev_env,
+                payload={"error": error_msg},
+            )
             return {"error": error_msg}
 
     # dict key can be path, but pygls fails to handle slashes in dict keys, use strings
@@ -574,6 +622,16 @@ async def run_action(
         for fmt, result in result_by_format.items()
     }
     result_str = json.dumps(converted_result_by_format, cls=CustomJSONEncoder)
+    er_wal.emit_run_event(
+        global_state.wal_writer,
+        event_type=er_wal.ErWalEventType.RUN_COMPLETED,
+        wal_run_id=wal_run_id,
+        action_name=action_name,
+        project_path=project_path,
+        trigger=trigger,
+        dev_env=dev_env,
+        payload={"status": status, "return_code": response.return_code},
+    )
     return {
         "status": status,
         "resultByFormat": result_str,

finecode_extension_runner/src/finecode_extension_runner/lsp_server.py

@@ -58,6 +58,7 @@ class RunActionRequest(BaseSchema):
 
 @dataclass
 class RunActionOptions(BaseSchema):
+    wal_run_id: str
     meta: code_action.RunActionMeta
     partial_result_token: int | str | None = None
     progress_token: int | str | None = None