diff --git a/axhub-gateway/build.gradle b/axhub-gateway/build.gradle index f1534de..c725784 100644 --- a/axhub-gateway/build.gradle +++ b/axhub-gateway/build.gradle @@ -24,6 +24,12 @@ dependencies { // implementation 'com.networknt:json-schema-validator:1.4.0' // Spring AI 내장 버전과 충돌 방지를 위해 주석 처리 implementation 'org.springdoc:springdoc-openapi-starter-webmvc-ui:2.5.0' + + testImplementation 'org.springframework.boot:spring-boot-starter-test' +} + +tasks.named('test') { + useJUnitPlatform() } dependencies { diff --git a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/config/AgentResponseBudgetProperties.java b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/config/AgentResponseBudgetProperties.java new file mode 100644 index 0000000..e7c24ae --- /dev/null +++ b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/config/AgentResponseBudgetProperties.java @@ -0,0 +1,26 @@ +package io.shinhanlife.axhub.biz.mcp.gateway.config; + +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.stereotype.Component; + +@Component +@ConfigurationProperties(prefix = "mcp.agent-response") +public class AgentResponseBudgetProperties { + private Integer maxPreviewItems; + private Integer maxFieldsPerItem; + private Long maxItemBytes; + private Long maxTotalBytes; + private Boolean includeSummary; + + public void setMaxPreviewItems(Integer maxPreviewItems) { this.maxPreviewItems = maxPreviewItems; } + public void setMaxFieldsPerItem(Integer maxFieldsPerItem) { this.maxFieldsPerItem = maxFieldsPerItem; } + public void setMaxItemBytes(Long maxItemBytes) { this.maxItemBytes = maxItemBytes; } + public void setMaxTotalBytes(Long maxTotalBytes) { this.maxTotalBytes = maxTotalBytes; } + public void setIncludeSummary(Boolean includeSummary) { this.includeSummary = includeSummary; } + + public int maxPreviewItems() { return maxPreviewItems == null || maxPreviewItems < 1 ? 20 : maxPreviewItems; } + public int maxFieldsPerItem() { return maxFieldsPerItem == null || maxFieldsPerItem < 1 ? 10 : maxFieldsPerItem; } + public long maxItemBytes() { return maxItemBytes == null || maxItemBytes < 1 ? 4_096 : maxItemBytes; } + public long maxTotalBytes() { return maxTotalBytes == null || maxTotalBytes < 1 ? 65_536 : maxTotalBytes; } + public boolean includeSummary() { return includeSummary == null || includeSummary; } +} diff --git a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/dto/ToolMetadata.java b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/dto/ToolMetadata.java index f8ec565..cf939c4 100644 --- a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/dto/ToolMetadata.java +++ b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/dto/ToolMetadata.java @@ -9,7 +9,7 @@ import lombok.Data; import lombok.NoArgsConstructor; import java.util.Map; -1import java.util.Set; +import java.util.Set; import java.util.List; import java.util.HashSet; /** diff --git a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/service/ExecuteService.java b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/service/ExecuteService.java index 091f10e..cf914a8 100644 --- a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/service/ExecuteService.java +++ b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/service/ExecuteService.java @@ -24,8 +24,11 @@ import lombok.extern.slf4j.Slf4j; import org.springframework.http.MediaType; import org.springframework.stereotype.Service; import org.springframework.web.client.RestClient; +import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; +import io.shinhanlife.axhub.biz.mcp.gateway.tool.large.LargeToolResponseService; +import io.shinhanlife.axhub.biz.mcp.gateway.tool.large.PaginationRequestValidator; @Slf4j @Service @@ -42,6 +45,8 @@ public class ExecuteService { private final ToolAuthorizationService authorizationService; private final RedisToolTraceService redisTrace; private final McpGatewayProperties properties; + private final LargeToolResponseService largeResponses; + private final PaginationRequestValidator paginationValidator; private final RestClient restClient = RestClient.create(); private final ExecutorService executor; @@ -55,7 +60,9 @@ public class ExecuteService { CircuitBreakerService circuitBreakerService, ToolAuthorizationService authorizationService, RedisToolTraceService redisTrace, - McpGatewayProperties properties) { + McpGatewayProperties properties, + LargeToolResponseService largeResponses, + PaginationRequestValidator paginationValidator) { this.planner = planner; this.killSwitchService = killSwitchService; this.objectMapper = objectMapper; @@ -67,6 +74,8 @@ public class ExecuteService { this.authorizationService = authorizationService; this.redisTrace = redisTrace; this.properties = properties; + this.largeResponses = largeResponses; + this.paginationValidator = paginationValidator; this.executor = new ThreadPoolExecutor( properties.toolExecutorCorePoolSize(), @@ -145,7 +154,7 @@ public class ExecuteService { try { redisTrace.attemptStarted(context, metadata, arguments, attempt, retryPolicy.maxAttempts()); circuitBreaker.beforeCall(); - Object result = executeOnce(metadata, payload); + Object result = executeOnce(context, metadata, arguments, payload); circuitBreaker.recordSuccess(); return result; } catch (ToolExecutionException error) { @@ -165,7 +174,7 @@ public class ExecuteService { : lastError; } - private Object executeOnce(ToolMetadata metadata, Map payload) { + private Object executeOnce(McpRequestContext context, ToolMetadata metadata, ObjectNode arguments, Map payload) { CompletableFuture future; try { @@ -177,23 +186,48 @@ public class ExecuteService { } String executeApiUrl = targetUrl + "/mcp/api/v1/tools/call"; - try { - log.info(" [ExecuteService] 요청 페이로드(마스킹 적용): {}", objectMapper.writeValueAsString(dataMasker.mask(objectMapper.valueToTree(payload)))); - } catch (Exception ignore) {} - - try { - return executeWithUrl(payload, executeApiUrl); - } catch (Exception e) { - if (executeApiUrl.contains("http://tool-")) { - String fallbackUrl = executeApiUrl.replaceAll("http://tool-[a-zA-Z0-9-]+", "http://localhost"); - log.warn(" [ExecuteService] 호스트를 찾을 수 없어 localhost로 재시도합니다: {}", fallbackUrl); - try { - return executeWithUrl(payload, fallbackUrl); - } catch (Exception ex) { - throw new ToolExecutionException(FailureType.SERVER_ERROR, "Tool Pod 호출 실패 (localhost 재시도 포함): " + ex.getMessage()); + ObjectNode pageArguments = paginationValidator.normalize(arguments); + LargeToolResponseService.Collector collector = largeResponses.newCollector(metadata.getName(), context.requestId()); + + while (true) { + Map pagePayload = new java.util.HashMap<>(payload); + if (pagePayload.containsKey("params")) { + Map params = new java.util.HashMap<>((Map) pagePayload.get("params")); + params.put("arguments", objectMapper.convertValue(pageArguments, Map.class)); + pagePayload.put("params", params); + } else { + pagePayload.put("arguments", objectMapper.convertValue(pageArguments, Map.class)); + } + + try { + log.info(" [ExecuteService] 요청 페이로드(마스킹 적용): {}", objectMapper.writeValueAsString(dataMasker.mask(objectMapper.valueToTree(pagePayload)))); + } catch (Exception ignore) {} + + JsonNode data = null; + try { + Object httpResult = executeWithUrl(pagePayload, executeApiUrl); + data = extractData(objectMapper.valueToTree(httpResult)); + } catch (Exception e) { + if (executeApiUrl.contains("http://tool-")) { + String fallbackUrl = executeApiUrl.replaceAll("http://tool-[a-zA-Z0-9-]+", "http://localhost"); + log.warn(" [ExecuteService] 호스트를 찾을 수 없어 localhost로 재시도합니다: {}", fallbackUrl); + try { + Object httpResult = executeWithUrl(pagePayload, fallbackUrl); + data = extractData(objectMapper.valueToTree(httpResult)); + } catch (Exception ex) { + throw new ToolExecutionException(FailureType.SERVER_ERROR, "Tool Pod 호출 실패 (localhost 재시도 포함): " + ex.getMessage()); + } + } else { + throw new ToolExecutionException(FailureType.SERVER_ERROR, "Tool Pod 호출 실패: " + e.getMessage()); } } - throw new ToolExecutionException(FailureType.SERVER_ERROR, "Tool Pod 호출 실패: " + e.getMessage()); + + collector.accept(data); + if (!collector.shouldFetchNextPage()) { + return collector.finish(); + } + pageArguments.put("cursor", collector.nextCursor()); + pageArguments.put("pageSize", collector.pageSize()); } } catch (ToolExecutionException error) { throw error; @@ -240,6 +274,13 @@ public class ExecuteService { .body(Object.class); } + private JsonNode extractData(JsonNode root) { + if (!root.path("success").asBoolean(true)) { + throw new ToolExecutionException(FailureType.BUSINESS_ERROR, "Tool 서버 업무 오류: " + root.path("error").asText()); + } + return root.has("data") ? root.get("data") : root; + } + private RetryPolicy retryPolicy(ToolMetadata metadata, ObjectNode arguments) { if (!retryAllowedByOperation(metadata, arguments)) { return RetryPolicy.disabled(); diff --git a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/AgentResponseBudgetService.java b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/AgentResponseBudgetService.java new file mode 100644 index 0000000..acf4112 --- /dev/null +++ b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/AgentResponseBudgetService.java @@ -0,0 +1,229 @@ +package io.shinhanlife.axhub.biz.mcp.gateway.tool.large; + +import io.shinhanlife.axhub.biz.mcp.gateway.config.AgentResponseBudgetProperties; +import io.shinhanlife.axhub.biz.mcp.gateway.guardrail.SensitiveDataMasker; +import org.springframework.stereotype.Service; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; + +@Service +public class AgentResponseBudgetService { + private final AgentResponseBudgetProperties properties; + private final ObjectMapper json; + private final SensitiveDataMasker masker; + + public AgentResponseBudgetService(AgentResponseBudgetProperties properties, ObjectMapper json, SensitiveDataMasker masker) { + this.properties = properties; + this.json = json; + this.masker = masker; + } + + public ObjectNode apply(ObjectNode response) { + ArrayNode originalItems = response.path("previewItems").isArray() + ? (ArrayNode) response.path("previewItems") + : json.createArrayNode(); + int originalCount = originalItems.size(); + BudgetStats stats = new BudgetStats(); + ArrayNode budgetItems = json.createArrayNode(); + + for (JsonNode item : originalItems) { + if (budgetItems.size() >= properties.maxPreviewItems()) { + stats.omittedItems++; + stats.limited = true; + continue; + } + JsonNode budgetItem = budgetItem(item, stats); + if (jsonBytes(budgetItem) > properties.maxItemBytes()) { + budgetItem = fallbackTruncatedItem(budgetItem); + stats.truncatedItems++; + stats.limited = true; + } + budgetItems.add(budgetItem); + } + + ObjectNode budgeted = copyWithout(response, "previewItems", "budget", "summary"); + budgeted.set("previewItems", budgetItems); + if (properties.includeSummary()) { + budgeted.put("summary", summary(budgeted, originalCount, budgetItems.size())); + } + + while (jsonBytes(budgeted) > properties.maxTotalBytes() && budgetItems.size() > 0) { + budgetItems.remove(budgetItems.size() - 1); + stats.omittedItems++; + stats.limited = true; + budgeted.set("previewItems", budgetItems); + if (properties.includeSummary()) { + budgeted.put("summary", summary(budgeted, originalCount, budgetItems.size())); + } + } + + ObjectNode budget = json.createObjectNode(); + updateBudget(budget, stats, originalCount, budgetItems.size()); + budgeted.put("previewCount", budgetItems.size()); + budgeted.set("budget", budget); + while (jsonBytes(budgeted) > properties.maxTotalBytes() && budgetItems.size() > 0) { + budgetItems.remove(budgetItems.size() - 1); + stats.omittedItems++; + stats.limited = true; + budgeted.set("previewItems", budgetItems); + budgeted.put("previewCount", budgetItems.size()); + if (properties.includeSummary()) { + budgeted.put("summary", summary(budgeted, originalCount, budgetItems.size())); + } + updateBudget(budget, stats, originalCount, budgetItems.size()); + budgeted.set("budget", budget); + } + return budgeted; + } + + private void updateBudget(ObjectNode budget, BudgetStats stats, int originalCount, int previewCount) { + budget.put("limited", stats.limited || stats.omittedItems > 0 || stats.truncatedItems > 0); + budget.put("maxPreviewItems", properties.maxPreviewItems()); + budget.put("maxFieldsPerItem", properties.maxFieldsPerItem()); + budget.put("maxItemBytes", properties.maxItemBytes()); + budget.put("maxTotalBytes", properties.maxTotalBytes()); + budget.put("originalPreviewCount", originalCount); + budget.put("previewCount", previewCount); + budget.put("omittedItems", stats.omittedItems); + budget.put("truncatedItems", stats.truncatedItems); + } + + private JsonNode budgetItem(JsonNode item, BudgetStats stats) { + JsonNode masked = masker.mask(item); + if (!masked.isObject()) { + return truncateByBytes(masked, stats); + } + ObjectNode limited = json.createObjectNode(); + int copied = 0; + int originalFields = 0; + + java.util.Iterator> fields = masked.fields(); + while (fields.hasNext()) { + java.util.Map.Entry entry = fields.next(); + originalFields++; + if (copied >= properties.maxFieldsPerItem()) { + continue; + } + limited.set(entry.getKey(), entry.getValue()); + copied++; + } + + if (originalFields > copied) { + limited.put("truncated", true); + limited.put("omittedFieldCount", originalFields - copied); + stats.truncatedItems++; + stats.limited = true; + } + return truncateByBytes(limited, stats); + } + + private JsonNode truncateByBytes(JsonNode item, BudgetStats stats) { + if (jsonBytes(item) <= properties.maxItemBytes()) { + return item; + } + if (!item.isObject()) { + stats.truncatedItems++; + stats.limited = true; + ObjectNode truncated = json.createObjectNode(); + truncated.put("truncated", true); + truncated.put("originalBytes", jsonBytes(item)); + truncated.put("maxItemBytes", properties.maxItemBytes()); + return truncated; + } + ObjectNode source = (ObjectNode) item; + ObjectNode limited = json.createObjectNode(); + + java.util.Iterator> fields = source.fields(); + while (fields.hasNext()) { + java.util.Map.Entry entry = fields.next(); + JsonNode value = entry.getValue(); + if (value.isTextual()) { + limited.put(entry.getKey(), truncateText(value.asText())); + } else { + limited.set(entry.getKey(), value); + } + } + + limited.put("truncated", true); + limited.put("originalBytes", jsonBytes(item)); + limited.put("maxItemBytes", properties.maxItemBytes()); + stats.truncatedItems++; + stats.limited = true; + if (jsonBytes(limited) <= properties.maxItemBytes()) { + return limited; + } + return fallbackTruncatedItem(item); + } + + private ObjectNode fallbackTruncatedItem(JsonNode item) { + ObjectNode fallback = json.createObjectNode(); + fallback.put("truncated", true); + fallback.put("originalBytes", jsonBytes(item)); + fallback.put("maxItemBytes", properties.maxItemBytes()); + if (item.isObject()) { + ArrayNode fieldNames = json.createArrayNode(); + java.util.Iterator fieldNamesIter = item.fieldNames(); + while (fieldNamesIter.hasNext()) { + fieldNames.add(fieldNamesIter.next()); + } + fallback.set("fieldNames", fieldNames); + } + return fallback; + } + + private String truncateText(String value) { + if (value == null) { + return ""; + } + int maxChars = Math.min(value.length(), 256); + String truncated = value.substring(0, maxChars); + while (truncated.length() > 16 && truncated.getBytes(java.nio.charset.StandardCharsets.UTF_8).length > properties.maxItemBytes() / 2) { + truncated = truncated.substring(0, truncated.length() / 2); + } + return truncated + "..."; + } + + private ObjectNode copyWithout(ObjectNode source, String... excludedFields) { + ObjectNode copy = json.createObjectNode(); + + java.util.Iterator> fields = source.fields(); + while (fields.hasNext()) { + java.util.Map.Entry entry = fields.next(); + if (!excluded(entry.getKey(), excludedFields)) { + copy.set(entry.getKey(), entry.getValue()); + } + } + + return copy; + } + + private boolean excluded(String field, String[] excludedFields) { + for (String excluded : excludedFields) { + if (excluded.equals(field)) { + return true; + } + } + return false; + } + + private String summary(ObjectNode response, int originalCount, int previewCount) { + long totalCount = response.path("totalCount").asLong(originalCount); + return "Returned " + previewCount + " preview item(s) out of " + totalCount + " total result(s)."; + } + + private long jsonBytes(JsonNode node) { + try { + return json.writeValueAsBytes(node).length; + } catch (Exception error) { + return Long.MAX_VALUE; + } + } + + private static final class BudgetStats { + private boolean limited; + private int omittedItems; + private int truncatedItems; + } +} diff --git a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/LargeToolResponseService.java b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/LargeToolResponseService.java new file mode 100644 index 0000000..4651cf5 --- /dev/null +++ b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/LargeToolResponseService.java @@ -0,0 +1,347 @@ +package io.shinhanlife.axhub.biz.mcp.gateway.tool.large; + +import io.shinhanlife.axhub.biz.mcp.gateway.config.McpGatewayProperties; +import io.shinhanlife.axhub.biz.mcp.gateway.guardrail.SensitiveDataMasker; +import io.shinhanlife.axhub.biz.mcp.gateway.resilience.FailureType; +import io.shinhanlife.axhub.biz.mcp.gateway.resilience.ToolExecutionException; +import org.springframework.stereotype.Service; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; + +import java.time.Duration; +import java.time.Instant; + +@Service +public class LargeToolResponseService { + private final McpGatewayProperties properties; + private final ObjectMapper json; + private final SensitiveDataMasker masker; + private final AgentResponseBudgetService agentBudget; + + public LargeToolResponseService(McpGatewayProperties properties, + ObjectMapper json, + SensitiveDataMasker masker, + AgentResponseBudgetService agentBudget) { + this.properties = properties; + this.json = json; + this.masker = masker; + this.agentBudget = agentBudget; + } + + public Collector newCollector(String toolName, String requestId) { + return new Collector(toolName, requestId); + } + + public ObjectNode tooLargeResult(String toolName, String requestId, long receivedBytes, String nextCursor) { + ObjectNode response = basePartialResponse(toolName, requestId); + response.put("failureType", FailureType.TOO_LARGE_RESULT.name()); + response.put("errorCode", "TOOL_RESULT_TOO_LARGE"); + response.put("message", "Tool 응답이 허용 크기를 초과했습니다. 원문 body는 저장하지 않았고 cursor/resultRef 기반 재조회가 필요합니다."); + response.put("receivedBytes", receivedBytes); + response.put("maxResponseBytesFromTool", properties.maxResponseBytesFromTool()); + response.put("hasMore", true); + response.put("nextCursor", safe(nextCursor)); + response.put("resultRef", resultRef(toolName, requestId, nextCursor)); + response.put("truncated", true); + response.put("omittedCount", 0); + response.set("previewItems", json.createArrayNode()); + response.set("streamEvents", streamEvents(toolName, 0, 0, true, nextCursor)); + attachPageInfoAndStructuredContent(response); + return agentBudget.apply(response); + } + + private ObjectNode basePartialResponse(String toolName, String requestId) { + ObjectNode response = json.createObjectNode(); + response.put("success", true); + response.put("toolName", toolName); + response.put("requestId", requestId); + response.put("status", "PARTIAL_RESULT"); + response.put("pageSize", pageSize()); + response.put("pageCount", 0); + response.put("returnedCount", 0); + response.put("totalCount", 0); + response.put("hasMore", false); + response.put("nextCursor", ""); + response.put("truncated", false); + response.put("omittedCount", 0); + response.put("maxStreamBytes", properties.maxStreamBytes()); + response.put("maxItemBytes", properties.largeResponseMaxItemBytes()); + return response; + } + + private ArrayNode streamEvents(String toolName, int pageCount, int returnedCount, boolean hasMore, String nextCursor) { + ArrayNode events = json.createArrayNode(); + ObjectNode start = json.createObjectNode(); + start.put("type", "start"); + start.put("tool", toolName); + events.add(start); + + ObjectNode summary = json.createObjectNode(); + summary.put("type", "summary"); + summary.put("pageCount", pageCount); + summary.put("returnedCount", returnedCount); + summary.put("message", returnedCount + "건의 preview를 생성했습니다."); + events.add(summary); + + ObjectNode done = json.createObjectNode(); + done.put("type", "done"); + done.put("hasMore", hasMore); + done.put("nextCursor", safe(nextCursor)); + events.add(done); + return events; + } + + private String resultRef(String toolName, String requestId, String nextCursor) { + if (nextCursor == null || nextCursor.isBlank()) { + return "tool-result:" + toolName + ":" + requestId; + } + return "tool-result:" + toolName + ":" + requestId + ":cursor:" + nextCursor; + } + + private int pageSize() { + return Math.min(properties.defaultPageSize(), properties.maxPageSize()); + } + + private String safe(String value) { + return value == null ? "" : value; + } + + public class Collector { + private final String toolName; + private final String requestId; + private final Instant startedAt = Instant.now(); + private final ArrayNode previewItems = json.createArrayNode(); + private final ArrayNode streamEvents = json.createArrayNode(); + private int pageCount; + private int returnedCount; + private long totalCount; + private boolean paginated; + private boolean hasMore; + private String nextCursor = ""; + private String stopReason = ""; + private JsonNode normalData; + private boolean truncated; + private int omittedCount; + + private Collector(String toolName, String requestId) { + this.toolName = toolName; + this.requestId = requestId; + ObjectNode start = json.createObjectNode(); + start.put("type", "start"); + start.put("tool", toolName); + streamEvents.add(start); + } + + public void accept(JsonNode data) { + if (data == null || data.isMissingNode() || data.isNull()) { + throw new ToolExecutionException(FailureType.BUSINESS_ERROR, "Tool response data is empty. tool=" + toolName); + } + Page page = pageFrom(data); + if (!page.paginated() && pageCount == 0 && count(page.items()) <= pageSize()) { + JsonNode masked = masker.mask(data); + normalData = masked; + pageCount = 1; + returnedCount = count(page.items()); + totalCount = returnedCount; + return; + } + + paginated = true; + pageCount++; + totalCount = page.totalCount() >= 0 ? page.totalCount() : totalCount; + hasMore = page.hasMore(); + nextCursor = page.nextCursor(); + + int pageItemCount = 0; + for (JsonNode item : page.items()) { + pageItemCount++; + JsonNode preview = previewItem(item); + if (previewItems.size() < maxPreviewItems() && canAddPreviewItem(preview)) { + previewItems.add(preview); + returnedCount++; + } else { + omittedCount++; + truncated = true; + } + } + + ObjectNode summary = json.createObjectNode(); + summary.put("type", "summary"); + summary.put("page", pageCount); + summary.put("count", pageItemCount); + summary.put("returnedCount", returnedCount); + summary.put("omittedCount", omittedCount); + summary.put("nextCursor", nextCursor()); + summary.put("message", returnedCount + " items previewed"); + addStreamEvent(summary); + + if (pageCount >= properties.maxPagesPerCall() && hasMore) { + stopReason = "MAX_PAGES_PER_CALL"; + truncated = true; + } + if (Duration.between(startedAt, Instant.now()).toSeconds() >= properties.maxDurationSeconds() && hasMore) { + stopReason = "MAX_DURATION_SECONDS"; + truncated = true; + } + } + + public boolean shouldFetchNextPage() { + return paginated + && hasMore + && nextCursor != null + && !nextCursor.isBlank() + && pageCount < properties.maxPagesPerCall() + && Duration.between(startedAt, Instant.now()).toSeconds() < properties.maxDurationSeconds() + && currentResponseBytes() < properties.maxStreamBytes(); + } + + public String nextCursor() { + return safe(nextCursor); + } + + public int pageSize() { + return LargeToolResponseService.this.pageSize(); + } + + public ObjectNode finish() { + if (normalData != null) { + ObjectNode wrap = json.createObjectNode(); + wrap.put("success", true); + wrap.put("toolName", toolName); + wrap.put("requestId", requestId); + wrap.set("data", normalData); + return wrap; + } + + ObjectNode done = json.createObjectNode(); + done.put("type", "done"); + done.put("totalCount", totalCount); + done.put("hasMore", hasMore); + done.put("nextCursor", nextCursor()); + addStreamEvent(done); + + ObjectNode response = basePartialResponse(toolName, requestId); + response.put("pageCount", pageCount); + response.put("returnedCount", returnedCount); + response.put("totalCount", totalCount); + response.put("hasMore", hasMore); + response.put("nextCursor", nextCursor()); + response.put("resultRef", resultRef(toolName, requestId, nextCursor())); + response.put("truncated", truncated); + response.put("omittedCount", omittedCount); + response.put("message", stopReason.isBlank() + ? "대량 결과라 일부 preview만 반환했습니다." + : "대량 결과라 " + stopReason + " 조건에서 중단하고 일부 preview만 반환했습니다."); + response.set("previewItems", previewItems); + response.set("streamEvents", streamEvents); + attachPageInfoAndStructuredContent(response); + return agentBudget.apply(response); + } + + private Page pageFrom(JsonNode data) { + if (data.isObject() && data.path("items").isArray()) { + return new Page(true, data.path("items"), data.path("nextCursor").asText(""), + data.path("hasMore").asBoolean(false), data.path("totalCount").asLong(-1)); + } + if (data.isArray()) { + return new Page(false, data, "", false, count(data)); + } + ArrayNode one = json.createArrayNode(); + one.add(data); + return new Page(false, one, "", false, 1); + } + + private JsonNode previewItem(JsonNode item) { + JsonNode masked = masker.mask(item); + long itemBytes = jsonBytes(masked); + if (itemBytes <= properties.largeResponseMaxItemBytes()) { + return masked; + } + truncated = true; + ObjectNode preview = json.createObjectNode(); + preview.put("truncated", true); + preview.put("originalBytes", itemBytes); + preview.put("maxItemBytes", properties.largeResponseMaxItemBytes()); + if (masked.isObject()) { + ArrayNode fieldNames = json.createArrayNode(); + java.util.Iterator fieldNamesIter = masked.fieldNames(); + while (fieldNamesIter.hasNext()) { + fieldNames.add(fieldNamesIter.next()); + } + preview.set("fieldNames", fieldNames); + } + return preview; + } + + private boolean canAddPreviewItem(JsonNode item) { + return jsonBytes(item) <= properties.largeResponseMaxItemBytes() + && currentResponseBytes() + jsonBytes(item) <= properties.maxStreamBytes(); + } + + private void addStreamEvent(ObjectNode event) { + if (currentResponseBytes() + jsonBytes(event) <= properties.maxStreamBytes()) { + streamEvents.add(event); + return; + } + truncated = true; + } + + private int maxPreviewItems() { + return Math.min(pageSize() * properties.maxPagesPerCall(), properties.largeResponseMaxPreviewItems()); + } + + private long currentResponseBytes() { + ObjectNode estimate = json.createObjectNode(); + estimate.set("previewItems", previewItems); + estimate.set("streamEvents", streamEvents); + return jsonBytes(estimate); + } + + private long jsonBytes(JsonNode node) { + try { + return json.writeValueAsBytes(node).length; + } catch (Exception error) { + return Long.MAX_VALUE; + } + } + + private int count(JsonNode node) { + int count = 0; + java.util.Iterator iter = node.elements(); + while(iter.hasNext()) { + iter.next(); + count++; + } + return count; + } + } + + private record Page(boolean paginated, JsonNode items, String nextCursor, boolean hasMore, long totalCount) {} + + private void attachPageInfoAndStructuredContent(ObjectNode response) { + ObjectNode pageInfo = json.createObjectNode(); + pageInfo.put("pageSize", response.path("pageSize").asInt(pageSize())); + pageInfo.put("pageCount", response.path("pageCount").asInt(0)); + pageInfo.put("returnedCount", response.path("returnedCount").asInt(0)); + pageInfo.put("totalCount", response.path("totalCount").asLong(0)); + pageInfo.put("hasMore", response.path("hasMore").asBoolean(false)); + pageInfo.put("nextCursor", response.path("nextCursor").asText("")); + response.set("pageInfo", pageInfo); + + ObjectNode summary = json.createObjectNode(); + summary.put("status", response.path("status").asText("PARTIAL_RESULT")); + summary.put("toolName", response.path("toolName").asText("")); + summary.put("message", response.path("message").asText("")); + summary.put("truncated", response.path("truncated").asBoolean(false)); + summary.put("omittedCount", response.path("omittedCount").asInt(0)); + summary.put("resultRef", response.path("resultRef").asText("")); + + ObjectNode structuredContent = json.createObjectNode(); + structuredContent.set("summary", summary); + structuredContent.set("pageInfo", pageInfo.deepCopy()); + structuredContent.set("previewItems", response.path("previewItems").deepCopy()); + response.set("structuredContent", structuredContent); + } +} diff --git a/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/PaginationRequestValidator.java b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/PaginationRequestValidator.java new file mode 100644 index 0000000..667b8b5 --- /dev/null +++ b/axhub-gateway/src/main/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/PaginationRequestValidator.java @@ -0,0 +1,101 @@ +package io.shinhanlife.axhub.biz.mcp.gateway.tool.large; + +import io.shinhanlife.axhub.biz.mcp.gateway.config.McpGatewayProperties; +import io.shinhanlife.axhub.biz.mcp.gateway.resilience.FailureType; +import io.shinhanlife.axhub.biz.mcp.gateway.resilience.ToolExecutionException; +import org.springframework.stereotype.Component; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; + +/** + * Tool 서버로 전달할 pagination 요청값을 MCP 정책 기준으로 검증하고 보정합니다. + */ +@Component +public class PaginationRequestValidator { + private static final int MAX_CURSOR_LENGTH = 2048; + + private final McpGatewayProperties properties; + private final ObjectMapper json; + + public PaginationRequestValidator(McpGatewayProperties properties, ObjectMapper json) { + this.properties = properties; + this.json = json; + } + + /** + * pageSize/cursor를 검증한 뒤 Tool 서버에 넘길 안전한 arguments 복사본을 만듭니다. + */ + public ObjectNode normalize(ObjectNode arguments) { + try { + ObjectNode normalized = arguments == null + ? json.createObjectNode() + : (ObjectNode) json.readTree(json.writeValueAsString(arguments)); + normalizePageSize(normalized); + validateCursor(normalized); + return normalized; + } catch (ToolExecutionException error) { + throw error; + } catch (Exception error) { + throw standardError("INVALID_PAGINATION_REQUEST", + "Pagination 요청값을 검증하는 중 오류가 발생했습니다.", + detail("cause", error.getMessage())); + } + } + + private void normalizePageSize(ObjectNode normalized) { + int pageSize = normalized.has("pageSize") + ? normalized.path("pageSize").asInt(-1) + : properties.defaultPageSize(); + if (pageSize < 1) { + throw standardError("INVALID_PAGINATION_REQUEST", + "pageSize는 1 이상이어야 합니다.", + detail("requestedPageSize", pageSize)); + } + if (pageSize > properties.maxPageSize()) { + ObjectNode detail = detail("requestedPageSize", pageSize); + detail.put("maxPageSize", properties.maxPageSize()); + detail.put("suggestedPageSize", properties.defaultPageSize()); + throw standardError("PAGE_SIZE_EXCEEDED", + "요청한 pageSize가 MCP 최대 허용값을 초과했습니다.", + detail); + } + normalized.put("pageSize", pageSize); + } + + private void validateCursor(ObjectNode normalized) { + if (!normalized.has("cursor")) { + return; + } + String cursor = normalized.path("cursor").asText(""); + if (cursor.length() > MAX_CURSOR_LENGTH) { + ObjectNode detail = detail("cursorLength", cursor.length()); + detail.put("maxCursorLength", MAX_CURSOR_LENGTH); + throw standardError("INVALID_CURSOR", + "cursor 값이 너무 깁니다.", + detail); + } + } + + private ToolExecutionException standardError(String code, String message, ObjectNode detail) { + ObjectNode error = json.createObjectNode(); + error.put("isError", true); + error.put("errorCode", code); + error.put("message", message); + error.set("details", detail); + try { + return new ToolExecutionException(FailureType.CLIENT_ERROR, json.writeValueAsString(error)); + } catch (Exception serializeError) { + return new ToolExecutionException(FailureType.CLIENT_ERROR, code + ": " + message); + } + } + + private ObjectNode detail(String name, Object value) { + ObjectNode detail = json.createObjectNode(); + if (value instanceof Number number) { + detail.put(name, number.longValue()); + } else { + detail.put(name, value == null ? "" : String.valueOf(value)); + } + return detail; + } +} diff --git a/axhub-gateway/src/test/java/io/shinhanlife/axhub/biz/mcp/gateway/AxHubGatewayApplicationTests.java b/axhub-gateway/src/test/java/io/shinhanlife/axhub/biz/mcp/gateway/AxHubGatewayApplicationTests.java new file mode 100644 index 0000000..df18da4 --- /dev/null +++ b/axhub-gateway/src/test/java/io/shinhanlife/axhub/biz/mcp/gateway/AxHubGatewayApplicationTests.java @@ -0,0 +1,27 @@ +package io.shinhanlife.axhub.biz.mcp.gateway; + +import org.junit.jupiter.api.Test; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.test.context.ActiveProfiles; +import org.springframework.test.context.bean.override.mockito.MockitoBean; +import io.shinhanlife.axhub.biz.mcp.gateway.audit.AuditLogService; +import io.shinhanlife.axhub.biz.mcp.gateway.redis.RedisToolTraceService; + +@SpringBootTest +@ActiveProfiles("test") +class AxHubGatewayApplicationTests { + + // Mock components that might require external dependencies (like Redis/DB) to pass the context load + @MockitoBean + private RedisToolTraceService redisToolTraceService; + + @MockitoBean + private AuditLogService auditLogService; + + @Test + void contextLoads() { + // This test ensures that all Spring beans, @ConfigurationProperties, and dependencies are correctly wired. + // It validates that the ported classes (ExecuteService, AgentResponseBudgetService, LargeToolResponseService, etc.) + // have no @Autowired or initialization errors. + } +} diff --git a/axhub-gateway/src/test/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/PaginationRequestValidatorTest.java b/axhub-gateway/src/test/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/PaginationRequestValidatorTest.java new file mode 100644 index 0000000..ec495f8 --- /dev/null +++ b/axhub-gateway/src/test/java/io/shinhanlife/axhub/biz/mcp/gateway/tool/large/PaginationRequestValidatorTest.java @@ -0,0 +1,67 @@ +package io.shinhanlife.axhub.biz.mcp.gateway.tool.large; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import io.shinhanlife.axhub.biz.mcp.gateway.config.McpGatewayProperties; +import io.shinhanlife.axhub.biz.mcp.gateway.resilience.ToolExecutionException; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.*; + +class PaginationRequestValidatorTest { + + private PaginationRequestValidator validator; + private ObjectMapper json; + private McpGatewayProperties properties; + + @BeforeEach + void setUp() { + properties = new McpGatewayProperties(); + properties.setDefaultPageSize(100); + properties.setMaxPageSize(500); + json = new ObjectMapper(); + validator = new PaginationRequestValidator(properties, json); + } + + @Test + void shouldSetDefaultPageSizeIfMissing() { + ObjectNode arguments = json.createObjectNode(); + ObjectNode normalized = validator.normalize(arguments); + + assertEquals(100, normalized.get("pageSize").asInt()); + } + + @Test + void shouldKeepValidPageSize() { + ObjectNode arguments = json.createObjectNode(); + arguments.put("pageSize", 200); + ObjectNode normalized = validator.normalize(arguments); + + assertEquals(200, normalized.get("pageSize").asInt()); + } + + @Test + void shouldThrowIfPageSizeTooLarge() { + ObjectNode arguments = json.createObjectNode(); + arguments.put("pageSize", 600); + + ToolExecutionException exception = assertThrows(ToolExecutionException.class, () -> { + validator.normalize(arguments); + }); + + assertTrue(exception.getMessage().contains("PAGE_SIZE_EXCEEDED")); + } + + @Test + void shouldThrowIfCursorTooLong() { + ObjectNode arguments = json.createObjectNode(); + arguments.put("cursor", "a".repeat(3000)); + + ToolExecutionException exception = assertThrows(ToolExecutionException.class, () -> { + validator.normalize(arguments); + }); + + assertTrue(exception.getMessage().contains("INVALID_CURSOR")); + } +}