AgentRunStore.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.SQLException;
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.om.Insight;
import prerna.query.querystruct.SelectQueryStruct;
import prerna.query.querystruct.filters.SimpleQueryFilter;
import prerna.query.querystruct.selectors.QueryColumnSelector;
import prerna.reactor.agent.AgentRunContext;
import prerna.util.ConnectionUtils;
import prerna.util.QueryExecutionUtility;
import prerna.util.SystemEngineRegistry;
import prerna.util.Utility;

/**
 * All SQL against the {@code AGENT_RUN} table: inserts, status transitions,
 * queue queries, and the user-facing reads.
 *
 * <p>
 * The table is the durable record of a run, so some of the concurrency rules
 * are expressed here as SQL rather than as Java.
 * {@link #markRunningIfSubmitted} is the conditional update that decides which
 * node gets to execute a run, and {@link #isOldestSubmittedRunForRoom} is the
 * per-room FIFO order the worker honors before claiming anything.
 *
 * <p>
 * Reads made on a user's behalf are scoped by {@code USER_ID} so one user
 * cannot read another's runs. The two queue queries the worker uses are
 * deliberately unscoped, because the worker serves every user on the node and
 * has no caller identity of its own.
 */
public final class AgentRunStore {

	private static final Gson GSON = new Gson();
	private static final String ACTIVITY_LOG_COLUMNS = "ar.RUN_ID, ar.PARENT_RUN_ID, ar.ROOM_ID, ar.WORKSPACE_ID, ar.MODEL_ID, "
			+ "ar.HARNESS_TYPE, ar.JOB_ID, ar.STATUS, ar.INPUT, ar.INPUT_MESSAGE_ID, ar.FINAL_OUTPUT, ar.FINAL_OUTPUT_MESSAGE_ID, "
			+ "ar.ERROR_MESSAGE, ar.PROGRESS_JSON, ar.DATE_CREATED, ar.STARTED_AT, ar.COMPLETED_AT, ar.USER_ID, r.ROOM_NAME";
	// Rooms are keyed per user, so the name join must match on both columns.
	private static final String ACTIVITY_LOG_FROM = "FROM AGENT_RUN ar "
			+ "LEFT JOIN ROOM r ON ar.ROOM_ID = r.ROOM_ID AND ar.USER_ID = r.USER_ID";

	private AgentRunStore() {

	}

	public static void insertSubmitted(String runId, AgentRunRequest request, String userId) {
		insert(runId, request, userId, AgentRunStatus.SUBMITTED);
	}

	// Human children start waiting; no worker ever picks them up.
	static void insertInputRequired(String runId, AgentRunRequest request, String userId) {
		insert(runId, request, userId, AgentRunStatus.INPUT_REQUIRED);
	}

	private static void insert(String runId, AgentRunRequest request, String userId, AgentRunStatus status) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			String query = "INSERT INTO AGENT_RUN (RUN_ID, PARENT_RUN_ID, ROOM_ID, WORKSPACE_ID, MODEL_ID, HARNESS_TYPE, JOB_ID, "
					+ "STATUS, INPUT, REQUEST_JSON, INPUT_MESSAGE_ID, FINAL_OUTPUT, FINAL_OUTPUT_MESSAGE_ID, ERROR_MESSAGE, "
					+ "DATE_CREATED, STARTED_AT, COMPLETED_AT, USER_ID) "
					+ "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)";
			ps = db.getPreparedStatement(query);
			int idx = 1;
			ps.setString(idx++, runId);
			setNullableString(ps, idx++, request.getParentRunId());
			setNullableString(ps, idx++, request.getRoomId());
			setNullableString(ps, idx++, request.getWorkspaceId());
			setNullableString(ps, idx++, request.getEngineIdFallback());
			setNullableString(ps, idx++, request.getHarnessType());
			ps.setString(idx++, runId);
			ps.setString(idx++, status.name());
			setClob(db, ps, idx++, request.getInput());
			setClob(db, ps, idx++, GSON.toJson(request.toPersistedMap()));
			ps.setNull(idx++, Types.VARCHAR);
			ps.setNull(idx++, Types.CLOB);
			ps.setNull(idx++, Types.VARCHAR);
			ps.setNull(idx++, Types.CLOB);
			ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			ps.setNull(idx++, Types.TIMESTAMP);
			ps.setNull(idx++, Types.TIMESTAMP);
			setNullableString(ps, idx++, userId);
			ps.executeUpdate();
			commitIfNeeded(ps);
		} catch (Exception e) {
			throw new IllegalStateException("Failed to create AGENT_RUN row for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	public static AgentRunRecord getRun(String runId, Insight insight) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String userId = resolveInsightUserId(insight);
			String query = "SELECT RUN_ID, PARENT_RUN_ID, ROOM_ID, WORKSPACE_ID, MODEL_ID, HARNESS_TYPE, STATUS, INPUT, REQUEST_JSON, "
					+ "USER_ID, JOB_ID FROM AGENT_RUN WHERE RUN_ID = ? AND USER_ID = ?";
			ps = db.getPreparedStatement(query);
			ps.setString(1, runId);
			ps.setString(2, userId);
			rs = ps.executeQuery();
			if (!rs.next()) {
				return null;
			}
			AgentRunRequest request = requestFromRow(rs, insight);
			return new AgentRunRecord(rs.getString("RUN_ID"), rs.getString("ROOM_ID"),
					parseRunStatus(rs.getString("STATUS")), request, rs.getString("USER_ID"), rs.getString("JOB_ID"));
		} catch (Exception e) {
			if (e instanceof SecurityException) {
				throw (SecurityException) e;
			}
			throw new IllegalStateException("Failed to load AGENT_RUN row for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/**
	 * Loads an agent run without applying the owning-user predicate.
	 *
	 * <p>
	 * This method is reserved for Automation APIs that have already verified project
	 * access and the exact persisted Automation run, node, and agent-run trace.
	 * Generic agent APIs must use {@link #getRun(String, Insight)}.
	 *
	 * @param runId agent run identifier
	 * @param insight current Automation editor insight used to reconstruct the request
	 * @return matching run, or {@code null} when it does not exist
	 */
	public static AgentRunRecord getRunForAutomation(String runId, Insight insight) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		try {
			SelectQueryStruct qs = new SelectQueryStruct();
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__RUN_ID", "runId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__PARENT_RUN_ID", "parentRunId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__ROOM_ID", "roomId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__WORKSPACE_ID", "workspaceId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__MODEL_ID", "modelId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__HARNESS_TYPE", "harnessType"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__STATUS", "status"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__INPUT", "input"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__REQUEST_JSON", "requestJson"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__USER_ID", "userId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__JOB_ID", "jobId"));
			qs.addExplicitFilter(SimpleQueryFilter.makeColToValFilter("AGENT_RUN__RUN_ID", "==", runId));

			List<Map<String, Object>> rows = QueryExecutionUtility.flushRsToMap(db, qs);
			if (rows.isEmpty()) {
				return null;
			}
			Map<String, Object> row = rows.get(0);
			AgentRunRequest request = requestFromMap(row, insight);
			return new AgentRunRecord(stringValue(row.get("runId")), stringValue(row.get("roomId")),
					parseRunStatus(stringValue(row.get("status"))), request, stringValue(row.get("userId")),
					stringValue(row.get("jobId")));
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load Automation AGENT_RUN row for runId=" + runId, e);
		}
	}

	public static Map<String, Object> getRunMap(String runId, Insight insight) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String userId = resolveInsightUserId(insight);
			String query = "SELECT " + ACTIVITY_LOG_COLUMNS + ", ar.REQUEST_JSON " + ACTIVITY_LOG_FROM
					+ " WHERE ar.RUN_ID = ? AND ar.USER_ID = ?";
			ps = db.getPreparedStatement(query);
			ps.setString(1, runId);
			ps.setString(2, userId);
			rs = ps.executeQuery();
			if (!rs.next()) {
				return null;
			}
			return runMapFromRow(rs);
		} catch (Exception e) {
			if (e instanceof SecurityException) {
				throw (SecurityException) e;
			}
			throw new IllegalStateException("Failed to load AGENT_RUN details for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/**
	 * Loads agent-run details without applying the owning-user predicate.
	 *
	 * <p>
	 * This is the detail-map equivalent of {@link #getRunForAutomation(String, Insight)}
	 * and is only valid after Automation trace authorization succeeds.
	 *
	 * @param runId agent run identifier
	 * @return matching run details, or {@code null} when it does not exist
	 */
	public static Map<String, Object> getRunMapForAutomation(String runId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		try {
			SelectQueryStruct qs = new SelectQueryStruct();
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__RUN_ID", "runId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__PARENT_RUN_ID", "parentRunId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__ROOM_ID", "roomId"));
			qs.addSelector(new QueryColumnSelector("ROOM__ROOM_NAME", "roomName"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__WORKSPACE_ID", "workspaceId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__MODEL_ID", "modelId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__HARNESS_TYPE", "harnessType"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__JOB_ID", "jobId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__STATUS", "status"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__INPUT", "input"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__INPUT_MESSAGE_ID", "inputMessageId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__FINAL_OUTPUT", "finalText"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__FINAL_OUTPUT_MESSAGE_ID", "finalOutputMessageId"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__ERROR_MESSAGE", "errorMessage"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__DATE_CREATED", "dateCreated"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__STARTED_AT", "startedAt"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__COMPLETED_AT", "completedAt"));
			qs.addSelector(new QueryColumnSelector("AGENT_RUN__USER_ID", "userId"));
			qs.addRelation("AGENT_RUN__ROOM_ID", "ROOM__ROOM_ID", "left.outer.join");
			qs.addRelation("AGENT_RUN__USER_ID", "ROOM__USER_ID", "left.outer.join");
			qs.addExplicitFilter(SimpleQueryFilter.makeColToValFilter("AGENT_RUN__RUN_ID", "==", runId));

			List<Map<String, Object>> rows = QueryExecutionUtility.flushRsToMap(db, qs);
			if (rows.isEmpty()) {
				return null;
			}
			Map<String, Object> run = rows.get(0);
			run.put("dateCreated", stringValue(run.get("dateCreated")));
			run.put("startedAt", stringValue(run.get("startedAt")));
			run.put("completedAt", stringValue(run.get("completedAt")));
			run.put("artifacts", new ArrayList<>());
			return run;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load Automation AGENT_RUN details for runId=" + runId, e);
		}
	}

	/**
	 * Return one page of runs for an agent owned by the current user, newest first.
	 *
	 * @param insight current request insight, used to scope the query to its user
	 * @param agentId agent workspace identifier used to scope the runs
	 * @param limit   maximum number of runs to return
	 * @param offset  number of matching runs to skip
	 * @return the requested page of agent runs
	 */
	public static List<Map<String, Object>> getActivityLog(Insight insight, String agentId, long limit, long offset) {
		if (agentId == null || agentId.trim().isEmpty()) {
			throw new IllegalArgumentException("agentId is required");
		}
		if (limit <= 0) {
			throw new IllegalArgumentException("limit must be greater than 0");
		}
		if (offset < 0) {
			throw new IllegalArgumentException("offset must be greater than or equal to 0");
		}

		String userId = resolveInsightUserId(insight);
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			StringBuilder query = new StringBuilder("SELECT " + ACTIVITY_LOG_COLUMNS + " " + ACTIVITY_LOG_FROM
					+ " WHERE ar.USER_ID = ? AND ar.WORKSPACE_ID = ? ORDER BY ar.DATE_CREATED DESC, ar.RUN_ID DESC");
			db.getQueryUtil().addLimitOffsetToQuery(query, limit, offset);

			ps = db.getPreparedStatement(query.toString());
			ps.setString(1, userId);
			ps.setString(2, agentId.trim());
			rs = ps.executeQuery();

			List<Map<String, Object>> runs = new ArrayList<>();
			while (rs.next()) {
				runs.add(runMapFromRow(rs));
			}
			return runs;
		} catch (Exception e) {
			if (e instanceof SecurityException) {
				throw (SecurityException) e;
			}
			throw new IllegalStateException("Failed to load AGENT_RUN activity log", e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/**
	 * Return every run in a room owned by the current user, newest first.
	 *
	 * @param insight current request insight, used to scope the query to its user
	 * @param roomId  room whose complete run history should be returned
	 * @return all current-user runs for the room
	 */
	public static List<Map<String, Object>> getRunsForRoom(Insight insight, String roomId) {
		if (roomId == null || roomId.trim().isEmpty()) {
			throw new IllegalArgumentException("roomId is required");
		}

		String userId = resolveInsightUserId(insight);
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT " + ACTIVITY_LOG_COLUMNS + " " + ACTIVITY_LOG_FROM
					+ " WHERE ar.USER_ID = ? AND ar.ROOM_ID = ? ORDER BY ar.DATE_CREATED DESC, ar.RUN_ID DESC";
			ps = db.getPreparedStatement(query);
			ps.setString(1, userId);
			ps.setString(2, roomId.trim());
			rs = ps.executeQuery();

			List<Map<String, Object>> runs = new ArrayList<>();
			while (rs.next()) {
				runs.add(runMapFromRow(rs));
			}
			return runs;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load AGENT_RUN rows for roomId=" + roomId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/**
	 * Return every direct subagent run for a parent run owned by the current user,
	 * newest first.
	 *
	 * @param insight     current request insight, used to scope the query to its
	 *                    user
	 * @param parentRunId parent run whose direct subagent runs should be returned
	 * @return all current-user runs with the supplied parent run ID
	 */
	public static List<Map<String, Object>> getSubagentRuns(Insight insight, String parentRunId) {
		if (parentRunId == null || parentRunId.trim().isEmpty()) {
			throw new IllegalArgumentException("runId is required");
		}

		String userId = resolveInsightUserId(insight);
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT " + ACTIVITY_LOG_COLUMNS + ", ar.REQUEST_JSON " + ACTIVITY_LOG_FROM
					+ " WHERE ar.USER_ID = ? AND ar.PARENT_RUN_ID = ? ORDER BY ar.DATE_CREATED DESC, ar.RUN_ID DESC";
			ps = db.getPreparedStatement(query);
			ps.setString(1, userId);
			ps.setString(2, parentRunId.trim());
			rs = ps.executeQuery();

			List<Map<String, Object>> runs = new ArrayList<>();
			while (rs.next()) {
				Map<String, Object> run = runMapFromRow(rs);
				// Lets child cards tell a person's task from an agent's.
				AgentRunRequest request = requestFromJson(rs.getString("REQUEST_JSON"));
				boolean human = request != null && request.isHumanExecutor();
				run.put("executorType", human ? "HUMAN" : "AGENT");
				run.put("executorLabel", human ? request.getHumanExecutorLabel() : null);
				runs.add(run);
			}
			return runs;
		} catch (Exception e) {
			if (e instanceof SecurityException) {
				throw (SecurityException) e;
			}
			throw new IllegalStateException("Failed to load subagent AGENT_RUN rows for parentRunId=" + parentRunId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	/** Internal durable inputs for detached-child delivery and repair. */
	static List<Map<String, Object>> getTerminalChildCompletions(String childRunId, String parentRunId,
			String parentRoomId, String userId, long limit) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			StringBuilder query = new StringBuilder(
					"SELECT child.RUN_ID AS CHILD_RUN_ID, child.PARENT_RUN_ID, child.STATUS, child.FINAL_OUTPUT, "
							+ "child.ERROR_MESSAGE, child.REQUEST_JSON, child.USER_ID, parent.ROOM_ID AS PARENT_ROOM_ID, "
							+ "parent.REQUEST_JSON AS PARENT_REQUEST_JSON "
							+ "FROM AGENT_RUN child JOIN AGENT_RUN parent ON child.PARENT_RUN_ID = parent.RUN_ID "
							+ "AND child.USER_ID = parent.USER_ID WHERE child.PARENT_RUN_ID IS NOT NULL "
							+ "AND child.STATUS IN (?, ?, ?)");
			if (childRunId != null && !childRunId.isBlank()) {
				query.append(" AND child.RUN_ID = ?");
			}
			if (parentRunId != null && !parentRunId.isBlank()) {
				query.append(" AND child.PARENT_RUN_ID = ?");
			}
			if (parentRoomId != null && !parentRoomId.isBlank()) {
				query.append(" AND parent.ROOM_ID = ?");
			}
			if (userId != null && !userId.isBlank()) {
				query.append(" AND child.USER_ID = ?");
			}
			query.append(" ORDER BY child.COMPLETED_AT DESC, child.RUN_ID DESC");
			if (limit > 0) {
				db.getQueryUtil().addLimitOffsetToQuery(query, limit, 0);
			}

			ps = db.getPreparedStatement(query.toString());
			int idx = 1;
			ps.setString(idx++, AgentRunStatus.COMPLETED.name());
			ps.setString(idx++, AgentRunStatus.FAILED.name());
			ps.setString(idx++, AgentRunStatus.CANCELLED.name());
			if (childRunId != null && !childRunId.isBlank()) {
				ps.setString(idx++, childRunId.trim());
			}
			if (parentRunId != null && !parentRunId.isBlank()) {
				ps.setString(idx++, parentRunId.trim());
			}
			if (parentRoomId != null && !parentRoomId.isBlank()) {
				ps.setString(idx++, parentRoomId.trim());
			}
			if (userId != null && !userId.isBlank()) {
				ps.setString(idx++, userId.trim());
			}

			rs = ps.executeQuery();
			List<Map<String, Object>> completions = new ArrayList<>();
			while (rs.next()) {
				Map<String, Object> completion = new HashMap<>();
				completion.put("childRunId", rs.getString("CHILD_RUN_ID"));
				completion.put("parentRunId", rs.getString("PARENT_RUN_ID"));
				completion.put("parentRoomId", rs.getString("PARENT_ROOM_ID"));
				completion.put("status", rs.getString("STATUS"));
				completion.put("finalText", rs.getString("FINAL_OUTPUT"));
				completion.put("errorMessage", rs.getString("ERROR_MESSAGE"));
				completion.put("requestJson", rs.getString("REQUEST_JSON"));
				completion.put("parentRequestJson", rs.getString("PARENT_REQUEST_JSON"));
				completion.put("userId", rs.getString("USER_ID"));
				completions.add(completion);
			}
			return completions;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load terminal child completions", e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	// Unscoped: the assignee answering a delegation is not the owner of either run.
	static String getParentRoomId(String childRunId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			ps = db.getPreparedStatement("SELECT parent.ROOM_ID FROM AGENT_RUN child JOIN AGENT_RUN parent "
					+ "ON child.PARENT_RUN_ID = parent.RUN_ID AND child.USER_ID = parent.USER_ID WHERE child.RUN_ID = ?");
			ps.setString(1, childRunId);
			rs = ps.executeQuery();
			return rs.next() ? rs.getString(1) : null;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to load parent room for runId=" + childRunId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	static boolean runExists(String runId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT 1 FROM AGENT_RUN WHERE RUN_ID = ?";
			ps = db.getPreparedStatement(query);
			ps.setString(1, runId);
			rs = ps.executeQuery();
			return rs.next();
		} catch (Exception e) {
			throw new IllegalStateException("Failed to inspect AGENT_RUN existence for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	public static List<AgentRunRecord> getSubmittedRuns(int limit, Insight insight) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT RUN_ID, PARENT_RUN_ID, ROOM_ID, WORKSPACE_ID, MODEL_ID, HARNESS_TYPE, STATUS, INPUT, REQUEST_JSON, "
					+ "USER_ID, JOB_ID FROM AGENT_RUN WHERE STATUS = ? ORDER BY DATE_CREATED ASC, RUN_ID ASC";
			ps = db.getPreparedStatement(query);
			ps.setString(1, AgentRunStatus.SUBMITTED.name());
			rs = ps.executeQuery();
			List<AgentRunRecord> records = new ArrayList<>();
			while (rs.next() && (limit <= 0 || records.size() < limit)) {
				AgentRunRequest request = requestFromRow(rs, insight);
				records.add(new AgentRunRecord(rs.getString("RUN_ID"), rs.getString("ROOM_ID"),
						parseRunStatus(rs.getString("STATUS")), request, rs.getString("USER_ID"),
						rs.getString("JOB_ID")));
			}
			return records;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to inspect submitted AGENT_RUN rows", e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	public static boolean isOldestSubmittedRunForRoom(String runId, String roomId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		ResultSet rs = null;
		try {
			String query = "SELECT RUN_ID FROM AGENT_RUN WHERE ROOM_ID = ? AND STATUS = ? "
					+ "ORDER BY DATE_CREATED ASC, RUN_ID ASC";
			ps = db.getPreparedStatement(query);
			setNullableString(ps, 1, roomId);
			ps.setString(2, AgentRunStatus.SUBMITTED.name());
			rs = ps.executeQuery();
			if (!rs.next()) {
				return false;
			}
			return runId.equals(rs.getString(1));
		} catch (Exception e) {
			throw new IllegalStateException("Failed to inspect AGENT_RUN queue for roomId=" + roomId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, rs);
		}
	}

	public static void markRunning(String runId, String jobId) {
		updateStatus(runId, AgentRunStatus.RUNNING, jobId, null, null, true, false);
	}

	public static boolean markRunningIfSubmitted(String runId, String jobId) {
		return updateStatusIfCurrent(runId, AgentRunStatus.SUBMITTED, AgentRunStatus.RUNNING, jobId, null, null, true,
				false);
	}

	public static void markInputMessage(String runId, String messageId) {
		updateMessageId(runId, "INPUT_MESSAGE_ID", messageId);
	}

	public static void markFinalOutputMessage(String runId, String messageId) {
		updateMessageId(runId, "FINAL_OUTPUT_MESSAGE_ID", messageId);
	}

	public static void markCompleted(String runId, String jobId, String finalOutput) {
		updateStatus(runId, AgentRunStatus.COMPLETED, jobId, finalOutput, null, false, true);
	}

	public static void markIncomplete(String runId, String jobId, String finalOutput, String errorMessage) {
        updateStatus(runId, AgentRunStatus.FAILED, jobId, finalOutput, errorMessage, false, true);
    }

	public static void markFailed(String runId, String jobId, String errorMessage) {
		updateStatus(runId, AgentRunStatus.FAILED, jobId, null, errorMessage, false, true);
	}

	public static boolean markFailedIfSubmitted(String runId, String jobId, String errorMessage) {
		return updateStatusIfCurrent(runId, AgentRunStatus.SUBMITTED, AgentRunStatus.FAILED, jobId, null,
				errorMessage, false, true);
	}

	public static void markCancelled(String runId, String jobId, String errorMessage) {
		updateStatus(runId, AgentRunStatus.CANCELLED, jobId, null, errorMessage, false, true);
	}

	/**
	 * Transition a run from RUNNING to INPUT_REQUIRED. Called by the worker when
	 * the harness pauses on {@code SMSS_MCP_EXECUTION=ask} tools.
	 */
	public static void markInputRequired(String runId, String jobId) {
		updateStatusIfCurrent(runId, AgentRunStatus.RUNNING, AgentRunStatus.INPUT_REQUIRED, jobId, null, null, false,
				false);
	}

	/**
	 * Transition a run from INPUT_REQUIRED back to SUBMITTED so the worker picks it
	 * up and resumes. Called by {@code RunMCPToolReactor} after the user's decision
	 * has been written to the room.
	 */
	public static boolean markResumed(String runId, String jobId) {
		return updateStatusIfCurrent(runId, AgentRunStatus.INPUT_REQUIRED, AgentRunStatus.SUBMITTED, jobId, null, null,
				false, false);
	}

	static boolean markCompletedIfInputRequired(String runId, String finalOutput) {
		return updateStatusIfCurrent(runId, AgentRunStatus.INPUT_REQUIRED, AgentRunStatus.COMPLETED, runId,
				finalOutput, null, false, true);
	}

	static boolean markFailedIfInputRequired(String runId, String errorMessage) {
		return updateStatusIfCurrent(runId, AgentRunStatus.INPUT_REQUIRED, AgentRunStatus.FAILED, runId, null,
				errorMessage, false, true);
	}

	public static boolean markCancelledIfNotTerminal(String runId, String jobId, String errorMessage) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			String query = "UPDATE AGENT_RUN SET STATUS = ?, JOB_ID = ?, ERROR_MESSAGE = ?, COMPLETED_AT = ? "
					+ "WHERE RUN_ID = ? AND STATUS IN (?, ?, ?)";
			ps = db.getPreparedStatement(query);
			int idx = 1;
			ps.setString(idx++, AgentRunStatus.CANCELLED.name());
			setNullableString(ps, idx++, jobId);
			setClob(db, ps, idx++, errorMessage);
			ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			ps.setString(idx++, runId);
			ps.setString(idx++, AgentRunStatus.SUBMITTED.name());
			ps.setString(idx++, AgentRunStatus.RUNNING.name());
			ps.setString(idx++, AgentRunStatus.INPUT_REQUIRED.name());
			int updated = ps.executeUpdate();
			commitIfNeeded(ps);
			return updated > 0;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to cancel AGENT_RUN runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	private static void updateStatus(String runId, AgentRunStatus status, String jobId, String finalOutput,
			String errorMessage, boolean setStartedAt, boolean setCompletedAt) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			StringBuilder query = new StringBuilder("UPDATE AGENT_RUN SET STATUS = ?, JOB_ID = ?");
			if (finalOutput != null) {
				query.append(", FINAL_OUTPUT = ?");
			}
			if (errorMessage != null) {
				query.append(", ERROR_MESSAGE = ?");
			}
			if (setStartedAt) {
				query.append(", STARTED_AT = ?");
			}
			if (setCompletedAt) {
				query.append(", COMPLETED_AT = ?");
			}
			query.append(" WHERE RUN_ID = ?");

			ps = db.getPreparedStatement(query.toString());
			int idx = 1;
			ps.setString(idx++, status.name());
			setNullableString(ps, idx++, jobId);
			if (finalOutput != null) {
				setClob(db, ps, idx++, finalOutput);
			}
			if (errorMessage != null) {
				setClob(db, ps, idx++, errorMessage);
			}
			if (setStartedAt) {
				ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			}
			if (setCompletedAt) {
				ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			}
			ps.setString(idx++, runId);
			ps.executeUpdate();
			commitIfNeeded(ps);
		} catch (Exception e) {
			throw new IllegalStateException("Failed to update AGENT_RUN status for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	private static boolean updateStatusIfCurrent(String runId, AgentRunStatus currentStatus, AgentRunStatus nextStatus,
			String jobId, String finalOutput, String errorMessage, boolean setStartedAt, boolean setCompletedAt) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			StringBuilder query = new StringBuilder("UPDATE AGENT_RUN SET STATUS = ?, JOB_ID = ?");
			if (finalOutput != null) {
				query.append(", FINAL_OUTPUT = ?");
			}
			if (errorMessage != null) {
				query.append(", ERROR_MESSAGE = ?");
			}
			if (setStartedAt) {
				query.append(", STARTED_AT = ?");
			}
			if (setCompletedAt) {
				query.append(", COMPLETED_AT = ?");
			}
			query.append(" WHERE RUN_ID = ? AND STATUS = ?");

			ps = db.getPreparedStatement(query.toString());
			int idx = 1;
			ps.setString(idx++, nextStatus.name());
			setNullableString(ps, idx++, jobId);
			if (finalOutput != null) {
				setClob(db, ps, idx++, finalOutput);
			}
			if (errorMessage != null) {
				setClob(db, ps, idx++, errorMessage);
			}
			if (setStartedAt) {
				ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			}
			if (setCompletedAt) {
				ps.setTimestamp(idx++, Utility.getCurrentSqlTimestampUTC());
			}
			ps.setString(idx++, runId);
			ps.setString(idx++, currentStatus.name());
			int updated = ps.executeUpdate();
			commitIfNeeded(ps);
			return updated > 0;
		} catch (Exception e) {
			throw new IllegalStateException("Failed to update AGENT_RUN status for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

    /** Store progress independently of the request and final outcome, including failed runs. */
    public static void updateProgress(String runId, Map<String, Object> progress) {
        IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
        PreparedStatement ps = null;
        try {
            ps = db.getPreparedStatement("UPDATE AGENT_RUN SET PROGRESS_JSON = ? WHERE RUN_ID = ?");
            setClob(db, ps, 1, GSON.toJson(progress));
            ps.setString(2, runId);
            ps.executeUpdate();
            commitIfNeeded(ps);
        } catch (Exception e) {
            throw new IllegalStateException("Failed to persist progress for runId=" + runId, e);
        } finally {
            ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
        }
    }

	private static void updateMessageId(String runId, String columnName, String messageId) {
		IRDBMSEngine db = SystemEngineRegistry.getModelInferenceLogsDb();
		PreparedStatement ps = null;
		try {
			String safeColumnName = safeMessageIdColumn(columnName);
			String query = "UPDATE AGENT_RUN SET " + safeColumnName + " = ? WHERE RUN_ID = ?";
			ps = db.getPreparedStatement(query);
			setNullableString(ps, 1, messageId);
			ps.setString(2, runId);
			ps.executeUpdate();
			commitIfNeeded(ps);
		} catch (Exception e) {
			throw new IllegalStateException("Failed to update AGENT_RUN " + columnName + " for runId=" + runId, e);
		} finally {
			ConnectionUtils.closeAllConnectionsIfPooling(db, null, ps, null);
		}
	}

	private static String safeMessageIdColumn(String columnName) {
		if ("INPUT_MESSAGE_ID".equals(columnName)) {
			return "INPUT_MESSAGE_ID";
		}
		if ("FINAL_OUTPUT_MESSAGE_ID".equals(columnName)) {
			return "FINAL_OUTPUT_MESSAGE_ID";
		}
		throw new IllegalArgumentException("Unsupported AGENT_RUN message-id column: " + columnName);
	}

	@SuppressWarnings("unchecked")
	private static AgentRunRequest requestFromRow(ResultSet rs, Insight insight) throws Exception {
		String requestJson = rs.getString("REQUEST_JSON");
		if (requestJson != null && !requestJson.trim().isEmpty()) {
			Map<String, Object> persisted = GSON.fromJson(requestJson, Map.class);
			AgentRunRequest request = AgentRunRequest.fromPersistedMap(persisted, insight);
			if (request != null) {
				String parentRunId = rs.getString("PARENT_RUN_ID");
				return request.getParentRunId() == null && parentRunId != null ? request.withParentRunId(parentRunId)
						: request;
			}
		}
		return new AgentRunRequest(rs.getString("ROOM_ID"), rs.getString("INPUT"), rs.getString("MODEL_ID"),
				rs.getString("HARNESS_TYPE"), rs.getString("WORKSPACE_ID"), AgentRunContext.DEFAULT_MAX_TURNS,
				AgentRunContext.DEFAULT_MAX_REFLECTIONS, null, null, null, null, insight)
				.withParentRunId(rs.getString("PARENT_RUN_ID"));
	}

	@SuppressWarnings("unchecked")
	private static AgentRunRequest requestFromMap(Map<String, Object> row, Insight insight) {
		String requestJson = stringValue(row.get("requestJson"));
		if (requestJson != null) {
			Map<String, Object> persisted = GSON.fromJson(requestJson, Map.class);
			AgentRunRequest request = AgentRunRequest.fromPersistedMap(persisted, insight);
			if (request != null) {
				String parentRunId = stringValue(row.get("parentRunId"));
				return request.getParentRunId() == null && parentRunId != null ? request.withParentRunId(parentRunId)
						: request;
			}
		}
		return new AgentRunRequest(stringValue(row.get("roomId")), stringValue(row.get("input")),
				stringValue(row.get("modelId")), stringValue(row.get("harnessType")),
				stringValue(row.get("workspaceId")), AgentRunContext.DEFAULT_MAX_TURNS,
				AgentRunContext.DEFAULT_MAX_REFLECTIONS, null, null, null, null, insight)
				.withParentRunId(stringValue(row.get("parentRunId")));
	}

	@SuppressWarnings("unchecked")
	private static AgentRunRequest requestFromJson(String json) {
		try {
			return json == null ? null : AgentRunRequest.fromPersistedMap(GSON.fromJson(json, Map.class), null);
		} catch (RuntimeException e) {
			return null;
		}
	}

	private static AgentRunStatus parseRunStatus(String value) {
		if (value == null || value.trim().isEmpty()) {
			return null;
		}
		return AgentRunStatus.valueOf(value.trim());
	}

	private static String resolveInsightUserId(Insight insight) {
		if (insight == null) {
			throw new SecurityException("AGENT_RUN access requires an authenticated user");
		}
		String insightUserId = insight.getUserId();
		if (insightUserId == null || insightUserId.trim().isEmpty() || "-1".equals(insightUserId)) {
			throw new SecurityException("AGENT_RUN access requires an authenticated user");
		}
		return insightUserId;
	}

	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();
		}
	}

	private static String stringValue(Object value) {
		return value == null ? null : String.valueOf(value);
	}

	private static Map<String, Object> runMapFromRow(ResultSet rs) throws SQLException {
		Map<String, Object> map = new HashMap<>();
		map.put("runId", rs.getString("RUN_ID"));
		map.put("parentRunId", rs.getString("PARENT_RUN_ID"));
		map.put("roomId", rs.getString("ROOM_ID"));
		map.put("roomName", rs.getString("ROOM_NAME"));
		map.put("workspaceId", rs.getString("WORKSPACE_ID"));
		map.put("modelId", rs.getString("MODEL_ID"));
		map.put("harnessType", rs.getString("HARNESS_TYPE"));
		map.put("jobId", rs.getString("JOB_ID"));
		map.put("status", rs.getString("STATUS"));
		map.put("input", rs.getString("INPUT"));
		map.put("inputMessageId", rs.getString("INPUT_MESSAGE_ID"));
		map.put("finalText", rs.getString("FINAL_OUTPUT"));
		map.put("finalOutputMessageId", rs.getString("FINAL_OUTPUT_MESSAGE_ID"));
		map.put("errorMessage", rs.getString("ERROR_MESSAGE"));
        String progressJson = rs.getString("PROGRESS_JSON");
        if (progressJson != null && !progressJson.isBlank()) map.put("progress", GSON.fromJson(progressJson, Map.class));
        Map<String, Object> liveProgress = prerna.reactor.agent.runtime.AgentRunProgress.activeSnapshot(rs.getString("RUN_ID"));
        if (liveProgress != null) map.put("progress", liveProgress);
		map.put("dateCreated", stringValue(rs.getTimestamp("DATE_CREATED")));
		map.put("startedAt", stringValue(rs.getTimestamp("STARTED_AT")));
		map.put("completedAt", stringValue(rs.getTimestamp("COMPLETED_AT")));
		map.put("userId", rs.getString("USER_ID"));
		map.put("artifacts", new ArrayList<>());
		addDeliveryMetadata(map);
		return map;
	}

    /** Project code-owned delivery evidence independently of the run's execution status. */
    static void addDeliveryMetadata(Map<String, Object> run) {
        if (!(run.get("progress") instanceof Map<?, ?> progress)
                || !(progress.get("workflow") instanceof Map<?, ?> workflow)) return;
        if (workflow.get("reviewOutcome") instanceof Map<?, ?> review) run.put("reviewOutcome", new HashMap<>(review));
        if (workflow.get("warning") instanceof String warning && !warning.isBlank()) run.put("warnings", List.of(warning));
        if (workflow.get("artifact") instanceof Map<?, ?> artifact && "available".equals(artifact.get("status")))
            run.put("artifacts", List.of(new HashMap<>(artifact)));
    }
}