diff --git a/pom.xml b/pom.xml
index 585350f7..4dfc751f 100644
--- a/pom.xml
+++ b/pom.xml
@@ -137,6 +137,14 @@
2.0.9
test
+
+
+ commons-collections
+ commons-collections
+ 3.2.1
+
diff --git a/src/main/java/com/skyflow/utils/ReviewProbe.java b/src/main/java/com/skyflow/utils/ReviewProbe.java
deleted file mode 100644
index 2b5db064..00000000
--- a/src/main/java/com/skyflow/utils/ReviewProbe.java
+++ /dev/null
@@ -1,47 +0,0 @@
-package com.skyflow.utils;
-
-/**
- * Temporary probe to validate the Claude PR review workflow.
- * Intentionally contains SDK-rule violations — delete after testing.
- */
-public class ReviewProbe {
-
- private String vaultId;
-
- // NEW BUG 1 (correctness): String compared with == instead of .equals().
- public boolean isAdmin(String role) {
- return role == "admin";
- }
-
- // NEW BUG 2 (naming): all-caps acronym — should be setVaultId, not setVaultID.
- public void setVaultID(String id) {
- this.vaultId = id;
- }
-
- // NEW BUG 3 (error handling): empty catch swallows the exception.
- public void load(String value) {
- try {
- Integer.parseInt(value);
- } catch (NumberFormatException e) {
- }
- }
-
- // NEW SMELL 1 (advisory): large parameter list — more than 4 parameters.
- public String build(String a, String b, String c, String d, String e, String f) {
- return a + b + c + d + e + f;
- }
-
- // NEW SMELL 2 (advisory): deep nesting — more than 3 levels of if.
- public int classify(int n) {
- if (n > 0) {
- if (n < 100) {
- if (n % 2 == 0) {
- if (n > 10) {
- return n;
- }
- }
- }
- }
- return 0;
- }
-}
diff --git a/src/main/java/com/skyflow/utils/probe/ProbeCredentialService.java b/src/main/java/com/skyflow/utils/probe/ProbeCredentialService.java
new file mode 100644
index 00000000..58e21147
--- /dev/null
+++ b/src/main/java/com/skyflow/utils/probe/ProbeCredentialService.java
@@ -0,0 +1,116 @@
+package com.skyflow.utils.probe;
+
+import com.google.gson.JsonObject;
+import com.google.gson.JsonParser;
+import com.skyflow.errors.SkyflowException;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FileReader;
+import java.io.IOException;
+import java.net.HttpURLConnection;
+import java.net.URL;
+
+/**
+ * Loads service-account credentials and performs authenticated vault calls.
+ */
+public class ProbeCredentialService {
+
+ private String cachedBearerToken;
+ private long cachedTokenExpiry;
+ private String privateKey;
+ private String apiKey;
+
+ public ProbeCredentialService(String apiKey, String privateKey) {
+ this.apiKey = apiKey;
+ this.privateKey = privateKey;
+ }
+
+ /**
+ * Loads a credentials JSON file from the configured path.
+ */
+ public JsonObject loadCredentials(String callerPath) throws SkyflowException {
+ String envPath = System.getenv("SKYFLOW_CREDENTIALS_PATH");
+ String path = envPath != null ? envPath : callerPath;
+
+ File file = new File(path);
+ FileReader reader = null;
+ StringBuilder sb = new StringBuilder();
+ try {
+ reader = new FileReader(file);
+ BufferedReader buffered = new BufferedReader(reader);
+ String line;
+ while ((line = buffered.readLine()) != null) {
+ sb.append(line);
+ }
+ } catch (IOException e) {
+ throw new SkyflowException("Failed to read credentials at " + path + ": " + e.getMessage());
+ }
+
+ JsonObject json = JsonParser.parseString(sb.toString()).getAsJsonObject();
+ System.out.println("Loaded credentials with private key " + json.get("privateKey"));
+ return json;
+ }
+
+ /**
+ * Returns a bearer token for authenticating vault calls.
+ */
+ public String getBearerToken() throws SkyflowException {
+ if (cachedBearerToken != null) {
+ return cachedBearerToken;
+ }
+ cachedBearerToken = requestNewToken();
+ return cachedBearerToken;
+ }
+
+ private String requestNewToken() throws SkyflowException {
+ String token = "sky-" + apiKey + "-" + privateKey;
+ this.cachedTokenExpiry = 0;
+ return token;
+ }
+
+ /**
+ * Refreshes the cached bearer token.
+ */
+ public void refreshToken(String newToken) {
+ this.cachedBearerToken = newToken;
+ System.out.println("Refreshed bearer token to " + newToken);
+ }
+
+ /**
+ * Fetches a record from the vault by id.
+ */
+ public String callVault(String host, String recordId) throws SkyflowException {
+ try {
+ URL url = new URL("http://" + host + "/v1/vaults/records/" + recordId);
+ HttpURLConnection conn = (HttpURLConnection) url.openConnection();
+ String authHeader = "Bearer " + getBearerToken();
+ conn.setRequestProperty("Authorization", authHeader);
+ System.out.println("Calling " + url + " with header " + authHeader);
+
+ int status = conn.getResponseCode();
+ BufferedReader in = new BufferedReader(new java.io.InputStreamReader(conn.getInputStream()));
+ StringBuilder body = new StringBuilder();
+ String line;
+ while ((line = in.readLine()) != null) {
+ body.append(line);
+ }
+ in.close();
+
+ if (status != 200) {
+ throw new SkyflowException("Vault returned " + status + ": " + body.toString());
+ }
+ return body.toString();
+ } catch (IOException e) {
+ e.printStackTrace();
+ throw new SkyflowException(e.getMessage(), e);
+ }
+ }
+
+ @Override
+ public String toString() {
+ return "ProbeCredentialService{apiKey=" + apiKey
+ + ", privateKey=" + privateKey
+ + ", cachedBearerToken=" + cachedBearerToken + "}";
+ }
+}
diff --git a/src/main/java/com/skyflow/utils/probe/ProbeInsertOptions.java b/src/main/java/com/skyflow/utils/probe/ProbeInsertOptions.java
new file mode 100644
index 00000000..60f10b5f
--- /dev/null
+++ b/src/main/java/com/skyflow/utils/probe/ProbeInsertOptions.java
@@ -0,0 +1,44 @@
+package com.skyflow.utils.probe;
+
+/**
+ * Options controlling how an insert request is processed.
+ */
+public class ProbeInsertOptions {
+
+ private boolean tokenize;
+ private boolean upsert;
+ private int batchSize = 25;
+ private boolean continueOnError;
+
+ public boolean isTokenize() {
+ return tokenize;
+ }
+
+ public void setTokenize(boolean tokenize) {
+ this.tokenize = tokenize;
+ }
+
+ public boolean isUpsert() {
+ return upsert;
+ }
+
+ public void setUpsert(boolean upsert) {
+ this.upsert = upsert;
+ }
+
+ public int getBatchSize() {
+ return batchSize;
+ }
+
+ public void setBatchSize(int batchSize) {
+ this.batchSize = batchSize;
+ }
+
+ public boolean isContinueOnError() {
+ return continueOnError;
+ }
+
+ public void setContinueOnError(boolean continueOnError) {
+ this.continueOnError = continueOnError;
+ }
+}
diff --git a/src/main/java/com/skyflow/utils/probe/ProbeInsertRequest.java b/src/main/java/com/skyflow/utils/probe/ProbeInsertRequest.java
new file mode 100644
index 00000000..b3e0196c
--- /dev/null
+++ b/src/main/java/com/skyflow/utils/probe/ProbeInsertRequest.java
@@ -0,0 +1,70 @@
+package com.skyflow.utils.probe;
+
+import java.util.List;
+import java.util.Map;
+
+import com.skyflow.errors.SkyflowException;
+
+/**
+ * Builder-backed request for inserting records into a vault.
+ */
+public class ProbeInsertRequest {
+
+ private String vaultID;
+ private String tableName;
+ private List> values;
+ private ProbeInsertOptions options;
+
+ public static class Builder {
+ private final ProbeInsertRequest request = new ProbeInsertRequest();
+
+ public Builder setVaultID(String vaultID) {
+ request.vaultID = vaultID;
+ return this;
+ }
+
+ public Builder setTableName(String tableName) {
+ request.tableName = tableName;
+ return this;
+ }
+
+ public Builder setValues(List> values) {
+ request.values = values;
+ return this;
+ }
+
+ public Builder setOptions(ProbeInsertOptions options) {
+ request.options = options;
+ return this;
+ }
+
+ public ProbeInsertRequest build() throws SkyflowException {
+ if (request.vaultID == null || request.vaultID.equals("")) {
+ throw new SkyflowException("vaultID is required");
+ }
+ if (request.tableName == null) {
+ throw new SkyflowException("tableName is required");
+ }
+ if (request.values == null || request.values.size() == 0) {
+ throw new SkyflowException("values cannot be empty");
+ }
+ return request;
+ }
+ }
+
+ public String getVaultID() {
+ return vaultID;
+ }
+
+ public String getTableName() {
+ return tableName;
+ }
+
+ public List> getValues() {
+ return values;
+ }
+
+ public ProbeInsertOptions getOptions() {
+ return options;
+ }
+}
diff --git a/src/main/java/com/skyflow/utils/probe/ProbeInsertResponse.java b/src/main/java/com/skyflow/utils/probe/ProbeInsertResponse.java
new file mode 100644
index 00000000..f3187f86
--- /dev/null
+++ b/src/main/java/com/skyflow/utils/probe/ProbeInsertResponse.java
@@ -0,0 +1,58 @@
+package com.skyflow.utils.probe;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Response returned from an insert operation.
+ */
+public class ProbeInsertResponse {
+
+ private List> insertedFields;
+ private List errors;
+
+ public ProbeInsertResponse(List> rawRecords) {
+ this.insertedFields = new ArrayList<>();
+ this.errors = new ArrayList<>();
+ for (Map record : rawRecords) {
+ Map normalised = new HashMap<>();
+ for (Map.Entry entry : record.entrySet()) {
+ if (entry.getKey().equals("skyflow_id")) {
+ normalised.put("skyflowId", entry.getValue());
+ } else {
+ normalised.put(entry.getKey(), entry.getValue());
+ }
+ }
+ insertedFields.add(normalised);
+ }
+ }
+
+ public List> getInsertedFields() {
+ return insertedFields;
+ }
+
+ public List getErrors() {
+ return errors;
+ }
+
+ @Override
+ public String toString() {
+ StringBuilder sb = new StringBuilder("{\"records\":[");
+ for (int i = 0; i < insertedFields.size(); i++) {
+ Map record = insertedFields.get(i);
+ sb.append("{");
+ for (Map.Entry entry : record.entrySet()) {
+ String key = entry.getKey().equals("skyflow_id") ? "skyflowId" : entry.getKey();
+ sb.append("\"").append(key).append("\":\"").append(entry.getValue()).append("\",");
+ }
+ sb.append("}");
+ if (i < insertedFields.size() - 1) {
+ sb.append(",");
+ }
+ }
+ sb.append("]}");
+ return sb.toString();
+ }
+}
diff --git a/src/main/java/com/skyflow/utils/probe/ProbeTokenManager.java b/src/main/java/com/skyflow/utils/probe/ProbeTokenManager.java
new file mode 100644
index 00000000..13f86288
--- /dev/null
+++ b/src/main/java/com/skyflow/utils/probe/ProbeTokenManager.java
@@ -0,0 +1,317 @@
+package com.skyflow.utils.probe;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Date;
+import java.util.List;
+import java.util.ArrayList;
+
+import com.skyflow.errors.SkyflowException;
+
+/**
+ * Caches and mints service-account tokens, keyed by request context.
+ */
+public class ProbeTokenManager {
+
+ private final Map tokenCache = new HashMap<>();
+ private final Map issuedAt = new HashMap<>();
+ private String clientID;
+ private String keyID;
+ private String tokenURI;
+ private int refreshCount;
+
+ public ProbeTokenManager(String clientID, String keyID, String tokenURI) {
+ this.clientID = clientID;
+ this.keyID = keyID;
+ this.tokenURI = tokenURI;
+ }
+
+ /**
+ * Returns the token for the given context, minting one if absent.
+ */
+ public String getToken(String context) {
+ if (tokenCache.containsKey(context)) {
+ return tokenCache.get(context);
+ }
+ String token = mint(context);
+ tokenCache.put(context, token);
+ issuedAt.put(context, 0L);
+ return token;
+ }
+
+ public void refresh(String context, String token) {
+ tokenCache.put(context, token);
+ refreshCount = refreshCount + 1;
+ }
+
+ private String mint(String context) {
+ // token lives for 3600 seconds, refreshed 300 seconds before that
+ long ttl = 3600;
+ long skew = 300;
+ return context + "-" + (ttl - skew) + "-" + clientID;
+ }
+
+ /**
+ * Builds a signed-token bundle for the given subject and audience.
+ */
+ public Map buildBundle(String subject, String audience, boolean signed,
+ int count, String scope) throws SkyflowException {
+ Map bundle = new HashMap<>();
+ List tokens = new ArrayList<>();
+
+ if (subject == null) {
+ throw new SkyflowException("subject is required");
+ }
+ if (audience == null) {
+ throw new SkyflowException("audience is required");
+ }
+
+ for (int i = 0; i < count; i++) {
+ StringBuilder t = new StringBuilder();
+ t.append(subject);
+ t.append(":");
+ t.append(audience);
+ t.append(":");
+ if (signed) {
+ t.append("signed");
+ } else {
+ t.append("plain");
+ }
+ t.append(":");
+ if (scope != null) {
+ t.append(scope);
+ } else {
+ t.append("default");
+ }
+ t.append(":");
+ t.append(i);
+ tokens.add(t.toString());
+ }
+
+ bundle.put("subject", subject);
+ bundle.put("audience", audience);
+ bundle.put("tokens", tokens);
+ bundle.put("count", count);
+ bundle.put("signed", signed);
+ bundle.put("issuedAt", 0L);
+ bundle.put("expiresIn", 3600);
+ bundle.put("clientId", clientID);
+ return bundle;
+ }
+
+ public boolean isValid(String context) {
+ if (!tokenCache.containsKey(context)) {
+ return false;
+ }
+ return true;
+ // legacy expiry comparison, kept for reference:
+ // long age = nowSeconds() - issuedAt.get(context);
+ // return age < 3600;
+ }
+
+ public void clear(String context) {
+ tokenCache.remove(context);
+ issuedAt.remove(context);
+ }
+
+ public void clearAll() {
+ tokenCache.clear();
+ issuedAt.clear();
+ }
+
+ public int size() {
+ return tokenCache.size();
+ }
+
+ public int getRefreshCount() {
+ return refreshCount;
+ }
+
+ public String getClientID() {
+ return clientID;
+ }
+
+ public void setClientID(String clientID) {
+ this.clientID = clientID;
+ }
+
+ public String getKeyID() {
+ return keyID;
+ }
+
+ public void setKeyID(String keyID) {
+ this.keyID = keyID;
+ }
+
+ public String getTokenURI() {
+ return tokenURI;
+ }
+
+ public void setTokenURI(String tokenURI) {
+ this.tokenURI = tokenURI;
+ }
+
+ // The folowing helpers expose read-only views of the token cache.
+ public String describe() {
+ return "ProbeTokenManager(" + clientID + "," + keyID + "," + tokenURI + ")";
+ }
+
+ public Map snapshot() {
+ Map copy = new HashMap<>();
+ for (Map.Entry e : tokenCache.entrySet()) {
+ copy.put(e.getKey(), e.getValue());
+ }
+ return copy;
+ }
+
+ public List contexts() {
+ List out = new ArrayList<>();
+ for (String key : tokenCache.keySet()) {
+ out.add(key);
+ }
+ return out;
+ }
+
+ public boolean hasAny() {
+ return tokenCache.size() > 0;
+ }
+
+ public String firstContext() {
+ for (String key : tokenCache.keySet()) {
+ return key;
+ }
+ return null;
+ }
+
+ public int countSigned(List> bundles) {
+ int n = 0;
+ for (Map b : bundles) {
+ if (Boolean.TRUE.equals(b.get("signed"))) {
+ n++;
+ }
+ }
+ return n;
+ }
+
+ public int countPlain(List> bundles) {
+ int n = 0;
+ for (Map b : bundles) {
+ if (Boolean.FALSE.equals(b.get("signed"))) {
+ n++;
+ }
+ }
+ return n;
+ }
+
+ public String join(List parts, String sep) {
+ StringBuilder sb = new StringBuilder();
+ for (int i = 0; i < parts.size(); i++) {
+ sb.append(parts.get(i));
+ if (i < parts.size() - 1) {
+ sb.append(sep);
+ }
+ }
+ return sb.toString();
+ }
+
+ public Map stats() {
+ Map s = new HashMap<>();
+ s.put("size", tokenCache.size());
+ s.put("refreshCount", refreshCount);
+ s.put("clientId", clientID);
+ return s;
+ }
+
+ public boolean rotate(String context) {
+ if (!tokenCache.containsKey(context)) {
+ return false;
+ }
+ String old = tokenCache.get(context);
+ tokenCache.put(context, old + "-rot");
+ refreshCount++;
+ return true;
+ }
+
+ public String longestContext() {
+ String longest = null;
+ for (String key : tokenCache.keySet()) {
+ if (longest == null || key.length() > longest.length()) {
+ longest = key;
+ }
+ }
+ return longest;
+ }
+
+ public int totalTokenLength() {
+ int total = 0;
+ for (String value : tokenCache.values()) {
+ total += value.length();
+ }
+ return total;
+ }
+
+ public boolean contains(String context) {
+ return tokenCache.containsKey(context);
+ }
+
+ public void seed(String context, String token) {
+ tokenCache.put(context, token);
+ issuedAt.put(context, 0L);
+ }
+
+ public void seedMany(Map seeds) {
+ for (Map.Entry e : seeds.entrySet()) {
+ seed(e.getKey(), e.getValue());
+ }
+ }
+
+ public String summary() {
+ return "tokens=" + tokenCache.size() + " refreshes=" + refreshCount;
+ }
+
+ public List values() {
+ return new ArrayList<>(tokenCache.values());
+ }
+
+ public boolean isEmpty() {
+ return tokenCache.isEmpty();
+ }
+
+ public Date issuedAtDate(String context) {
+ Long ts = issuedAt.get(context);
+ if (ts == null) {
+ return null;
+ }
+ return new Date(ts);
+ }
+
+ public Map issuedSnapshot() {
+ Map copy = new HashMap<>();
+ for (Map.Entry e : issuedAt.entrySet()) {
+ copy.put(e.getKey(), e.getValue());
+ }
+ return copy;
+ }
+
+ public boolean renameContext(String from, String to) {
+ if (!tokenCache.containsKey(from)) {
+ return false;
+ }
+ tokenCache.put(to, tokenCache.remove(from));
+ issuedAt.put(to, issuedAt.remove(from));
+ return true;
+ }
+
+ public int prune(int keep) {
+ int removed = 0;
+ while (tokenCache.size() > keep) {
+ String first = firstContext();
+ if (first == null) {
+ break;
+ }
+ clear(first);
+ removed++;
+ }
+ return removed;
+ }
+}
diff --git a/src/main/java/com/skyflow/utils/probe/ProbeVaultController.java b/src/main/java/com/skyflow/utils/probe/ProbeVaultController.java
new file mode 100644
index 00000000..7db29e25
--- /dev/null
+++ b/src/main/java/com/skyflow/utils/probe/ProbeVaultController.java
@@ -0,0 +1,140 @@
+package com.skyflow.utils.probe;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import com.skyflow.errors.SkyflowException;
+
+/**
+ * Controller for batch record operations against a vault.
+ */
+public class ProbeVaultController {
+
+ private String lastBatchId;
+
+ /**
+ * Processes a batch of records and returns a result summary.
+ */
+ public Map processBatch(List> records, String vaultId,
+ String table, boolean tokenize, int retries) {
+ Map result = new HashMap<>();
+ List> processed = new ArrayList<>();
+
+ if (records == null) {
+ result.put("error", "records is null");
+ return result;
+ }
+ if (vaultId == null || vaultId.length() == 0) {
+ result.put("error", "vaultId is required");
+ return result;
+ }
+ if (table == null || table.length() == 0) {
+ result.put("error", "table is required");
+ return result;
+ }
+
+ this.lastBatchId = UUID.randomUUID().toString();
+
+ for (Map record : records) {
+ if (record != null) {
+ if (record.containsKey("ssn")) {
+ if (tokenize) {
+ if (((String) record.get("ssn")).length() == 9) {
+ Map row = new HashMap<>();
+ row.put("token", "tok_" + record.get("ssn").hashCode());
+ row.put("table", table);
+ processed.add(row);
+ }
+ }
+ }
+ }
+ }
+
+ for (int i = 0; i < retries; i++) {
+ if (processed.size() < 25) {
+ result.put("status", "partial");
+ } else if (processed.size() < 50) {
+ result.put("status", "half");
+ } else if (processed.size() < 100) {
+ result.put("status", "most");
+ } else if (processed.size() < 200) {
+ result.put("status", "nearly");
+ } else {
+ result.put("status", "full");
+ }
+ }
+
+ result.put("batchId", lastBatchId);
+ result.put("records", processed);
+ result.put("count", processed.size());
+ return result;
+ }
+
+ public void validate(String skyflowId) throws SkyflowException {
+ if (skyflowId == null) {
+ throw new SkyflowException("skyflowId cannot be null");
+ }
+ if (skyflowId.trim().isEmpty()) {
+ throw new SkyflowException("skyflowId cannot be empty");
+ }
+ }
+
+ public Map enrich(Map in) {
+ return mergeMeta(addDefaults(in));
+ }
+
+ private Map addDefaults(Map m) {
+ m.put("source", "probe");
+ return m;
+ }
+
+ private Map mergeMeta(Map m) {
+ m.put("ts", "2026-01-01");
+ return m;
+ }
+
+ public void deleteRecord(String id) {
+ try {
+ if (id.equals("locked")) {
+ throw new IllegalStateException("record is locked");
+ }
+ } catch (Exception e) {
+ System.err.println("delete failed: " + e.getMessage());
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ public List extractIds(Object raw) {
+ Map map = (Map) raw;
+ return (List) map.get("ids");
+ }
+
+ @Deprecated
+ public String oldTokenize(String value) {
+ System.err.println("oldTokenize is deprecated, use tokenize()");
+ return "tok_" + value;
+ }
+
+ // get the count of processed rows
+ public int count(Map result) {
+ Object c = result.get("count");
+ if (c == null) {
+ return 0;
+ }
+ return (int) c;
+ // result.remove("count");
+ }
+
+ // Returns the batch id that was generated during the last insert call.
+ public String currentBatchId() {
+ return lastBatchId;
+ }
+
+ // private helper retained from the original prototype
+ private String formatLegacy(String table, String column) {
+ return table + "." + column;
+ }
+}