AgentRunActionStore.java

/*******************************************************************************
 * Copyright 2015 Defense Health Agency (DHA)
 *
 * If your use of this software does not include any GPLv2 components:
 * 	Licensed under the Apache License, Version 2.0 (the "License");
 * 	you may not use this file except in compliance with the License.
 * 	You may obtain a copy of the License at
 *
 * 	  http://www.apache.org/licenses/LICENSE-2.0
 *
 * 	Unless required by applicable law or agreed to in writing, software
 * 	distributed under the License is distributed on an "AS IS" BASIS,
 * 	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * 	See the License for the specific language governing permissions and
 * 	limitations under the License.
 * ----------------------------------------------------------------------------
 * If your use of this software includes any GPLv2 components:
 * 	This program is free software; you can redistribute it and/or
 * 	modify it under the terms of the GNU General Public License
 * 	as published by the Free Software Foundation; either version 2
 * 	of the License, or (at your option) any later version.
 *
 * 	This program is distributed in the hope that it will be useful,
 * 	but WITHOUT ANY WARRANTY; without even the implied warranty of
 * 	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * 	GNU General Public License for more details.
 *******************************************************************************/
package prerna.reactor.agent.run;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Types;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import com.google.gson.Gson;

import prerna.engine.api.IRDBMSEngine;
import prerna.query.querystruct.SelectQueryStruct;
import prerna.query.querystruct.filters.SimpleQueryFilter;
import prerna.query.querystruct.selectors.QueryColumnSelector;
import prerna.util.ConnectionUtils;
import prerna.util.QueryExecutionUtility;
import prerna.util.SystemEngineRegistry;
import prerna.util.Utility;
import prerna.util.ValueUtils;

/**
 * Persistence layer for the {@code AGENT_RUN_ACTION} table.
 *
 * <p>
 * Each row represents a single MCP tool call that was paused because its
 * {@code SMSS_MCP_EXECUTION} is {@code "ask"}. The harness creates rows when it
 * transitions a run to {@code INPUT_REQUIRED}; {@code RunMCPToolReactor}
 * updates them when the user decides; {@code GetAgentRunReactor} reads them to
 * surface pending actions to the UI.
 */
public final class AgentRunActionStore {

	private static final Gson GSON = new Gson();

	private AgentRunActionStore() {

	}

	/**
	 * Insert a batch of pending actions for a single run. Called by the harness
	 * when it pauses on {@code ask} tools.
	 *
	 * @param runId   the agent run id
	 * @param roomId  the room id
	 * @param userId  the user id
	 * @param actions one map per pending tool call, each carrying at minimum:
	 *                {@code toolCallId}, {@code toolName}, {@code toolArgs},
	 *                {@code toolMeta}, {@code parentMessageId}, {@code hasUi},
	 *                {@code uiUrl}
	 */
	public static void insertPendingActions(String runId, String roomId, String userId,
			List<Map<String, Object>> actions) {
		if (actions == null || actions.isEmpty()) {
			return;
		}
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			String query = "INSERT INTO AGENT_RUN_ACTION (ACTION_ID, RUN_ID, ROOM_ID, PARENT_MESSAGE_ID, "
					+ "TOOL_CALL_ID, TOOL_NAME, TOOL_ARGS, TOOL_META, HAS_UI, UI_URL, STATUS, "
					+ "DATE_CREATED, USER_ID) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)";
			ps = db.getPreparedStatement(query);
			for (Map<String, Object> action : actions) {
				int idx = 1;
				String actionId = ValueUtils.trimToNull(action.get("actionId"));
				if (actionId == null) {
					throw new IllegalArgumentException("actionId is required for AGENT_RUN_ACTION rows");
				}
				ps.setString(idx++, actionId);
				ps.setString(idx++, runId);
				ps.setString(idx++, roomId);
				setNullableString(ps, idx++, ValueUtils.trimToNull(action.get("parentMessageId")));
				setNullableString(ps, idx++, ValueUtils.trimToNull(action.get("toolCallId")));
				setNullableString(ps, idx++, ValueUtils.trimToNull(action.get("toolName")));
				setClob(db, ps, idx++, toJson(action.get("toolArgs")));
				setClob(db, ps, idx++, toJson(action.get("toolMeta")));
				ps.setString(idx++, booleanValue(action.get("hasUi")) ? "true" : "false");
				setClob(db, ps, idx++, ValueUtils.trimToNull(action.get("uiUrl")));
				ps.setString(idx++, "PENDING");
				ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
				setNullableString(ps, idx++, userId);
				ps.addBatch();
			}
			ps.executeBatch();
			commitIfNeeded(ps);
		} catch (Exception e) {
			throw new IllegalStateException("Failed to insert AGENT_RUN_ACTION rows for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	/**
	 * Return all action rows for a run id, ordered by creation time. Used by
	 * {@code GetAgentRunReactor} to surface pending actions.
	 */
	public static List<Map<String, Object>> getActionsForRun(String runId) {
		return getActionsForRun(runId, null);
	}

	/**
	 * Return all action rows for a run id and, when supplied, the owning user.
	 * Owner-scoped read surfaces should always use this overload.
	 */
	public static List<Map<String, Object>> getActionsForRun(String runId, String userId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT ACTION_ID, RUN_ID, ROOM_ID, PARENT_MESSAGE_ID, TOOL_CALL_ID, TOOL_NAME, "
					+ "TOOL_ARGS, EDITED_ARGS, TOOL_META, HAS_UI, UI_URL, STATUS, "
					+ "RESULT, TOOL_STATUS, DATE_CREATED, DECIDED_AT, USER_ID "
					+ "FROM AGENT_RUN_ACTION WHERE RUN_ID = ?" + (userId != null ? " AND USER_ID = ?" : "")
					+ " ORDER BY DATE_CREATED ASC";
			ps = db.getPreparedStatement(query);
			ps.setString(1, runId);
			if (userId != null) {
				ps.setString(2, userId);
			}
			rs = ps.executeQuery();
			List<Map<String, Object>> results = new ArrayList<>();
			while (rs.next()) {
				results.add(rowToMap(rs));
			}
			return results;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load AGENT_RUN_ACTION rows for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/**
	 * Return only PENDING actions for a run.
	 */
	public static List<Map<String, Object>> getPendingActions(String runId) {
		List<Map<String, Object>> all = getActionsForRun(runId);
		List<Map<String, Object>> pending = new ArrayList<>();
		for (Map<String, Object> a : all) {
			// A delegation belongs to its assignee, not the run owner's approval UI.
			if ("PENDING".equals(a.get("status")) && !HumanDelegationService.isDelegationAction(a)) {
				pending.add(a);
			}
		}
		return pending;
	}

	/**
	 * Return one pending action row by action id and owner. ACTION_ID is a globally
	 * unique v7 UUID (the table PK), so scoping by actionId + userId is sufficient;
	 * runId is derivable from the row.
	 */
	public static Map<String, Object> getPendingActionById(String actionId, String userId) {
		Map<String, Object> action = getActionById(actionId, userId);
		if (action != null && "PENDING".equals(action.get("status"))) {
			return action;
		}
		return null;
	}

	/**
	 * Return one action row by action id and owner, regardless of status.
	 */
	public static Map<String, Object> getActionById(String actionId, String userId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT ACTION_ID, RUN_ID, ROOM_ID, PARENT_MESSAGE_ID, TOOL_CALL_ID, TOOL_NAME, "
					+ "TOOL_ARGS, EDITED_ARGS, TOOL_META, HAS_UI, UI_URL, STATUS, "
					+ "RESULT, TOOL_STATUS, DATE_CREATED, DECIDED_AT, USER_ID "
					+ "FROM AGENT_RUN_ACTION WHERE ACTION_ID = ? AND USER_ID = ?";
			ps = db.getPreparedStatement(query);
			ps.setString(1, actionId);
			ps.setString(2, userId);
			rs = ps.executeQuery();
			if (rs.next()) {
				return rowToMap(rs);
			}
			return null;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load AGENT_RUN_ACTION actionId=" + actionId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/**
	 * Loads an action without applying the owning-user predicate.
	 *
	 * <p>
	 * This method is reserved for Automation APIs that have already verified
	 * project edit access and the exact persisted Automation trace. Generic agent
	 * APIs must use {@link #getActionById(String, String)}.
	 *
	 * @param actionId agent action identifier
	 * @return matching action, or {@code null} when it does not exist
	 */
	public static Map<String, Object> getActionByIdForAutomation(String actionId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		try {
			SelectQueryStruct qs = new SelectQueryStruct();
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__ACTION_ID", "actionId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__RUN_ID", "runId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__ROOM_ID", "roomId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__PARENT_MESSAGE_ID", "parentMessageId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__TOOL_CALL_ID", "toolCallId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__TOOL_NAME", "toolName"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__TOOL_ARGS", "toolArgs"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__EDITED_ARGS", "editedArgs"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__TOOL_META", "toolMeta"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__HAS_UI", "hasUi"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__UI_URL", "uiUrl"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__STATUS", "status"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__RESULT", "result"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__TOOL_STATUS", "toolStatus"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__DATE_CREATED", "dateCreated"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__DECIDED_AT", "decidedAt"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN_ACTION__USER_ID", "userId"));
			qs.addExplicitFilter(SimpleQueryFilter.makeColToValFilter("AGENT_RUN_ACTION__ACTION_ID", "==", actionId));

			List<Map<String, Object>> rows = QueryExecutionUtility.flushRsToMap(db, qs);
			if (rows.isEmpty()) {
				return null;
			}
			Map<String, Object> action = rows.get(0);
			action.put("hasUi", booleanValue(action.get("hasUi")));
			action.put("toolArgs", ValueUtils.trimToNull(action.get("toolArgs")));
			action.put("editedArgs", ValueUtils.trimToNull(action.get("editedArgs")));
			action.put("toolMeta", ValueUtils.trimToNull(action.get("toolMeta")));
			action.put("uiUrl", ValueUtils.trimToNull(action.get("uiUrl")));
			action.put("result", ValueUtils.trimToNull(action.get("result")));
			action.put("dateCreated", ValueUtils.trimToNull(action.get("dateCreated")));
			action.put("decidedAt", ValueUtils.trimToNull(action.get("decidedAt")));
			return action;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load Automation AGENT_RUN_ACTION actionId=" + actionId, e);
		}
	}

	/**
	 * Claim a pending action before executing a side-effecting tool. Only one
	 * request can move the action from PENDING to EXECUTING.
	 */
	public static boolean claimForExecution(String actionId, String runId, String userId) {
		return updateStatus(actionId, runId, userId, "PENDING", "EXECUTING", false);
	}

	/**
	 * Return an EXECUTING action to PENDING when the tool was not executed to a
	 * durable result.
	 */
	public static boolean releaseExecutionClaim(String actionId, String runId, String userId) {
		return updateStatus(actionId, runId, userId, "EXECUTING", "PENDING", false);
	}

	/**
	 * Mark an action as decided. STATUS is the single lifecycle column: it moves
	 * from PENDING to one of APPROVED / EDITED / REJECTED / RESPONDED, which also
	 * encodes what the user chose (there is no separate DECISION column).
	 *
	 * @param actionId   the action id
	 * @param editedArgs final args when they differ from what the model proposed,
	 *                   else null
	 * @param result     the tool result (approve/edit), the user's response
	 *                   (respond), or the rejection message (reject)
	 * @param status     the decided status: APPROVED, EDITED, REJECTED, RESPONDED
	 * @param toolStatus execution status persisted for exact replay
	 */
	public static boolean markDecided(String actionId, String runId, String userId, Object editedArgs, String result,
			String status, String toolStatus) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			StringBuilder query = new StringBuilder("UPDATE AGENT_RUN_ACTION SET STATUS = ?, DECIDED_AT = ?");
			if (editedArgs != null) {
				query.append(", EDITED_ARGS = ?");
			}
			if (result != null) {
				query.append(", RESULT = ?");
			}
			if (toolStatus != null) {
				query.append(", TOOL_STATUS = ?");
			}
			query.append(" WHERE ACTION_ID = ? AND RUN_ID = ? AND USER_ID = ? AND STATUS IN ('PENDING', 'EXECUTING')");

			ps = db.getPreparedStatement(query.toString());
			int idx = 1;
			ps.setString(idx++, status);
			ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			if (editedArgs != null) {
				setClob(db, ps, idx++, toJson(editedArgs));
			}
			if (result != null) {
				setClob(db, ps, idx++, result);
			}
			if (toolStatus != null) {
				ps.setString(idx++, toolStatus);
			}
			ps.setString(idx++, actionId);
			ps.setString(idx++, runId);
			ps.setString(idx++, userId);
			int updated = ps.executeUpdate();
			commitIfNeeded(ps);
			return updated > 0;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to update AGENT_RUN_ACTION for actionId=" + actionId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	/**
	 * Actions of one tool type assigned to a user, newest first; status is
	 * optional.
	 */
	static List<Map<String, Object>> getAssignedActions(String userId, String toolName, String status, int limit) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			StringBuilder query = new StringBuilder("SELECT ACTION_ID, RUN_ID, ROOM_ID, PARENT_MESSAGE_ID, "
					+ "TOOL_CALL_ID, TOOL_NAME, TOOL_ARGS, EDITED_ARGS, TOOL_META, HAS_UI, UI_URL, STATUS, "
					+ "RESULT, TOOL_STATUS, DATE_CREATED, DECIDED_AT, USER_ID "
					+ "FROM AGENT_RUN_ACTION WHERE USER_ID = ? AND TOOL_NAME = ?");
			if (status != null) {
				query.append(" AND STATUS = ?");
			}
			query.append(" ORDER BY DATE_CREATED DESC");
			db.getQueryUtil().addLimitOffsetToQuery(query, limit, 0);
			ps = db.getPreparedStatement(query.toString());
			int idx = 1;
			ps.setString(idx++, userId);
			ps.setString(idx++, toolName);
			if (status != null) {
				ps.setString(idx++, status);
			}
			rs = ps.executeQuery();
			List<Map<String, Object>> results = new ArrayList<>();
			while (rs.next()) {
				results.add(rowToMap(rs));
			}
			return results;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load assigned AGENT_RUN_ACTION rows", e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/** Cancel every still-pending action on a run, whoever it is assigned to. */
	static void cancelPendingForRun(String runId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			ps = db.getPreparedStatement(
					"UPDATE AGENT_RUN_ACTION SET STATUS = ?, DECIDED_AT = ? " + "WHERE RUN_ID = ? AND STATUS = ?");
			ps.setString(1, "CANCELLED");
			ps.setTimestamp(2, Utility.getCurrentSqlTimestampUTC());
			ps.setString(3, runId);
			ps.setString(4, "PENDING");
			ps.executeUpdate();
			commitIfNeeded(ps);
		} catch (Exception e) {
			throw new IllegalStateException("Failed to cancel AGENT_RUN_ACTION rows for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	/**
	 * Return true if there is at least one action row for the given run. Used by
	 * the worker as a short-term marker that the run has entered the HITL resume
	 * flow without relying on persisted REQUEST_JSON.resumeMode.
	 */
	public static boolean hasAnyActions(String runId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT COUNT(*) FROM AGENT_RUN_ACTION WHERE RUN_ID = ?";
			ps = db.getPreparedStatement(query);
			ps.setString(1, runId);
			rs = ps.executeQuery();
			return rs.next() && rs.getInt(1) > 0;
		} catch (Exception e) {
			return false;
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/**
	 * Check if all actions for a run have been decided (no PENDING remaining).
	 */
	public static boolean allActionsDecided(String runId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT COUNT(*) FROM AGENT_RUN_ACTION WHERE RUN_ID = ? AND STATUS IN ('PENDING', 'EXECUTING')";
			ps = db.getPreparedStatement(query);
			ps.setString(1, runId);
			rs = ps.executeQuery();
			if (rs.next()) {
				return rs.getInt(1) == 0;
			}
			return true;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to count pending AGENT_RUN_ACTION rows for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	// --- helpers ---

	private static boolean updateStatus(String actionId, String runId, String userId, String currentStatus,
			String nextStatus, boolean setDecidedAt) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			StringBuilder query = new StringBuilder("UPDATE AGENT_RUN_ACTION SET STATUS = ?");
			if (setDecidedAt) {
				query.append(", DECIDED_AT = ?");
			}
			query.append(" WHERE ACTION_ID = ? AND RUN_ID = ? AND USER_ID = ? AND STATUS = ?");
			ps = db.getPreparedStatement(query.toString());
			int idx = 1;
			ps.setString(idx++, nextStatus);
			if (setDecidedAt) {
				ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			}
			ps.setString(idx++, actionId);
			ps.setString(idx++, runId);
			ps.setString(idx++, userId);
			ps.setString(idx++, currentStatus);
			int updated = ps.executeUpdate();
			commitIfNeeded(ps);
			return updated > 0;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to update AGENT_RUN_ACTION status for actionId=" + actionId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	private static Map<String, Object> rowToMap(ResultSet rs) throws Exception {
		Map<String, Object> map = new HashMap<>();
		map.put("actionId", rs.getString("ACTION_ID"));
		map.put("runId", rs.getString("RUN_ID"));
		map.put("roomId", rs.getString("ROOM_ID"));
		map.put("parentMessageId", rs.getString("PARENT_MESSAGE_ID"));
		map.put("toolCallId", rs.getString("TOOL_CALL_ID"));
		map.put("toolName", rs.getString("TOOL_NAME"));
		map.put("toolArgs", clobToString(rs, "TOOL_ARGS"));
		map.put("editedArgs", clobToString(rs, "EDITED_ARGS"));
		map.put("toolMeta", clobToString(rs, "TOOL_META"));
		map.put("hasUi", "true".equalsIgnoreCase(rs.getString("HAS_UI")));
		map.put("uiUrl", clobToString(rs, "UI_URL"));
		map.put("status", rs.getString("STATUS"));
		map.put("result", clobToString(rs, "RESULT"));
		map.put("toolStatus", rs.getString("TOOL_STATUS"));
		map.put("dateCreated", ValueUtils.trimToNull(rs.getTimestamp("DATE_CREATED")));
		map.put("decidedAt", ValueUtils.trimToNull(rs.getTimestamp("DECIDED_AT")));
		map.put("userId", rs.getString("USER_ID"));
		return map;
	}

	private static String clobToString(ResultSet rs, String colName) throws Exception {
		String val = rs.getString(colName);
		return val == null || val.trim().isEmpty() ? null : val;
	}

	private static boolean booleanValue(Object value) {
		if (value == null) {
			return false;
		}
		return Boolean.TRUE.equals(value) || "true".equalsIgnoreCase(String.valueOf(value));
	}

	private static String toJson(Object value) {
		if (value == null) {
			return null;
		}
		if (value instanceof String) {
			return (String) value;
		}
		return GSON.toJson(value);
	}

	private static void setNullableString(PreparedStatement ps, int idx, String value) throws Exception {
		if (value == null || value.trim().isEmpty()) {
			ps.setNull(idx, Types.VARCHAR);
		} else {
			ps.setString(idx, value);
		}
	}

	private static void setClob(IRDBMSEngine db, PreparedStatement ps, int idx, String value) throws Exception {
		if (value == null) {
			ps.setNull(idx, Types.CLOB);
		} else {
			db.getQueryUtil().handleInsertionOfClob(ps, value, idx, GSON);
		}
	}

	private static void commitIfNeeded(PreparedStatement ps) throws Exception {
		if (!ps.getConnection().getAutoCommit()) {
			ps.getConnection().commit();
		}
	}
}