From 5600c59b666c02d9e584173d6e83301472026e9c Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 11:42:57 +0800 Subject: [PATCH 1/5] perf(workflow): remove synchronous step storage waits --- flocks/ingest/kafka/manager.py | 30 ++-- flocks/ingest/syslog/manager.py | 29 ++-- flocks/server/routes/workflow.py | 101 ++---------- flocks/workflow/execution_store.py | 151 +++++------------ flocks/workflow/poller_manager.py | 19 ++- flocks/workflow/store.py | 136 +++++++++++++++- flocks/workflow/triggers/runtime.py | 30 +++- tests/ingest/test_kafka_manager.py | 47 +++--- .../test_syslog_manager_backpressure.py | 27 ++-- .../server/routes/test_workflow_run_route.py | 95 +++++++++-- .../workflow/test_execution_store_compact.py | 145 +++++++++++++++-- tests/workflow/test_poller_manager.py | 31 ++-- tests/workflow/test_trigger_runtime.py | 30 +++- tests/workflow/test_workflow_store.py | 153 +++++++++++++++++- 14 files changed, 703 insertions(+), 321 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 1b752be8a..893156c43 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -91,6 +91,7 @@ def _worker_count_for_trigger(trigger: TriggerDefinition) -> int: def _queue_size_for_trigger(trigger: TriggerDefinition) -> int: return min(_MAX_QUEUE_SIZE, max(1, int(trigger.concurrency.queueSize))) + _KAFKA_STORAGE_LIST_KEYS = DEFAULT_LARGE_LIST_KEYS | frozenset( { "duplicate_alerts", @@ -446,10 +447,13 @@ async def restart_workflow( err = "workflow_not_found" if startup: self._status[workflow_id] = {"state": "stopped", "error": err} - log.info("kafka.workflow_not_found_on_start", { - "workflow_id": workflow_id, - "action": "stale_config_skipped", - }) + log.info( + "kafka.workflow_not_found_on_start", + { + "workflow_id": workflow_id, + "action": "stale_config_skipped", + }, + ) return {"state": "stopped", "error": err} self._status[workflow_id] = {"state": "failed", "error": err} log.warning("kafka.workflow_not_found", {"workflow_id": workflow_id}) @@ -689,9 +693,7 @@ async def _worker_loop( generation_cancel_event: Optional[threading.Event] = None, ) -> None: run_cancel_event = ( - generation_cancel_event - or self._generation_cancel_events.get(workflow_id) - or threading.Event() + generation_cancel_event or self._generation_cancel_events.get(workflow_id) or threading.Event() ) while not abort.is_set(): try: @@ -765,17 +767,15 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, + persist=False, ) exec_id = exec_data["id"] - loop = asyncio.get_running_loop() start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) trigger_input_keys = list((trigger.mapping or {}).keys()) or [input_key] step_recorder = ExecutionStepRecorder( exec_id=exec_id, - loop=loop, - logger=log, - log_event="kafka.execution_step.write_failed", + capture_steps=False, step_compactor=lambda step: _compact_step_for_kafka_storage( step, input_key=input_key, @@ -845,9 +845,15 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: } ) finally: + steps = step_recorder.take_steps() await cleanup_workflow_tool_context(tool_context) try: - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=steps, + ) except Exception as exc: log.warning("kafka.exec_record_failed", {"exec_id": exec_id, "error": str(exc)}) return exec_data diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 0c1439b94..699fa53cb 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -335,10 +335,13 @@ async def restart_workflow( err = "workflow_not_found" if startup: self._listener_status[workflow_id] = {"state": "stopped", "error": err} - log.info("syslog.workflow_not_found_on_start", { - "workflow_id": workflow_id, - "action": "stale_config_skipped", - }) + log.info( + "syslog.workflow_not_found_on_start", + { + "workflow_id": workflow_id, + "action": "stale_config_skipped", + }, + ) return {"state": "stopped", "error": err} self._listener_status[workflow_id] = {"state": "failed", "error": err} log.warning("syslog.workflow_not_found", {"workflow_id": workflow_id}) @@ -553,9 +556,7 @@ async def _worker_loop( of in-flight workflow runs is exactly ``_MAX_CONCURRENT_EXECUTIONS``. """ run_cancel_event = ( - generation_cancel_event - or self._generation_cancel_events.get(workflow_id) - or threading.Event() + generation_cancel_event or self._generation_cancel_events.get(workflow_id) or threading.Event() ) while not abort.is_set(): try: @@ -618,14 +619,12 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, + persist=False, ) exec_id = exec_data["id"] - loop = asyncio.get_running_loop() step_recorder = ExecutionStepRecorder( exec_id=exec_id, - loop=loop, - logger=log, - log_event="syslog.execution_step.write_failed", + capture_steps=False, ) start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) @@ -692,9 +691,15 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: } ) finally: + steps = step_recorder.take_steps() await cleanup_workflow_tool_context(tool_context) try: - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=steps, + ) except Exception as exc: log.warning("syslog.exec_record_failed", {"exec_id": exec_id, "error": str(exc)}) return exec_data diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index b838c543a..928b6a900 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -57,7 +57,6 @@ derive_loop_progress, load_execution_steps, normalize_execution_status as _normalize_execution_status, - record_execution_step, record_execution_result as _record_execution_result, resolve_execution_outcome as _resolve_execution_outcome, workflow_execution_key as _workflow_execution_key, @@ -99,7 +98,6 @@ webhook_router = APIRouter() log = Log.create(service="workflow-routes") -_PROGRESS_FLUSH_EVERY_STEPS = 5 _WORKFLOW_LIST_ENRICH_CONCURRENCY = 8 _WORKFLOW_API_HEALTH_INTERVAL_S = 5.0 _WORKFLOW_API_HEALTH_PROBE_CONCURRENCY = 4 @@ -1127,7 +1125,6 @@ async def _run_workflow_execution_task( """Execute a workflow in the background and keep the execution record updated.""" start_time = time.time() step_count = 0 - loop = asyncio.get_running_loop() pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None execution_summary: Dict[str, Any] = { @@ -1150,21 +1147,6 @@ async def _run_workflow_execution_task( } ) - def _write_progress(update_fields: Dict[str, Any]) -> None: - try: - execution_summary.update(update_fields) - asyncio.run_coroutine_threadsafe( - WorkflowStore.upsert_execution(compact_execution_summary(execution_summary)), loop - ).result(timeout=5) - except Exception as exc: - log.warning( - "workflow.step_progress.write_failed", - { - "exec_id": exec_id, - "error": str(exc), - }, - ) - def _on_step_start(_run_id, step_index, node, _inputs): nonlocal pending_step_index, pending_step node_id = getattr(node, "id", None) @@ -1185,7 +1167,7 @@ def _on_step_start(_run_id, step_index, node, _inputs): "error": "Run cancelled before node completed", } ) - _write_progress( + execution_summary.update( { "currentNodeId": node_id, "currentNodeType": node_type, @@ -1220,47 +1202,6 @@ def _on_step_complete(step_result) -> None: "updatedAt": int(time.time() * 1000), } ) - try: - asyncio.run_coroutine_threadsafe( - record_execution_step(exec_id, step_count, step_dict), - loop, - ).result(timeout=5) - except Exception as exc: - log.warning( - "workflow.execution_step.write_failed", - { - "exec_id": exec_id, - "step_index": step_count, - "error": str(exc), - }, - ) - if step_count % _PROGRESS_FLUSH_EVERY_STEPS == 0: - _write_progress( - { - "stepCount": step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) - - async def _flush_pending_step() -> None: - if pending_step_index is None or pending_step is None: - return - try: - await record_execution_step(exec_id, pending_step_index, pending_step) - except Exception as exc: - log.warning( - "workflow.pending_step.write_failed", - { - "exec_id": exec_id, - "step_index": pending_step_index, - "error": str(exc), - }, - ) try: result: RunWorkflowResult = await asyncio.to_thread( @@ -1284,10 +1225,9 @@ async def _flush_pending_step() -> None: # ``record_execution_result`` backfills this compacted history into # append-only step rows, then stores only the summary row. final_history = compact_history_for_storage(result.history) - if status_value == "cancelled" and not final_history: - await _flush_pending_step() final_steps = result.steps - if pending_step_index is not None: + if pending_step_index is not None and pending_step is not None: + final_history.append(pending_step) final_steps = max(final_steps, pending_step_index) current_data.update( { @@ -1319,15 +1259,18 @@ async def _flush_pending_step() -> None: except Exception as exc: duration = time.time() - start_time current_data = dict(execution_summary) + final_history = [pending_step] if pending_step is not None else [] + final_steps = max(step_count, pending_step_index or 0) current_data.update( { "status": "cancelled" if cancel_event.is_set() else "error", "finishedAt": int(time.time() * 1000), "duration": duration, "errorMessage": str(exc), - "executionLog": [], - "stepCount": step_count, + "executionLog": final_history, + "stepCount": final_steps, "currentPhase": "cancelled" if cancel_event.is_set() else "error", + "currentStepIndex": final_steps, "updatedAt": int(time.time() * 1000), } ) @@ -1521,19 +1464,13 @@ async def create_workflow(req: WorkflowCreateRequest): if strict_mapping_errors: raise HTTPException( status_code=400, - detail=( - "Workflow strict edge mapping failed: " - f"{strict_mapping_errors[:5]}" - ), + detail=(f"Workflow strict edge mapping failed: {strict_mapping_errors[:5]}"), ) schema_errors = _schema_lint_errors(workflow_model) if schema_errors: raise HTTPException( status_code=400, - detail=( - "Workflow schema lint failed: " - f"{schema_errors[:5]}" - ), + detail=(f"Workflow schema lint failed: {schema_errors[:5]}"), ) workflow_id = str(uuid.uuid4()) @@ -1632,19 +1569,13 @@ async def update_workflow(workflow_id: str, req: WorkflowUpdateRequest): if strict_mapping_errors: raise HTTPException( status_code=400, - detail=( - "Workflow strict edge mapping failed: " - f"{strict_mapping_errors[:5]}" - ), + detail=(f"Workflow strict edge mapping failed: {strict_mapping_errors[:5]}"), ) schema_errors = _schema_lint_errors(workflow_model) if schema_errors: raise HTTPException( status_code=400, - detail=( - "Workflow schema lint failed: " - f"{schema_errors[:5]}" - ), + detail=(f"Workflow schema lint failed: {schema_errors[:5]}"), ) workflow_json = req.workflow_json except Exception as e: @@ -2455,9 +2386,7 @@ async def refresh_workflow_api_health_cache() -> Dict[str, int]: active_workflow_ids = [ str(service.get("workflowId") or _workflow_id_from_api_service_key(key)) for key, service in zip(keys, services) - if isinstance(service, dict) - and service - and _summarize_capability_state(service.get("status")) == "running" + if isinstance(service, dict) and service and _summarize_capability_state(service.get("status")) == "running" ] semaphore = asyncio.Semaphore(_WORKFLOW_API_HEALTH_PROBE_CONCURRENCY) @@ -2542,9 +2471,7 @@ async def _get_workflow_integration_status( set_workflow_json_triggers(workflow_data.get("workflowJson") or {}, triggers), ) statuses_by_id = { - item.get("triggerId"): item - for item in statuses - if isinstance(item, dict) and item.get("triggerId") + item.get("triggerId"): item for item in statuses if isinstance(item, dict) and item.get("triggerId") } trigger_items: List[WorkflowTriggerStatusItemResponse] = [] for trigger in triggers: diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index 9d0ae1c79..da7c189c2 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -342,24 +342,6 @@ def derive_loop_progress( # for the same workflow serialize instead of skipping cleanup. _trim_locks: Dict[str, asyncio.Lock] = {} -# Per-workflow lock to serialize read-modify-write of stats. Concurrent -# executions of the same workflow (e.g. syslog-triggered runs with -# semaphore=8) would otherwise race on ``Storage.read → mutate → write`` -# and silently lose counter increments. -_stats_locks: Dict[str, asyncio.Lock] = {} - - -def _get_stats_lock(workflow_id: str) -> asyncio.Lock: - lock = _stats_locks.get(workflow_id) - if lock is None: - lock = asyncio.Lock() - _stats_locks[workflow_id] = lock - return lock - - -def _workflow_stats_key(workflow_id: str) -> str: - return f"workflow/{workflow_id}/stats" - def _get_trim_lock(workflow_id: str) -> asyncio.Lock: lock = _trim_locks.get(workflow_id) @@ -369,37 +351,6 @@ def _get_trim_lock(workflow_id: str) -> asyncio.Lock: return lock -_DEFAULT_STATS: Dict[str, Any] = { - "callCount": 0, - "successCount": 0, - "errorCount": 0, - "totalRuntime": 0.0, - "avgRuntime": 0.0, - "thumbsUp": 0, - "thumbsDown": 0, -} - - -async def _update_workflow_stats(workflow_id: str, success: bool, duration: float) -> None: - """Increment workflow call/success/error counters and update avgRuntime. - - Serialised per workflow to keep concurrent updates from clobbering each - other (read → mutate → write race). - """ - lock = _get_stats_lock(workflow_id) - async with lock: - try: - await WorkflowStore.increment_stats(workflow_id, success=success, duration=duration) - except Exception as exc: - log.warning( - "workflow.stats.update_failed", - { - "workflow_id": workflow_id, - "error": str(exc), - }, - ) - - def workflow_execution_key(exec_id: str) -> str: """Return the storage key for one workflow execution.""" return f"workflow_execution/{exec_id}" @@ -453,26 +404,21 @@ async def record_execution_step( class ExecutionStepRecorder: - """Bridge synchronous workflow step callbacks to append-only step rows.""" + """Collect compact workflow steps without blocking the runner thread.""" def __init__( self, *, exec_id: str, - loop: asyncio.AbstractEventLoop, - logger: Any = None, - log_event: str = "workflow.execution_step.write_failed", + capture_steps: bool = True, step_compactor: Callable[[Any], Dict[str, Any]] = compact_step_for_storage, - write_timeout_s: float = 5.0, ) -> None: self.exec_id = exec_id - self.loop = loop - self.logger = logger or log - self.log_event = log_event + self.capture_steps = capture_steps self.step_compactor = step_compactor - self.write_timeout_s = write_timeout_s self.step_count = 0 self.summary: Dict[str, Any] = {} + self._pending_steps: List[Tuple[int, Dict[str, Any]]] = [] def on_step_complete(self, step_result: Any) -> None: raw_step = step_result.model_dump(mode="json") if hasattr(step_result, "model_dump") else step_result @@ -498,48 +444,25 @@ def on_step_complete(self, step_result: Any) -> None: "updatedAt": int(time.time() * 1000), } ) - try: - asyncio.run_coroutine_threadsafe( - record_execution_step(self.exec_id, self.step_count, step_dict), - self.loop, - ).result(timeout=self.write_timeout_s) - except Exception as exc: - self.logger.warning( - self.log_event, - { - "exec_id": self.exec_id, - "step_index": self.step_count, - "error": str(exc), - }, - ) + if self.capture_steps: + self._pending_steps.append((self.step_count, step_dict)) + def take_steps(self) -> List[Tuple[int, Dict[str, Any]]]: + """Return buffered steps for the final execution transaction.""" + pending_steps = self._pending_steps + self._pending_steps = [] + return pending_steps -async def _backfill_execution_steps( - exec_id: str, - execution_log: Any, -) -> int: - """Persist legacy inline executionLog entries as append-only step rows.""" - if not isinstance(execution_log, list): - return 0 - written = 0 - for step_index, step in enumerate(execution_log, start=1): - step_payload = compact_step_for_storage(step) - if not isinstance(step_payload, dict): - continue - try: - await WorkflowStore.record_step(exec_id, step_index, step_payload) - written += 1 - except Exception as exc: - log.warning( - "workflow.execution_step.backfill_failed", - { - "exec_id": exec_id, - "step_index": step_index, - "error": str(exc), - }, - ) - return written +def _prepare_execution_steps(execution_log: Any) -> List[Tuple[int, Dict[str, Any]]]: + """Compact an inline execution log for one final batch transaction.""" + if not isinstance(execution_log, list): + return [] + return [ + (step_index, compact_step_for_storage(step)) + for step_index, step in enumerate(execution_log, start=1) + if isinstance(step, dict) + ] async def load_execution_steps( @@ -621,8 +544,9 @@ async def create_execution_record( *, input_params: Optional[Dict[str, Any]] = None, exec_id: Optional[str] = None, + persist: bool = True, ) -> Dict[str, Any]: - """Create and persist a running workflow execution record. + """Build a running workflow execution record and optionally persist it. *input_params* is passed through ``compact_outputs_for_storage`` before writing to SQLite so that batch HTTP calls whose inputs contain a key in @@ -637,7 +561,8 @@ async def create_execution_record( input_params=compacted_params, exec_id=exec_id, ) - await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) + if persist: + await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) return exec_data @@ -645,18 +570,23 @@ async def record_execution_result( workflow_id: str, exec_id: str, exec_data: Dict[str, Any], + *, + steps: Optional[Iterable[Tuple[int, Dict[str, Any]]]] = None, ) -> None: - """Persist the final execution record, audit trail, and workflow stats.""" + """Persist the final execution record, step batch, audit trail, and stats.""" summary_data = dict(exec_data) - backfilled_steps = await _backfill_execution_steps(exec_id, summary_data.get("executionLog")) + prepared_steps = ( + list(steps) + if steps is not None + else _prepare_execution_steps(summary_data.get("executionLog")) + ) + persisted_step_count = len(prepared_steps) existing_step_count = _as_positive_int(summary_data.get("stepCount")) - if backfilled_steps and (existing_step_count is None or existing_step_count < backfilled_steps): - summary_data["stepCount"] = backfilled_steps - - await WorkflowStore.upsert_execution(compact_execution_summary(summary_data)) + if persisted_step_count and ( + existing_step_count is None or existing_step_count < persisted_step_count + ): + summary_data["stepCount"] = persisted_step_count - # Update call/success/error counters so all trigger paths (HTTP, syslog, etc.) - # are reflected in the UI stats panel. status = summary_data.get("status", "error") success = status == "success" duration = summary_data.get("duration") @@ -664,7 +594,12 @@ async def record_execution_result( started_at = summary_data.get("startedAt", 0) finished_at = summary_data.get("finishedAt", int(time.time() * 1000)) duration = max(0.0, (finished_at - started_at) / 1000.0) - await _update_workflow_stats(workflow_id, success, float(duration)) + await WorkflowStore.complete_execution( + compact_execution_summary(summary_data), + prepared_steps, + success=success, + duration=float(duration), + ) # Recorder writes to its own SQLite tables and can be slow under load. # Run it as a background task so the syslog/HTTP dispatcher can release the diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index 9db0bdd4b..d596daa80 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -448,14 +448,15 @@ async def _execute_run( cancel_events = self._run_cancel_events.setdefault(workflow_id, set()) cancel_events.add(cancel_event) inputs = self._build_inputs(config) - exec_data = await create_execution_record(workflow_id, input_params=inputs) + exec_data = await create_execution_record( + workflow_id, + input_params=inputs, + persist=False, + ) exec_id = str(exec_data["id"]) - loop = asyncio.get_running_loop() step_recorder = ExecutionStepRecorder( exec_id=exec_id, - loop=loop, - logger=log, - log_event="poller.execution_step.write_failed", + capture_steps=False, ) current = self._status.get(workflow_id) or self._base_status(workflow_id) current["lastRunAt"] = started_at_ms @@ -555,9 +556,15 @@ async def _execute_run( self._status[workflow_id] = current log.warning("poller.run_failed", {"workflow_id": workflow_id, "error": str(exc)}) finally: + steps = step_recorder.take_steps() await cleanup_workflow_tool_context(tool_context) try: - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=steps, + ) except Exception as exc: log.warning( "poller.exec_record_failed", diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index 3c238c236..4a74e6401 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -8,7 +8,7 @@ import sqlite3 from datetime import UTC, datetime from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, Iterable, List, Optional, Tuple import aiosqlite @@ -46,6 +46,7 @@ class WorkflowStore: _conn: Optional[aiosqlite.Connection] = None _init_pid: Optional[int] = None _db_path: Optional[Path] = None + _completion_lock: Optional[asyncio.Lock] = None @classmethod def get_db_path(cls) -> Path: @@ -93,6 +94,7 @@ async def _open_and_migrate() -> None: cls._initialized = True cls._init_pid = current_pid cls._db_path = db_path + cls._completion_lock = asyncio.Lock() await cls._migrate_legacy_kv() try: @@ -124,6 +126,7 @@ async def close(cls) -> None: cls._initialized = False cls._init_pid = None cls._db_path = None + cls._completion_lock = None @classmethod async def _db(cls) -> aiosqlite.Connection: @@ -415,13 +418,58 @@ async def record_step( step_index: int, step_payload: Dict[str, Any], ) -> None: + await cls.record_steps(exec_id, [(step_index, step_payload)]) + + @classmethod + async def record_steps( + cls, + exec_id: str, + steps: Iterable[Tuple[int, Dict[str, Any]]], + ) -> None: + rows = [ + ( + exec_id, + int(step_index), + step_payload.get("node_id"), + step_payload.get("node_type") or step_payload.get("type"), + cls._json_dumps(step_payload.get("inputs") or {}), + cls._json_dumps(step_payload.get("outputs") or {}), + step_payload.get("error"), + cls._json_dumps(step_payload), + ) + for step_index, step_payload in steps + ] + if not rows: + return db = await cls._db() - await db.execute( + await db.executemany( """ INSERT OR REPLACE INTO workflow_execution_steps (exec_id, step_index, node_id, node_type, inputs, outputs, error, payload) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, + rows, + ) + await db.commit() + + @classmethod + async def complete_execution( + cls, + exec_data: Dict[str, Any], + steps: Iterable[Tuple[int, Dict[str, Any]]], + *, + success: bool, + duration: float, + ) -> None: + """Persist one completed execution and its stats in one transaction.""" + db = await cls._db() + payload = dict(exec_data) + exec_id = str(payload.get("id") or "") + workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") + if not exec_id or not workflow_id: + raise ValueError("workflow execution requires id and workflowId") + + step_rows = [ ( exec_id, int(step_index), @@ -431,9 +479,87 @@ async def record_step( cls._json_dumps(step_payload.get("outputs") or {}), step_payload.get("error"), cls._json_dumps(step_payload), - ), - ) - await db.commit() + ) + for step_index, step_payload in steps + ] + runtime = float(duration) + success_delta = 1 if success else 0 + error_delta = 0 if success else 1 + lock = cls._completion_lock + if lock is None: + lock = asyncio.Lock() + cls._completion_lock = lock + + async with lock: + try: + if step_rows: + await db.executemany( + """ + INSERT OR REPLACE INTO workflow_execution_steps + (exec_id, step_index, node_id, node_type, inputs, outputs, error, payload) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + step_rows, + ) + await db.execute( + """ + INSERT OR REPLACE INTO workflow_executions + (id, workflow_id, status, current_phase, current_node_id, current_node_type, + current_step_index, step_count, input_params, output_results, error_message, + trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + exec_id, + workflow_id, + str(payload.get("status") or "running"), + payload.get("currentPhase"), + payload.get("currentNodeId"), + payload.get("currentNodeType"), + cls._as_int(payload.get("currentStepIndex")), + cls._as_int(payload.get("stepCount")) or 0, + cls._json_dumps(payload.get("inputParams") or {}), + cls._json_dumps(payload.get("outputResults") or {}), + payload.get("errorMessage"), + payload.get("triggerId"), + payload.get("triggerType"), + cls._as_int(payload.get("startedAt")) or cls._now_ms(), + cls._as_int(payload.get("finishedAt")), + cls._as_float(payload.get("duration")), + cls._as_int(payload.get("updatedAt")) or cls._now_ms(), + cls._json_dumps(payload), + ), + ) + await db.execute( + """ + INSERT INTO workflow_stats ( + workflow_id, call_count, success_count, error_count, + total_runtime, avg_runtime, thumbs_up, thumbs_down, updated_at + ) + VALUES (?, 1, ?, ?, ?, ?, 0, 0, ?) + ON CONFLICT(workflow_id) DO UPDATE SET + call_count = workflow_stats.call_count + 1, + success_count = workflow_stats.success_count + excluded.success_count, + error_count = workflow_stats.error_count + excluded.error_count, + total_runtime = workflow_stats.total_runtime + excluded.total_runtime, + avg_runtime = ( + workflow_stats.total_runtime + excluded.total_runtime + ) / (workflow_stats.call_count + 1), + updated_at = excluded.updated_at + """, + ( + workflow_id, + success_delta, + error_delta, + runtime, + runtime, + cls._now_ms(), + ), + ) + await db.commit() + except Exception: + await db.rollback() + raise @classmethod async def list_steps( diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index 5e1a5b3da..2a3d2aedf 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -11,7 +11,7 @@ from flocks.hooks.pipeline import HookPipeline from flocks.utils.log import Log from flocks.workflow.execution_store import ( - compact_history_for_storage, + ExecutionStepRecorder, compact_outputs_for_storage, create_execution_record, record_execution_result, @@ -248,8 +248,13 @@ async def _execute_workflow_effect( exec_data = await create_execution_record( workflow_id, input_params=mapped_inputs, + persist=False, ) exec_id = exec_data["id"] + step_recorder = ExecutionStepRecorder( + exec_id=exec_id, + capture_steps=False, + ) started_at = time.time() tool_context = None try: @@ -266,10 +271,15 @@ async def _execute_workflow_effect( run_workflow, workflow=workflow_json, inputs=mapped_inputs, + run_id=exec_id, trace=False, + execution_profile="high_frequency", + on_step_complete=step_recorder.on_step_complete, tool_context=tool_context, ) status_value, error_message = resolve_execution_outcome(result) + step_count = step_recorder.step_count or result.steps + exec_data.update(step_recorder.summary) exec_data.update( { "status": status_value, @@ -277,10 +287,11 @@ async def _execute_workflow_effect( "finishedAt": _now_ms(), "duration": time.time() - started_at, "errorMessage": error_message, - "executionLog": compact_history_for_storage(result.history), + "executionLog": [], + "stepCount": step_count, "currentNodeId": result.last_node_id, "currentPhase": status_value, - "currentStepIndex": result.steps, + "currentStepIndex": step_count, "triggerId": trigger.id, "triggerType": trigger.type, "deliveryId": mapped_inputs.get("_flocks", {}).get("trigger", {}).get("deliveryId"), @@ -289,12 +300,18 @@ async def _execute_workflow_effect( } ) except Exception as exc: + step_count = step_recorder.step_count + exec_data.update(step_recorder.summary) exec_data.update( { "status": "error", "finishedAt": _now_ms(), "duration": time.time() - started_at, "errorMessage": str(exc), + "executionLog": [], + "stepCount": step_count, + "currentPhase": "error", + "currentStepIndex": step_count, "triggerId": trigger.id, "triggerType": trigger.type, "deliveryId": mapped_inputs.get("_flocks", {}).get("trigger", {}).get("deliveryId"), @@ -304,7 +321,12 @@ async def _execute_workflow_effect( ) finally: await cleanup_workflow_tool_context(tool_context) - await record_execution_result(workflow_id, exec_id, exec_data) + await record_execution_result( + workflow_id, + exec_id, + exec_data, + steps=step_recorder.take_steps(), + ) return exec_data async def dispatch_event( diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index a3c2b0432..6e494b566 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -24,7 +24,6 @@ import pytest from flocks.ingest.kafka import manager as kafka_manager -from flocks.workflow import execution_store from flocks.workflow.triggers.models import TriggerDefinition @@ -241,10 +240,7 @@ def test_trigger_concurrency_config_is_honored_with_safety_caps() -> None: "concurrency": {"maxParallel": 999, "queueSize": 999_999}, } ) - assert ( - kafka_manager._worker_count_for_trigger(oversized) - == kafka_manager._MAX_CONCURRENT_EXECUTIONS - ) + assert kafka_manager._worker_count_for_trigger(oversized) == kafka_manager._MAX_CONCURRENT_EXECUTIONS assert kafka_manager._queue_size_for_trigger(oversized) == kafka_manager._MAX_QUEUE_SIZE @@ -548,18 +544,20 @@ async def test_trigger_workflow_compacts_kafka_execution_record( captured_input_params: dict = {} captured_exec_data: dict = {} captured_run_kwargs: dict = {} - recorded_steps: list[tuple[str, int, dict]] = [] + captured_steps: list[tuple[int, dict]] = [] - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False captured_input_params.update(input_params or {}) return {"id": "exec-compact", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): captured_exec_data.update(exec_data) - - async def _fake_record_execution_step(exec_id, step_index, step): # noqa: ANN001 - recorded_steps.append((exec_id, step_index, step)) - return step + captured_steps.extend(steps or []) def _fake_run_workflow(**kwargs): # noqa: ANN003 captured_run_kwargs.update(kwargs) @@ -597,7 +595,6 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 monkeypatch.setattr(kafka_manager, "create_execution_record", _fake_create_execution_record) monkeypatch.setattr(kafka_manager, "record_execution_result", _fake_record_execution_result) monkeypatch.setattr(kafka_manager, "run_workflow", _fake_run_workflow) - monkeypatch.setattr(execution_store, "record_execution_step", _fake_record_execution_step) await manager._trigger_workflow( "wf-compact", @@ -619,11 +616,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 } assert captured_exec_data["executionLog"] == [] assert captured_exec_data["stepCount"] == 2 - assert recorded_steps[0][0] == "exec-compact" - assert recorded_steps[0][1] == 1 - assert recorded_steps[0][2]["outputs"] == {"_raw_alerts_count": 1} - assert recorded_steps[1][1] == 2 - assert recorded_steps[1][2]["inputs"] == {"_filtered_alerts_count": 1} + assert captured_steps == [] assert len(json.dumps(captured_exec_data, ensure_ascii=False)) < 10_000 @@ -635,11 +628,16 @@ async def test_trigger_workflow_merges_configured_inputs_with_consumed_message( captured_run_kwargs: dict = {} recorded_input_params: dict = {} - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False recorded_input_params.update(input_params or {}) return {"id": "exec-merge", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): return None def _fake_run_workflow(**kwargs): # noqa: ANN003 @@ -691,10 +689,15 @@ async def test_trigger_workflow_applies_mapping_and_filter( captured_run_kwargs: dict = {} recorded_exec_data: dict = {} - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False return {"id": "exec-filter", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): recorded_exec_data.update(exec_data) def _fake_run_workflow(**kwargs): # noqa: ANN003 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 1982caf84..3ad929893 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -25,7 +25,6 @@ import pytest from flocks.ingest.syslog import manager as syslog_manager -from flocks.workflow import execution_store from flocks.workflow.triggers.models import TriggerDefinition @@ -148,10 +147,7 @@ def test_trigger_concurrency_config_is_honored_with_safety_caps() -> None: "concurrency": {"maxParallel": 999, "queueSize": 999_999}, } ) - assert ( - syslog_manager._worker_count_for_trigger(oversized) - == syslog_manager._MAX_CONCURRENT_EXECUTIONS - ) + assert syslog_manager._worker_count_for_trigger(oversized) == syslog_manager._MAX_CONCURRENT_EXECUTIONS assert syslog_manager._queue_size_for_trigger(oversized) == syslog_manager._MAX_QUEUE_SIZE @@ -352,17 +348,19 @@ async def test_trigger_workflow_applies_mapping_and_filter( manager = syslog_manager.SyslogManager() captured_run_kwargs: dict = {} recorded_exec_data: dict = {} - recorded_steps: list[tuple[str, int, dict]] = [] + recorded_steps: list[tuple[int, dict]] = [] - async def _fake_create_execution_record(workflow_id, *, input_params=None, exec_id=None): # noqa: ANN001 + async def _fake_create_execution_record( # noqa: ANN001 + workflow_id, *, input_params=None, exec_id=None, persist=True + ): + assert persist is False return {"id": "exec-syslog", "workflowId": workflow_id, "inputParams": input_params} - async def _fake_record_execution_result(workflow_id, exec_id, exec_data): # noqa: ANN001 + async def _fake_record_execution_result( # noqa: ANN001 + workflow_id, exec_id, exec_data, *, steps=None + ): recorded_exec_data.update(exec_data) - - async def _fake_record_execution_step(exec_id, step_index, step): # noqa: ANN001 - recorded_steps.append((exec_id, step_index, step)) - return step + recorded_steps.extend(steps or []) def _fake_run_workflow(**kwargs): # noqa: ANN003 captured_run_kwargs.update(kwargs) @@ -392,7 +390,6 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 monkeypatch.setattr(syslog_manager, "create_execution_record", _fake_create_execution_record) monkeypatch.setattr(syslog_manager, "record_execution_result", _fake_record_execution_result) monkeypatch.setattr(syslog_manager, "run_workflow", _fake_run_workflow) - monkeypatch.setattr(execution_store, "record_execution_step", _fake_record_execution_step) trigger = TriggerDefinition.model_validate( { @@ -430,9 +427,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 action_name="trigger:syslog", ) trigger_tool_context.cleanup.assert_awaited_once_with(trigger_tool_context.context) - assert recorded_steps[0][0] == "exec-syslog" - assert recorded_steps[0][1] == 1 - assert recorded_steps[0][2]["node_id"] == "receive_alert" + assert recorded_steps == [] assert recorded_exec_data["triggerId"] == "syslog-alerts" assert recorded_exec_data["triggerSource"] == "udp://0.0.0.0:5514" assert recorded_exec_data["executionLog"] == [] diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index 588fe227b..5684a5cef 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -104,9 +104,7 @@ async def test_create_workflow_rejects_unmapped_edges_after_strict_default( req = workflow_module.WorkflowCreateRequest( name="new workflow", - workflowJson=_two_node_workflow_json( - {"from": "prepare_message", "to": "transform_message", "order": 0} - ), + workflowJson=_two_node_workflow_json({"from": "prepare_message", "to": "transform_message", "order": 0}), ) with pytest.raises(workflow_module.HTTPException) as exc_info: @@ -207,9 +205,7 @@ async def test_update_workflow_rejects_unmapped_edges_when_strict( req = workflow_module.WorkflowUpdateRequest( workflowJson={ - **_two_node_workflow_json( - {"from": "prepare_message", "to": "transform_message", "order": 0} - ), + **_two_node_workflow_json({"from": "prepare_message", "to": "transform_message", "order": 0}), "metadata": {"runtime": {"strict_edge_mapping": True, "dataflow_mode": "vertex_cache"}}, } ) @@ -228,15 +224,33 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( monkeypatch: pytest.MonkeyPatch, ) -> None: init_mock = AsyncMock() - run_mock = Mock( - return_value=SimpleNamespace( + step_result = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="tool"), + {}, + ) + kwargs["on_step_complete"](step_result) + return SimpleNamespace( outputs={"ok": True}, - history=[], + history=[step_result.model_dump(mode="json")], last_node_id="node-1", steps=1, ) - ) + + run_mock = Mock(side_effect=run_workflow_mock) record_result = AsyncMock(return_value=None) + upsert_execution = AsyncMock(return_value=None) storage_read = AsyncMock( return_value={ "id": "exec-1", @@ -248,6 +262,7 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( monkeypatch.setattr(MCP, "init", init_mock) monkeypatch.setattr(workflow_module, "run_workflow", run_mock) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) monkeypatch.setattr(workflow_module.Storage, "read", storage_read) @@ -270,7 +285,67 @@ async def test_run_workflow_execution_task_reuses_existing_mcp_without_reinit( init_mock.assert_not_awaited() run_mock.assert_called_once() assert run_mock.call_args.kwargs["tool_context"] is tool_context + upsert_execution.assert_not_awaited() record_result.assert_awaited_once() + assert record_result.await_args.args[2]["executionLog"] == [ + { + "node_id": "node-1", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + ] + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_batches_cancelled_pending_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="tool"), + {"message": "hello"}, + ) + return SimpleNamespace( + run_id="run-1", + outputs={}, + history=[], + last_node_id="node-1", + steps=0, + ) + + record_result = AsyncMock(return_value=None) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr(workflow_module, "compact_outputs_for_storage", lambda value: value) + monkeypatch.setattr(workflow_module, "compact_history_for_storage", lambda value: value) + + cancel_event = workflow_module.threading.Event() + cancel_event.set() + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={"message": "hello"}, trace=False), + exec_id="exec-cancelled", + cancel_event=cancel_event, + ) + + record_result.assert_awaited_once() + final_data = record_result.await_args.args[2] + assert final_data["status"] == "cancelled" + assert final_data["stepCount"] == 1 + assert final_data["executionLog"] == [ + { + "node_id": "node-1", + "node_type": "tool", + "inputs": {"message": "hello"}, + "outputs": {}, + "error": "Run cancelled before node completed", + } + ] @pytest.mark.asyncio diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index f5d02a555..de43af374 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -16,6 +16,8 @@ """ from __future__ import annotations + +import asyncio from typing import Any, Dict, List from unittest.mock import AsyncMock, patch @@ -25,11 +27,13 @@ DEFAULT_GENERIC_SEQUENCE_THRESHOLD, DEFAULT_LARGE_LIST_KEYS, DEFAULT_MAX_INLINE_COLLECTION_BYTES, + ExecutionStepRecorder, _trim_execution_history, compact_history_for_storage, compact_execution_summary, compact_outputs_for_storage, compact_step_for_storage, + create_execution_record, record_execution_result, workflow_execution_step_key, ) @@ -300,10 +304,82 @@ def test_workflow_execution_step_key_is_append_only_namespaced() -> None: @pytest.mark.asyncio -async def test_record_execution_result_backfills_execution_log_steps() -> None: - record_step = AsyncMock(return_value=None) +async def test_create_execution_record_can_skip_initial_database_write() -> None: upsert_execution = AsyncMock(return_value=None) - update_stats = AsyncMock(return_value=None) + + with patch.object(WorkflowStore, "upsert_execution", upsert_execution): + record = await create_execution_record( + "wf-trigger", + input_params={"message": "hello"}, + exec_id="exec-trigger", + persist=False, + ) + + assert record["id"] == "exec-trigger" + assert record["currentPhase"] == "queued" + upsert_execution.assert_not_awaited() + + +def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: + record_step = AsyncMock(return_value=None) + record_steps = AsyncMock(return_value=None) + recorder = ExecutionStepRecorder(exec_id="exec-batch") + + with ( + patch.object(WorkflowStore, "record_step", record_step), + patch.object(WorkflowStore, "record_steps", record_steps), + ): + recorder.on_step_complete({"node_id": "n1", "outputs": {"ok": 1}}) + recorder.on_step_complete({"node_id": "n2", "outputs": {"ok": 2}}) + + assert recorder.step_count == 2 + assert recorder.summary["currentNodeId"] == "n2" + assert recorder.take_steps() == [ + (1, {"node_id": "n1", "outputs": {"ok": 1}}), + (2, {"node_id": "n2", "outputs": {"ok": 2}}), + ] + assert recorder.take_steps() == [] + record_step.assert_not_awaited() + record_steps.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_four_trigger_workers_keep_step_history_disabled() -> None: + """Four trigger threads track progress without retaining step history.""" + record_step = AsyncMock(return_value=None) + record_steps = AsyncMock(return_value=None) + recorders = [ + ExecutionStepRecorder( + exec_id=f"exec-trigger-{worker}", + capture_steps=False, + ) + for worker in range(4) + ] + + def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: + for step in range(7): + recorder.on_step_complete( + {"node_id": f"node-{step}", "outputs": {"ok": True}} + ) + + with ( + patch.object(WorkflowStore, "record_step", record_step), + patch.object(WorkflowStore, "record_steps", record_steps), + ): + await asyncio.gather( + *(asyncio.to_thread(_run_seven_steps, recorder) for recorder in recorders) + ) + + batches = [recorder.take_steps() for recorder in recorders] + assert batches == [[], [], [], []] + assert [recorder.step_count for recorder in recorders] == [7, 7, 7, 7] + record_step.assert_not_awaited() + record_steps.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_record_execution_result_backfills_execution_log_steps() -> None: + complete_execution = AsyncMock(return_value=None) exec_data = { "id": "exec-1", "workflowId": "wf", @@ -320,24 +396,67 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 raise RuntimeError with ( - patch.object(WorkflowStore, "record_step", record_step), - patch.object(WorkflowStore, "upsert_execution", upsert_execution), - patch("flocks.workflow.execution_store._update_workflow_stats", update_stats), + patch.object(WorkflowStore, "complete_execution", complete_execution), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), ): await record_execution_result("wf", "exec-1", exec_data) - step_calls = record_step.await_args_list - assert step_calls[0].args[:2] == ("exec-1", 1) - assert step_calls[0].args[2]["outputs"] == {"_raw_alerts_count": 150} - assert step_calls[1].args[:2] == ("exec-1", 2) - assert step_calls[1].args[2]["inputs"] == {"_filtered_alerts_count": 150} - upsert_execution.assert_awaited_once() - summary = upsert_execution.await_args.args[0] + complete_execution.assert_awaited_once() + summary, steps = complete_execution.await_args.args + assert steps[0][0] == 1 + assert steps[0][1]["outputs"] == {"_raw_alerts_count": 150} + assert steps[1][0] == 2 + assert steps[1][1]["inputs"] == {"_filtered_alerts_count": 150} assert summary["executionLog"] == [] assert summary["stepCount"] == 2 + assert complete_execution.await_args.kwargs == { + "success": True, + "duration": 1.0, + } + + +@pytest.mark.asyncio +async def test_record_execution_result_accepts_explicit_step_batch() -> None: + complete_execution = AsyncMock(return_value=None) + explicit_steps = [ + (1, {"node_id": "step-1", "outputs": {"ok": True}}), + (2, {"node_id": "step-2", "outputs": {"ok": True}}), + ] + exec_data = { + "id": "exec-trigger", + "workflowId": "wf-trigger", + "status": "success", + "duration": 0.01, + "executionLog": [], + "stepCount": 2, + } + + def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + with ( + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), + ): + await record_execution_result( + "wf-trigger", + "exec-trigger", + exec_data, + steps=explicit_steps, + ) + + summary, persisted_steps = complete_execution.await_args.args + assert summary["executionLog"] == [] + assert persisted_steps == explicit_steps + assert complete_execution.await_args.kwargs == { + "success": True, + "duration": 0.01, + } def test_compact_history_compacts_each_step_inputs() -> None: diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index 7da394808..ae116df5c 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -8,7 +8,6 @@ import pytest from flocks.workflow import poller_manager -from flocks.workflow import execution_store from flocks.workflow.runner import RunWorkflowResult @@ -100,7 +99,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -141,7 +140,7 @@ async def test_run_once_records_execution_and_normalizes_business_failure( manager = poller_manager.WorkflowPollerManager() created_records: list[dict[str, Any]] = [] recorded_results: list[dict[str, Any]] = [] - recorded_steps: list[tuple[str, int, dict[str, Any]]] = [] + recorded_steps: list[tuple[int, dict[str, Any]]] = [] async def _fake_get_config(_workflow_id: str, *, kind: str) -> dict[str, Any]: return { @@ -155,7 +154,9 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, + persist: bool = True, ) -> dict[str, Any]: + assert persist is False record = { "id": exec_id or "exec-1", "workflowId": workflow_id, @@ -173,17 +174,12 @@ async def _fake_record_execution_result( workflow_id: str, exec_id: str, exec_data: dict[str, Any], + *, + steps: list[tuple[int, dict[str, Any]]] | None = None, ) -> None: _ = workflow_id, exec_id recorded_results.append(dict(exec_data)) - - async def _fake_record_execution_step( - exec_id: str, - step_index: int, - step: dict[str, Any], - ) -> dict[str, Any]: - recorded_steps.append((exec_id, step_index, step)) - return step + recorded_steps.extend(steps or []) def _fake_run_workflow( # noqa: ANN001 *, @@ -242,7 +238,6 @@ def _fake_run_workflow( # noqa: ANN001 ) monkeypatch.setattr(poller_manager, "create_execution_record", _fake_create_execution_record) monkeypatch.setattr(poller_manager, "record_execution_result", _fake_record_execution_result) - monkeypatch.setattr(execution_store, "record_execution_step", _fake_record_execution_step) monkeypatch.setattr(poller_manager, "run_workflow", _fake_run_workflow) status = await manager.run_once("wf-business-failure") @@ -255,9 +250,7 @@ def _fake_run_workflow( # noqa: ANN001 assert recorded_results[0]["executionLog"] == [] assert recorded_results[0]["stepCount"] == 1 assert recorded_results[0]["loopProgress"]["total_iterations"] == 2 - assert recorded_steps[0][0] == "exec-1" - assert recorded_steps[0][1] == 1 - assert recorded_steps[0][2]["node_id"] == "load" + assert recorded_steps == [] assert status["lastStatus"] == "error" assert status["lastError"] == "business rule blocked" assert status["selectedCount"] == 9 @@ -301,7 +294,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -356,7 +349,9 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, + persist: bool = True, ) -> dict[str, Any]: + assert persist is False _ = input_params return { "id": exec_id or f"exec-{workflow_id}", @@ -372,8 +367,10 @@ async def _fake_record_execution_result( workflow_id: str, exec_id: str, exec_data: dict[str, Any], + *, + steps: list[tuple[int, dict[str, Any]]] | None = None, ) -> None: - _ = workflow_id, exec_id, exec_data + _ = workflow_id, exec_id, exec_data, steps def _fake_run_workflow( # noqa: ANN001 *, diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index dabf9bb7e..dd34bec84 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -25,11 +25,21 @@ async def test_trigger_execution_builds_tool_context_for_workflow_tools( def _fake_run_workflow(**kwargs): # noqa: ANN003 missing_context = kwargs.get("tool_context") is None + kwargs["on_step_complete"]( + SimpleNamespace( + model_dump=lambda mode="json": { + "node_id": "notify", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + ) return SimpleNamespace( status="FAILED" if missing_context else "SUCCEEDED", outputs={}, error="Parent session not found" if missing_context else None, - history=[], + history=[{"node_id": "notify", "outputs": {"ok": True}}], last_node_id="notify", steps=1, ) @@ -41,12 +51,10 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 ) monkeypatch.setattr(runtime_module, "cleanup_workflow_tool_context", cleanup_context) monkeypatch.setattr(runtime_module, "run_workflow", Mock(side_effect=_fake_run_workflow)) - monkeypatch.setattr( - runtime_module, - "create_execution_record", - AsyncMock(return_value={"id": "exec-1"}), - ) - monkeypatch.setattr(runtime_module, "record_execution_result", AsyncMock()) + create_record = AsyncMock(return_value={"id": "exec-1"}) + record_result = AsyncMock() + monkeypatch.setattr(runtime_module, "create_execution_record", create_record) + monkeypatch.setattr(runtime_module, "record_execution_result", record_result) trigger = TriggerDefinition.model_validate({"id": "webhook-trigger", "type": "custom_webhook"}) runtime = runtime_module.TriggerRuntime() @@ -64,6 +72,14 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 action_name="trigger:custom_webhook", ) assert runtime_module.run_workflow.call_args.kwargs["tool_context"] is tool_context + assert runtime_module.run_workflow.call_args.kwargs["run_id"] == "exec-1" + assert runtime_module.run_workflow.call_args.kwargs["execution_profile"] == "high_frequency" + assert callable(runtime_module.run_workflow.call_args.kwargs["on_step_complete"]) + assert create_record.await_args.kwargs["persist"] is False + assert result["executionLog"] == [] + assert result["stepCount"] == 1 + record_result.assert_awaited_once() + assert record_result.await_args.kwargs["steps"] == [] cleanup_context.assert_awaited_once_with(tool_context) diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index bca4f3ada..5c904b844 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -20,6 +20,7 @@ def _reset_state() -> None: WorkflowStore._conn = None WorkflowStore._init_pid = None WorkflowStore._db_path = None + WorkflowStore._completion_lock = None @pytest.fixture(autouse=True) @@ -74,8 +75,13 @@ async def test_workflow_store_records_execution_steps_config_and_kv() -> None: ) assert [row["id"] for row in filtered] == ["exec-1"] - await WorkflowStore.record_step("exec-1", 1, {"node_id": "n1", "outputs": {"ok": 1}}) - await WorkflowStore.record_step("exec-1", 2, {"node_id": "n2", "outputs": {"ok": 2}}) + await WorkflowStore.record_steps( + "exec-1", + [ + (1, {"node_id": "n1", "outputs": {"ok": 1}}), + (2, {"node_id": "n2", "outputs": {"ok": 2}}), + ], + ) steps, total = await WorkflowStore.list_steps("exec-1", offset=1, limit=1) assert total == 2 assert steps == [{"node_id": "n2", "outputs": {"ok": 2}}] @@ -108,3 +114,146 @@ async def test_workflow_store_increment_stats_is_atomic_for_concurrent_updates() assert stats["errorCount"] == sum(1 for success, _ in updates if not success) assert stats["totalRuntime"] == pytest.approx(60.0) assert stats["avgRuntime"] == pytest.approx(1.0) + + +@pytest.mark.asyncio +async def test_complete_execution_writes_steps_summary_and_stats_with_one_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_db() + commit_count = 0 + original_commit = db.commit + + async def counted_commit() -> None: + nonlocal commit_count + commit_count += 1 + await original_commit() + + monkeypatch.setattr(db, "commit", counted_commit) + + await WorkflowStore.complete_execution( + { + "id": "exec-complete", + "workflowId": "wf-complete", + "status": "success", + "startedAt": 100, + "finishedAt": 350, + "duration": 0.25, + "executionLog": [], + }, + steps=[ + (1, {"node_id": "n1", "outputs": {"ok": 1}}), + (2, {"node_id": "n2", "outputs": {"ok": 2}}), + ], + success=True, + duration=0.25, + ) + + assert commit_count == 1 + execution = await WorkflowStore.get_execution("exec-complete") + assert execution is not None + assert execution["status"] == "success" + steps, total = await WorkflowStore.list_steps("exec-complete") + assert total == 2 + assert [step["node_id"] for step in steps] == ["n1", "n2"] + stats = await WorkflowStore.get_stats("wf-complete") + assert stats is not None + assert stats["callCount"] == 1 + assert stats["successCount"] == 1 + assert stats["totalRuntime"] == pytest.approx(0.25) + + +@pytest.mark.asyncio +async def test_complete_execution_reduces_28_step_writes_to_four_commits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_db() + commit_count = 0 + original_commit = db.commit + + async def counted_commit() -> None: + nonlocal commit_count + commit_count += 1 + await original_commit() + + monkeypatch.setattr(db, "commit", counted_commit) + steps = [ + (index, {"node_id": f"node-{index}", "outputs": {"ok": True}}) + for index in range(1, 8) + ] + + await asyncio.gather( + *( + WorkflowStore.complete_execution( + { + "id": f"exec-{index}", + "workflowId": "wf-trigger", + "status": "success", + "startedAt": index + 1, + "finishedAt": index + 2, + "duration": 0.01, + "executionLog": [], + "stepCount": 7, + }, + steps, + success=True, + duration=0.01, + ) + for index in range(4) + ) + ) + + assert commit_count == 4 + executions = await WorkflowStore.list_executions("wf-trigger", limit=50) + assert len(executions) == 4 + assert all(execution["executionLog"] == [] for execution in executions) + for index in range(4): + persisted_steps, total = await WorkflowStore.list_steps(f"exec-{index}") + assert total == 7 + assert [step["node_id"] for step in persisted_steps] == [ + f"node-{step_index}" for step_index in range(1, 8) + ] + stats = await WorkflowStore.get_stats("wf-trigger") + assert stats is not None + assert stats["callCount"] == 4 + assert stats["successCount"] == 4 + + +@pytest.mark.asyncio +async def test_complete_execution_rolls_back_partial_transaction( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_db() + original_commit = db.commit + + async def fail_commit() -> None: + raise RuntimeError("commit failed") + + monkeypatch.setattr(db, "commit", fail_commit) + + with pytest.raises(RuntimeError, match="commit failed"): + await WorkflowStore.complete_execution( + { + "id": "exec-rollback", + "workflowId": "wf-rollback", + "status": "success", + "startedAt": 1, + "finishedAt": 2, + "duration": 0.01, + "executionLog": [], + "stepCount": 1, + }, + [(1, {"node_id": "node-1", "outputs": {"ok": True}})], + success=True, + duration=0.01, + ) + + monkeypatch.setattr(db, "commit", original_commit) + assert await WorkflowStore.get_execution("exec-rollback") is None + persisted_steps, total = await WorkflowStore.list_steps("exec-rollback") + assert persisted_steps == [] + assert total == 0 + assert await WorkflowStore.get_stats("wf-rollback") is None From c01cb2530d1bf840a7fd7da4517e35ad7f9ef895 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 13:25:39 +0800 Subject: [PATCH 2/5] fix(workflow): finalize atomic step persistence --- flocks/workflow/execution_store.py | 33 +++++----- flocks/workflow/store.py | 61 +++++++++++++++++-- .../workflow/test_execution_store_compact.py | 22 +++---- tests/workflow/test_poller_manager.py | 12 +++- tests/workflow/test_workflow_store.py | 55 ++++++++++++++++- 5 files changed, 143 insertions(+), 40 deletions(-) diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index da7c189c2..c1c7cb443 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -594,12 +594,15 @@ async def record_execution_result( started_at = summary_data.get("startedAt", 0) finished_at = summary_data.get("finishedAt", int(time.time() * 1000)) duration = max(0.0, (finished_at - started_at) / 1000.0) - await WorkflowStore.complete_execution( + trimmed_exec_ids = await WorkflowStore.complete_execution( compact_execution_summary(summary_data), prepared_steps, success=success, duration=float(duration), + history_limit=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, ) + if not isinstance(trimmed_exec_ids, list): + trimmed_exec_ids = [] # Recorder writes to its own SQLite tables and can be slow under load. # Run it as a background task so the syslog/HTTP dispatcher can release the @@ -621,6 +624,19 @@ async def _record_audit() -> None: "error": str(exc), }, ) + for trimmed_exec_id in trimmed_exec_ids: + try: + record_path = Recorder.paths().workflow_dir / f"{trimmed_exec_id}.jsonl" + await asyncio.to_thread(record_path.unlink, missing_ok=True) + except Exception as exc: + log.warning( + "workflow.history.trim_delete_failed", + { + "workflow_id": workflow_id, + "exec_id": trimmed_exec_id, + "error": str(exc), + }, + ) asyncio.create_task(_record_audit(), name=f"audit-{exec_id}") except RuntimeError: @@ -634,21 +650,6 @@ async def _record_audit() -> None: except Exception: pass - # Prune old execution records when the per-workflow limit is exceeded. - # This is awaited so a successful completion does not silently leave the - # workflow above its retention cap. - try: - await _trim_execution_history(workflow_id) - except Exception as exc: - log.error( - "workflow.history.trim_failed", - { - "workflow_id": workflow_id, - "exec_id": exec_id, - "error": str(exc), - }, - ) - async def _delete_execution_history_record( execution_key: str, diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index 4a74e6401..db049f3a2 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -44,6 +44,7 @@ class WorkflowStore: _initialized = False _conn: Optional[aiosqlite.Connection] = None + _completion_conn: Optional[aiosqlite.Connection] = None _init_pid: Optional[int] = None _db_path: Optional[Path] = None _completion_lock: Optional[asyncio.Lock] = None @@ -73,7 +74,10 @@ async def init(cls) -> None: ) if cls._conn: await cls._conn.close() + if cls._completion_conn: + await cls._completion_conn.close() cls._conn = None + cls._completion_conn = None cls._initialized = False cls._init_pid = None @@ -91,6 +95,12 @@ async def _open_and_migrate() -> None: for stmt in _INDEX_STMTS: await cls._conn.execute(stmt) await cls._conn.commit() + cls._completion_conn = await aiosqlite.connect( + db_path, + timeout=Storage._sqlite_timeout_s, + ) + cls._completion_conn.row_factory = aiosqlite.Row + await Storage.configure_connection(cls._completion_conn) cls._initialized = True cls._init_pid = current_pid cls._db_path = db_path @@ -103,7 +113,10 @@ async def _open_and_migrate() -> None: except Exception as exc: if cls._conn: await cls._conn.close() + if cls._completion_conn: + await cls._completion_conn.close() cls._conn = None + cls._completion_conn = None cls._initialized = False cls._init_pid = None cls._db_path = None @@ -122,7 +135,10 @@ async def _open_and_migrate() -> None: async def close(cls) -> None: if cls._conn: await cls._conn.close() + if cls._completion_conn: + await cls._completion_conn.close() cls._conn = None + cls._completion_conn = None cls._initialized = False cls._init_pid = None cls._db_path = None @@ -140,6 +156,17 @@ async def _db(cls) -> aiosqlite.Connection: async def raw_db(cls) -> aiosqlite.Connection: return await cls._db() + @classmethod + async def _completion_db(cls) -> aiosqlite.Connection: + if not cls._completion_conn or not cls._initialized: + await cls.init() + return cls._completion_conn # type: ignore[return-value] + + @classmethod + async def raw_completion_db(cls) -> aiosqlite.Connection: + """Return the completion connection for transaction-level tests.""" + return await cls._completion_db() + @staticmethod def _json_dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False, default=str) @@ -357,7 +384,7 @@ async def list_executions( f""" SELECT payload FROM workflow_executions WHERE {" AND ".join(clauses)} - ORDER BY started_at DESC + ORDER BY started_at DESC, rowid DESC LIMIT ? """, tuple(params), @@ -399,7 +426,7 @@ async def trim_executions(cls, workflow_id: str, *, keep: int) -> List[str]: """ SELECT id FROM workflow_executions WHERE workflow_id = ? - ORDER BY started_at DESC + ORDER BY started_at DESC, rowid DESC LIMIT -1 OFFSET ? """, (workflow_id, max(int(keep), 0)), @@ -460,9 +487,10 @@ async def complete_execution( *, success: bool, duration: float, - ) -> None: - """Persist one completed execution and its stats in one transaction.""" - db = await cls._db() + history_limit: Optional[int] = None, + ) -> List[str]: + """Persist one completed execution, stats, and retention in one transaction.""" + db = await cls._completion_db() payload = dict(exec_data) exec_id = str(payload.get("id") or "") workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") @@ -492,6 +520,7 @@ async def complete_execution( async with lock: try: + await db.execute("BEGIN IMMEDIATE") if step_rows: await db.executemany( """ @@ -556,7 +585,29 @@ async def complete_execution( cls._now_ms(), ), ) + trimmed_exec_ids: List[str] = [] + if history_limit is not None: + async with db.execute( + """ + SELECT id FROM workflow_executions + WHERE workflow_id = ? + ORDER BY started_at DESC, rowid DESC + LIMIT -1 OFFSET ? + """, + (workflow_id, max(int(history_limit), 0)), + ) as cur: + trimmed_exec_ids = [str(row["id"]) for row in await cur.fetchall()] + for trimmed_exec_id in trimmed_exec_ids: + await db.execute( + "DELETE FROM workflow_execution_steps WHERE exec_id = ?", + (trimmed_exec_id,), + ) + await db.execute( + "DELETE FROM workflow_executions WHERE id = ?", + (trimmed_exec_id,), + ) await db.commit() + return trimmed_exec_ids except Exception: await db.rollback() raise diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index de43af374..9245eea3a 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -344,17 +344,11 @@ def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: @pytest.mark.asyncio -async def test_four_trigger_workers_keep_step_history_disabled() -> None: - """Four trigger threads track progress without retaining step history.""" +async def test_four_trigger_workers_collect_steps_without_storage() -> None: + """Four trigger threads collect complete batches without callback SQL.""" record_step = AsyncMock(return_value=None) record_steps = AsyncMock(return_value=None) - recorders = [ - ExecutionStepRecorder( - exec_id=f"exec-trigger-{worker}", - capture_steps=False, - ) - for worker in range(4) - ] + recorders = [ExecutionStepRecorder(exec_id=f"exec-trigger-{worker}") for worker in range(4)] def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: for step in range(7): @@ -371,7 +365,7 @@ def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: ) batches = [recorder.take_steps() for recorder in recorders] - assert batches == [[], [], [], []] + assert [len(batch) for batch in batches] == [7, 7, 7, 7] assert [recorder.step_count for recorder in recorders] == [7, 7, 7, 7] record_step.assert_not_awaited() record_steps.assert_not_awaited() @@ -379,7 +373,7 @@ def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: @pytest.mark.asyncio async def test_record_execution_result_backfills_execution_log_steps() -> None: - complete_execution = AsyncMock(return_value=None) + complete_execution = AsyncMock(return_value=[]) exec_data = { "id": "exec-1", "workflowId": "wf", @@ -399,7 +393,6 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 patch.object(WorkflowStore, "complete_execution", complete_execution), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), - patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), ): await record_execution_result("wf", "exec-1", exec_data) @@ -414,12 +407,13 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 assert complete_execution.await_args.kwargs == { "success": True, "duration": 1.0, + "history_limit": 30, } @pytest.mark.asyncio async def test_record_execution_result_accepts_explicit_step_batch() -> None: - complete_execution = AsyncMock(return_value=None) + complete_execution = AsyncMock(return_value=[]) explicit_steps = [ (1, {"node_id": "step-1", "outputs": {"ok": True}}), (2, {"node_id": "step-2", "outputs": {"ok": True}}), @@ -441,7 +435,6 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 patch.object(WorkflowStore, "complete_execution", complete_execution), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), - patch("flocks.workflow.execution_store._trim_execution_history", AsyncMock(return_value=None)), ): await record_execution_result( "wf-trigger", @@ -456,6 +449,7 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 assert complete_execution.await_args.kwargs == { "success": True, "duration": 0.01, + "history_limit": 30, } diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index ae116df5c..82ee82cf0 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -406,7 +406,10 @@ def _fake_run_workflow( # noqa: ANN001 assert manager.get_status("wf-stop")["activeRuns"] == 1 release_run.set() - await asyncio.sleep(0.05) + for _ in range(100): + if manager.get_status("wf-stop")["activeRuns"] == 0: + break + await asyncio.sleep(0.01) assert manager.get_status("wf-stop")["activeRuns"] == 0 @@ -421,7 +424,12 @@ async def _fake_list_configs(*, kind: str) -> list[tuple[str, dict[str, Any]]]: ("wf-disabled", {"enabled": False}), ] - async def _fake_restart(workflow_id: str) -> dict[str, Any]: + async def _fake_restart( + workflow_id: str, + *, + startup: bool = False, + ) -> dict[str, Any]: + assert startup is True restarted.append(workflow_id) return {"workflowId": workflow_id, "state": "running"} diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index 5c904b844..d699d0f9b 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -18,6 +18,7 @@ def _reset_state() -> None: Storage._init_pid = None WorkflowStore._initialized = False WorkflowStore._conn = None + WorkflowStore._completion_conn = None WorkflowStore._init_pid = None WorkflowStore._db_path = None WorkflowStore._completion_lock = None @@ -121,7 +122,7 @@ async def test_complete_execution_writes_steps_summary_and_stats_with_one_commit monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() - db = await WorkflowStore.raw_db() + db = await WorkflowStore.raw_completion_db() commit_count = 0 original_commit = db.commit @@ -169,7 +170,7 @@ async def test_complete_execution_reduces_28_step_writes_to_four_commits( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() - db = await WorkflowStore.raw_db() + db = await WorkflowStore.raw_completion_db() commit_count = 0 original_commit = db.commit @@ -226,7 +227,7 @@ async def test_complete_execution_rolls_back_partial_transaction( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() - db = await WorkflowStore.raw_db() + db = await WorkflowStore.raw_completion_db() original_commit = db.commit async def fail_commit() -> None: @@ -257,3 +258,51 @@ async def fail_commit() -> None: assert persisted_steps == [] assert total == 0 assert await WorkflowStore.get_stats("wf-rollback") is None + + +@pytest.mark.asyncio +async def test_complete_execution_applies_retention_before_single_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + await WorkflowStore.init() + db = await WorkflowStore.raw_completion_db() + commit_count = 0 + original_commit = db.commit + + async def counted_commit() -> None: + nonlocal commit_count + commit_count += 1 + await original_commit() + + monkeypatch.setattr(db, "commit", counted_commit) + trimmed: list[str] = [] + for index in range(4): + trimmed = await WorkflowStore.complete_execution( + { + "id": f"exec-retain-{index}", + "workflowId": "wf-retain", + "status": "success", + "startedAt": index + 1, + "finishedAt": index + 2, + "duration": 0.01, + "executionLog": [], + "stepCount": 1, + }, + [(1, {"node_id": f"node-{index}", "outputs": {"ok": True}})], + success=True, + duration=0.01, + history_limit=3, + ) + + assert commit_count == 4 + assert trimmed == ["exec-retain-0"] + assert await WorkflowStore.get_execution("exec-retain-0") is None + old_steps, old_total = await WorkflowStore.list_steps("exec-retain-0") + assert old_steps == [] + assert old_total == 0 + executions = await WorkflowStore.list_executions("wf-retain", limit=10) + assert [execution["id"] for execution in executions] == [ + "exec-retain-3", + "exec-retain-2", + "exec-retain-1", + ] From 76f30996f8dfcc292f83242d94cb067a45720382 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 17:01:57 +0800 Subject: [PATCH 3/5] fix(workflow): preserve steps without callback waits Keep workflow callbacks storage-free while batching complete step history and serializing interactive progress writes. Isolate terminal persistence from stats and retention failures. Co-Authored-By: Claude Opus 4.6 --- flocks/ingest/kafka/manager.py | 1 - flocks/ingest/syslog/manager.py | 5 +- flocks/server/routes/workflow.py | 230 +++++----- flocks/tool/task/run_workflow.py | 258 ++++++----- flocks/workflow/execution_store.py | 213 ++++++---- flocks/workflow/poller_manager.py | 5 +- flocks/workflow/store.py | 121 ++---- flocks/workflow/triggers/runtime.py | 5 +- tests/ingest/test_kafka_manager.py | 10 +- .../test_syslog_manager_backpressure.py | 12 +- .../server/routes/test_workflow_run_route.py | 400 +++++++++++++++++- .../workflow/test_execution_store_compact.py | 346 +++++++++++---- tests/workflow/test_poller_manager.py | 12 +- tests/workflow/test_tool_run_workflow.py | 269 +++++++++++- tests/workflow/test_trigger_runtime.py | 12 +- tests/workflow/test_workflow_store.py | 93 ++-- 16 files changed, 1411 insertions(+), 581 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index 893156c43..a44c28eec 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -775,7 +775,6 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: trigger_input_keys = list((trigger.mapping or {}).keys()) or [input_key] step_recorder = ExecutionStepRecorder( exec_id=exec_id, - capture_steps=False, step_compactor=lambda step: _compact_step_for_kafka_storage( step, input_key=input_key, diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 699fa53cb..ab15ad23a 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -622,10 +622,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: persist=False, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder( - exec_id=exec_id, - capture_steps=False, - ) + step_recorder = ExecutionStepRecorder(exec_id=exec_id) start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) tool_context = None diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index 928b6a900..0a35c4eab 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -55,6 +55,8 @@ compact_step_for_storage, create_execution_record, derive_loop_progress, + ExecutionProgressWriter, + ExecutionStepRecorder, load_execution_steps, normalize_execution_status as _normalize_execution_status, record_execution_result as _record_execution_result, @@ -144,6 +146,7 @@ class ActiveWorkflowExecution: workflow_id: str task: asyncio.Task[Any] cancel_event: threading.Event + progress_writer: ExecutionProgressWriter _active_workflow_executions: Dict[str, ActiveWorkflowExecution] = {} @@ -1120,11 +1123,12 @@ async def _run_workflow_execution_task( req: WorkflowRunRequest, exec_id: str, cancel_event: threading.Event, + progress_writer: ExecutionProgressWriter, tool_context: Optional[ToolContext] = None, ) -> None: """Execute a workflow in the background and keep the execution record updated.""" start_time = time.time() - step_count = 0 + step_recorder = ExecutionStepRecorder(exec_id=exec_id) pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None execution_summary: Dict[str, Any] = { @@ -1167,122 +1171,123 @@ def _on_step_start(_run_id, step_index, node, _inputs): "error": "Run cancelled before node completed", } ) - execution_summary.update( - { - "currentNodeId": node_id, - "currentNodeType": node_type, - "currentPhase": "running", - "currentStepIndex": step_index, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) + progress_update = { + "currentNodeId": node_id, + "currentNodeType": node_type, + "currentPhase": "cancelling" if cancel_event.is_set() else "running", + "currentStepIndex": step_index, + "loopProgress": loop_progress, + "updatedAt": int(time.time() * 1000), + } + execution_summary.update(progress_update) + progress_writer.submit(progress_update) return step_index def _on_step_complete(step_result) -> None: - nonlocal step_count, pending_step_index, pending_step - step_dict = compact_step_for_storage(step_result.model_dump(mode="json")) - step_count += 1 + nonlocal pending_step_index, pending_step + step_recorder.on_step_complete(step_result) pending_step_index = None pending_step = None - loop_progress = derive_loop_progress( - node_id=step_dict.get("node_id"), - global_step_index=step_count, - inputs=step_dict.get("inputs"), - outputs=step_dict.get("outputs"), - ) - execution_summary.update( - { - "stepCount": step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) - + progress_update = dict(step_recorder.summary) + if cancel_event.is_set(): + progress_update["currentPhase"] = "cancelling" + execution_summary.update(progress_update) + progress_writer.submit(progress_update) + + result: Optional[RunWorkflowResult] = None + execution_error: Optional[Exception] = None try: - result: RunWorkflowResult = await asyncio.to_thread( - run_workflow, - workflow=workflow_json, - inputs=req.inputs or {}, - timeout_s=req.timeout_s, - trace=req.trace, - on_step_start=_on_step_start, - on_step_complete=_on_step_complete, - cancel=cancel_event.is_set, - tool_context=tool_context, - ) + try: + result = await asyncio.to_thread( + run_workflow, + workflow=workflow_json, + inputs=req.inputs or {}, + timeout_s=req.timeout_s, + trace=req.trace, + on_step_start=_on_step_start, + on_step_complete=_on_step_complete, + cancel=cancel_event.is_set, + tool_context=tool_context, + ) + except Exception as exc: + execution_error = exc duration = time.time() - start_time current_data = dict(execution_summary) - status_value, error_message = _resolve_execution_outcome(result) - if cancel_event.is_set() and status_value == "success": - status_value = "cancelled" - error_message = error_message or f"Run cancelled: run_id={result.run_id or exec_id}" - # ``record_execution_result`` backfills this compacted history into - # append-only step rows, then stores only the summary row. - final_history = compact_history_for_storage(result.history) - final_steps = result.steps + final_step_batch = step_recorder.take_steps() if pending_step_index is not None and pending_step is not None: - final_history.append(pending_step) - final_steps = max(final_steps, pending_step_index) - current_data.update( - { - "outputResults": compact_outputs_for_storage(result.outputs), - "status": status_value, - "finishedAt": int(time.time() * 1000), - "duration": duration, - "executionLog": final_history, - "stepCount": final_steps, - "errorMessage": error_message, - "currentNodeId": result.last_node_id, - "currentNodeType": current_data.get("currentNodeType"), - "currentPhase": status_value, - "currentStepIndex": final_steps, - "updatedAt": int(time.time() * 1000), - } + final_step_batch.append((pending_step_index, pending_step)) + final_step_batch.sort(key=lambda item: item[0]) + final_steps = max( + max((step_index for step_index, _ in final_step_batch), default=0), + pending_step_index or 0, ) + final_history = [step for _, step in final_step_batch] + + if execution_error is None: + assert result is not None + status_value, error_message = _resolve_execution_outcome(result) + if cancel_event.is_set() and status_value == "success": + status_value = "cancelled" + error_message = error_message or f"Run cancelled: run_id={result.run_id or exec_id}" + final_steps = max(result.steps, final_steps) + current_data.update( + { + "outputResults": compact_outputs_for_storage(result.outputs), + "status": status_value, + "finishedAt": int(time.time() * 1000), + "duration": duration, + "executionLog": final_history, + "stepCount": final_steps, + "errorMessage": error_message, + "currentNodeId": result.last_node_id, + "currentNodeType": current_data.get("currentNodeType"), + "currentPhase": status_value, + "currentStepIndex": final_steps, + "updatedAt": int(time.time() * 1000), + } + ) + else: + current_data.update( + { + "status": "cancelled" if cancel_event.is_set() else "error", + "finishedAt": int(time.time() * 1000), + "duration": duration, + "errorMessage": str(execution_error), + "executionLog": final_history, + "stepCount": final_steps, + "currentPhase": "cancelled" if cancel_event.is_set() else "error", + "currentStepIndex": final_steps, + "updatedAt": int(time.time() * 1000), + } + ) - await _record_execution_result(workflow_id, exec_id, current_data) - log.info( - "workflow.executed", - { - "id": workflow_id, - "exec_id": exec_id, - "status": status_value, - "duration": duration, - }, - ) - except Exception as exc: - duration = time.time() - start_time - current_data = dict(execution_summary) - final_history = [pending_step] if pending_step is not None else [] - final_steps = max(step_count, pending_step_index or 0) - current_data.update( - { - "status": "cancelled" if cancel_event.is_set() else "error", - "finishedAt": int(time.time() * 1000), - "duration": duration, - "errorMessage": str(exc), - "executionLog": final_history, - "stepCount": final_steps, - "currentPhase": "cancelled" if cancel_event.is_set() else "error", - "currentStepIndex": final_steps, - "updatedAt": int(time.time() * 1000), - } - ) - await _record_execution_result(workflow_id, exec_id, current_data) - log.error( - "workflow.execute.error", - { - "id": workflow_id, - "exec_id": exec_id, - "error": str(exc), - }, + await progress_writer.close_and_drain() + await _record_execution_result( + workflow_id, + exec_id, + current_data, + steps=final_step_batch, ) + if execution_error is None: + log.info( + "workflow.executed", + { + "id": workflow_id, + "exec_id": exec_id, + "status": current_data["status"], + "duration": duration, + }, + ) + else: + log.error( + "workflow.execute.error", + { + "id": workflow_id, + "exec_id": exec_id, + "error": str(execution_error), + }, + ) finally: _active_workflow_executions.pop(exec_id, None) @@ -1692,6 +1697,7 @@ async def run_workflow_endpoint(workflow_id: str, req: WorkflowRunRequest): ) await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) exec_id = str(exec_data["id"]) + progress_writer = ExecutionProgressWriter(exec_data) cancel_event = threading.Event() task = asyncio.create_task( @@ -1701,6 +1707,7 @@ async def run_workflow_endpoint(workflow_id: str, req: WorkflowRunRequest): req=req, exec_id=exec_id, cancel_event=cancel_event, + progress_writer=progress_writer, tool_context=tool_context, ), name=f"workflow-run-{exec_id}", @@ -1709,6 +1716,7 @@ async def run_workflow_endpoint(workflow_id: str, req: WorkflowRunRequest): workflow_id=workflow_id, task=task, cancel_event=cancel_event, + progress_writer=progress_writer, ) # Guarantee cleanup of the registry entry even when the task is @@ -1757,13 +1765,13 @@ async def cancel_workflow_execution(workflow_id: str, exec_id: str): raise HTTPException(status_code=404, detail="Execution not found for this workflow") active.cancel_event.set() - exec_data.update( - { - "currentPhase": "cancelling", - "errorMessage": exec_data.get("errorMessage") or "Cancellation requested", - } - ) - await WorkflowStore.upsert_execution(exec_data) + progress_update = { + "currentPhase": "cancelling", + "errorMessage": exec_data.get("errorMessage") or "Cancellation requested", + "updatedAt": int(time.time() * 1000), + } + exec_data.update(progress_update) + await active.progress_writer.update(progress_update) log.info( "workflow.execution.cancel_requested", { diff --git a/flocks/tool/task/run_workflow.py b/flocks/tool/task/run_workflow.py index 9b78e5123..ab5fcbf10 100644 --- a/flocks/tool/task/run_workflow.py +++ b/flocks/tool/task/run_workflow.py @@ -11,32 +11,29 @@ import time from pathlib import Path from types import SimpleNamespace -from typing import Optional, Dict, Any, Union +from typing import Optional, Dict, Any, Union, List, Tuple from flocks.tool.registry import ToolRegistry, ToolCategory, ToolParameter, ParameterType, ToolResult, ToolContext from flocks.utils.log import Log from flocks.session.recorder import Recorder from flocks.workflow.execution_store import ( compact_history_for_storage, - compact_execution_summary, compact_outputs_for_storage, compact_step_for_storage, create_execution_record, derive_loop_progress, + ExecutionProgressWriter, + ExecutionStepRecorder, normalize_execution_status, - record_execution_step, record_execution_result, resolve_execution_outcome, ) from flocks.workflow.fs_store import read_workflow_from_fs, resolve_workflow_id_from_source -from flocks.workflow.store import WorkflowStore from flocks.tool.truncation import truncate_output log = Log.create(service="tool.run_workflow") -_PROGRESS_FLUSH_EVERY_STEPS = 5 - # Lazy import to avoid circular import (flocks.tool <-> flocks.workflow) _WORKFLOW_AVAILABLE: Optional[bool] = None RequirementsInstaller = None @@ -573,33 +570,17 @@ async def run_workflow_tool( canonical_workflow_id = registered_workflow_id or resolve_workflow_id_from_source(workflow_source) display_workflow_id = canonical_workflow_id or workflow_id tracked_execution: Optional[Dict[str, Any]] = None - tracked_step_count = 0 + step_recorder: Optional[ExecutionStepRecorder] = None + progress_writer: Optional[ExecutionProgressWriter] = None + callback_step_count = 0 pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None + final_step_batch: Optional[List[Tuple[int, Dict[str, Any]]]] = None loop = asyncio.get_running_loop() def _emit_metadata(metadata: Dict[str, Any]) -> None: loop.call_soon_threadsafe(ctx.metadata, metadata) - def _update_execution_progress(update_fields: Dict[str, Any]) -> None: - try: - if tracked_execution is None: - return - tracked_execution.update(update_fields) - asyncio.run_coroutine_threadsafe( - WorkflowStore.upsert_execution(compact_execution_summary(tracked_execution)), - loop, - ).result(timeout=5) - except Exception as exc: - log.warning( - "run_workflow.execution_progress.write_failed", - { - "workflow_id": display_workflow_id, - "exec_id": tracked_execution["id"] if tracked_execution else None, - "error": str(exc), - }, - ) - def _on_step_start( _run_id: Optional[str], step_index: int, @@ -625,17 +606,19 @@ def _on_step_start( "error": "Run cancelled before node completed", } ) + current_phase = "cancelling" if ctx.abort.is_set() else "running" + progress_update = { + "currentNodeId": current_node_id, + "currentNodeType": current_node_type, + "currentPhase": current_phase, + "currentStepIndex": step_index, + "loopProgress": loop_progress, + "updatedAt": int(time.time() * 1000), + } if tracked_execution is not None: - _update_execution_progress( - { - "currentNodeId": current_node_id, - "currentNodeType": current_node_type, - "currentPhase": "running", - "currentStepIndex": step_index, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) + tracked_execution.update(progress_update) + if progress_writer is not None: + progress_writer.submit(progress_update) _emit_metadata( { "title": f"Running workflow: {workflow_name}", @@ -645,7 +628,7 @@ def _on_step_start( "total_nodes": workflow_total_nodes, "workflow_execution_id": tracked_execution["id"] if tracked_execution else None, "status": "running", - "phase": "running", + "phase": current_phase, "current_node_id": current_node_id, "current_node_type": current_node_type, "step_index": step_index, @@ -656,64 +639,43 @@ def _on_step_start( return step_index def _on_step_complete(step_result: Any) -> None: - nonlocal tracked_step_count, pending_step_index, pending_step - if hasattr(step_result, "model_dump"): - step_dict = step_result.model_dump(mode="json") - elif isinstance(step_result, dict): - step_dict = dict(step_result) + nonlocal callback_step_count, pending_step_index, pending_step + if step_recorder is not None: + step_recorder.on_step_complete(step_result) + callback_step_count = step_recorder.step_count + progress_update = dict(step_recorder.summary) else: - step_dict = {"node_id": None, "outputs": {}, "error": str(step_result)} - step_index = tracked_step_count + 1 - compacted_step = compact_step_for_storage(step_dict) + if hasattr(step_result, "model_dump"): + step_dict = step_result.model_dump(mode="json") + elif isinstance(step_result, dict): + step_dict = dict(step_result) + else: + step_dict = {"node_id": None, "outputs": {}, "error": str(step_result)} + callback_step_count += 1 + compacted_step = compact_step_for_storage(step_dict) + progress_update = { + "stepCount": callback_step_count, + "currentNodeId": compacted_step.get("node_id"), + "currentNodeType": compacted_step.get("node_type") + or compacted_step.get("type"), + "currentPhase": "running", + "currentStepIndex": callback_step_count, + "loopProgress": derive_loop_progress( + node_id=compacted_step.get("node_id"), + global_step_index=callback_step_count, + inputs=compacted_step.get("inputs"), + outputs=compacted_step.get("outputs"), + ), + "updatedAt": int(time.time() * 1000), + } pending_step_index = None pending_step = None - loop_progress = derive_loop_progress( - node_id=step_dict.get("node_id"), - global_step_index=step_index, - inputs=step_dict.get("inputs"), - outputs=step_dict.get("outputs"), - ) - tracked_step_count = step_index + if ctx.abort.is_set(): + progress_update["currentPhase"] = "cancelling" if tracked_execution is not None: - tracked_execution.update( - { - "stepCount": tracked_step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": tracked_step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) - if tracked_execution is not None: - try: - asyncio.run_coroutine_threadsafe( - record_execution_step(tracked_execution["id"], step_index, compacted_step), - loop, - ).result(timeout=5) - except Exception as exc: - log.warning( - "run_workflow.execution_step.write_failed", - { - "workflow_id": display_workflow_id, - "exec_id": tracked_execution["id"], - "step_index": step_index, - "error": str(exc), - }, - ) - if tracked_step_count % _PROGRESS_FLUSH_EVERY_STEPS == 0: - _update_execution_progress( - { - "stepCount": tracked_step_count, - "currentNodeId": step_dict.get("node_id"), - "currentNodeType": step_dict.get("node_type") or step_dict.get("type"), - "currentPhase": "running", - "currentStepIndex": tracked_step_count, - "loopProgress": loop_progress, - "updatedAt": int(time.time() * 1000), - } - ) + tracked_execution.update(progress_update) + if progress_writer is not None: + progress_writer.submit(progress_update) _emit_metadata( { "title": f"Running workflow: {workflow_name}", @@ -723,36 +685,24 @@ def _on_step_complete(step_result: Any) -> None: "total_nodes": workflow_total_nodes, "workflow_execution_id": tracked_execution["id"] if tracked_execution else None, "status": "running", - "phase": "running", - "current_node_id": step_dict.get("node_id"), - "current_node_type": step_dict.get("node_type") or step_dict.get("type"), - "step_index": tracked_step_count, - "step_count": tracked_step_count, - "loop_progress": loop_progress, + "phase": progress_update["currentPhase"], + "current_node_id": progress_update.get("currentNodeId"), + "current_node_type": progress_update.get("currentNodeType"), + "step_index": callback_step_count, + "step_count": callback_step_count, + "loop_progress": progress_update.get("loopProgress"), }, } ) - return - async def _flush_pending_step() -> None: - if tracked_execution is None or pending_step_index is None or pending_step is None: - return - try: - await record_execution_step( - tracked_execution["id"], - pending_step_index, - pending_step, - ) - except Exception as exc: - log.warning( - "run_workflow.pending_step.write_failed", - { - "workflow_id": display_workflow_id, - "exec_id": tracked_execution["id"], - "step_index": pending_step_index, - "error": str(exc), - }, - ) + def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: + nonlocal final_step_batch + if final_step_batch is None: + final_step_batch = step_recorder.take_steps() if step_recorder is not None else [] + if pending_step_index is not None and pending_step is not None: + final_step_batch.append((pending_step_index, pending_step)) + final_step_batch.sort(key=lambda item: item[0]) + return final_step_batch await ctx.ask( permission="run_workflow", @@ -771,6 +721,8 @@ async def _flush_pending_step() -> None: canonical_workflow_id, input_params=workflow_inputs, ) + step_recorder = ExecutionStepRecorder(exec_id=tracked_execution["id"]) + progress_writer = ExecutionProgressWriter(tracked_execution) # Update metadata to show workflow is running _emit_metadata( @@ -890,7 +842,8 @@ async def _flush_pending_step() -> None: result_dict = {"status": "UNKNOWN", "output": str(result)} status = result_dict.get("status", "UNKNOWN") - success = status == "SUCCEEDED" + status_value = normalize_execution_status(status) + success = status_value == "success" error = result_dict.get("error") output, output_truncated, output_path = _format_workflow_result_for_tool(result_dict) @@ -905,19 +858,25 @@ async def _flush_pending_step() -> None: }, ) - # Append-only recording for audit/replay - await _record_workflow_tool_result(display_workflow_id, result_dict) - - status_value = normalize_execution_status(status) compacted_history = compact_history_for_storage(result_dict.get("history")) - history_count = len(compacted_history) - if status_value == "cancelled" and not compacted_history: - await _flush_pending_step() + tracked_steps = _take_final_step_batch() if tracked_execution is not None else [] + final_history = ( + [step for _, step in tracked_steps] + if tracked_execution is not None + else compacted_history + ) + history_count = len(final_history) final_step_count = result_dict.get("steps") if not isinstance(final_step_count, int): - final_step_count = tracked_step_count - if pending_step_index is not None: - final_step_count = max(final_step_count, pending_step_index) + final_step_count = callback_step_count + final_step_count = max( + final_step_count, + max((step_index for step_index, _ in tracked_steps), default=0), + ) + + if tracked_execution is None: + await _record_workflow_tool_result(display_workflow_id, result_dict) + if tracked_execution and canonical_workflow_id: current_data = dict(tracked_execution) outcome_result = result @@ -928,13 +887,20 @@ async def _flush_pending_step() -> None: error=result_dict.get("error"), ) status_value, error_message = resolve_execution_outcome(outcome_result) # type: ignore[arg-type] + if ctx.abort.is_set() and status_value == "success": + status_value = "cancelled" + error_message = error_message or ( + f"Run cancelled: run_id={result_dict.get('run_id') or tracked_execution['id']}" + ) + error = error or error_message + success = status_value == "success" current_data.update( { "outputResults": compact_outputs_for_storage(result_dict.get("outputs")), "status": status_value, "finishedAt": int(time.time() * 1000), "duration": time.time() - execution_started_at, - "executionLog": compacted_history, + "executionLog": final_history, "stepCount": final_step_count, "errorMessage": error_message, "currentNodeId": result_dict.get("last_node_id"), @@ -943,10 +909,13 @@ async def _flush_pending_step() -> None: "updatedAt": int(time.time() * 1000), } ) + if progress_writer is not None: + await progress_writer.close_and_drain() await record_execution_result( canonical_workflow_id, tracked_execution["id"], current_data, + steps=tracked_steps, ) _emit_metadata( { @@ -979,7 +948,7 @@ async def _flush_pending_step() -> None: total_nodes=workflow_total_nodes, workflow_execution_id=tracked_execution["id"] if tracked_execution else None, status=status_value, - steps=result_dict.get("steps", 0), + steps=final_step_count, last_node_id=result_dict.get("last_node_id"), outputs=result_dict.get("outputs"), history_count=history_count, @@ -1018,24 +987,34 @@ async def _flush_pending_step() -> None: "error": error_msg, }, ) + terminal_status = "cancelled" if ctx.abort.is_set() else "error" + final_step_count = callback_step_count if tracked_execution and canonical_workflow_id: + tracked_steps = _take_final_step_batch() + final_step_count = max( + final_step_count, + max((step_index for step_index, _ in tracked_steps), default=0), + ) current_data = dict(tracked_execution) current_data.update( { - "status": "error", + "status": terminal_status, "finishedAt": int(time.time() * 1000), "errorMessage": error_msg, - "executionLog": [], - "stepCount": tracked_step_count, - "currentPhase": "error", - "currentStepIndex": tracked_step_count, + "executionLog": [step for _, step in tracked_steps], + "stepCount": final_step_count, + "currentPhase": terminal_status, + "currentStepIndex": final_step_count, "updatedAt": int(time.time() * 1000), } ) + if progress_writer is not None: + await progress_writer.close_and_drain() await record_execution_result( canonical_workflow_id, tracked_execution["id"], current_data, + steps=tracked_steps, ) _emit_metadata( { @@ -1045,9 +1024,9 @@ async def _flush_pending_step() -> None: "workflow_name": workflow_name, "total_nodes": workflow_total_nodes, "workflow_execution_id": tracked_execution["id"], - "status": "error", - "phase": "error", - "step_index": tracked_step_count, + "status": terminal_status, + "phase": terminal_status, + "step_index": final_step_count, }, } ) @@ -1061,6 +1040,7 @@ async def _flush_pending_step() -> None: workflow_name=workflow_name, total_nodes=workflow_total_nodes, workflow_execution_id=tracked_execution["id"] if tracked_execution else None, - status="FAILED", + status="CANCELLED" if terminal_status == "cancelled" else "FAILED", + steps=final_step_count, ), ) diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index c1c7cb443..3dbc59cd4 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -5,6 +5,7 @@ import asyncio from itertools import islice import sys +import threading import time import uuid from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple @@ -337,18 +338,6 @@ def derive_loop_progress( # Keep this intentionally small so high-frequency workflows do not keep # inflating the SQLite row set and matching JSONL audit files indefinitely. _MAX_EXECUTION_HISTORY_PER_WORKFLOW = 30 -# Per-workflow trim lock. Trims are awaited by the writer so the retention cap -# is enforced before ``record_execution_result`` returns, while concurrent runs -# for the same workflow serialize instead of skipping cleanup. -_trim_locks: Dict[str, asyncio.Lock] = {} - - -def _get_trim_lock(workflow_id: str) -> asyncio.Lock: - lock = _trim_locks.get(workflow_id) - if lock is None: - lock = asyncio.Lock() - _trim_locks[workflow_id] = lock - return lock def workflow_execution_key(exec_id: str) -> str: @@ -410,11 +399,9 @@ def __init__( self, *, exec_id: str, - capture_steps: bool = True, step_compactor: Callable[[Any], Dict[str, Any]] = compact_step_for_storage, ) -> None: self.exec_id = exec_id - self.capture_steps = capture_steps self.step_compactor = step_compactor self.step_count = 0 self.summary: Dict[str, Any] = {} @@ -444,8 +431,7 @@ def on_step_complete(self, step_result: Any) -> None: "updatedAt": int(time.time() * 1000), } ) - if self.capture_steps: - self._pending_steps.append((self.step_count, step_dict)) + self._pending_steps.append((self.step_count, step_dict)) def take_steps(self) -> List[Tuple[int, Dict[str, Any]]]: """Return buffered steps for the final execution transaction.""" @@ -454,6 +440,94 @@ def take_steps(self) -> List[Tuple[int, Dict[str, Any]]]: return pending_steps +class ExecutionProgressWriter: + """Coalesce nonblocking execution-summary updates onto one SQLite writer.""" + + def __init__(self, execution_summary: Dict[str, Any]) -> None: + self._loop = asyncio.get_running_loop() + self._summary = compact_execution_summary(execution_summary) + self._pending_summary: Optional[Dict[str, Any]] = None + self._pending_waiters: List[asyncio.Future[None]] = [] + self._writer_task: Optional[asyncio.Task[None]] = None + self._submission_lock = threading.Lock() + self._closed = False + + def submit(self, update: Dict[str, Any]) -> None: + """Queue an update from any thread without waiting for persistence.""" + with self._submission_lock: + if self._closed: + return + self._loop.call_soon_threadsafe(self._merge_update, dict(update), None) + + async def update(self, update: Dict[str, Any]) -> None: + """Queue and await an owner-loop update, preserving submission order.""" + if asyncio.get_running_loop() is not self._loop: + raise RuntimeError("ExecutionProgressWriter.update must run on its owner loop") + + waiter = self._loop.create_future() + with self._submission_lock: + if self._closed: + return + self._loop.call_soon(self._merge_update, dict(update), waiter) + await waiter + + async def close_and_drain(self) -> None: + """Reject new updates and flush every update accepted before closing.""" + if asyncio.get_running_loop() is not self._loop: + raise RuntimeError("ExecutionProgressWriter.close_and_drain must run on its owner loop") + + barrier = self._loop.create_future() + with self._submission_lock: + self._closed = True + self._loop.call_soon(barrier.set_result, None) + await barrier + + writer_task = self._writer_task + if writer_task is not None: + await asyncio.shield(writer_task) + + def _merge_update( + self, + update: Dict[str, Any], + waiter: Optional[asyncio.Future[None]], + ) -> None: + self._summary.update(update) + self._pending_summary = compact_execution_summary(self._summary) + if waiter is not None: + self._pending_waiters.append(waiter) + if self._writer_task is None: + exec_id = str(self._summary.get("id") or "unknown") + self._writer_task = self._loop.create_task( + self._flush(), + name=f"workflow-progress-{exec_id}", + ) + + async def _flush(self) -> None: + try: + while self._pending_summary is not None: + summary = self._pending_summary + waiters = self._pending_waiters + self._pending_summary = None + self._pending_waiters = [] + try: + await WorkflowStore.upsert_execution(summary) + except Exception as exc: + log.warning( + "workflow.progress.update_failed", + { + "workflow_id": summary.get("workflowId"), + "exec_id": summary.get("id"), + "error": str(exc), + }, + ) + finally: + for waiter in waiters: + if not waiter.done(): + waiter.set_result(None) + finally: + self._writer_task = None + + def _prepare_execution_steps(execution_log: Any) -> List[Tuple[int, Dict[str, Any]]]: """Compact an inline execution log for one final batch transaction.""" if not isinstance(execution_log, list): @@ -594,15 +668,48 @@ async def record_execution_result( started_at = summary_data.get("startedAt", 0) finished_at = summary_data.get("finishedAt", int(time.time() * 1000)) duration = max(0.0, (finished_at - started_at) / 1000.0) - trimmed_exec_ids = await WorkflowStore.complete_execution( + + await WorkflowStore.complete_execution( compact_execution_summary(summary_data), prepared_steps, - success=success, - duration=float(duration), - history_limit=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, ) - if not isinstance(trimmed_exec_ids, list): - trimmed_exec_ids = [] + + try: + await WorkflowStore.increment_stats( + workflow_id, + success=success, + duration=float(duration), + ) + except Exception as exc: + log.warning( + "workflow.stats.update_failed", + { + "workflow_id": workflow_id, + "exec_id": exec_id, + "error": str(exc), + }, + ) + + trimmed_exec_ids: List[str] = [] + try: + trimmed_exec_ids = await WorkflowStore.trim_executions( + workflow_id, + keep=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, + ) + except Exception as exc: + log.error( + "workflow.history.trim_failed", + { + "workflow_id": workflow_id, + "exec_id": exec_id, + "error": str(exc), + }, + ) + + audit_data = dict(exec_data) + audit_data["executionLog"] = [ + step for _, step in sorted(prepared_steps, key=lambda item: item[0]) + ] # Recorder writes to its own SQLite tables and can be slow under load. # Run it as a background task so the syslog/HTTP dispatcher can release the @@ -614,7 +721,7 @@ async def _record_audit() -> None: await Recorder.record_workflow_execution( exec_id=exec_id, workflow_id=workflow_id, - run_result=exec_data, + run_result=audit_data, ) except Exception as exc: log.debug( @@ -645,65 +752,7 @@ async def _record_audit() -> None: await Recorder.record_workflow_execution( exec_id=exec_id, workflow_id=workflow_id, - run_result=exec_data, + run_result=audit_data, ) except Exception: pass - - -async def _delete_execution_history_record( - execution_key: str, - *, - index_key: Optional[str] = None, -) -> None: - exec_id = execution_key.rsplit("/", 1)[-1] - deleted_steps = await WorkflowStore.clear_steps(exec_id) - removed_execution = await WorkflowStore.delete_execution(exec_id) - record_path = Recorder.paths().workflow_dir / f"{exec_id}.jsonl" - await asyncio.to_thread(record_path.unlink, missing_ok=True) - log.debug( - "workflow.history.trim_deleted", - { - "exec_id": exec_id, - "execution_key": execution_key, - "steps": deleted_steps, - "removed_execution": removed_execution, - }, - ) - - -async def _trim_execution_history(workflow_id: str) -> None: - """Delete the oldest execution records once the per-workflow cap is exceeded. - - New records carry a per-workflow ``workflow_execution_index`` key, so hot - trims avoid scanning unrelated workflows. This path is intentionally - index-only: if an old execution has no index key, it is outside the hot - retention path and should be handled by a separate migration/GC task. - - A per-workflow lock serializes concurrent trims. Cleanup is awaited by - ``record_execution_result`` so the retention cap is enforced synchronously - instead of being an opportunistic background task. - """ - lock = _get_trim_lock(workflow_id) - async with lock: - failures: List[str] = [] - for exec_id in await WorkflowStore.trim_executions( - workflow_id, - keep=_MAX_EXECUTION_HISTORY_PER_WORKFLOW, - ): - try: - record_path = Recorder.paths().workflow_dir / f"{exec_id}.jsonl" - await asyncio.to_thread(record_path.unlink, missing_ok=True) - except Exception as exc: - failures.append(f"{exec_id}: {exc}") - log.warning( - "workflow.history.trim_delete_failed", - { - "workflow_id": workflow_id, - "exec_id": exec_id, - "error": str(exc), - }, - ) - - if failures: - raise RuntimeError("Failed to trim workflow execution history: " + "; ".join(failures[:3])) diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index d596daa80..7fdb7a101 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -454,10 +454,7 @@ async def _execute_run( persist=False, ) exec_id = str(exec_data["id"]) - step_recorder = ExecutionStepRecorder( - exec_id=exec_id, - capture_steps=False, - ) + step_recorder = ExecutionStepRecorder(exec_id=exec_id) current = self._status.get(workflow_id) or self._base_status(workflow_id) current["lastRunAt"] = started_at_ms current["activeRuns"] = self._cleanup_done_runs(workflow_id) diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index db049f3a2..b4d9b129f 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -158,6 +158,8 @@ async def raw_db(cls) -> aiosqlite.Connection: @classmethod async def _completion_db(cls) -> aiosqlite.Connection: + if cls._initialized and cls._init_pid is not None and cls._init_pid != os.getpid(): + await cls.init() if not cls._completion_conn or not cls._initialized: await cls.init() return cls._completion_conn # type: ignore[return-value] @@ -439,21 +441,12 @@ async def trim_executions(cls, workflow_id: str, *, keep: int) -> List[str]: return exec_ids @classmethod - async def record_step( - cls, - exec_id: str, - step_index: int, - step_payload: Dict[str, Any], - ) -> None: - await cls.record_steps(exec_id, [(step_index, step_payload)]) - - @classmethod - async def record_steps( + def _step_rows( cls, exec_id: str, steps: Iterable[Tuple[int, Dict[str, Any]]], - ) -> None: - rows = [ + ) -> List[Tuple[Any, ...]]: + return [ ( exec_id, int(step_index), @@ -466,6 +459,23 @@ async def record_steps( ) for step_index, step_payload in steps ] + + @classmethod + async def record_step( + cls, + exec_id: str, + step_index: int, + step_payload: Dict[str, Any], + ) -> None: + await cls.record_steps(exec_id, [(step_index, step_payload)]) + + @classmethod + async def record_steps( + cls, + exec_id: str, + steps: Iterable[Tuple[int, Dict[str, Any]]], + ) -> None: + rows = cls._step_rows(exec_id, steps) if not rows: return db = await cls._db() @@ -484,12 +494,8 @@ async def complete_execution( cls, exec_data: Dict[str, Any], steps: Iterable[Tuple[int, Dict[str, Any]]], - *, - success: bool, - duration: float, - history_limit: Optional[int] = None, - ) -> List[str]: - """Persist one completed execution, stats, and retention in one transaction.""" + ) -> None: + """Atomically persist one final execution summary and its step batch.""" db = await cls._completion_db() payload = dict(exec_data) exec_id = str(payload.get("id") or "") @@ -497,22 +503,7 @@ async def complete_execution( if not exec_id or not workflow_id: raise ValueError("workflow execution requires id and workflowId") - step_rows = [ - ( - exec_id, - int(step_index), - step_payload.get("node_id"), - step_payload.get("node_type") or step_payload.get("type"), - cls._json_dumps(step_payload.get("inputs") or {}), - cls._json_dumps(step_payload.get("outputs") or {}), - step_payload.get("error"), - cls._json_dumps(step_payload), - ) - for step_index, step_payload in steps - ] - runtime = float(duration) - success_delta = 1 if success else 0 - error_delta = 0 if success else 1 + step_rows = cls._step_rows(exec_id, steps) lock = cls._completion_lock if lock is None: lock = asyncio.Lock() @@ -559,57 +550,19 @@ async def complete_execution( cls._json_dumps(payload), ), ) - await db.execute( - """ - INSERT INTO workflow_stats ( - workflow_id, call_count, success_count, error_count, - total_runtime, avg_runtime, thumbs_up, thumbs_down, updated_at - ) - VALUES (?, 1, ?, ?, ?, ?, 0, 0, ?) - ON CONFLICT(workflow_id) DO UPDATE SET - call_count = workflow_stats.call_count + 1, - success_count = workflow_stats.success_count + excluded.success_count, - error_count = workflow_stats.error_count + excluded.error_count, - total_runtime = workflow_stats.total_runtime + excluded.total_runtime, - avg_runtime = ( - workflow_stats.total_runtime + excluded.total_runtime - ) / (workflow_stats.call_count + 1), - updated_at = excluded.updated_at - """, - ( - workflow_id, - success_delta, - error_delta, - runtime, - runtime, - cls._now_ms(), - ), - ) - trimmed_exec_ids: List[str] = [] - if history_limit is not None: - async with db.execute( - """ - SELECT id FROM workflow_executions - WHERE workflow_id = ? - ORDER BY started_at DESC, rowid DESC - LIMIT -1 OFFSET ? - """, - (workflow_id, max(int(history_limit), 0)), - ) as cur: - trimmed_exec_ids = [str(row["id"]) for row in await cur.fetchall()] - for trimmed_exec_id in trimmed_exec_ids: - await db.execute( - "DELETE FROM workflow_execution_steps WHERE exec_id = ?", - (trimmed_exec_id,), - ) - await db.execute( - "DELETE FROM workflow_executions WHERE id = ?", - (trimmed_exec_id,), - ) await db.commit() - return trimmed_exec_ids - except Exception: - await db.rollback() + except BaseException: + try: + await db.rollback() + except BaseException as rollback_exc: + log.error( + "workflow.store.completion_rollback_failed", + { + "workflow_id": workflow_id, + "exec_id": exec_id, + "error": str(rollback_exc), + }, + ) raise @classmethod diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index 2a3d2aedf..d41aa012d 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -251,10 +251,7 @@ async def _execute_workflow_effect( persist=False, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder( - exec_id=exec_id, - capture_steps=False, - ) + step_recorder = ExecutionStepRecorder(exec_id=exec_id) started_at = time.time() tool_context = None try: diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 6e494b566..2134fa9c9 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -616,7 +616,15 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 } assert captured_exec_data["executionLog"] == [] assert captured_exec_data["stepCount"] == 2 - assert captured_steps == [] + assert [step_index for step_index, _ in captured_steps] == [1, 2] + assert [step["node_id"] for _, step in captured_steps] == [ + "receive_alert", + "dedup_and_write", + ] + assert captured_steps[0][1]["inputs"]["kafka_message"]["_type"] == "dict" + assert captured_steps[0][1]["outputs"] == {"_raw_alerts_count": 1} + assert captured_steps[1][1]["inputs"] == {"_filtered_alerts_count": 1} + assert captured_steps[1][1]["outputs"] == {"_enriched_alerts_count": 1} assert len(json.dumps(captured_exec_data, ensure_ascii=False)) < 10_000 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 3ad929893..0ce283747 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -427,7 +427,17 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 action_name="trigger:syslog", ) trigger_tool_context.cleanup.assert_awaited_once_with(trigger_tool_context.context) - assert recorded_steps == [] + assert recorded_steps == [ + ( + 1, + { + "node_id": "receive_alert", + "node_type": "python", + "inputs": {"message": "demo"}, + "outputs": {"ok": True}, + }, + ) + ] assert recorded_exec_data["triggerId"] == "syslog-alerts" assert recorded_exec_data["triggerSource"] == "udp://0.0.0.0:5514" assert recorded_exec_data["executionLog"] == [] diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index 5684a5cef..f0891bafa 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -1,3 +1,4 @@ +import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock, Mock @@ -243,7 +244,7 @@ def run_workflow_mock(**kwargs): kwargs["on_step_complete"](step_result) return SimpleNamespace( outputs={"ok": True}, - history=[step_result.model_dump(mode="json")], + history=[], last_node_id="node-1", steps=1, ) @@ -272,6 +273,14 @@ def run_workflow_mock(**kwargs): req = workflow_module.WorkflowRunRequest(inputs={"ip": "8.8.8.8"}, trace=False) tool_context = ToolContext(session_id="session-1", message_id="message-1", agent="rex") + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-1", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) await workflow_module._run_workflow_execution_task( workflow_id="wf-1", @@ -279,22 +288,25 @@ def run_workflow_mock(**kwargs): req=req, exec_id="exec-1", cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, tool_context=tool_context, ) init_mock.assert_not_awaited() run_mock.assert_called_once() assert run_mock.call_args.kwargs["tool_context"] is tool_context - upsert_execution.assert_not_awaited() + assert upsert_execution.await_count >= 1 + assert all(call.args[0]["executionLog"] == [] for call in upsert_execution.await_args_list) + assert upsert_execution.await_args.args[0]["currentNodeId"] == "node-1" record_result.assert_awaited_once() - assert record_result.await_args.args[2]["executionLog"] == [ - { - "node_id": "node-1", - "node_type": "tool", - "inputs": {}, - "outputs": {"ok": True}, - } - ] + expected_step = { + "node_id": "node-1", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + } + assert record_result.await_args.args[2]["executionLog"] == [expected_step] + assert record_result.await_args.kwargs["steps"] == [(1, expected_step)] @pytest.mark.asyncio @@ -320,34 +332,392 @@ def run_workflow_mock(**kwargs): monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr( + workflow_module.WorkflowStore, + "upsert_execution", + AsyncMock(return_value=None), + ) monkeypatch.setattr(workflow_module, "compact_outputs_for_storage", lambda value: value) monkeypatch.setattr(workflow_module, "compact_history_for_storage", lambda value: value) cancel_event = workflow_module.threading.Event() cancel_event.set() + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-cancelled", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, req=workflow_module.WorkflowRunRequest(inputs={"message": "hello"}, trace=False), exec_id="exec-cancelled", cancel_event=cancel_event, + progress_writer=progress_writer, ) record_result.assert_awaited_once() final_data = record_result.await_args.args[2] assert final_data["status"] == "cancelled" assert final_data["stepCount"] == 1 - assert final_data["executionLog"] == [ + pending_step = { + "node_id": "node-1", + "node_type": "tool", + "inputs": {"message": "hello"}, + "outputs": {}, + "error": "Run cancelled before node completed", + } + assert final_data["executionLog"] == [pending_step] + assert record_result.await_args.kwargs["steps"] == [(1, pending_step)] + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_keeps_completed_and_pending_step_indices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cancel_event = workflow_module.threading.Event() + completed_step = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "python", + "inputs": {"value": 1}, + "outputs": {"value": 2}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {"value": 1}, + ) + kwargs["on_step_complete"](completed_step) + cancel_event.set() + kwargs["on_step_start"]( + "run-1", + 2, + SimpleNamespace(id="node-2", type="tool"), + {"message": "hello"}, + ) + return SimpleNamespace( + run_id="run-1", + outputs={"value": 2}, + history=[], + last_node_id="node-2", + steps=1, + ) + + record_result = AsyncMock(return_value=None) + upsert_execution = AsyncMock(return_value=None) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) + + progress_writer = workflow_module.ExecutionProgressWriter( { + "id": "exec-partial-cancel", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={"value": 1}, trace=False), + exec_id="exec-partial-cancel", + cancel_event=cancel_event, + progress_writer=progress_writer, + ) + + steps = record_result.await_args.kwargs["steps"] + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + assert steps[1][1]["error"] == "Run cancelled before node completed" + assert record_result.await_args.args[2]["status"] == "cancelled" + assert upsert_execution.await_args.args[0]["currentPhase"] == "cancelling" + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_keeps_steps_when_runner_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + completed_step = SimpleNamespace( + model_dump=lambda mode: { "node_id": "node-1", - "node_type": "tool", - "inputs": {"message": "hello"}, - "outputs": {}, - "error": "Run cancelled before node completed", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"](completed_step) + kwargs["on_step_start"]( + "run-1", + 2, + SimpleNamespace(id="node-2", type="tool"), + {"message": "hello"}, + ) + raise RuntimeError("runner failed") + + record_result = AsyncMock(return_value=None) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr( + workflow_module.WorkflowStore, + "upsert_execution", + AsyncMock(return_value=None), + ) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-runner-error", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={}, trace=False), + exec_id="exec-runner-error", + cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, + ) + + final_data = record_result.await_args.args[2] + steps = record_result.await_args.kwargs["steps"] + assert final_data["status"] == "error" + assert final_data["errorMessage"] == "runner failed" + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + + +@pytest.mark.asyncio +async def test_run_workflow_execution_task_does_not_reclassify_persistence_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_result = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"](step_result) + return SimpleNamespace( + outputs={"ok": True}, + history=[], + last_node_id="node-1", + steps=1, + ) + + record_result = AsyncMock(side_effect=RuntimeError("storage failed")) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) + monkeypatch.setattr( + workflow_module.WorkflowStore, + "upsert_execution", + AsyncMock(return_value=None), + ) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-storage-error", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], } + ) + + with pytest.raises(RuntimeError, match="storage failed"): + await workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={}, trace=False), + exec_id="exec-storage-error", + cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, + ) + + record_result.assert_awaited_once() + final_data = record_result.await_args.args[2] + assert final_data["status"] == "success" + assert record_result.await_args.kwargs["steps"] == [ + ( + 1, + { + "node_id": "node-1", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + }, + ) ] +@pytest.mark.asyncio +async def test_run_workflow_callbacks_do_not_wait_for_blocked_progress_write( + monkeypatch: pytest.MonkeyPatch, +) -> None: + write_started = asyncio.Event() + release_write = asyncio.Event() + runner_finished = workflow_module.threading.Event() + write_order: list[str] = [] + step_result = SimpleNamespace( + model_dump=lambda mode: { + "node_id": "node-1", + "node_type": "python", + "inputs": {}, + "outputs": {"ok": True}, + } + ) + + async def blocked_upsert(_summary): + write_order.append("progress-start") + write_started.set() + await release_write.wait() + write_order.append("progress-end") + + async def record_result(*args, **kwargs): # noqa: ANN002, ANN003 + write_order.append("final") + + def run_workflow_mock(**kwargs): + kwargs["on_step_start"]( + "run-1", + 1, + SimpleNamespace(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"](step_result) + runner_finished.set() + return SimpleNamespace( + outputs={"ok": True}, + history=[], + last_node_id="node-1", + steps=1, + ) + + record_result_mock = AsyncMock(side_effect=record_result) + monkeypatch.setattr(workflow_module, "run_workflow", Mock(side_effect=run_workflow_mock)) + monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) + monkeypatch.setattr(workflow_module, "_record_execution_result", record_result_mock) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", blocked_upsert) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-blocked-progress", + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + ) + task = asyncio.create_task( + workflow_module._run_workflow_execution_task( + workflow_id="wf-1", + workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, + req=workflow_module.WorkflowRunRequest(inputs={}, trace=False), + exec_id="exec-blocked-progress", + cancel_event=workflow_module.threading.Event(), + progress_writer=progress_writer, + ) + ) + + await write_started.wait() + assert await asyncio.to_thread(runner_finished.wait, 0.1) + record_result_mock.assert_not_awaited() + release_write.set() + await task + + record_result_mock.assert_awaited_once() + assert write_order[-2:] == ["progress-end", "final"] + + +@pytest.mark.asyncio +async def test_cancel_workflow_execution_uses_active_progress_writer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + persisted_summaries: list[dict] = [] + + async def capture_upsert(summary): + persisted_summaries.append(dict(summary)) + + monkeypatch.setattr( + workflow_module.WorkflowStore, + "get_execution", + AsyncMock( + return_value={ + "id": "exec-cancel-route", + "workflowId": "wf-1", + "status": "running", + "currentPhase": "running", + "executionLog": [], + } + ), + ) + monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", capture_upsert) + + progress_writer = workflow_module.ExecutionProgressWriter( + { + "id": "exec-cancel-route", + "workflowId": "wf-1", + "status": "running", + "currentPhase": "queued", + "executionLog": [], + } + ) + progress_writer.submit({"currentPhase": "running", "currentNodeId": "node-1"}) + cancel_event = workflow_module.threading.Event() + current_task = asyncio.current_task() + assert current_task is not None + workflow_module._active_workflow_executions["exec-cancel-route"] = ( + workflow_module.ActiveWorkflowExecution( + workflow_id="wf-1", + task=current_task, + cancel_event=cancel_event, + progress_writer=progress_writer, + ) + ) + + try: + response = await workflow_module.cancel_workflow_execution( + "wf-1", + "exec-cancel-route", + ) + await progress_writer.close_and_drain() + finally: + workflow_module._active_workflow_executions.pop("exec-cancel-route", None) + + assert response["status"] == "accepted" + assert cancel_event.is_set() + assert persisted_summaries[-1]["currentPhase"] == "cancelling" + assert persisted_summaries[-1]["currentNodeId"] == "node-1" + assert persisted_summaries[-1]["errorMessage"] == "Cancellation requested" + + @pytest.mark.asyncio async def test_workflow_tool_context_preserves_current_opaque_extension_context( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index 9245eea3a..0ceb57849 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +from types import SimpleNamespace from typing import Any, Dict, List from unittest.mock import AsyncMock, patch @@ -27,8 +28,8 @@ DEFAULT_GENERIC_SEQUENCE_THRESHOLD, DEFAULT_LARGE_LIST_KEYS, DEFAULT_MAX_INLINE_COLLECTION_BYTES, + ExecutionProgressWriter, ExecutionStepRecorder, - _trim_execution_history, compact_history_for_storage, compact_execution_summary, compact_outputs_for_storage, @@ -371,9 +372,130 @@ def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: record_steps.assert_not_awaited() +@pytest.mark.asyncio +async def test_progress_writer_submits_without_waiting_and_coalesces_updates() -> None: + write_started = asyncio.Event() + release_write = asyncio.Event() + writes: List[Dict[str, Any]] = [] + active_writes = 0 + max_active_writes = 0 + + async def blocked_upsert(summary: Dict[str, Any]) -> None: + nonlocal active_writes, max_active_writes + active_writes += 1 + max_active_writes = max(max_active_writes, active_writes) + writes.append(dict(summary)) + try: + if len(writes) == 1: + write_started.set() + await release_write.wait() + finally: + active_writes -= 1 + + writer = ExecutionProgressWriter( + { + "id": "exec-progress", + "workflowId": "wf-progress", + "status": "running", + "executionLog": [{"node_id": "ignored"}], + } + ) + + with patch.object(WorkflowStore, "upsert_execution", side_effect=blocked_upsert): + writer.submit({"currentNodeId": "node-1", "currentStepIndex": 1}) + await write_started.wait() + await asyncio.wait_for( + asyncio.to_thread( + writer.submit, + {"currentNodeId": "node-2", "currentStepIndex": 2}, + ), + timeout=0.1, + ) + await asyncio.to_thread( + writer.submit, + {"currentNodeId": "node-3", "currentStepIndex": 3}, + ) + release_write.set() + await writer.close_and_drain() + + assert max_active_writes == 1 + assert len(writes) == 2 + assert writes[0]["currentNodeId"] == "node-1" + assert writes[-1]["currentNodeId"] == "node-3" + assert writes[-1]["currentStepIndex"] == 3 + assert writes[-1]["executionLog"] == [] + + +@pytest.mark.asyncio +async def test_progress_writer_awaited_update_is_ordered_and_close_rejects_late_updates() -> None: + writes: List[Dict[str, Any]] = [] + + async def capture_upsert(summary: Dict[str, Any]) -> None: + writes.append(dict(summary)) + + writer = ExecutionProgressWriter( + { + "id": "exec-cancelling", + "workflowId": "wf-cancelling", + "status": "running", + "currentPhase": "queued", + "executionLog": [], + } + ) + + with patch.object(WorkflowStore, "upsert_execution", side_effect=capture_upsert): + await asyncio.to_thread( + writer.submit, + {"currentPhase": "running", "currentNodeId": "node-1"}, + ) + await writer.update({"currentPhase": "cancelling"}) + await writer.close_and_drain() + writer.submit({"currentPhase": "running", "currentNodeId": "late-node"}) + await asyncio.sleep(0) + + assert writes[-1]["currentPhase"] == "cancelling" + assert writes[-1]["currentNodeId"] == "node-1" + assert all(write.get("currentNodeId") != "late-node" for write in writes) + + +@pytest.mark.asyncio +async def test_progress_writer_logs_write_failures_without_raising() -> None: + writer = ExecutionProgressWriter( + { + "id": "exec-write-failure", + "workflowId": "wf-write-failure", + "status": "running", + "executionLog": [], + } + ) + + with patch.object( + WorkflowStore, + "upsert_execution", + AsyncMock(side_effect=RuntimeError("database locked")), + ): + await writer.update({"currentPhase": "running"}) + await writer.close_and_drain() + + @pytest.mark.asyncio async def test_record_execution_result_backfills_execution_log_steps() -> None: - complete_execution = AsyncMock(return_value=[]) + calls: List[str] = [] + + async def complete_execution(*args, **kwargs): # noqa: ANN002, ANN003 + calls.append("complete") + + async def increment_stats(*args, **kwargs): # noqa: ANN002, ANN003 + calls.append("stats") + + async def trim_executions(*args, **kwargs): # noqa: ANN002, ANN003 + calls.append("trim") + return [] + + complete_execution_mock = AsyncMock(side_effect=complete_execution) + increment_stats_mock = AsyncMock(side_effect=increment_stats) + trim_executions_mock = AsyncMock(side_effect=trim_executions) + record_audit = AsyncMock(return_value=None) exec_data = { "id": "exec-1", "workflowId": "wf", @@ -390,33 +512,39 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 raise RuntimeError with ( - patch.object(WorkflowStore, "complete_execution", complete_execution), - patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch.object(WorkflowStore, "complete_execution", complete_execution_mock), + patch.object(WorkflowStore, "increment_stats", increment_stats_mock), + patch.object(WorkflowStore, "trim_executions", trim_executions_mock), + patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), ): await record_execution_result("wf", "exec-1", exec_data) - complete_execution.assert_awaited_once() - summary, steps = complete_execution.await_args.args + assert calls == ["complete", "stats", "trim"] + complete_execution_mock.assert_awaited_once() + summary, steps = complete_execution_mock.await_args.args + assert complete_execution_mock.await_args.kwargs == {} assert steps[0][0] == 1 assert steps[0][1]["outputs"] == {"_raw_alerts_count": 150} assert steps[1][0] == 2 assert steps[1][1]["inputs"] == {"_filtered_alerts_count": 150} assert summary["executionLog"] == [] assert summary["stepCount"] == 2 - assert complete_execution.await_args.kwargs == { - "success": True, - "duration": 1.0, - "history_limit": 30, - } + increment_stats_mock.assert_awaited_once_with("wf", success=True, duration=1.0) + trim_executions_mock.assert_awaited_once_with("wf", keep=30) + audit_data = record_audit.await_args.kwargs["run_result"] + assert audit_data["executionLog"] == [step for _, step in steps] @pytest.mark.asyncio async def test_record_execution_result_accepts_explicit_step_batch() -> None: - complete_execution = AsyncMock(return_value=[]) + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(return_value=None) + trim_executions = AsyncMock(return_value=[]) + record_audit = AsyncMock(return_value=None) explicit_steps = [ - (1, {"node_id": "step-1", "outputs": {"ok": True}}), (2, {"node_id": "step-2", "outputs": {"ok": True}}), + (1, {"node_id": "step-1", "outputs": {"ok": True}}), ] exec_data = { "id": "exec-trigger", @@ -433,7 +561,9 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 with ( patch.object(WorkflowStore, "complete_execution", complete_execution), - patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), ): await record_execution_result( @@ -444,13 +574,88 @@ def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 ) summary, persisted_steps = complete_execution.await_args.args + assert complete_execution.await_args.kwargs == {} assert summary["executionLog"] == [] assert persisted_steps == explicit_steps - assert complete_execution.await_args.kwargs == { - "success": True, - "duration": 0.01, - "history_limit": 30, - } + increment_stats.assert_awaited_once_with("wf-trigger", success=True, duration=0.01) + trim_executions.assert_awaited_once_with("wf-trigger", keep=30) + audit_data = record_audit.await_args.kwargs["run_result"] + assert [step["node_id"] for step in audit_data["executionLog"]] == [ + "step-1", + "step-2", + ] + + +@pytest.mark.asyncio +async def test_record_execution_result_stats_failure_does_not_block_retention() -> None: + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(side_effect=RuntimeError("stats locked")) + trim_executions = AsyncMock(return_value=[]) + + def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + with ( + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + ): + await record_execution_result( + "wf-stats-failure", + "exec-stats-failure", + { + "id": "exec-stats-failure", + "workflowId": "wf-stats-failure", + "status": "success", + "duration": 0.5, + "executionLog": [], + }, + ) + + complete_execution.assert_awaited_once() + increment_stats.assert_awaited_once() + trim_executions.assert_awaited_once_with("wf-stats-failure", keep=30) + + +@pytest.mark.asyncio +async def test_record_execution_result_retention_failure_keeps_committed_execution() -> None: + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(return_value=None) + trim_executions = AsyncMock(side_effect=RuntimeError("retention locked")) + + def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + with ( + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + ): + await record_execution_result( + "wf-retention-failure", + "exec-retention-failure", + { + "id": "exec-retention-failure", + "workflowId": "wf-retention-failure", + "status": "error", + "duration": 0.5, + "executionLog": [], + }, + ) + + complete_execution.assert_awaited_once() + increment_stats.assert_awaited_once_with( + "wf-retention-failure", + success=False, + duration=0.5, + ) + trim_executions.assert_awaited_once_with("wf-retention-failure", keep=30) def test_compact_history_compacts_each_step_inputs() -> None: @@ -538,62 +743,55 @@ def test_compact_outputs_covers_raw_alerts_in_input_params() -> None: @pytest.mark.asyncio -async def test_trim_execution_history_keeps_only_30_and_deletes_matching_jsonl( - tmp_path, -) -> None: - workflow_id = "wf-trim" - for idx in range(32): - exec_id = f"exec-{idx:02d}" - workflow_record = tmp_path / "workflow" / f"{exec_id}.jsonl" - workflow_record.parent.mkdir(parents=True, exist_ok=True) - workflow_record.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") - - # Another workflow's record should be ignored entirely because the trim - # only reads workflow_execution_index//. - other_record = tmp_path / "workflow" / "other-exec.jsonl" - other_record.parent.mkdir(parents=True, exist_ok=True) - other_record.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") - - trim_mock = AsyncMock(return_value=["exec-00", "exec-01"]) +async def test_record_execution_result_deletes_jsonl_for_trimmed_executions(tmp_path) -> None: + workflow_dir = tmp_path / "workflow" + workflow_dir.mkdir(parents=True, exist_ok=True) + trimmed_paths = [workflow_dir / "exec-00.jsonl", workflow_dir / "exec-01.jsonl"] + retained_path = workflow_dir / "exec-02.jsonl" + for record_path in [*trimmed_paths, retained_path]: + record_path.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") + + complete_execution = AsyncMock(return_value=None) + increment_stats = AsyncMock(return_value=None) + trim_executions = AsyncMock(return_value=["exec-00", "exec-01"]) + record_audit = AsyncMock(return_value=None) + created_tasks: List[asyncio.Task[None]] = [] + real_create_task = asyncio.create_task + + def capture_create_task(coro, *args, **kwargs): # noqa: ANN001 + task = real_create_task(coro, *args, **kwargs) + created_tasks.append(task) + return task with ( - patch.object(WorkflowStore, "trim_executions", trim_mock), - patch("flocks.session.recorder._record_dir", return_value=tmp_path), - ): - await _trim_execution_history(workflow_id) - - trim_mock.assert_awaited_once_with(workflow_id, keep=30) - assert not (tmp_path / "workflow" / "exec-00.jsonl").exists() - assert not (tmp_path / "workflow" / "exec-01.jsonl").exists() - assert (tmp_path / "workflow" / "exec-02.jsonl").exists() - assert other_record.exists() - - -@pytest.mark.asyncio -async def test_trim_execution_history_uses_index_without_full_scan(tmp_path) -> None: - workflow_id = "wf-indexed" - for idx in range(32): - exec_id = f"exec-{idx:02d}" - workflow_record = tmp_path / "workflow" / f"{exec_id}.jsonl" - workflow_record.parent.mkdir(parents=True, exist_ok=True) - workflow_record.write_text('{"type":"workflow.summary"}\n', encoding="utf-8") - - trim_mock = AsyncMock(return_value=["exec-00", "exec-01"]) - - with ( - patch.object(WorkflowStore, "trim_executions", trim_mock), - patch("flocks.session.recorder._record_dir", return_value=tmp_path), + patch.object(WorkflowStore, "complete_execution", complete_execution), + patch.object(WorkflowStore, "increment_stats", increment_stats), + patch.object(WorkflowStore, "trim_executions", trim_executions), + patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), + patch( + "flocks.workflow.execution_store.Recorder.paths", + return_value=SimpleNamespace(workflow_dir=workflow_dir), + ), + patch( + "flocks.workflow.execution_store.asyncio.create_task", + side_effect=capture_create_task, + ), ): - await _trim_execution_history(workflow_id) - - trim_mock.assert_awaited_once_with(workflow_id, keep=30) - assert not (tmp_path / "workflow" / "exec-00.jsonl").exists() - assert not (tmp_path / "workflow" / "exec-01.jsonl").exists() - + await record_execution_result( + "wf-trim", + "exec-32", + { + "id": "exec-32", + "workflowId": "wf-trim", + "status": "success", + "duration": 0.25, + "executionLog": [], + }, + steps=[(1, {"node_id": "node-1", "outputs": {"ok": True}})], + ) + await asyncio.gather(*created_tasks) -@pytest.mark.asyncio -async def test_trim_execution_history_surfaces_delete_failures() -> None: - workflow_id = "wf-trim-fail" - with patch.object(WorkflowStore, "trim_executions", AsyncMock(side_effect=RuntimeError("locked"))): - with pytest.raises(RuntimeError, match="locked"): - await _trim_execution_history(workflow_id) + trim_executions.assert_awaited_once_with("wf-trim", keep=30) + record_audit.assert_awaited_once() + assert all(not path.exists() for path in trimmed_paths) + assert retained_path.exists() diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index 82ee82cf0..f4069d1bd 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -250,7 +250,17 @@ def _fake_run_workflow( # noqa: ANN001 assert recorded_results[0]["executionLog"] == [] assert recorded_results[0]["stepCount"] == 1 assert recorded_results[0]["loopProgress"]["total_iterations"] == 2 - assert recorded_steps == [] + assert recorded_steps == [ + ( + 1, + { + "node_id": "load", + "node_type": "python", + "inputs": {"iteration": 1, "total_iterations": 2}, + "outputs": {"load_stats": {"record_count": 9}}, + }, + ) + ] assert status["lastStatus"] == "error" assert status["lastError"] == "business rule blocked" assert status["selectedCount"] == 9 diff --git a/tests/workflow/test_tool_run_workflow.py b/tests/workflow/test_tool_run_workflow.py index 223d4397f..0842f90e5 100644 --- a/tests/workflow/test_tool_run_workflow.py +++ b/tests/workflow/test_tool_run_workflow.py @@ -10,6 +10,8 @@ """ import asyncio +import threading + import pytest from unittest.mock import AsyncMock, Mock, patch, MagicMock from typing import Dict, Any @@ -27,6 +29,7 @@ import flocks.tool.task.run_workflow as run_workflow_module from flocks.mcp.client import McpClient from flocks.workflow.runner import RunWorkflowResult, run_workflow +from flocks.workflow.store import WorkflowStore class FakeRunWorkflowResult: @@ -285,7 +288,12 @@ async def test_run_workflow_success(self, tool_context_with_permission, simple_w } ) mock_run = Mock(name="run_workflow", return_value=fake) - with patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)): + direct_audit = AsyncMock(return_value=None) + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value=None), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), + ): result = await ToolRegistry.execute( "run_workflow", ctx=tool_context_with_permission, workflow=simple_workflow, inputs={} ) @@ -297,6 +305,7 @@ async def test_run_workflow_success(self, tool_context_with_permission, simple_w assert result.metadata["status"] == "success" assert result.metadata["steps"] == 1 assert "run_id" not in result.metadata + direct_audit.assert_awaited_once_with("test-workflow-001", fake.__dict__) # Check that permission was requested assert len(tool_context_with_permission._permissions_requested) > 0 @@ -326,7 +335,7 @@ def run_side_effect(**kwargs): steps=1, last_node_id="node-1", outputs={"message": "ok"}, - history=[{"node_id": "node-1", "node_type": "python", "outputs": {"message": "ok"}}], + history=[], error=None, ) @@ -343,6 +352,7 @@ def run_side_effect(**kwargs): ) upsert_execution = AsyncMock(return_value=None) record_result = AsyncMock(return_value=None) + direct_audit = AsyncMock(return_value=None) with ( patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), @@ -353,8 +363,9 @@ def run_side_effect(**kwargs): ), patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), patch.object(run_workflow_module, "create_execution_record", create_execution), - patch.object(run_workflow_module.WorkflowStore, "upsert_execution", upsert_execution), + patch.object(WorkflowStore, "upsert_execution", upsert_execution), patch.object(run_workflow_module, "record_execution_result", record_result), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), ): result = await ToolRegistry.execute( "run_workflow", @@ -368,9 +379,242 @@ def run_side_effect(**kwargs): assert "run_id" not in result.metadata create_execution.assert_awaited_once() record_result.assert_awaited_once() + expected_step = { + "node_id": "node-1", + "node_type": "python", + "outputs": {"message": "ok"}, + } + assert record_result.await_args.kwargs["steps"] == [(1, expected_step)] + assert record_result.await_args.args[2]["executionLog"] == [expected_step] assert upsert_execution.await_count >= 1 + assert all(call.args[0]["executionLog"] == [] for call in upsert_execution.await_args_list) + direct_audit.assert_not_awaited() assert any(update.get("workflow_execution_id") == "exec-registered" for update in metadata_updates) + @pytest.mark.anyio + async def test_run_workflow_registered_callbacks_do_not_wait_for_progress_storage( + self, + tool_context_with_permission, + simple_workflow, + ): + write_started = asyncio.Event() + release_write = asyncio.Event() + runner_finished = threading.Event() + write_order: list[str] = [] + + async def blocked_upsert(_summary): + write_order.append("progress-start") + write_started.set() + await release_write.wait() + write_order.append("progress-end") + + async def record_result(*args, **kwargs): # noqa: ANN002, ANN003 + write_order.append("final") + + def run_side_effect(**kwargs): + kwargs["on_step_start"]( + kwargs["run_id"], + 1, + MagicMock(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"]( + { + "node_id": "node-1", + "node_type": "python", + "outputs": {"ok": True}, + } + ) + runner_finished.set() + return FakeRunWorkflowResult( + status="SUCCEEDED", + run_id=kwargs["run_id"], + steps=1, + last_node_id="node-1", + outputs={"ok": True}, + history=[], + error=None, + ) + + mock_run = Mock(name="run_workflow", side_effect=run_side_effect) + create_execution = AsyncMock( + return_value={ + "id": "exec-blocked", + "workflowId": "test-workflow-001", + "status": "running", + "startedAt": 1, + "executionLog": [], + } + ) + record_result_mock = AsyncMock(side_effect=record_result) + + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), + patch.object(run_workflow_module, "create_execution_record", create_execution), + patch.object(WorkflowStore, "upsert_execution", blocked_upsert), + patch.object(run_workflow_module, "record_execution_result", record_result_mock), + patch.object(run_workflow_module, "_record_workflow_tool_result", AsyncMock(return_value=None)), + ): + task = asyncio.create_task( + ToolRegistry.execute( + "run_workflow", + ctx=tool_context_with_permission, + workflow=simple_workflow, + inputs={}, + ) + ) + await write_started.wait() + assert await asyncio.to_thread(runner_finished.wait, 0.1) + record_result_mock.assert_not_awaited() + release_write.set() + result = await task + + assert result.success is True + record_result_mock.assert_awaited_once() + assert write_order[-2:] == ["progress-end", "final"] + + @pytest.mark.anyio + async def test_run_workflow_registered_cancellation_keeps_completed_and_pending_steps( + self, + tool_context_with_permission, + simple_workflow, + ): + persisted_summaries: list[dict[str, Any]] = [] + + async def capture_upsert(summary): + persisted_summaries.append(dict(summary)) + + def run_side_effect(**kwargs): + kwargs["on_step_start"]( + kwargs["run_id"], + 1, + MagicMock(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"]( + { + "node_id": "node-1", + "node_type": "python", + "outputs": {"ok": True}, + } + ) + tool_context_with_permission.abort.set() + kwargs["on_step_start"]( + kwargs["run_id"], + 2, + MagicMock(id="node-2", type="tool"), + {"message": "hello"}, + ) + return FakeRunWorkflowResult( + status="SUCCEEDED", + run_id=kwargs["run_id"], + steps=1, + last_node_id="node-2", + outputs={"ok": True}, + history=[], + error=None, + ) + + mock_run = Mock(name="run_workflow", side_effect=run_side_effect) + create_execution = AsyncMock( + return_value={ + "id": "exec-cancelled", + "workflowId": "test-workflow-001", + "status": "running", + "startedAt": 1, + "executionLog": [], + } + ) + record_result = AsyncMock(return_value=None) + direct_audit = AsyncMock(return_value=None) + + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), + patch.object(run_workflow_module, "create_execution_record", create_execution), + patch.object(WorkflowStore, "upsert_execution", capture_upsert), + patch.object(run_workflow_module, "record_execution_result", record_result), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), + ): + result = await ToolRegistry.execute( + "run_workflow", + ctx=tool_context_with_permission, + workflow=simple_workflow, + inputs={}, + ) + + steps = record_result.await_args.kwargs["steps"] + assert result.success is False + assert result.metadata["status"] == "cancelled" + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + assert steps[1][1]["error"] == "Run cancelled before node completed" + assert record_result.await_args.args[2]["status"] == "cancelled" + assert persisted_summaries[-1]["currentPhase"] == "cancelling" + direct_audit.assert_not_awaited() + + @pytest.mark.anyio + async def test_run_workflow_registered_failure_keeps_callback_steps( + self, + tool_context_with_permission, + simple_workflow, + ): + def run_side_effect(**kwargs): + kwargs["on_step_start"]( + kwargs["run_id"], + 1, + MagicMock(id="node-1", type="python"), + {}, + ) + kwargs["on_step_complete"]( + { + "node_id": "node-1", + "node_type": "python", + "outputs": {"ok": True}, + } + ) + kwargs["on_step_start"]( + kwargs["run_id"], + 2, + MagicMock(id="node-2", type="tool"), + {"message": "hello"}, + ) + raise RuntimeError("runner failed") + + mock_run = Mock(name="run_workflow", side_effect=run_side_effect) + create_execution = AsyncMock( + return_value={ + "id": "exec-failed", + "workflowId": "test-workflow-001", + "status": "running", + "startedAt": 1, + "executionLog": [], + } + ) + record_result = AsyncMock(return_value=None) + + with ( + patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), + patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), + patch.object(run_workflow_module, "create_execution_record", create_execution), + patch.object(WorkflowStore, "upsert_execution", AsyncMock(return_value=None)), + patch.object(run_workflow_module, "record_execution_result", record_result), + ): + result = await ToolRegistry.execute( + "run_workflow", + ctx=tool_context_with_permission, + workflow=simple_workflow, + inputs={}, + ) + + steps = record_result.await_args.kwargs["steps"] + assert result.success is False + assert "runner failed" in result.error + assert [step_index for step_index, _ in steps] == [1, 2] + assert [step["node_id"] for _, step in steps] == ["node-1", "node-2"] + assert record_result.await_args.args[2]["executionLog"] == [step for _, step in steps] + @pytest.mark.anyio async def test_run_workflow_registered_id_overrides_missing_workflow_json_id( self, @@ -416,7 +660,7 @@ def run_side_effect(**kwargs): return_value={"id": "wf-directory-id", "workflowJson": workflow_without_id}, ), patch.object(run_workflow_module, "create_execution_record", create_execution), - patch.object(run_workflow_module.WorkflowStore, "upsert_execution", AsyncMock(return_value=None)), + patch.object(WorkflowStore, "upsert_execution", AsyncMock(return_value=None)), patch.object(run_workflow_module, "record_execution_result", AsyncMock(return_value=None)), ): result = await ToolRegistry.execute( @@ -476,17 +720,16 @@ def run_side_effect(**kwargs): } ) upsert_execution = AsyncMock(return_value=None) - record_step = AsyncMock(return_value=None) record_result = AsyncMock(return_value=None) + direct_audit = AsyncMock(return_value=None) with ( patch.object(run_workflow_module, "_get_workflow_runtime", return_value=_runtime_tuple(run_fn=mock_run)), patch.object(run_workflow_module, "resolve_workflow_id_from_source", return_value="test-workflow-001"), patch.object(run_workflow_module, "create_execution_record", create_execution), - patch.object(run_workflow_module.WorkflowStore, "upsert_execution", upsert_execution), - patch.object(run_workflow_module, "record_execution_step", record_step), + patch.object(WorkflowStore, "upsert_execution", upsert_execution), patch.object(run_workflow_module, "record_execution_result", record_result), - patch.object(run_workflow_module, "_record_workflow_tool_result", AsyncMock(return_value=None)), + patch.object(run_workflow_module, "_record_workflow_tool_result", direct_audit), ): result = await ToolRegistry.execute( "run_workflow", @@ -496,8 +739,9 @@ def run_side_effect(**kwargs): ) assert result.success is True - record_step.assert_awaited() - step_payload = record_step.await_args.args[2] + steps = record_result.await_args.kwargs["steps"] + assert [step_index for step_index, _ in steps] == [1] + step_payload = steps[0][1] assert step_payload["inputs"] == { "_raw_alerts_count": 150, "source": "syslog", @@ -506,18 +750,19 @@ def run_side_effect(**kwargs): "_raw_alerts_count": 150, "message": "ok", } + direct_audit.assert_not_awaited() assert result.metadata["has_output"] is True assert result.metadata["output_keys"] == ["enriched_alerts", "message"] assert "outputs" not in result.metadata assert "history" not in result.metadata - assert result.metadata["history_count"] == 0 + assert result.metadata["history_count"] == 1 final_exec_data = record_result.await_args.args[2] assert final_exec_data["outputResults"] == { "_enriched_alerts_count": 150, "message": "done", } - assert final_exec_data["executionLog"] == [] + assert final_exec_data["executionLog"] == [step_payload] assert final_exec_data["stepCount"] == 1 assert any(update.get("workflow_execution_id") == "exec-compacted" for update in metadata_updates) diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index dd34bec84..6456a9227 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -79,7 +79,17 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert result["executionLog"] == [] assert result["stepCount"] == 1 record_result.assert_awaited_once() - assert record_result.await_args.kwargs["steps"] == [] + assert record_result.await_args.kwargs["steps"] == [ + ( + 1, + { + "node_id": "notify", + "node_type": "tool", + "inputs": {}, + "outputs": {"ok": True}, + }, + ) + ] cleanup_context.assert_awaited_once_with(tool_context) diff --git a/tests/workflow/test_workflow_store.py b/tests/workflow/test_workflow_store.py index d699d0f9b..6608db1b3 100644 --- a/tests/workflow/test_workflow_store.py +++ b/tests/workflow/test_workflow_store.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import os from pathlib import Path import pytest @@ -118,7 +119,7 @@ async def test_workflow_store_increment_stats_is_atomic_for_concurrent_updates() @pytest.mark.asyncio -async def test_complete_execution_writes_steps_summary_and_stats_with_one_commit( +async def test_complete_execution_writes_steps_and_summary_with_one_commit( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() @@ -147,8 +148,6 @@ async def counted_commit() -> None: (1, {"node_id": "n1", "outputs": {"ok": 1}}), (2, {"node_id": "n2", "outputs": {"ok": 2}}), ], - success=True, - duration=0.25, ) assert commit_count == 1 @@ -158,11 +157,7 @@ async def counted_commit() -> None: steps, total = await WorkflowStore.list_steps("exec-complete") assert total == 2 assert [step["node_id"] for step in steps] == ["n1", "n2"] - stats = await WorkflowStore.get_stats("wf-complete") - assert stats is not None - assert stats["callCount"] == 1 - assert stats["successCount"] == 1 - assert stats["totalRuntime"] == pytest.approx(0.25) + assert await WorkflowStore.get_stats("wf-complete") is None @pytest.mark.asyncio @@ -199,8 +194,6 @@ async def counted_commit() -> None: "stepCount": 7, }, steps, - success=True, - duration=0.01, ) for index in range(4) ) @@ -216,10 +209,7 @@ async def counted_commit() -> None: assert [step["node_id"] for step in persisted_steps] == [ f"node-{step_index}" for step_index in range(1, 8) ] - stats = await WorkflowStore.get_stats("wf-trigger") - assert stats is not None - assert stats["callCount"] == 4 - assert stats["successCount"] == 4 + assert await WorkflowStore.get_stats("wf-trigger") is None @pytest.mark.asyncio @@ -248,8 +238,6 @@ async def fail_commit() -> None: "stepCount": 1, }, [(1, {"node_id": "node-1", "outputs": {"ok": True}})], - success=True, - duration=0.01, ) monkeypatch.setattr(db, "commit", original_commit) @@ -261,48 +249,59 @@ async def fail_commit() -> None: @pytest.mark.asyncio -async def test_complete_execution_applies_retention_before_single_commit( +async def test_complete_execution_rolls_back_cancelled_transaction( monkeypatch: pytest.MonkeyPatch, ) -> None: await WorkflowStore.init() db = await WorkflowStore.raw_completion_db() - commit_count = 0 original_commit = db.commit - async def counted_commit() -> None: - nonlocal commit_count - commit_count += 1 - await original_commit() + async def cancel_commit() -> None: + raise asyncio.CancelledError - monkeypatch.setattr(db, "commit", counted_commit) - trimmed: list[str] = [] - for index in range(4): - trimmed = await WorkflowStore.complete_execution( + monkeypatch.setattr(db, "commit", cancel_commit) + + with pytest.raises(asyncio.CancelledError): + await WorkflowStore.complete_execution( { - "id": f"exec-retain-{index}", - "workflowId": "wf-retain", + "id": "exec-cancelled-commit", + "workflowId": "wf-cancelled-commit", "status": "success", - "startedAt": index + 1, - "finishedAt": index + 2, - "duration": 0.01, + "startedAt": 1, + "finishedAt": 2, "executionLog": [], "stepCount": 1, }, - [(1, {"node_id": f"node-{index}", "outputs": {"ok": True}})], - success=True, - duration=0.01, - history_limit=3, + [(1, {"node_id": "node-1", "outputs": {"ok": True}})], ) - assert commit_count == 4 - assert trimmed == ["exec-retain-0"] - assert await WorkflowStore.get_execution("exec-retain-0") is None - old_steps, old_total = await WorkflowStore.list_steps("exec-retain-0") - assert old_steps == [] - assert old_total == 0 - executions = await WorkflowStore.list_executions("wf-retain", limit=10) - assert [execution["id"] for execution in executions] == [ - "exec-retain-3", - "exec-retain-2", - "exec-retain-1", - ] + monkeypatch.setattr(db, "commit", original_commit) + await WorkflowStore.complete_execution( + { + "id": "exec-after-cancel", + "workflowId": "wf-cancelled-commit", + "status": "success", + "startedAt": 3, + "finishedAt": 4, + "executionLog": [], + "stepCount": 1, + }, + [(1, {"node_id": "node-2", "outputs": {"ok": True}})], + ) + + assert await WorkflowStore.get_execution("exec-cancelled-commit") is None + assert await WorkflowStore.get_execution("exec-after-cancel") is not None + + +@pytest.mark.asyncio +async def test_completion_connection_reinitializes_after_pid_change() -> None: + await WorkflowStore.init() + original_connection = await WorkflowStore.raw_completion_db() + original_lock = WorkflowStore._completion_lock + WorkflowStore._init_pid = -1 + + refreshed_connection = await WorkflowStore.raw_completion_db() + + assert refreshed_connection is not original_connection + assert WorkflowStore._completion_lock is not original_lock + assert WorkflowStore._init_pid == os.getpid() From 0908fd356fc72ebc063f7a52f3c68d348b632c2f Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 17:12:34 +0800 Subject: [PATCH 4/5] fix(workflow): persist queued trigger executions --- flocks/ingest/kafka/manager.py | 2 +- flocks/ingest/syslog/manager.py | 2 +- flocks/workflow/poller_manager.py | 2 +- flocks/workflow/triggers/runtime.py | 2 +- tests/ingest/test_kafka_manager.py | 6 +++--- tests/ingest/test_syslog_manager_backpressure.py | 2 +- tests/workflow/test_poller_manager.py | 4 ++-- tests/workflow/test_trigger_runtime.py | 2 +- 8 files changed, 11 insertions(+), 11 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index a44c28eec..d5bd88f44 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -767,7 +767,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=False, + persist=True, ) exec_id = exec_data["id"] start_time = time.time() diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index ab15ad23a..88a2a68cb 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -619,7 +619,7 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=False, + persist=True, ) exec_id = exec_data["id"] step_recorder = ExecutionStepRecorder(exec_id=exec_id) diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index 7fdb7a101..5e774fba1 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -451,7 +451,7 @@ async def _execute_run( exec_data = await create_execution_record( workflow_id, input_params=inputs, - persist=False, + persist=True, ) exec_id = str(exec_data["id"]) step_recorder = ExecutionStepRecorder(exec_id=exec_id) diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index d41aa012d..f7121a106 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -248,7 +248,7 @@ async def _execute_workflow_effect( exec_data = await create_execution_record( workflow_id, input_params=mapped_inputs, - persist=False, + persist=True, ) exec_id = exec_data["id"] step_recorder = ExecutionStepRecorder(exec_id=exec_id) diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 2134fa9c9..19afec5d1 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -549,7 +549,7 @@ async def test_trigger_workflow_compacts_kafka_execution_record( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True captured_input_params.update(input_params or {}) return {"id": "exec-compact", "workflowId": workflow_id, "inputParams": input_params} @@ -639,7 +639,7 @@ async def test_trigger_workflow_merges_configured_inputs_with_consumed_message( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True recorded_input_params.update(input_params or {}) return {"id": "exec-merge", "workflowId": workflow_id, "inputParams": input_params} @@ -700,7 +700,7 @@ async def test_trigger_workflow_applies_mapping_and_filter( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True return {"id": "exec-filter", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 0ce283747..1b8a459a2 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -353,7 +353,7 @@ async def test_trigger_workflow_applies_mapping_and_filter( async def _fake_create_execution_record( # noqa: ANN001 workflow_id, *, input_params=None, exec_id=None, persist=True ): - assert persist is False + assert persist is True return {"id": "exec-syslog", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index f4069d1bd..f2f5f6914 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -156,7 +156,7 @@ async def _fake_create_execution_record( exec_id: str | None = None, persist: bool = True, ) -> dict[str, Any]: - assert persist is False + assert persist is True record = { "id": exec_id or "exec-1", "workflowId": workflow_id, @@ -361,7 +361,7 @@ async def _fake_create_execution_record( exec_id: str | None = None, persist: bool = True, ) -> dict[str, Any]: - assert persist is False + assert persist is True _ = input_params return { "id": exec_id or f"exec-{workflow_id}", diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index 6456a9227..279ade3b0 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -75,7 +75,7 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert runtime_module.run_workflow.call_args.kwargs["run_id"] == "exec-1" assert runtime_module.run_workflow.call_args.kwargs["execution_profile"] == "high_frequency" assert callable(runtime_module.run_workflow.call_args.kwargs["on_step_complete"]) - assert create_record.await_args.kwargs["persist"] is False + assert create_record.await_args.kwargs["persist"] is True assert result["executionLog"] == [] assert result["stepCount"] == 1 record_result.assert_awaited_once() From afc664994688ff40ef7b63e82f9a7317184817e0 Mon Sep 17 00:00:00 2001 From: xiami762 <> Date: Fri, 21 Aug 2026 17:51:01 +0800 Subject: [PATCH 5/5] refactor(workflow): simplify step persistence plumbing Remove obsolete persistence options and duplicated step/row handling while preserving atomic completion and nonblocking progress behavior. Co-Authored-By: Claude Opus 4.6 --- flocks/ingest/kafka/manager.py | 2 - flocks/ingest/syslog/manager.py | 3 +- flocks/server/routes/workflow.py | 2 +- flocks/tool/task/run_workflow.py | 46 ++-------- flocks/workflow/execution_store.py | 10 +-- flocks/workflow/poller_manager.py | 3 +- flocks/workflow/store.py | 67 +++++---------- flocks/workflow/triggers/runtime.py | 3 +- tests/ingest/test_kafka_manager.py | 9 +- .../test_syslog_manager_backpressure.py | 3 +- .../server/routes/test_workflow_run_route.py | 86 ++++--------------- .../workflow/test_execution_store_compact.py | 77 +++-------------- tests/workflow/test_poller_manager.py | 8 +- tests/workflow/test_trigger_runtime.py | 5 +- 14 files changed, 76 insertions(+), 248 deletions(-) diff --git a/flocks/ingest/kafka/manager.py b/flocks/ingest/kafka/manager.py index d5bd88f44..dea17a7e9 100644 --- a/flocks/ingest/kafka/manager.py +++ b/flocks/ingest/kafka/manager.py @@ -767,14 +767,12 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=True, ) exec_id = exec_data["id"] start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) trigger_input_keys = list((trigger.mapping or {}).keys()) or [input_key] step_recorder = ExecutionStepRecorder( - exec_id=exec_id, step_compactor=lambda step: _compact_step_for_kafka_storage( step, input_key=input_key, diff --git a/flocks/ingest/syslog/manager.py b/flocks/ingest/syslog/manager.py index 88a2a68cb..63e1d1b18 100644 --- a/flocks/ingest/syslog/manager.py +++ b/flocks/ingest/syslog/manager.py @@ -619,10 +619,9 @@ async def _executor(mapped_inputs: Dict[str, Any]) -> Dict[str, Any]: exec_data = await create_execution_record( workflow_id, input_params=summarized_inputs, - persist=True, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() start_time = time.time() trigger_meta = mapped_inputs.get("_flocks", {}).get("trigger", {}) tool_context = None diff --git a/flocks/server/routes/workflow.py b/flocks/server/routes/workflow.py index 0a35c4eab..b1ca85243 100644 --- a/flocks/server/routes/workflow.py +++ b/flocks/server/routes/workflow.py @@ -1128,7 +1128,7 @@ async def _run_workflow_execution_task( ) -> None: """Execute a workflow in the background and keep the execution record updated.""" start_time = time.time() - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None execution_summary: Dict[str, Any] = { diff --git a/flocks/tool/task/run_workflow.py b/flocks/tool/task/run_workflow.py index ab5fcbf10..7c50452b7 100644 --- a/flocks/tool/task/run_workflow.py +++ b/flocks/tool/task/run_workflow.py @@ -570,9 +570,8 @@ async def run_workflow_tool( canonical_workflow_id = registered_workflow_id or resolve_workflow_id_from_source(workflow_source) display_workflow_id = canonical_workflow_id or workflow_id tracked_execution: Optional[Dict[str, Any]] = None - step_recorder: Optional[ExecutionStepRecorder] = None + step_recorder = ExecutionStepRecorder() progress_writer: Optional[ExecutionProgressWriter] = None - callback_step_count = 0 pending_step_index: Optional[int] = None pending_step: Optional[Dict[str, Any]] = None final_step_batch: Optional[List[Tuple[int, Dict[str, Any]]]] = None @@ -639,35 +638,9 @@ def _on_step_start( return step_index def _on_step_complete(step_result: Any) -> None: - nonlocal callback_step_count, pending_step_index, pending_step - if step_recorder is not None: - step_recorder.on_step_complete(step_result) - callback_step_count = step_recorder.step_count - progress_update = dict(step_recorder.summary) - else: - if hasattr(step_result, "model_dump"): - step_dict = step_result.model_dump(mode="json") - elif isinstance(step_result, dict): - step_dict = dict(step_result) - else: - step_dict = {"node_id": None, "outputs": {}, "error": str(step_result)} - callback_step_count += 1 - compacted_step = compact_step_for_storage(step_dict) - progress_update = { - "stepCount": callback_step_count, - "currentNodeId": compacted_step.get("node_id"), - "currentNodeType": compacted_step.get("node_type") - or compacted_step.get("type"), - "currentPhase": "running", - "currentStepIndex": callback_step_count, - "loopProgress": derive_loop_progress( - node_id=compacted_step.get("node_id"), - global_step_index=callback_step_count, - inputs=compacted_step.get("inputs"), - outputs=compacted_step.get("outputs"), - ), - "updatedAt": int(time.time() * 1000), - } + nonlocal pending_step_index, pending_step + step_recorder.on_step_complete(step_result) + progress_update = dict(step_recorder.summary) pending_step_index = None pending_step = None if ctx.abort.is_set(): @@ -688,8 +661,8 @@ def _on_step_complete(step_result: Any) -> None: "phase": progress_update["currentPhase"], "current_node_id": progress_update.get("currentNodeId"), "current_node_type": progress_update.get("currentNodeType"), - "step_index": callback_step_count, - "step_count": callback_step_count, + "step_index": step_recorder.step_count, + "step_count": step_recorder.step_count, "loop_progress": progress_update.get("loopProgress"), }, } @@ -698,7 +671,7 @@ def _on_step_complete(step_result: Any) -> None: def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: nonlocal final_step_batch if final_step_batch is None: - final_step_batch = step_recorder.take_steps() if step_recorder is not None else [] + final_step_batch = step_recorder.take_steps() if pending_step_index is not None and pending_step is not None: final_step_batch.append((pending_step_index, pending_step)) final_step_batch.sort(key=lambda item: item[0]) @@ -721,7 +694,6 @@ def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: canonical_workflow_id, input_params=workflow_inputs, ) - step_recorder = ExecutionStepRecorder(exec_id=tracked_execution["id"]) progress_writer = ExecutionProgressWriter(tracked_execution) # Update metadata to show workflow is running @@ -868,7 +840,7 @@ def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: history_count = len(final_history) final_step_count = result_dict.get("steps") if not isinstance(final_step_count, int): - final_step_count = callback_step_count + final_step_count = step_recorder.step_count final_step_count = max( final_step_count, max((step_index for step_index, _ in tracked_steps), default=0), @@ -988,7 +960,7 @@ def _take_final_step_batch() -> List[Tuple[int, Dict[str, Any]]]: }, ) terminal_status = "cancelled" if ctx.abort.is_set() else "error" - final_step_count = callback_step_count + final_step_count = step_recorder.step_count if tracked_execution and canonical_workflow_id: tracked_steps = _take_final_step_batch() final_step_count = max( diff --git a/flocks/workflow/execution_store.py b/flocks/workflow/execution_store.py index 3dbc59cd4..6503ee981 100644 --- a/flocks/workflow/execution_store.py +++ b/flocks/workflow/execution_store.py @@ -372,7 +372,7 @@ def workflow_execution_step_prefix(exec_id: str) -> str: def compact_execution_summary(exec_data: Dict[str, Any]) -> Dict[str, Any]: """Return an execution record safe to keep in the hot summary row. - Step details are stored separately under ``workflow_execution_step`` keys. + Step details are stored separately in ``workflow_execution_steps`` rows. Keeping ``executionLog`` out of the summary row avoids rewriting an ever-growing JSON blob on every progress update. """ @@ -398,10 +398,8 @@ class ExecutionStepRecorder: def __init__( self, *, - exec_id: str, step_compactor: Callable[[Any], Dict[str, Any]] = compact_step_for_storage, ) -> None: - self.exec_id = exec_id self.step_compactor = step_compactor self.step_count = 0 self.summary: Dict[str, Any] = {} @@ -618,9 +616,8 @@ async def create_execution_record( *, input_params: Optional[Dict[str, Any]] = None, exec_id: Optional[str] = None, - persist: bool = True, ) -> Dict[str, Any]: - """Build a running workflow execution record and optionally persist it. + """Build and persist a running workflow execution record. *input_params* is passed through ``compact_outputs_for_storage`` before writing to SQLite so that batch HTTP calls whose inputs contain a key in @@ -635,8 +632,7 @@ async def create_execution_record( input_params=compacted_params, exec_id=exec_id, ) - if persist: - await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) + await WorkflowStore.upsert_execution(compact_execution_summary(exec_data)) return exec_data diff --git a/flocks/workflow/poller_manager.py b/flocks/workflow/poller_manager.py index 5e774fba1..9d348d5db 100644 --- a/flocks/workflow/poller_manager.py +++ b/flocks/workflow/poller_manager.py @@ -451,10 +451,9 @@ async def _execute_run( exec_data = await create_execution_record( workflow_id, input_params=inputs, - persist=True, ) exec_id = str(exec_data["id"]) - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() current = self._status.get(workflow_id) or self._base_status(workflow_id) current["lastRunAt"] = started_at_ms current["activeRuns"] = self._cleanup_done_runs(workflow_id) diff --git a/flocks/workflow/store.py b/flocks/workflow/store.py index b4d9b129f..f7efbd0a1 100644 --- a/flocks/workflow/store.py +++ b/flocks/workflow/store.py @@ -37,6 +37,13 @@ "workflow_syslog_config/", ) _WORKFLOW_PREFIXES = _WORKFLOW_KV_PREFIXES + _WORKFLOW_TABLE_PREFIXES +_EXECUTION_UPSERT_SQL = """ + INSERT OR REPLACE INTO workflow_executions + (id, workflow_id, status, current_phase, current_node_id, current_node_type, + current_step_index, step_count, input_params, output_results, error_message, + trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +""" class WorkflowStore: @@ -312,21 +319,18 @@ async def _migrate_legacy_kv(cls) -> None: log.info("workflow.store.legacy_kv_migrated", counts) @classmethod - async def upsert_execution(cls, exec_data: Dict[str, Any]) -> None: - db = await cls._db() + def _execution_row( + cls, + exec_data: Dict[str, Any], + ) -> Tuple[str, str, Tuple[Any, ...]]: payload = dict(exec_data) exec_id = str(payload.get("id") or "") workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") if not exec_id or not workflow_id: raise ValueError("workflow execution requires id and workflowId") - await db.execute( - """ - INSERT OR REPLACE INTO workflow_executions - (id, workflow_id, status, current_phase, current_node_id, current_node_type, - current_step_index, step_count, input_params, output_results, error_message, - trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, + return ( + exec_id, + workflow_id, ( exec_id, workflow_id, @@ -348,6 +352,12 @@ async def upsert_execution(cls, exec_data: Dict[str, Any]) -> None: cls._json_dumps(payload), ), ) + + @classmethod + async def upsert_execution(cls, exec_data: Dict[str, Any]) -> None: + db = await cls._db() + _, _, row = cls._execution_row(exec_data) + await db.execute(_EXECUTION_UPSERT_SQL, row) await db.commit() @classmethod @@ -497,12 +507,7 @@ async def complete_execution( ) -> None: """Atomically persist one final execution summary and its step batch.""" db = await cls._completion_db() - payload = dict(exec_data) - exec_id = str(payload.get("id") or "") - workflow_id = str(payload.get("workflowId") or payload.get("workflow_id") or "") - if not exec_id or not workflow_id: - raise ValueError("workflow execution requires id and workflowId") - + exec_id, workflow_id, execution_row = cls._execution_row(exec_data) step_rows = cls._step_rows(exec_id, steps) lock = cls._completion_lock if lock is None: @@ -521,35 +526,7 @@ async def complete_execution( """, step_rows, ) - await db.execute( - """ - INSERT OR REPLACE INTO workflow_executions - (id, workflow_id, status, current_phase, current_node_id, current_node_type, - current_step_index, step_count, input_params, output_results, error_message, - trigger_id, trigger_type, started_at, finished_at, duration, updated_at, payload) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - exec_id, - workflow_id, - str(payload.get("status") or "running"), - payload.get("currentPhase"), - payload.get("currentNodeId"), - payload.get("currentNodeType"), - cls._as_int(payload.get("currentStepIndex")), - cls._as_int(payload.get("stepCount")) or 0, - cls._json_dumps(payload.get("inputParams") or {}), - cls._json_dumps(payload.get("outputResults") or {}), - payload.get("errorMessage"), - payload.get("triggerId"), - payload.get("triggerType"), - cls._as_int(payload.get("startedAt")) or cls._now_ms(), - cls._as_int(payload.get("finishedAt")), - cls._as_float(payload.get("duration")), - cls._as_int(payload.get("updatedAt")) or cls._now_ms(), - cls._json_dumps(payload), - ), - ) + await db.execute(_EXECUTION_UPSERT_SQL, execution_row) await db.commit() except BaseException: try: diff --git a/flocks/workflow/triggers/runtime.py b/flocks/workflow/triggers/runtime.py index f7121a106..2d88fb1e4 100644 --- a/flocks/workflow/triggers/runtime.py +++ b/flocks/workflow/triggers/runtime.py @@ -248,10 +248,9 @@ async def _execute_workflow_effect( exec_data = await create_execution_record( workflow_id, input_params=mapped_inputs, - persist=True, ) exec_id = exec_data["id"] - step_recorder = ExecutionStepRecorder(exec_id=exec_id) + step_recorder = ExecutionStepRecorder() started_at = time.time() tool_context = None try: diff --git a/tests/ingest/test_kafka_manager.py b/tests/ingest/test_kafka_manager.py index 19afec5d1..2b9b6eb09 100644 --- a/tests/ingest/test_kafka_manager.py +++ b/tests/ingest/test_kafka_manager.py @@ -547,9 +547,8 @@ async def test_trigger_workflow_compacts_kafka_execution_record( captured_steps: list[tuple[int, dict]] = [] async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True captured_input_params.update(input_params or {}) return {"id": "exec-compact", "workflowId": workflow_id, "inputParams": input_params} @@ -637,9 +636,8 @@ async def test_trigger_workflow_merges_configured_inputs_with_consumed_message( recorded_input_params: dict = {} async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True recorded_input_params.update(input_params or {}) return {"id": "exec-merge", "workflowId": workflow_id, "inputParams": input_params} @@ -698,9 +696,8 @@ async def test_trigger_workflow_applies_mapping_and_filter( recorded_exec_data: dict = {} async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True return {"id": "exec-filter", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/ingest/test_syslog_manager_backpressure.py b/tests/ingest/test_syslog_manager_backpressure.py index 1b8a459a2..29750f055 100644 --- a/tests/ingest/test_syslog_manager_backpressure.py +++ b/tests/ingest/test_syslog_manager_backpressure.py @@ -351,9 +351,8 @@ async def test_trigger_workflow_applies_mapping_and_filter( recorded_steps: list[tuple[int, dict]] = [] async def _fake_create_execution_record( # noqa: ANN001 - workflow_id, *, input_params=None, exec_id=None, persist=True + workflow_id, *, input_params=None, exec_id=None ): - assert persist is True return {"id": "exec-syslog", "workflowId": workflow_id, "inputParams": input_params} async def _fake_record_execution_result( # noqa: ANN001 diff --git a/tests/server/routes/test_workflow_run_route.py b/tests/server/routes/test_workflow_run_route.py index f0891bafa..69c213656 100644 --- a/tests/server/routes/test_workflow_run_route.py +++ b/tests/server/routes/test_workflow_run_route.py @@ -41,6 +41,17 @@ def _two_node_workflow_json(edge): } +def _progress_writer(exec_id: str, **updates): + summary = { + "id": exec_id, + "workflowId": "wf-1", + "status": "running", + "executionLog": [], + } + summary.update(updates) + return workflow_module.ExecutionProgressWriter(summary) + + @pytest.mark.asyncio async def test_create_workflow_applies_vertex_cache_runtime_defaults(monkeypatch: pytest.MonkeyPatch) -> None: writes: list[dict] = [] @@ -252,35 +263,17 @@ def run_workflow_mock(**kwargs): run_mock = Mock(side_effect=run_workflow_mock) record_result = AsyncMock(return_value=None) upsert_execution = AsyncMock(return_value=None) - storage_read = AsyncMock( - return_value={ - "id": "exec-1", - "workflowId": "wf-1", - "currentNodeType": "tool", - "executionLog": [], - } - ) - monkeypatch.setattr(MCP, "init", init_mock) monkeypatch.setattr(workflow_module, "run_workflow", run_mock) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) monkeypatch.setattr(workflow_module, "_resolve_execution_outcome", lambda _result: ("success", None)) monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) - monkeypatch.setattr(workflow_module.Storage, "read", storage_read) - monkeypatch.setattr(workflow_module.Storage, "write", AsyncMock(return_value=None)) monkeypatch.setattr(workflow_module, "compact_outputs_for_storage", lambda value: value) monkeypatch.setattr(workflow_module, "compact_history_for_storage", lambda value: value) req = workflow_module.WorkflowRunRequest(inputs={"ip": "8.8.8.8"}, trace=False) tool_context = ToolContext(session_id="session-1", message_id="message-1", agent="rex") - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-1", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-1") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", @@ -342,14 +335,7 @@ def run_workflow_mock(**kwargs): cancel_event = workflow_module.threading.Event() cancel_event.set() - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-cancelled", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-cancelled") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, @@ -418,14 +404,7 @@ def run_workflow_mock(**kwargs): monkeypatch.setattr(workflow_module, "_record_execution_result", record_result) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", upsert_execution) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-partial-cancel", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-partial-cancel") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, @@ -481,14 +460,7 @@ def run_workflow_mock(**kwargs): AsyncMock(return_value=None), ) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-runner-error", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-runner-error") await workflow_module._run_workflow_execution_task( workflow_id="wf-1", workflow_json={"id": "wf-1", "start": "node-1", "nodes": [], "edges": []}, @@ -544,14 +516,7 @@ def run_workflow_mock(**kwargs): AsyncMock(return_value=None), ) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-storage-error", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-storage-error") with pytest.raises(RuntimeError, match="storage failed"): await workflow_module._run_workflow_execution_task( @@ -627,14 +592,7 @@ def run_workflow_mock(**kwargs): monkeypatch.setattr(workflow_module, "_record_execution_result", record_result_mock) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", blocked_upsert) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-blocked-progress", - "workflowId": "wf-1", - "status": "running", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-blocked-progress") task = asyncio.create_task( workflow_module._run_workflow_execution_task( workflow_id="wf-1", @@ -680,15 +638,7 @@ async def capture_upsert(summary): ) monkeypatch.setattr(workflow_module.WorkflowStore, "upsert_execution", capture_upsert) - progress_writer = workflow_module.ExecutionProgressWriter( - { - "id": "exec-cancel-route", - "workflowId": "wf-1", - "status": "running", - "currentPhase": "queued", - "executionLog": [], - } - ) + progress_writer = _progress_writer("exec-cancel-route", currentPhase="queued") progress_writer.submit({"currentPhase": "running", "currentNodeId": "node-1"}) cancel_event = workflow_module.threading.Event() current_task = asyncio.current_task() diff --git a/tests/workflow/test_execution_store_compact.py b/tests/workflow/test_execution_store_compact.py index 0ceb57849..ea8c3d5cd 100644 --- a/tests/workflow/test_execution_store_compact.py +++ b/tests/workflow/test_execution_store_compact.py @@ -34,7 +34,6 @@ compact_execution_summary, compact_outputs_for_storage, compact_step_for_storage, - create_execution_record, record_execution_result, workflow_execution_step_key, ) @@ -45,6 +44,11 @@ def _make_alerts(n: int) -> List[Dict[str, Any]]: return [{"sip": f"1.2.3.{i % 256}", "url": f"/p/{i}"} for i in range(n)] +def _raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 + coro.close() + raise RuntimeError + + # ── compact_outputs_for_storage ─────────────────────────────────────────────── @@ -304,27 +308,10 @@ def test_workflow_execution_step_key_is_append_only_namespaced() -> None: assert workflow_execution_step_key("exec-1", 12) == "workflow_execution_step/exec-1/00000012" -@pytest.mark.asyncio -async def test_create_execution_record_can_skip_initial_database_write() -> None: - upsert_execution = AsyncMock(return_value=None) - - with patch.object(WorkflowStore, "upsert_execution", upsert_execution): - record = await create_execution_record( - "wf-trigger", - input_params={"message": "hello"}, - exec_id="exec-trigger", - persist=False, - ) - - assert record["id"] == "exec-trigger" - assert record["currentPhase"] == "queued" - upsert_execution.assert_not_awaited() - - def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: record_step = AsyncMock(return_value=None) record_steps = AsyncMock(return_value=None) - recorder = ExecutionStepRecorder(exec_id="exec-batch") + recorder = ExecutionStepRecorder() with ( patch.object(WorkflowStore, "record_step", record_step), @@ -344,34 +331,6 @@ def test_execution_step_recorder_collects_steps_without_storage_calls() -> None: record_steps.assert_not_awaited() -@pytest.mark.asyncio -async def test_four_trigger_workers_collect_steps_without_storage() -> None: - """Four trigger threads collect complete batches without callback SQL.""" - record_step = AsyncMock(return_value=None) - record_steps = AsyncMock(return_value=None) - recorders = [ExecutionStepRecorder(exec_id=f"exec-trigger-{worker}") for worker in range(4)] - - def _run_seven_steps(recorder: ExecutionStepRecorder) -> None: - for step in range(7): - recorder.on_step_complete( - {"node_id": f"node-{step}", "outputs": {"ok": True}} - ) - - with ( - patch.object(WorkflowStore, "record_step", record_step), - patch.object(WorkflowStore, "record_steps", record_steps), - ): - await asyncio.gather( - *(asyncio.to_thread(_run_seven_steps, recorder) for recorder in recorders) - ) - - batches = [recorder.take_steps() for recorder in recorders] - assert [len(batch) for batch in batches] == [7, 7, 7, 7] - assert [recorder.step_count for recorder in recorders] == [7, 7, 7, 7] - record_step.assert_not_awaited() - record_steps.assert_not_awaited() - - @pytest.mark.asyncio async def test_progress_writer_submits_without_waiting_and_coalesces_updates() -> None: write_started = asyncio.Event() @@ -507,16 +466,12 @@ async def trim_executions(*args, **kwargs): # noqa: ANN002, ANN003 ], } - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution_mock), patch.object(WorkflowStore, "increment_stats", increment_stats_mock), patch.object(WorkflowStore, "trim_executions", trim_executions_mock), patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result("wf", "exec-1", exec_data) @@ -555,16 +510,12 @@ async def test_record_execution_result_accepts_explicit_step_batch() -> None: "stepCount": 2, } - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution), patch.object(WorkflowStore, "increment_stats", increment_stats), patch.object(WorkflowStore, "trim_executions", trim_executions), patch("flocks.session.recorder.Recorder.record_workflow_execution", record_audit), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result( "wf-trigger", @@ -592,16 +543,12 @@ async def test_record_execution_result_stats_failure_does_not_block_retention() increment_stats = AsyncMock(side_effect=RuntimeError("stats locked")) trim_executions = AsyncMock(return_value=[]) - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution), patch.object(WorkflowStore, "increment_stats", increment_stats), patch.object(WorkflowStore, "trim_executions", trim_executions), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result( "wf-stats-failure", @@ -626,16 +573,12 @@ async def test_record_execution_result_retention_failure_keeps_committed_executi increment_stats = AsyncMock(return_value=None) trim_executions = AsyncMock(side_effect=RuntimeError("retention locked")) - def raise_create_task(coro, *args, **kwargs): # noqa: ANN001, ARG001 - coro.close() - raise RuntimeError - with ( patch.object(WorkflowStore, "complete_execution", complete_execution), patch.object(WorkflowStore, "increment_stats", increment_stats), patch.object(WorkflowStore, "trim_executions", trim_executions), patch("flocks.session.recorder.Recorder.record_workflow_execution", AsyncMock(return_value=None)), - patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=raise_create_task), + patch("flocks.workflow.execution_store.asyncio.create_task", side_effect=_raise_create_task), ): await record_execution_result( "wf-retention-failure", diff --git a/tests/workflow/test_poller_manager.py b/tests/workflow/test_poller_manager.py index f2f5f6914..cb10353ec 100644 --- a/tests/workflow/test_poller_manager.py +++ b/tests/workflow/test_poller_manager.py @@ -99,7 +99,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -154,9 +154,7 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, - persist: bool = True, ) -> dict[str, Any]: - assert persist is True record = { "id": exec_id or "exec-1", "workflowId": workflow_id, @@ -304,7 +302,7 @@ def _fake_run_workflow( # noqa: ANN001 monkeypatch.setattr( poller_manager, "create_execution_record", - lambda workflow_id, *, input_params=None, exec_id=None, persist=True: asyncio.sleep( + lambda workflow_id, *, input_params=None, exec_id=None: asyncio.sleep( 0, result={ "id": exec_id or f"exec-{workflow_id}", @@ -359,9 +357,7 @@ async def _fake_create_execution_record( *, input_params: dict[str, Any] | None = None, exec_id: str | None = None, - persist: bool = True, ) -> dict[str, Any]: - assert persist is True _ = input_params return { "id": exec_id or f"exec-{workflow_id}", diff --git a/tests/workflow/test_trigger_runtime.py b/tests/workflow/test_trigger_runtime.py index 279ade3b0..7f05ce780 100644 --- a/tests/workflow/test_trigger_runtime.py +++ b/tests/workflow/test_trigger_runtime.py @@ -75,7 +75,10 @@ def _fake_run_workflow(**kwargs): # noqa: ANN003 assert runtime_module.run_workflow.call_args.kwargs["run_id"] == "exec-1" assert runtime_module.run_workflow.call_args.kwargs["execution_profile"] == "high_frequency" assert callable(runtime_module.run_workflow.call_args.kwargs["on_step_complete"]) - assert create_record.await_args.kwargs["persist"] is True + create_record.assert_awaited_once_with( + "wf-trigger", + input_params={"message": "hello"}, + ) assert result["executionLog"] == [] assert result["stepCount"] == 1 record_result.assert_awaited_once()