From 03f6d3086d339457d6867da98b121fa1bb540563 Mon Sep 17 00:00:00 2001 From: Ark2307 Date: Mon, 13 Apr 2026 16:16:00 +0530 Subject: [PATCH 1/3] Added query data to send to db --- .../agentic_sessions/AgentQueryDataDao.java | 76 +++++++++++++++++++ .../dto/agentic_sessions/AgentQueryData.java | 72 ++++++++++++++++++ .../java/com/akto/data_actor/ClientActor.java | 15 ++++ .../java/com/akto/data_actor/DataActor.java | 3 + .../java/com/akto/data_actor/DbActor.java | 6 ++ .../java/com/akto/data_actor/DbLayer.java | 8 +- 6 files changed, 179 insertions(+), 1 deletion(-) create mode 100644 libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java create mode 100644 libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java diff --git a/libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java b/libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java new file mode 100644 index 0000000000..050ba5dc3a --- /dev/null +++ b/libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java @@ -0,0 +1,76 @@ +package com.akto.dao.agentic_sessions; + +import com.akto.dao.AccountsContextDao; +import com.akto.dao.MCollection; +import com.akto.dto.HttpResponseParams; +import com.akto.dto.agentic_sessions.AgentQueryData; +import com.mongodb.client.model.IndexOptions; +import com.mongodb.client.model.Indexes; +import org.bson.conversions.Bson; + +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; + +public class AgentQueryDataDao extends AccountsContextDao { + + public static final AgentQueryDataDao instance = new AgentQueryDataDao(); + + private static final long TTL_DAYS = 14; + + public void createIndicesIfAbsent() { + // TTL index: auto-delete records older than 14 days + Bson ttlIndex = Indexes.ascending(AgentQueryData.TIME_STAMP); + IndexOptions ttlOptions = new IndexOptions() + .name("timeStamp_ttl") + .expireAfter(TTL_DAYS * 24 * 60 * 60, TimeUnit.SECONDS); + MCollection.createIndexIfAbsent(getDBName(), getCollName(), ttlIndex, ttlOptions); + + // Compound index for cron queries: fetch by user within a time window + String[] compoundFields = new String[]{ + AgentQueryData.SERVICE_ID, + AgentQueryData.UNIQUE_USER_ID, + AgentQueryData.TIME_STAMP + }; + MCollection.createIndexIfAbsent(getDBName(), getCollName(), compoundFields, false); + } + + public AgentQueryData createAgentQueryDataFromHttpResponseParams(HttpResponseParams httpResponseParam) { + AgentQueryData agentQueryData = new AgentQueryData(); + + Map> requestHeaders = httpResponseParam.getRequestParams().getHeaders(); + + agentQueryData.setServiceId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_INSTALLER)); + agentQueryData.setUniqueUserId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_USER_EMAIL)); + agentQueryData.setSessionIdentifier(getFirstHeader(requestHeaders, AgentQueryData.HEADER_SESSION_ID)); + agentQueryData.setConversationId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_CONVERSATION_ID)); + agentQueryData.setGenerationId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_GENERATION_ID)); + agentQueryData.setModel(getFirstHeader(requestHeaders, AgentQueryData.HEADER_MODEL)); + agentQueryData.setTranscriptPath(getFirstHeader(requestHeaders, AgentQueryData.HEADER_TRANSCRIPT_PATH)); + agentQueryData.setCwd(getFirstHeader(requestHeaders, AgentQueryData.HEADER_CWD)); + agentQueryData.setPermissionMode(getFirstHeader(requestHeaders, AgentQueryData.HEADER_PERMISSION_MODE)); + agentQueryData.setHookEventName(getFirstHeader(requestHeaders, AgentQueryData.HEADER_HOOK_EVENT_NAME)); + + agentQueryData.setQueryPayload(httpResponseParam.getRequestParams().getPayload()); + agentQueryData.setResponsePayload(httpResponseParam.getPayload()); + agentQueryData.setTimeStamp(httpResponseParam.getTimeOrNow()); + + return agentQueryData; + } + + private static String getFirstHeader(Map> headers, String name) { + if (headers == null) return null; + List values = headers.get(name); + return (values != null && !values.isEmpty()) ? values.get(0) : null; + } + + @Override + public String getCollName() { + return "agent_query_data"; + } + + @Override + public Class getClassT() { + return AgentQueryData.class; + } +} diff --git a/libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java b/libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java new file mode 100644 index 0000000000..77f9cdf947 --- /dev/null +++ b/libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java @@ -0,0 +1,72 @@ +package com.akto.dto.agentic_sessions; + +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; + +@AllArgsConstructor +@NoArgsConstructor +@Getter +@Setter + +// use this in future for learning the user type, build the flow graph of the agent +public class AgentQueryData { + + // DB field name constants + public static final String SERVICE_ID = "serviceId"; + public static final String UNIQUE_USER_ID = "uniqueUserId"; + public static final String SESSION_IDENTIFIER = "sessionIdentifier"; + public static final String QUERY_PAYLOAD = "queryPayload"; + public static final String TIME_STAMP = "timeStamp"; + public static final String INPUT_TOKENS = "inputTokens"; + public static final String OUTPUT_TOKENS = "outputTokens"; + + // Request header name constants + public static final String HEADER_INSTALLER = "x-akto-installer"; + public static final String HEADER_USER_EMAIL = "user_email"; + public static final String HEADER_SESSION_ID = "session_id"; + public static final String HEADER_CONVERSATION_ID = "conversation_id"; + public static final String HEADER_GENERATION_ID = "generation_id"; + public static final String HEADER_MODEL = "model"; + public static final String HEADER_TRANSCRIPT_PATH = "transcript_path"; + public static final String HEADER_CWD = "cwd"; + public static final String HEADER_PERMISSION_MODE = "permission_mode"; + public static final String HEADER_HOOK_EVENT_NAME = "hook_event_name"; + + // from x-akto-installer + private String serviceId; + + // from user_email + private String uniqueUserId; + + // from session_id + private String sessionIdentifier; + + // from conversation_id + private String conversationId; + + // from generation_id + private String generationId; + + // from model + private String model; + + // from transcript_path + private String transcriptPath; + + // from cwd + private String cwd; + + // from permission_mode + private String permissionMode; + + // from hook_event_name + private String hookEventName; + + private String queryPayload; + private String responsePayload; + private long timeStamp; + private int inputTokens; + private int outputTokens; +} \ No newline at end of file diff --git a/libs/utils/src/main/java/com/akto/data_actor/ClientActor.java b/libs/utils/src/main/java/com/akto/data_actor/ClientActor.java index 5a99632062..0c27ef63ed 100644 --- a/libs/utils/src/main/java/com/akto/data_actor/ClientActor.java +++ b/libs/utils/src/main/java/com/akto/data_actor/ClientActor.java @@ -15,6 +15,7 @@ import com.akto.database_abstractor_authenticator.JwtAuthenticator; import com.akto.dto.*; import com.akto.dto.ApiInfo.ApiInfoKey; +import com.akto.dto.agentic_sessions.AgentQueryData; import com.akto.dto.billing.Organization; import com.akto.dto.billing.Tokens; import com.akto.dto.bulk_updates.BulkUpdates; @@ -4460,4 +4461,18 @@ public void storeTestingRunWebhook(TestingRunWebhook testingRunWebhook) { return; } } + + @Override + public void storeAgentQueryData(AgentQueryData agentQueryData) { + Map> headers = buildHeaders(); + BasicDBObject obj = new BasicDBObject(); + obj.put("agentQueryData", agentQueryData); + OriginalHttpRequest request = new OriginalHttpRequest(url + "/storeAgentQueryData", "", "POST", obj.toString(), headers, ""); + try { + OriginalHttpResponse response = ApiExecutor.sendRequest(request, true, null, false, null); + } catch (Exception e) { + loggerMaker.errorAndAddToDb("error in storeAgentQueryData" + e, LoggerMaker.LogDb.RUNTIME); + return; + } + } } diff --git a/libs/utils/src/main/java/com/akto/data_actor/DataActor.java b/libs/utils/src/main/java/com/akto/data_actor/DataActor.java index 3e976fff0d..1e6e858fbb 100644 --- a/libs/utils/src/main/java/com/akto/data_actor/DataActor.java +++ b/libs/utils/src/main/java/com/akto/data_actor/DataActor.java @@ -1,6 +1,7 @@ package com.akto.data_actor; import com.akto.dto.*; +import com.akto.dto.agentic_sessions.AgentQueryData; import com.akto.dto.billing.Organization; import com.akto.dto.billing.Tokens; import com.akto.dto.dependency_flow.Node; @@ -402,4 +403,6 @@ public void updateNewNodesInBatches(List updateNodes) { public abstract void storeSpans(List spans); public abstract void storeTestingRunWebhook(TestingRunWebhook testingRunWebhook); + + public abstract void storeAgentQueryData(AgentQueryData agentQueryData); } diff --git a/libs/utils/src/main/java/com/akto/data_actor/DbActor.java b/libs/utils/src/main/java/com/akto/data_actor/DbActor.java index 0bd762c02c..c76e5d77d0 100644 --- a/libs/utils/src/main/java/com/akto/data_actor/DbActor.java +++ b/libs/utils/src/main/java/com/akto/data_actor/DbActor.java @@ -2,6 +2,7 @@ import com.akto.dto.*; import com.akto.dto.ApiInfo.ApiInfoKey; +import com.akto.dto.agentic_sessions.AgentQueryData; import com.akto.dto.billing.Organization; import com.akto.dto.billing.Tokens; import com.akto.dto.dependency_flow.Node; @@ -744,4 +745,9 @@ public void storeSpans(List spans) { public void storeTestingRunWebhook(TestingRunWebhook testingRunWebhook) { DbLayer.storeTestingRunWebhook(testingRunWebhook); } + + @Override + public void storeAgentQueryData(AgentQueryData agentQueryData) { + DbLayer.storeAgentQueryData(agentQueryData); + } } \ No newline at end of file diff --git a/libs/utils/src/main/java/com/akto/data_actor/DbLayer.java b/libs/utils/src/main/java/com/akto/data_actor/DbLayer.java index 7faa8635d7..c0dde84852 100644 --- a/libs/utils/src/main/java/com/akto/data_actor/DbLayer.java +++ b/libs/utils/src/main/java/com/akto/data_actor/DbLayer.java @@ -21,6 +21,7 @@ import com.akto.bulk_update_util.ApiInfoBulkUpdate; import com.akto.dao.*; +import com.akto.dao.agentic_sessions.AgentQueryDataDao; import com.akto.dao.filter.MergedUrlsDao; import com.akto.dao.graph.SvcToSvcGraphEdgesDao; import com.akto.dao.graph.SvcToSvcGraphNodesDao; @@ -76,6 +77,7 @@ import com.akto.dao.traffic_metrics.TrafficMetricsDao; import com.akto.dao.upload.FileUploadsDao; import com.akto.dto.ApiInfo.ApiInfoKey; +import com.akto.dto.agentic_sessions.AgentQueryData; import com.akto.dto.billing.Organization; import com.akto.dto.billing.Tokens; import com.akto.dto.dependency_flow.Node; @@ -2762,4 +2764,8 @@ public static void bulkUpsertAgenticSessionContext(List session SessionDocumentDao.instance.getMCollection().bulkWrite(bulkUpdates); } } -} \ No newline at end of file + + public static void storeAgentQueryData(AgentQueryData agentQueryData) { + AgentQueryDataDao.instance.insertOne(agentQueryData); + } +} From bbcfa72a9f56d8a78146d8819dd6e9b73561a96c Mon Sep 17 00:00:00 2001 From: Ark2307 Date: Tue, 14 Apr 2026 10:31:32 +0530 Subject: [PATCH 2/3] Fixing query data --- .../agentic_sessions/AgentQueryDataDao.java | 17 +++------ .../dto/agentic_sessions/AgentQueryData.java | 38 ++----------------- 2 files changed, 10 insertions(+), 45 deletions(-) diff --git a/libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java b/libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java index 050ba5dc3a..ccdb178b5c 100644 --- a/libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java +++ b/libs/dao/src/main/java/com/akto/dao/agentic_sessions/AgentQueryDataDao.java @@ -2,6 +2,7 @@ import com.akto.dao.AccountsContextDao; import com.akto.dao.MCollection; +import com.akto.dao.context.Context; import com.akto.dto.HttpResponseParams; import com.akto.dto.agentic_sessions.AgentQueryData; import com.mongodb.client.model.IndexOptions; @@ -40,20 +41,14 @@ public AgentQueryData createAgentQueryDataFromHttpResponseParams(HttpResponsePar Map> requestHeaders = httpResponseParam.getRequestParams().getHeaders(); - agentQueryData.setServiceId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_INSTALLER)); - agentQueryData.setUniqueUserId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_USER_EMAIL)); - agentQueryData.setSessionIdentifier(getFirstHeader(requestHeaders, AgentQueryData.HEADER_SESSION_ID)); - agentQueryData.setConversationId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_CONVERSATION_ID)); - agentQueryData.setGenerationId(getFirstHeader(requestHeaders, AgentQueryData.HEADER_GENERATION_ID)); - agentQueryData.setModel(getFirstHeader(requestHeaders, AgentQueryData.HEADER_MODEL)); - agentQueryData.setTranscriptPath(getFirstHeader(requestHeaders, AgentQueryData.HEADER_TRANSCRIPT_PATH)); - agentQueryData.setCwd(getFirstHeader(requestHeaders, AgentQueryData.HEADER_CWD)); - agentQueryData.setPermissionMode(getFirstHeader(requestHeaders, AgentQueryData.HEADER_PERMISSION_MODE)); - agentQueryData.setHookEventName(getFirstHeader(requestHeaders, AgentQueryData.HEADER_HOOK_EVENT_NAME)); + agentQueryData.setServiceId(getFirstHeader(requestHeaders, "host")); + agentQueryData.setUserName(getFirstHeader(requestHeaders, (AgentQueryData.HEADER_PREFIX + AgentQueryData.HEADER_USER_EMAIL))); + agentQueryData.setSessionIdentifier(getFirstHeader(requestHeaders, (AgentQueryData.HEADER_PREFIX + AgentQueryData.HEADER_SESSION_ID))); + agentQueryData.setConversationId(getFirstHeader(requestHeaders, (AgentQueryData.HEADER_PREFIX + AgentQueryData.HEADER_CONVERSATION_ID))); agentQueryData.setQueryPayload(httpResponseParam.getRequestParams().getPayload()); agentQueryData.setResponsePayload(httpResponseParam.getPayload()); - agentQueryData.setTimeStamp(httpResponseParam.getTimeOrNow()); + agentQueryData.setTimeStamp(Context.now()); return agentQueryData; } diff --git a/libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java b/libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java index 77f9cdf947..c758c99cb1 100644 --- a/libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java +++ b/libs/dao/src/main/java/com/akto/dto/agentic_sessions/AgentQueryData.java @@ -23,47 +23,17 @@ public class AgentQueryData { public static final String OUTPUT_TOKENS = "outputTokens"; // Request header name constants - public static final String HEADER_INSTALLER = "x-akto-installer"; + public static final String HEADER_PREFIX = "x-akto-installer-"; public static final String HEADER_USER_EMAIL = "user_email"; public static final String HEADER_SESSION_ID = "session_id"; public static final String HEADER_CONVERSATION_ID = "conversation_id"; public static final String HEADER_GENERATION_ID = "generation_id"; - public static final String HEADER_MODEL = "model"; - public static final String HEADER_TRANSCRIPT_PATH = "transcript_path"; - public static final String HEADER_CWD = "cwd"; - public static final String HEADER_PERMISSION_MODE = "permission_mode"; - public static final String HEADER_HOOK_EVENT_NAME = "hook_event_name"; - - // from x-akto-installer + private String serviceId; - - // from user_email - private String uniqueUserId; - - // from session_id + private String deviceId; + private String userName; private String sessionIdentifier; - - // from conversation_id private String conversationId; - - // from generation_id - private String generationId; - - // from model - private String model; - - // from transcript_path - private String transcriptPath; - - // from cwd - private String cwd; - - // from permission_mode - private String permissionMode; - - // from hook_event_name - private String hookEventName; - private String queryPayload; private String responsePayload; private long timeStamp; From 34fe368bd751de016b850ba7a593e8022c7b2eb8 Mon Sep 17 00:00:00 2001 From: Ark2307 Date: Tue, 14 Apr 2026 15:28:56 +0530 Subject: [PATCH 3/3] Opening api to send query data --- .../src/main/java/com/akto/action/DbAction.java | 14 ++++++++++++++ .../src/main/resources/struts.xml | 10 ++++++++++ 2 files changed, 24 insertions(+) diff --git a/apps/database-abstractor/src/main/java/com/akto/action/DbAction.java b/apps/database-abstractor/src/main/java/com/akto/action/DbAction.java index eabb4db483..c69bf9b19a 100644 --- a/apps/database-abstractor/src/main/java/com/akto/action/DbAction.java +++ b/apps/database-abstractor/src/main/java/com/akto/action/DbAction.java @@ -43,6 +43,7 @@ import com.akto.dto.graph.SvcToSvcGraphNode; import com.akto.dto.metrics.MetricData; import com.akto.dto.monitoring.ModuleInfo; +import com.akto.dto.agentic_sessions.AgentQueryData; import com.akto.dto.agentic_sessions.SessionDocument; import com.akto.dto.notifications.SlackWebhook; import com.akto.dto.runtime_filters.RuntimeFilter; @@ -190,6 +191,9 @@ public class DbAction extends ActionSupport { @Getter @Setter private List sessionDocuments; + @Getter @Setter + private AgentQueryData agentQueryData; + private static final LoggerMaker loggerMaker = new LoggerMaker(DbAction.class, LogDb.DB_ABS); public List getWritesForTestingRunIssues() { @@ -3635,6 +3639,16 @@ public String storeSpans() { return Action.SUCCESS.toUpperCase(); } + public String storeAgentQueryData() { + try { + DbLayer.storeAgentQueryData(agentQueryData); + } catch (Exception e) { + loggerMaker.errorAndAddToDb(e, "Error in storeAgentQueryData " + e.toString()); + return Action.ERROR.toUpperCase(); + } + return Action.SUCCESS.toUpperCase(); + } + private TestingRunWebhook testingRunWebhook; public TestingRunWebhook getTestingRunWebhook() { diff --git a/apps/database-abstractor/src/main/resources/struts.xml b/apps/database-abstractor/src/main/resources/struts.xml index 04868530a5..ea03ed82db 100644 --- a/apps/database-abstractor/src/main/resources/struts.xml +++ b/apps/database-abstractor/src/main/resources/struts.xml @@ -2378,6 +2378,16 @@ ^actionErrors.* + + + + + + 422 + false + ^actionErrors.* + +