101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import timedelta
|
|
from temporalio import workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from ops_agent.temporal.activities import (
|
|
run_ops_diagnose_activity,
|
|
run_ops_remediate_activity
|
|
)
|
|
|
|
@workflow.defn
|
|
class OpsRemediateWorkflow:
|
|
"""Orchestrates OpsAgent diagnosis, waits for human operator approval signal, and executes remediation."""
|
|
|
|
def __init__(self) -> None:
|
|
self._approved = False
|
|
self._rejected = False
|
|
|
|
@workflow.signal
|
|
def approve(self) -> None:
|
|
self._approved = True
|
|
|
|
@workflow.signal
|
|
def reject(self) -> None:
|
|
self._rejected = True
|
|
|
|
@workflow.run
|
|
async def run(self, target_type: str, target_id: str) -> dict:
|
|
# Step 1: Diagnose issues
|
|
params = {"target_type": target_type, "target_id": target_id}
|
|
diagnose_res = await workflow.execute_activity(
|
|
run_ops_diagnose_activity,
|
|
params,
|
|
start_to_close_timeout=timedelta(minutes=5),
|
|
)
|
|
|
|
proposed_action = diagnose_res.get("proposed_action", "none")
|
|
|
|
# Step 2: Check if remediation is proposed. If 'none', exit.
|
|
if proposed_action == "none":
|
|
return {
|
|
"status": "completed",
|
|
"diagnosis": diagnose_res["diagnosis"],
|
|
"action_proposed": "none",
|
|
"action_executed": "none",
|
|
"result": "No problems detected. No action needed."
|
|
}
|
|
|
|
# Step 3: Approval Gate (wait for approval signal or reject signal or timeout)
|
|
# Timeout after 24 hours (or shorter in test/production depending on requirements)
|
|
approval_timeout = timedelta(hours=2)
|
|
|
|
workflow.logger.info(
|
|
f"Action proposed: {proposed_action}. Waiting for approval signal on workflow ID: {workflow.info().workflow_id}"
|
|
)
|
|
|
|
# Wait for approval condition
|
|
try:
|
|
is_condition_met = await workflow.wait_condition(
|
|
lambda: self._approved or self._rejected,
|
|
timeout=approval_timeout
|
|
)
|
|
except Exception:
|
|
# wait_condition times out
|
|
is_condition_met = False
|
|
|
|
if not is_condition_met:
|
|
return {
|
|
"status": "timed_out",
|
|
"diagnosis": diagnose_res["diagnosis"],
|
|
"action_proposed": proposed_action,
|
|
"action_executed": "none",
|
|
"result": "Timed out waiting for administrator approval."
|
|
}
|
|
|
|
if self._rejected:
|
|
return {
|
|
"status": "rejected",
|
|
"diagnosis": diagnose_res["diagnosis"],
|
|
"action_proposed": proposed_action,
|
|
"action_executed": "none",
|
|
"result": "Remediation action rejected by administrator."
|
|
}
|
|
|
|
# Step 4: Execute remediation
|
|
remed_params = {**diagnose_res, "is_approved": True}
|
|
exec_res = await workflow.execute_activity(
|
|
run_ops_remediate_activity,
|
|
remed_params,
|
|
start_to_close_timeout=timedelta(minutes=5),
|
|
)
|
|
|
|
return {
|
|
"status": "remediated",
|
|
"diagnosis": diagnose_res["diagnosis"],
|
|
"action_proposed": proposed_action,
|
|
"action_executed": proposed_action,
|
|
"result": exec_res["action_result"]
|
|
}
|