From ed392434fc7b41b40632073e5f99fe9cbdee2db2 Mon Sep 17 00:00:00 2001 From: Daijiro Wachi Date: Wed, 4 Nov 2026 19:12:52 +0900 Subject: [PATCH 1/4] Add an MCP server for AI agents, with traced and audited tool calls WebAPI can now serve the Model Context Protocol at POST /mcp (Spring AI 1.1.8, stateless Streamable HTTP), off by default behind spring.ai.mcp.server.enabled. It sits behind the REST API's security chain, so agents authenticate with a JWT or X-API-KEY and every tool runs under that user's existing @PreAuthorize rules. Tools cover the phenotype workflow: vocabulary search and hierarchy, concept set editing, cohort definition create/update, Circe checks plus SqlRender translation, cohort generation and status with attrition. Calls are made traceable on the server side rather than trusted to the client: - tools that change concept sets or cohorts require clinical_rationale and evidence_reference; - concept ids are re-checked against the vocabulary by id and name before anything is saved, so assumed ids are refused; - _meta.traceparent is continued with a new span per call and returned in the result's _meta; - each call is written to mcp_audit_log before it runs and completed with its outcome; a trigger makes the table append-only, and get_audit_trail reads a trace back as a decision log. The MCP server bean is excluded from lazy initialization; left lazy it is never created and /mcp cannot answer. --- README.md | 34 ++ pom.xml | 16 + .../org/ohdsi/webapi/mcp/AuditMcpTools.java | 98 ++++ .../org/ohdsi/webapi/mcp/CohortMcpTools.java | 419 ++++++++++++++++++ .../ohdsi/webapi/mcp/ConceptSetMcpTools.java | 240 ++++++++++ .../org/ohdsi/webapi/mcp/ConceptViews.java | 45 ++ .../org/ohdsi/webapi/mcp/McpServerConfig.java | 48 ++ .../org/ohdsi/webapi/mcp/McpToolFactory.java | 215 +++++++++ .../org/ohdsi/webapi/mcp/McpToolProvider.java | 9 + .../org/ohdsi/webapi/mcp/ReasoningStep.java | 14 + .../org/ohdsi/webapi/mcp/ToolArguments.java | 136 ++++++ .../webapi/mcp/ToolContractException.java | 12 + .../org/ohdsi/webapi/mcp/ToolDefinition.java | 156 +++++++ .../org/ohdsi/webapi/mcp/TraceContext.java | 66 +++ .../ohdsi/webapi/mcp/VocabularyMcpTools.java | 210 +++++++++ .../ohdsi/webapi/mcp/audit/McpAuditEntry.java | 138 ++++++ .../mcp/audit/McpAuditEntryRepository.java | 11 + .../webapi/mcp/audit/McpAuditService.java | 69 +++ src/main/resources/application.yaml | 23 + .../V3.0.1.20261104000000__mcp_audit_log.sql | 69 +++ .../ohdsi/webapi/mcp/McpToolFactoryTest.java | 153 +++++++ .../ohdsi/webapi/mcp/TraceContextTest.java | 57 +++ .../org/ohdsi/webapi/test/McpServerIT.java | 347 +++++++++++++++ 23 files changed, 2585 insertions(+) create mode 100644 src/main/java/org/ohdsi/webapi/mcp/AuditMcpTools.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/ConceptViews.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/McpServerConfig.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/McpToolProvider.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/ReasoningStep.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/ToolArguments.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/ToolContractException.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/TraceContext.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/VocabularyMcpTools.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java create mode 100644 src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java create mode 100644 src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql create mode 100644 src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java create mode 100644 src/test/java/org/ohdsi/webapi/mcp/TraceContextTest.java create mode 100644 src/test/java/org/ohdsi/webapi/test/McpServerIT.java diff --git a/README.md b/README.md index fd035ff17a..157a1c6298 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,40 @@ Acceptable values are `true` and `false`. Default parameters for each authentication provider are provided as examples in the embedded `application.yaml` file. All providers are disabled by default. +## MCP server for AI agents + +WebAPI can expose a [Model Context Protocol](https://modelcontextprotocol.io) server so AI agents can search vocabularies, build concept sets and cohort definitions, validate the cohort SQL and generate cohorts. It is off by default. To turn it on: + +```yaml +spring: + ai: + mcp: + server: + enabled: true +``` + +The server listens at `POST /mcp` (stateless Streamable HTTP; set the path with `spring.ai.mcp.server.streamable-http.mcp-endpoint`). It uses the same authentication as the REST API, so clients send a WebAPI JWT (`Authorization: Bearer ...`) or a personal API key (`X-API-KEY: ...`). Each tool runs with that user's permissions. For example, with Claude Code: + +```bash +claude mcp add --transport http webapi https://your-host/WebAPI/mcp --header "X-API-KEY: " +``` + +| Step | Tools | +|------|-------| +| 1 Concept exploration | `list_sources`, `search_concepts`, `get_concept_details`, `get_concept_hierarchy` | +| 2 Concept sets | `create_concept_set`, `get_concept_set`, `add_concept_to_set`, `remove_concept_from_set`, `count_concept_set_concepts` | +| 3 Cohort criteria | `create_cohort_definition`, `update_cohort_definition`, `get_cohort_definition` | +| 4 SQL validation | `validate_cohort_definition` (Circe checks, concept verification, SqlRender translation; the SQL is not run) | +| 5 Generation | `generate_cohort`, `get_cohort_generation_status` (status, counts, inclusion-rule attrition) | +| Audit | `get_audit_trail` | + +Traceability rules the server enforces: + +- **Rationale.** Tools that change concept sets or cohort definitions require `clinical_rationale` (at least `mcp.minRationaleLength` characters, default 20) and `evidence_reference`. A call that leaves them out is refused. +- **Verified concepts.** `add_concept_to_set`, `create_cohort_definition` and `update_cohort_definition` look up every concept id in the source vocabulary. If an id does not exist, or the name the caller gave does not match the vocabulary's name, the call is refused and nothing is saved. +- **Tracing.** Send a W3C `traceparent` in the call's `_meta`, plus an optional `reasoning_step`. WebAPI keeps the trace id, records your span as the parent, and returns the tool call's own span in the result's `_meta.traceparent`, along with `_meta.audit_entry_id`. A call without a traceparent starts a new trace, unless `mcp.requireTraceparent: true`. +- **Audit.** Every call is written to `mcp_audit_log` before it runs, then completed with its outcome: `SUCCESS`, `ERROR`, `REJECTED` or `DENIED`. A database trigger makes the table append-only. `get_audit_trail` returns a trace's entries in order, as the decision log for an audit manifest. + ## Geospatial support Instructions can be found at [webapi-component-geospatial](https://github.com/OHDSI/webapi-component-geospatial) diff --git a/pom.xml b/pom.xml index e10736be82..473447e02c 100644 --- a/pom.xml +++ b/pom.xml @@ -43,6 +43,9 @@ false + + 1.1.8 + @@ -322,6 +325,13 @@ pom import + + org.springframework.ai + spring-ai-bom + ${spring-ai.version} + pom + import + com.amazon.redshift redshift-jdbc42 @@ -448,6 +458,12 @@ spring-boot-starter-cache + + + org.springframework.ai + spring-ai-starter-mcp-server-webmvc + + org.springframework.boot diff --git a/src/main/java/org/ohdsi/webapi/mcp/AuditMcpTools.java b/src/main/java/org/ohdsi/webapi/mcp/AuditMcpTools.java new file mode 100644 index 0000000000..9134fb4420 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/AuditMcpTools.java @@ -0,0 +1,98 @@ +package org.ohdsi.webapi.mcp; + +import static org.ohdsi.webapi.mcp.ToolDefinition.stringParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.tool; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry; +import org.ohdsi.webapi.mcp.audit.McpAuditService; +import org.ohdsi.webapi.security.authz.AuthorizationService; +import org.ohdsi.webapi.security.identity.WebApiPrincipal; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +/** + * Reads back the audit trail of a trace, as the decision_log of an Audit Manifest. A user sees + * only the entries they made themselves. + */ +@Component +@ConditionalOnProperty(prefix = "spring.ai.mcp.server", name = "enabled", havingValue = "true") +public class AuditMcpTools implements McpToolProvider { + + private final McpAuditService audit; + private final AuthorizationService authorizationService; + private final ObjectMapper mapper; + + public AuditMcpTools(McpAuditService audit, AuthorizationService authorizationService, ObjectMapper mapper) { + this.audit = audit; + this.authorizationService = authorizationService; + this.mapper = mapper; + } + + @Override + public List tools() { + return List.of(tool("get_audit_trail") + .title("Get audit trail") + .description("Returns, in order, every tool call recorded for a trace id by the current user: " + + "span ids, reasoning step, arguments, clinical rationale, evidence and outcome. Use it to " + + "build the decision_log of the Audit Manifest from what the server actually recorded.") + .step(ReasoningStep.AUDIT) + .readOnly() + .requiredParam("trace_id", stringParam("The 32-hex-digit W3C trace id.")) + .handler(this::trail) + .build()); + } + + private Map trail(ToolArguments args) { + String traceId = args.string("trace_id").trim().toLowerCase(Locale.ROOT); + if (!traceId.matches("[0-9a-f]{32}")) { + throw new ToolContractException("'trace_id' must be 32 lowercase hex digits"); + } + WebApiPrincipal me = authorizationService.getAuthenticatedPrincipal(); + List> log = new ArrayList<>(); + for (McpAuditEntry e : audit.findByTrace(traceId)) { + if (!Objects.equals(e.getUserLogin(), me.getName())) { + continue; + } + Map m = new LinkedHashMap<>(); + m.put("audit_entry_id", e.getId()); + m.put("span_id", e.getSpanId()); + m.put("parent_span_id", e.getParentSpanId()); + m.put("reasoning_step", e.getReasoningStep()); + m.put("tool", e.getToolName()); + m.put("arguments", readTree(e.getArgumentsJson())); + m.put("clinical_rationale", e.getClinicalRationale()); + m.put("evidence_reference", e.getEvidenceReference()); + m.put("outcome", e.getOutcome().name()); + m.put("result", readTree(e.getResultJson())); + m.put("error_message", e.getErrorMessage()); + m.put("started_at", e.getStartedAt() == null ? null : e.getStartedAt().toString()); + m.put("finished_at", e.getFinishedAt() == null ? null : e.getFinishedAt().toString()); + log.add(m); + } + Map out = new LinkedHashMap<>(); + out.put("trace_id", traceId); + out.put("entry_count", log.size()); + out.put("decision_log", log); + return out; + } + + private Object readTree(String json) { + if (json == null) { + return null; + } + try { + return mapper.readTree(json); + } catch (JsonProcessingException e) { + // truncated results are no longer valid JSON; return them as text + return json; + } + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java b/src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java new file mode 100644 index 0000000000..cae98f5b90 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java @@ -0,0 +1,419 @@ +package org.ohdsi.webapi.mcp; + +import static org.ohdsi.webapi.mcp.ToolDefinition.booleanParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.integerParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.objectParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.stringParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.tool; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.ohdsi.circe.check.Checker; +import org.ohdsi.circe.check.Warning; +import org.ohdsi.circe.check.warnings.BaseWarning; +import org.ohdsi.circe.cohortdefinition.CohortExpression; +import org.ohdsi.circe.cohortdefinition.CohortExpressionQueryBuilder.BuildExpressionQueryOptions; +import org.ohdsi.circe.cohortdefinition.ConceptSet; +import org.ohdsi.circe.vocabulary.ConceptSetExpression; +import org.ohdsi.sql.SqlTranslate; +import org.ohdsi.vocabulary.Concept; +import org.ohdsi.webapi.cohortdefinition.CohortDefinitionService; +import org.ohdsi.webapi.cohortdefinition.ExpressionType; +import org.ohdsi.webapi.cohortdefinition.dto.CohortDTO; +import org.ohdsi.webapi.cohortdefinition.dto.CohortGenerationInfoDTO; +import org.ohdsi.webapi.cohortdefinition.dto.CohortRawDTO; +import org.ohdsi.webapi.cohortdefinition.InclusionRuleReport; +import org.ohdsi.webapi.job.JobExecutionResource; +import org.ohdsi.webapi.mcp.ToolDefinition.Rationale; +import org.ohdsi.webapi.service.VocabularyService; +import org.ohdsi.webapi.source.Source; +import org.ohdsi.webapi.source.SourceDaimon.DaimonType; +import org.ohdsi.webapi.source.SourceService; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +/** + * STEP 3-5 tools: compose cohort definitions, validate them and their SQL, generate them on a CDM + * source and read back the counts. + */ +@Component +@ConditionalOnProperty(prefix = "spring.ai.mcp.server", name = "enabled", havingValue = "true") +public class CohortMcpTools implements McpToolProvider { + + private final CohortDefinitionService cohortService; + private final VocabularyService vocabularyService; + private final SourceService sourceService; + private final VocabularyMcpTools vocabularyTools; + private final ObjectMapper mapper; + + public CohortMcpTools(CohortDefinitionService cohortService, VocabularyService vocabularyService, + SourceService sourceService, VocabularyMcpTools vocabularyTools, ObjectMapper mapper) { + this.cohortService = cohortService; + this.vocabularyService = vocabularyService; + this.sourceService = sourceService; + this.vocabularyTools = vocabularyTools; + this.mapper = mapper; + } + + @Override + public List tools() { + return List.of( + tool("create_cohort_definition") + .title("Create cohort definition") + .description("Saves a new cohort definition from an OHDSI (Circe) cohort expression. Every " + + "concept in the expression's ConceptSets is checked against the source vocabulary " + + "first; the call is refused if any id is unknown or its name does not match.") + .step(ReasoningStep.STEP_3_COHORT_CRITERIA_DEFINITION) + .rationale(Rationale.REQUIRED) + .requiredParam("name", stringParam("Cohort definition name; must be unique in WebAPI.")) + .param("description", stringParam("Plain-language summary of the phenotype logic.")) + .requiredParam("expression", objectParam("The Circe cohort expression (ConceptSets, " + + "PrimaryCriteria, InclusionRules, EndStrategy, ...), as JSON.")) + .requiredParam("source_key", stringParam("Vocabulary source used to verify the concepts.")) + .handler(this::create) + .build(), + tool("update_cohort_definition") + .title("Update cohort definition") + .description("Replaces the expression (and optionally name/description) of a cohort " + + "definition. The previous version is kept in the definition's version history. " + + "Concepts are verified as in create_cohort_definition.") + .step(ReasoningStep.STEP_3_COHORT_CRITERIA_DEFINITION) + .rationale(Rationale.REQUIRED) + .requiredParam("cohort_definition_id", integerParam("The cohort definition to change.")) + .requiredParam("expression", objectParam("The complete new Circe cohort expression.")) + .requiredParam("source_key", stringParam("Vocabulary source used to verify the concepts.")) + .param("name", stringParam("New name; unchanged if omitted.")) + .param("description", stringParam("New description; unchanged if omitted.")) + .handler(this::update) + .build(), + tool("get_cohort_definition") + .title("Get cohort definition") + .description("Returns a cohort definition's name, description and Circe expression.") + .step(ReasoningStep.STEP_3_COHORT_CRITERIA_DEFINITION) + .readOnly() + .requiredParam("cohort_definition_id", integerParam("The cohort definition id.")) + .handler(this::get) + .build(), + tool("validate_cohort_definition") + .title("Validate cohort definition and SQL") + .description("Validates a cohort expression (pass 'expression', or 'cohort_definition_id' for " + + "a saved one): runs the Circe logic checker, verifies every concept against the " + + "source vocabulary, builds the cohort SQL and translates it to the source's SQL " + + "dialect (or 'target_dialect'). Translation proves SqlRender can produce SQL for the " + + "dialect; the SQL is not executed. valid is false if any check reports an error.") + .step(ReasoningStep.STEP_4_SQL_VALIDATION) + .readOnly() + .param("cohort_definition_id", integerParam("A saved cohort definition to validate.")) + .param("expression", objectParam("An unsaved Circe cohort expression to validate.")) + .requiredParam("source_key", stringParam("Source whose vocabulary, schemas and dialect to use.")) + .param("target_dialect", stringParam("SqlRender dialect to translate to, e.g. 'postgresql', " + + "'sql server', 'oracle', 'redshift', 'spark', 'snowflake', 'bigquery'. Defaults to " + + "the source's dialect.")) + .param("include_sql", booleanParam("Return the translated SQL. Default false.")) + .handler(this::validate) + .build(), + tool("generate_cohort") + .title("Generate cohort") + .description("Queues generation of a saved cohort definition on a CDM source. This runs " + + "the cohort SQL against the CDM and replaces earlier results for that source. Poll " + + "get_cohort_generation_status for the outcome.") + .step(ReasoningStep.STEP_5_GENERATION_VERIFICATION) + .rationale(Rationale.OPTIONAL) + .requiredParam("cohort_definition_id", integerParam("The cohort definition to generate.")) + .requiredParam("source_key", stringParam("CDM source to generate on.")) + .handler(this::generate) + .build(), + tool("get_cohort_generation_status") + .title("Get cohort generation status") + .description("Returns the latest generation of a cohort definition on a source: status, " + + "person and record counts, failure message, and once complete the inclusion-rule " + + "attrition (how many people each rule keeps) for plausibility review.") + .step(ReasoningStep.STEP_5_GENERATION_VERIFICATION) + .readOnly() + .requiredParam("cohort_definition_id", integerParam("The cohort definition id.")) + .requiredParam("source_key", stringParam("CDM source the cohort was generated on.")) + .handler(this::status) + .build()); + } + + // ---- STEP 3 ------------------------------------------------------------------------- + + private Map create(ToolArguments args) { + String sourceKey = vocabularyTools.requireSource(args); + CohortExpression expression = parseExpression(args.raw("expression")); + requireVerifiedConcepts(sourceKey, expression); + + CohortDTO dto = new CohortDTO(); + dto.setName(args.string("name")); + dto.setDescription(args.optString("description")); + dto.setExpressionType(ExpressionType.SIMPLE_EXPRESSION); + dto.setExpression(expression); + CohortDTO created = cohortService.createCohortDefinition(dto); + return saved(created); + } + + private Map update(ToolArguments args) { + int id = args.intValue("cohort_definition_id"); + String sourceKey = vocabularyTools.requireSource(args); + CohortExpression expression = parseExpression(args.raw("expression")); + requireVerifiedConcepts(sourceKey, expression); + + CohortRawDTO current = cohortService.getCohortDefinitionRaw(id); // enforces read access + CohortDTO dto = new CohortDTO(); + dto.setName(args.has("name") ? args.string("name") : current.getName()); + dto.setDescription(args.has("description") ? args.optString("description") : current.getDescription()); + dto.setExpressionType(ExpressionType.SIMPLE_EXPRESSION); + dto.setExpression(expression); + return saved(cohortService.saveCohortDefinition(id, dto)); + } + + private Map get(ToolArguments args) { + CohortRawDTO raw = cohortService.getCohortDefinitionRaw(args.intValue("cohort_definition_id")); + Map out = new LinkedHashMap<>(); + out.put("cohort_definition_id", raw.getId()); + out.put("name", raw.getName()); + out.put("description", raw.getDescription()); + out.put("expression_type", raw.getExpressionType()); + out.put("expression", readTree(raw.getExpression())); + return out; + } + + private Map saved(CohortDTO dto) { + Map out = new LinkedHashMap<>(); + out.put("cohort_definition_id", dto.getId()); + out.put("name", dto.getName()); + out.put("concept_set_count", dto.getExpression() == null || dto.getExpression().conceptSets == null + ? 0 : dto.getExpression().conceptSets.length); + out.put("checker_warnings", check(dto.getExpression())); + return out; + } + + // ---- STEP 4 ------------------------------------------------------------------------- + + private Map validate(ToolArguments args) { + String sourceKey = vocabularyTools.requireSource(args); + Source source = sourceService.findBySourceKey(sourceKey); + Integer cohortId = args.has("cohort_definition_id") ? args.intValue("cohort_definition_id") : null; + CohortExpression expression; + if (args.has("expression")) { + expression = parseExpression(args.raw("expression")); + } else if (cohortId != null) { + expression = parseExpression(cohortService.getCohortDefinitionRaw(cohortId).getExpression()); + } else { + throw new ToolContractException("Pass 'expression' or 'cohort_definition_id'"); + } + + List> warnings = check(expression); + List conceptIssues = conceptIssues(sourceKey, expression); + boolean hasCheckerError = warnings.stream().anyMatch(w -> "CRITICAL".equals(w.get("severity"))); + + String dialect = args.has("target_dialect") ? args.string("target_dialect") : source.getSourceDialect(); + Map sql = new LinkedHashMap<>(); + sql.put("dialect", dialect); + String translated = null; + try { + BuildExpressionQueryOptions options = new BuildExpressionQueryOptions(); + options.cohortId = cohortId != null ? cohortId : 0; + options.cdmSchema = source.getTableQualifierOrNull(DaimonType.CDM); + options.vocabularySchema = source.getTableQualifierOrNull(DaimonType.Vocabulary); + options.resultSchema = source.getTableQualifierOrNull(DaimonType.Results); + options.targetTable = options.resultSchema != null ? options.resultSchema + ".cohort" : null; + options.generateStats = true; + CohortDefinitionService.GenerateSqlRequest request = new CohortDefinitionService.GenerateSqlRequest(); + request.expression = expression; + request.options = options; + String templateSql = cohortService.generateSql(request).templateSql; + translated = SqlTranslate.translateSql(templateSql, dialect); + sql.put("translated", true); + sql.put("length", translated.length()); + } catch (RuntimeException e) { + sql.put("translated", false); + sql.put("error", e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName()); + } + if (translated != null && args.optBool("include_sql", false)) { + sql.put("sql", translated); + } + + Map out = new LinkedHashMap<>(); + out.put("valid", !hasCheckerError && conceptIssues.isEmpty() && Boolean.TRUE.equals(sql.get("translated"))); + out.put("checker_warnings", warnings); + out.put("concept_issues", conceptIssues); + out.put("sql", sql); + return out; + } + + private List> check(CohortExpression expression) { + List> out = new ArrayList<>(); + if (expression == null) { + return out; + } + for (Warning w : new Checker().check(expression)) { + Map m = new LinkedHashMap<>(); + m.put("severity", w instanceof BaseWarning bw && bw.getSeverity() != null ? bw.getSeverity().name() : null); + m.put("message", w.toMessage()); + out.add(m); + } + return out; + } + + // ---- STEP 5 ------------------------------------------------------------------------- + + private Map generate(ToolArguments args) { + int id = args.intValue("cohort_definition_id"); + String sourceKey = vocabularyTools.requireSource(args); + JobExecutionResource job = cohortService.generateCohort(id, sourceKey, false); + Map out = new LinkedHashMap<>(); + out.put("cohort_definition_id", id); + out.put("source_key", sourceKey); + out.put("job_execution_id", job.getExecutionId()); + out.put("status", job.getStatus()); + return out; + } + + private Map status(ToolArguments args) { + int id = args.intValue("cohort_definition_id"); + String sourceKey = vocabularyTools.requireSource(args); + Source source = sourceService.findBySourceKey(sourceKey); + CohortGenerationInfoDTO info = cohortService.getInfo(id).stream() + .filter(i -> i.getId() != null && source.getSourceId() == i.getId().getSourceId()) + .findFirst() + .orElse(null); + + Map out = new LinkedHashMap<>(); + out.put("cohort_definition_id", id); + out.put("source_key", sourceKey); + if (info == null) { + out.put("status", "NOT_GENERATED"); + return out; + } + out.put("status", info.getStatus()); + out.put("is_valid", info.getIsValid()); + out.put("is_canceled", info.getIsCanceled()); + out.put("start_time", info.getStartTime()); + out.put("execution_duration_ms", info.getExecutionDuration()); + out.put("person_count", info.getPersonCount()); + out.put("record_count", info.getRecordCount()); + out.put("fail_message", info.getFailMessage()); + if (info.getIsValid() && "COMPLETE".equals(String.valueOf(info.getStatus()))) { + InclusionRuleReport report = cohortService.getInclusionRuleReport(id, sourceKey, 0); + Map attrition = new LinkedHashMap<>(); + if (report.summary != null) { + attrition.put("base_count", report.summary.baseCount); + attrition.put("final_count", report.summary.finalCount); + attrition.put("lost_count", report.summary.lostCount); + attrition.put("percent_matched", report.summary.percentMatched); + } + List> rules = new ArrayList<>(); + if (report.inclusionRuleStats != null) { + for (InclusionRuleReport.InclusionRuleStatistic s : report.inclusionRuleStats) { + Map r = new LinkedHashMap<>(); + r.put("rule_index", s.id); + r.put("name", s.name); + r.put("count_satisfying", s.countSatisfying); + r.put("percent_satisfying", s.percentSatisfying); + r.put("percent_excluded", s.percentExcluded); + rules.add(r); + } + } + attrition.put("inclusion_rules", rules); + out.put("attrition_all_events", attrition); + } + return out; + } + + // ---- helpers ------------------------------------------------------------------------ + + CohortExpression parseExpression(Object raw) { + if (raw == null) { + throw new ToolContractException("'expression' is required"); + } + String json; + try { + json = raw instanceof String s ? s : mapper.writeValueAsString(raw); + } catch (JsonProcessingException e) { + throw new ToolContractException("'expression' is not valid JSON: " + e.getOriginalMessage()); + } + try { + CohortExpression expression = CohortExpression.fromJson(json); + if (expression == null) { + throw new ToolContractException("'expression' is empty"); + } + return expression; + } catch (RuntimeException e) { + if (e instanceof ToolContractException) { + throw e; + } + throw new ToolContractException("'expression' is not a valid Circe cohort expression: " + e.getMessage()); + } + } + + private void requireVerifiedConcepts(String sourceKey, CohortExpression expression) { + List issues = conceptIssues(sourceKey, expression); + if (!issues.isEmpty()) { + throw new ToolContractException("The expression uses concepts that could not be verified against source '" + + sourceKey + "'; nothing was saved. " + String.join(" ", issues)); + } + } + + /** Every concept the expression's concept sets name, checked for existence and name against the vocabulary. */ + List conceptIssues(String sourceKey, CohortExpression expression) { + List issues = new ArrayList<>(); + Map claimed = new HashMap<>(); + Set ids = new LinkedHashSet<>(); + if (expression.conceptSets != null) { + for (ConceptSet cs : expression.conceptSets) { + if (cs.expression == null || cs.expression.items == null) { + continue; + } + for (ConceptSetExpression.ConceptSetItem item : cs.expression.items) { + if (item.concept == null || item.concept.conceptId == null) { + issues.add("Concept set " + cs.id + " ('" + cs.name + "') has an item without a conceptId."); + continue; + } + ids.add(item.concept.conceptId); + claimed.putIfAbsent(item.concept.conceptId, item.concept.conceptName); + } + } + } + if (ids.isEmpty()) { + return issues; + } + Collection found = vocabularyService.executeIdentifierLookup(sourceKey, + ids.stream().mapToLong(Long::longValue).toArray()); + Map byId = new HashMap<>(); + for (Concept c : found) { + byId.put(c.conceptId, c); + } + for (Long id : ids) { + Concept actual = byId.get(id); + String name = claimed.get(id); + if (actual == null) { + issues.add("Concept " + id + " does not exist in the vocabulary."); + } else if (name == null || name.isBlank()) { + issues.add("Concept " + id + " has no CONCEPT_NAME; it is '" + actual.conceptName + "'."); + } else if (!ConceptViews.sameName(name, actual.conceptName)) { + issues.add("Concept " + id + " is '" + actual.conceptName + "', not '" + name + "'."); + } + } + return issues; + } + + private Object readTree(String json) { + if (json == null) { + return null; + } + try { + return mapper.readTree(json); + } catch (JsonProcessingException e) { + return json; + } + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java b/src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java new file mode 100644 index 0000000000..0ef54c2d6c --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java @@ -0,0 +1,240 @@ +package org.ohdsi.webapi.mcp; + +import static org.ohdsi.webapi.mcp.ToolDefinition.booleanParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.integerParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.stringParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.tool; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.ohdsi.circe.vocabulary.ConceptSetExpression; +import org.ohdsi.vocabulary.Concept; +import org.ohdsi.webapi.conceptset.ConceptSetItem; +import org.ohdsi.webapi.conceptset.ConceptSetService; +import org.ohdsi.webapi.mcp.ToolDefinition.Rationale; +import org.ohdsi.webapi.service.VocabularyService; +import org.ohdsi.webapi.service.dto.ConceptSetDTO; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +/** + * STEP 2 tools: build concept sets one verified concept at a time. Every change carries a + * clinical rationale, and a concept is only added after its id and name are checked against the + * vocabulary, so a hallucinated or mistyped concept id is refused instead of saved. + */ +@Component +@ConditionalOnProperty(prefix = "spring.ai.mcp.server", name = "enabled", havingValue = "true") +public class ConceptSetMcpTools implements McpToolProvider { + + private final ConceptSetService conceptSetService; + private final VocabularyService vocabularyService; + private final VocabularyMcpTools vocabularyTools; + + public ConceptSetMcpTools(ConceptSetService conceptSetService, VocabularyService vocabularyService, + VocabularyMcpTools vocabularyTools) { + this.conceptSetService = conceptSetService; + this.vocabularyService = vocabularyService; + this.vocabularyTools = vocabularyTools; + } + + @Override + public List tools() { + return List.of( + tool("create_concept_set") + .title("Create concept set") + .description("Creates an empty, named concept set owned by the current user. Add concepts " + + "with add_concept_to_set.") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .rationale(Rationale.REQUIRED) + .requiredParam("name", stringParam("Concept set name; must be unique in WebAPI.")) + .param("description", stringParam("What the concept set is meant to capture.")) + .handler(this::create) + .build(), + tool("get_concept_set") + .title("Get concept set") + .description("Returns a concept set's items with their flags and the concept set expression " + + "JSON, ready to embed in a cohort definition's ConceptSets.") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .readOnly() + .requiredParam("concept_set_id", integerParam("The concept set id.")) + .param("source_key", stringParam("Vocabulary source used to fill in concept details. " + + "Defaults to the priority vocabulary source.")) + .handler(this::get) + .build(), + tool("add_concept_to_set") + .title("Add concept to concept set") + .description("Adds a concept to a concept set, or updates its flags if it is already there. " + + "The concept is looked up in the source vocabulary first and the call is refused " + + "if the id does not exist or concept_name does not match the vocabulary's name. " + + "Non-standard or invalid concepts are accepted but reported as warnings.") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .rationale(Rationale.REQUIRED) + .requiredParam("concept_set_id", integerParam("The concept set to change.")) + .requiredParam("source_key", stringParam("Vocabulary source used to verify the concept.")) + .requiredParam("concept_id", integerParam("The concept id, as found by search_concepts.")) + .requiredParam("concept_name", stringParam("The concept's name exactly as the vocabulary has it.")) + .requiredParam("include_descendants", booleanParam("Include all descendants of the concept.")) + .param("include_mapped", booleanParam("Include source concepts mapped to it. Default false.")) + .param("is_excluded", booleanParam("Exclude rather than include the concept. Default false.")) + .handler(this::addConcept) + .build(), + tool("remove_concept_from_set") + .title("Remove concept from concept set") + .description("Removes a concept from a concept set.") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .destructive() + .rationale(Rationale.REQUIRED) + .requiredParam("concept_set_id", integerParam("The concept set to change.")) + .requiredParam("concept_id", integerParam("The concept id to remove.")) + .handler(this::removeConcept) + .build(), + tool("count_concept_set_concepts") + .title("Count resolved concepts") + .description("Resolves the concept set against a source vocabulary (applying descendants, " + + "mapped concepts and exclusions) and returns how many concepts it includes.") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .readOnly() + .requiredParam("concept_set_id", integerParam("The concept set id.")) + .requiredParam("source_key", stringParam("Vocabulary source to resolve against.")) + .handler(this::count) + .build()); + } + + private Map create(ToolArguments args) { + ConceptSetDTO dto = new ConceptSetDTO(); + dto.setName(args.string("name")); + dto.setDescription(args.optString("description")); + ConceptSetDTO created = conceptSetService.createConceptSet(dto); + Map out = new LinkedHashMap<>(); + out.put("concept_set_id", created.getId()); + out.put("name", created.getName()); + return out; + } + + private Map get(ToolArguments args) { + int id = args.intValue("concept_set_id"); + ConceptSetDTO dto = conceptSetService.getConceptSet(id); + ConceptSetExpression expression = args.has("source_key") + ? conceptSetService.getConceptSetExpressionByIdAndSource(id, vocabularyTools.requireSource(args)) + : conceptSetService.getConceptSetExpressionById(id); + List> items = new ArrayList<>(); + if (expression.items != null) { + for (ConceptSetExpression.ConceptSetItem item : expression.items) { + Map m = ConceptViews.view(item.concept); + m.put("include_descendants", item.includeDescendants); + m.put("include_mapped", item.includeMapped); + m.put("is_excluded", item.isExcluded); + items.add(m); + } + } + Map out = new LinkedHashMap<>(); + out.put("concept_set_id", dto.getId()); + out.put("name", dto.getName()); + out.put("description", dto.getDescription()); + out.put("items", items); + out.put("expression", expression); + return out; + } + + private Map addConcept(ToolArguments args) { + int conceptSetId = args.intValue("concept_set_id"); + String sourceKey = vocabularyTools.requireSource(args); + long conceptId = args.longValue("concept_id"); + String claimedName = args.string("concept_name"); + boolean includeDescendants = args.bool("include_descendants"); + boolean includeMapped = args.optBool("include_mapped", false); + boolean excluded = args.optBool("is_excluded", false); + + Concept concept = verifiedConcept(sourceKey, conceptId, claimedName); + + List items = new ArrayList<>(); + conceptSetService.getConceptSetItems(conceptSetId).forEach(items::add); + ConceptSetItem existing = items.stream().filter(i -> i.getConceptId() == conceptId).findFirst().orElse(null); + ConceptSetItem item = existing != null ? existing : new ConceptSetItem(); + item.setConceptId(conceptId); + item.setIncludeDescendants(includeDescendants ? 1 : 0); + item.setIncludeMapped(includeMapped ? 1 : 0); + item.setIsExcluded(excluded ? 1 : 0); + if (existing == null) { + items.add(item); + } + conceptSetService.saveConceptSetItems(conceptSetId, items.toArray(new ConceptSetItem[0])); + + List warnings = new ArrayList<>(); + if (!excluded && !ConceptViews.isStandard(concept)) { + warnings.add("Concept " + conceptId + " is not a standard concept (standard_concept=" + + concept.standardConcept + "); cohort criteria match standard concepts, so it will only match " + + "via include_mapped or not at all."); + } + if (!ConceptViews.isValid(concept)) { + warnings.add("Concept " + conceptId + " is not valid (invalid_reason=" + concept.invalidReason + + "); consider its replacement concept."); + } + + Map out = new LinkedHashMap<>(); + out.put("concept_set_id", conceptSetId); + out.put("action", existing == null ? "added" : "updated"); + out.put("concept", ConceptViews.view(concept)); + out.put("include_descendants", includeDescendants); + out.put("include_mapped", includeMapped); + out.put("is_excluded", excluded); + out.put("item_count", items.size()); + out.put("warnings", warnings); + return out; + } + + private Map removeConcept(ToolArguments args) { + int conceptSetId = args.intValue("concept_set_id"); + long conceptId = args.longValue("concept_id"); + List items = new ArrayList<>(); + conceptSetService.getConceptSetItems(conceptSetId).forEach(items::add); + boolean removed = items.removeIf(i -> i.getConceptId() == conceptId); + if (!removed) { + throw new ToolContractException("Concept " + conceptId + " is not in concept set " + conceptSetId); + } + conceptSetService.saveConceptSetItems(conceptSetId, items.toArray(new ConceptSetItem[0])); + Map out = new LinkedHashMap<>(); + out.put("concept_set_id", conceptSetId); + out.put("action", "removed"); + out.put("concept_id", conceptId); + out.put("item_count", items.size()); + return out; + } + + private Map count(ToolArguments args) { + int id = args.intValue("concept_set_id"); + String sourceKey = vocabularyTools.requireSource(args); + ConceptSetExpression expression = conceptSetService.getConceptSetExpressionByIdAndSource(id, sourceKey); + Integer count = vocabularyService.countIncludedConceptSets(sourceKey, expression); + Map out = new LinkedHashMap<>(); + out.put("concept_set_id", id); + out.put("source_key", sourceKey); + out.put("included_concept_count", count); + return out; + } + + /** + * Looks the concept up in the vocabulary and refuses it unless the caller's name for it matches. + * This is the server-side guard against concept ids that were assumed rather than looked up. + */ + Concept verifiedConcept(String sourceKey, long conceptId, String claimedName) { + Concept concept; + try { + concept = vocabularyService.getConcept(sourceKey, conceptId); + } catch (org.springframework.web.server.ResponseStatusException e) { + if (e.getStatusCode().value() == 404) { + throw new ToolContractException("Concept " + conceptId + " does not exist in the vocabulary of source '" + + sourceKey + "'. Find the concept with search_concepts instead of assuming an id."); + } + throw e; + } + if (!ConceptViews.sameName(claimedName, concept.conceptName)) { + throw new ToolContractException("Concept " + conceptId + " is '" + concept.conceptName + "' in source '" + + sourceKey + "', not '" + claimedName + "'. Verify the id with search_concepts or " + + "get_concept_details before adding it."); + } + return concept; + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/ConceptViews.java b/src/main/java/org/ohdsi/webapi/mcp/ConceptViews.java new file mode 100644 index 0000000000..847fe57b7f --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/ConceptViews.java @@ -0,0 +1,45 @@ +package org.ohdsi.webapi.mcp; + +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; +import org.ohdsi.circe.vocabulary.Concept; + +/** Compact, snake_case views of vocabulary concepts for tool results. */ +final class ConceptViews { + + private ConceptViews() { + } + + static Map view(Concept c) { + Map m = new LinkedHashMap<>(); + m.put("concept_id", c.conceptId); + m.put("concept_name", c.conceptName); + m.put("domain_id", c.domainId); + m.put("vocabulary_id", c.vocabularyId); + m.put("concept_class_id", c.conceptClassId); + m.put("concept_code", c.conceptCode); + m.put("standard_concept", c.standardConcept); + m.put("invalid_reason", c.invalidReason); + return m; + } + + /** Vocabulary queries report a standard concept as 'S' and a non-standard one as 'N' (or null). */ + static boolean isStandard(Concept c) { + return "S".equals(c.standardConcept); + } + + /** Vocabulary queries report a valid concept as 'V' (or null). */ + static boolean isValid(Concept c) { + return c.invalidReason == null || "V".equals(c.invalidReason); + } + + /** Case- and whitespace-insensitive name comparison, used to catch concept ids that do not mean what the caller thinks. */ + static boolean sameName(String claimed, String actual) { + return Objects.equals(normalize(claimed), normalize(actual)); + } + + private static String normalize(String s) { + return s == null ? null : s.strip().replaceAll("\\s+", " ").toLowerCase(java.util.Locale.ROOT); + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/McpServerConfig.java b/src/main/java/org/ohdsi/webapi/mcp/McpServerConfig.java new file mode 100644 index 0000000000..00f5d49e69 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/McpServerConfig.java @@ -0,0 +1,48 @@ +package org.ohdsi.webapi.mcp; + +import io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncToolSpecification; +import io.modelcontextprotocol.server.McpStatelessSyncServer; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.springframework.boot.LazyInitializationExcludeFilter; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; + +/** + * Publishes WebAPI's tools on the MCP endpoint (Spring AI MCP server, stateless Streamable HTTP, + * {@code spring.ai.mcp.server.streamable-http.mcp-endpoint}, default {@code /mcp}). + * + *

The endpoint sits behind the same security filter chain as the REST API, so clients + * authenticate with a WebAPI JWT ({@code Authorization: Bearer}) or API key ({@code X-API-KEY}), + * and each tool runs with that user's permissions. + */ +@Configuration +@ConditionalOnProperty(prefix = "spring.ai.mcp.server", name = "enabled", havingValue = "true") +public class McpServerConfig { + + @Bean + public List webApiMcpTools(List providers, McpToolFactory factory) { + Set names = new HashSet<>(); + return providers.stream() + .flatMap(p -> p.tools().stream()) + .peek(def -> { + if (!names.add(def.name())) { + throw new IllegalStateException("Duplicate MCP tool name: " + def.name()); + } + }) + .map(factory::build) + .toList(); + } + + /** + * WebAPI runs with {@code spring.main.lazy-initialization=true}. The MCP server bean is never + * injected anywhere, and it is the bean that attaches the request handler to the transport, so + * left lazy it would never be created and every call to /mcp would fail. + */ + @Bean + public static LazyInitializationExcludeFilter mcpServerEagerInit() { + return LazyInitializationExcludeFilter.forBeanTypes(McpStatelessSyncServer.class); + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java b/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java new file mode 100644 index 0000000000..836404eb9a --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java @@ -0,0 +1,215 @@ +package org.ohdsi.webapi.mcp; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncToolSpecification; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpSchema.CallToolRequest; +import io.modelcontextprotocol.spec.McpSchema.CallToolResult; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.ohdsi.webapi.mcp.ToolDefinition.Rationale; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry.Outcome; +import org.ohdsi.webapi.mcp.audit.McpAuditService; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.security.access.AccessDeniedException; +import org.springframework.stereotype.Component; +import org.springframework.web.server.ResponseStatusException; + +/** + * Turns a {@link ToolDefinition} into an MCP tool whose every call is traced and audited. + * + *

For each call: the {@code _meta.traceparent} is continued (or a trace started), the audit row is + * written before the tool runs, the rationale contract is enforced, the tool runs under the + * caller's WebAPI identity (so the services' {@code @PreAuthorize} rules apply unchanged), and the + * audit row is completed with the outcome. The result carries the tool span's {@code traceparent} + * and the audit entry id in its {@code _meta}. + */ +@Component +@ConditionalOnProperty(prefix = "spring.ai.mcp.server", name = "enabled", havingValue = "true") +public class McpToolFactory { + + public static final String CLINICAL_RATIONALE = "clinical_rationale"; + public static final String EVIDENCE_REFERENCE = "evidence_reference"; + public static final String META_TRACEPARENT = "traceparent"; + public static final String META_REASONING_STEP = "reasoning_step"; + public static final String META_AUDIT_ENTRY_ID = "audit_entry_id"; + + private static final Logger log = LoggerFactory.getLogger(McpToolFactory.class); + + private final McpAuditService audit; + private final ObjectMapper mapper; + private final boolean requireTraceparent; + private final int minRationaleLength; + + public McpToolFactory(McpAuditService audit, ObjectMapper mapper, + @Value("${mcp.requireTraceparent:false}") boolean requireTraceparent, + @Value("${mcp.minRationaleLength:20}") int minRationaleLength) { + this.audit = audit; + this.mapper = mapper; + this.requireTraceparent = requireTraceparent; + this.minRationaleLength = minRationaleLength; + } + + public SyncToolSpecification build(ToolDefinition def) { + return SyncToolSpecification.builder() + .tool(describe(def)) + .callHandler((transportContext, request) -> call(def, request)) + .build(); + } + + McpSchema.Tool describe(ToolDefinition def) { + Map properties = new LinkedHashMap<>(def.properties); + List required = new ArrayList<>(def.required); + if (def.rationale != Rationale.NONE) { + properties.put(CLINICAL_RATIONALE, ToolDefinition.stringParam( + "Why this action is clinically correct for the target phenotype: why each concept fits, why " + + "descendants are or are not included, and which terminology logic was applied. Stored in the " + + "append-only audit log.")); + properties.put(EVIDENCE_REFERENCE, ToolDefinition.stringParam( + "Guidelines, terminology hierarchies or literature consulted (e.g. 'SNOMED-CT clinical finding " + + "hierarchy; ADA Standards of Care'). Stored in the audit log.")); + if (def.rationale == Rationale.REQUIRED) { + required.add(CLINICAL_RATIONALE); + required.add(EVIDENCE_REFERENCE); + } + } + String description = def.description + + "\n\nSend the caller's W3C traceparent in _meta.traceparent (and optionally _meta.reasoning_step); " + + "the result's _meta.traceparent is this call's span and _meta.audit_entry_id its audit record."; + return McpSchema.Tool.builder() + .name(def.name) + .title(def.title) + .description(description) + .inputSchema(new McpSchema.JsonSchema("object", properties, required, false, null, null)) + .annotations(new McpSchema.ToolAnnotations(def.title, def.readOnly, def.destructive, null, false, null)) + .build(); + } + + CallToolResult call(ToolDefinition def, CallToolRequest request) { + Map meta = request.meta() == null ? Map.of() : request.meta(); + ToolArguments args = new ToolArguments(request.arguments()); + + TraceContext trace; + String traceProblem = null; + Object traceparent = meta.get(META_TRACEPARENT); + if (traceparent instanceof String tp && !tp.isBlank()) { + try { + trace = TraceContext.continueFrom(tp); + } catch (IllegalArgumentException e) { + trace = TraceContext.newRoot(); + traceProblem = e.getMessage(); + } + } else { + trace = TraceContext.newRoot(); + if (requireTraceparent) { + traceProblem = "_meta.traceparent is required by this server"; + } + } + String step = meta.get(META_REASONING_STEP) instanceof String s && !s.isBlank() ? s : def.step.name(); + String rationale = args.raw(CLINICAL_RATIONALE) instanceof String s ? s : null; + String evidence = args.raw(EVIDENCE_REFERENCE) instanceof String s ? s : null; + + McpAuditEntry entry; + try { + entry = audit.start(trace, def.name, step, toJson(args.asMap()), rationale, evidence); + } catch (Exception e) { + log.error("MCP audit trail unavailable; refusing to run tool {}", def.name, e); + return error(trace, null, Outcome.ERROR, "The audit trail could not be written, so the tool was not run."); + } + + Outcome outcome; + Object result = null; + String message = null; + try { + if (traceProblem != null) { + throw new ToolContractException(traceProblem); + } + checkRationale(def, rationale, evidence); + result = def.handler.handle(args); + outcome = Outcome.SUCCESS; + } catch (ToolContractException e) { + outcome = Outcome.REJECTED; + message = e.getMessage(); + } catch (AccessDeniedException e) { + outcome = Outcome.DENIED; + message = "Access denied: the authenticated WebAPI user lacks the permission this tool needs."; + } catch (ResponseStatusException e) { + outcome = Outcome.ERROR; + message = e.getReason() != null ? e.getReason() : e.getStatusCode().toString(); + } catch (Exception e) { + log.warn("MCP tool {} failed", def.name, e); + outcome = Outcome.ERROR; + message = e.getMessage() != null ? e.getMessage() : e.getClass().getSimpleName(); + } + + String resultJson = outcome == Outcome.SUCCESS ? toJson(result) : null; + try { + audit.finish(entry, outcome, resultJson, message); + } catch (Exception e) { + // The action already happened and its STARTED row is on record; do not hide the result. + log.error("Could not record outcome {} of MCP audit entry {}", outcome, entry.getId(), e); + } + + if (outcome != Outcome.SUCCESS) { + return error(trace, entry.getId(), outcome, message); + } + return CallToolResult.builder() + .addTextContent(resultJson) + .isError(false) + .meta(resultMeta(trace, entry.getId(), step)) + .build(); + } + + private void checkRationale(ToolDefinition def, String rationale, String evidence) { + if (def.rationale != Rationale.REQUIRED) { + return; + } + if (rationale == null || rationale.strip().length() < minRationaleLength) { + throw new ToolContractException("'" + CLINICAL_RATIONALE + "' is required for " + def.name + + " and must be at least " + minRationaleLength + " characters: state why this action fits the " + + "target phenotype and why descendants are or are not included."); + } + if (evidence == null || evidence.isBlank()) { + throw new ToolContractException("'" + EVIDENCE_REFERENCE + "' is required for " + def.name + + ": name the terminology hierarchy or guideline consulted."); + } + } + + private CallToolResult error(TraceContext trace, Long auditId, Outcome outcome, String message) { + Map body = new LinkedHashMap<>(); + body.put("outcome", outcome.name()); + body.put("message", message); + return CallToolResult.builder() + .addTextContent(toJson(body)) + .isError(true) + .meta(resultMeta(trace, auditId, null)) + .build(); + } + + private static Map resultMeta(TraceContext trace, Long auditId, String step) { + Map meta = new LinkedHashMap<>(); + meta.put(META_TRACEPARENT, trace.traceparent()); + if (auditId != null) { + meta.put(META_AUDIT_ENTRY_ID, auditId); + } + if (step != null) { + meta.put(META_REASONING_STEP, step); + } + return meta; + } + + private String toJson(Object value) { + try { + return mapper.writeValueAsString(value); + } catch (JsonProcessingException e) { + throw new IllegalStateException("Could not serialize MCP tool payload", e); + } + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/McpToolProvider.java b/src/main/java/org/ohdsi/webapi/mcp/McpToolProvider.java new file mode 100644 index 0000000000..79447f12f6 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/McpToolProvider.java @@ -0,0 +1,9 @@ +package org.ohdsi.webapi.mcp; + +import java.util.List; + +/** A group of MCP tools. Every bean of this type is published on the MCP endpoint. */ +public interface McpToolProvider { + + List tools(); +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/ReasoningStep.java b/src/main/java/org/ohdsi/webapi/mcp/ReasoningStep.java new file mode 100644 index 0000000000..f8b4356213 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/ReasoningStep.java @@ -0,0 +1,14 @@ +package org.ohdsi.webapi.mcp; + +/** + * The phenotype-engineering step a tool belongs to. Recorded on every audit entry; a client may + * override it per call with {@code _meta.reasoning_step}. + */ +public enum ReasoningStep { + STEP_1_CONCEPT_EXPLORATION, + STEP_2_CONCEPT_SET_CONSOLIDATION, + STEP_3_COHORT_CRITERIA_DEFINITION, + STEP_4_SQL_VALIDATION, + STEP_5_GENERATION_VERIFICATION, + AUDIT +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/ToolArguments.java b/src/main/java/org/ohdsi/webapi/mcp/ToolArguments.java new file mode 100644 index 0000000000..359a1d4ef1 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/ToolArguments.java @@ -0,0 +1,136 @@ +package org.ohdsi.webapi.mcp; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * Typed access to the {@code arguments} of an MCP tool call. Any missing or mistyped argument + * fails with a {@link ToolContractException} naming the argument. + */ +public final class ToolArguments { + + private final Map args; + + public ToolArguments(Map args) { + this.args = args == null ? Collections.emptyMap() : args; + } + + public Map asMap() { + return args; + } + + public boolean has(String name) { + return args.get(name) != null; + } + + public Object raw(String name) { + return args.get(name); + } + + public String string(String name) { + String value = optString(name); + if (value == null || value.isBlank()) { + throw new ToolContractException("'" + name + "' is required"); + } + return value; + } + + public String optString(String name) { + Object value = args.get(name); + if (value == null) { + return null; + } + if (!(value instanceof String s)) { + throw new ToolContractException("'" + name + "' must be a string"); + } + return s; + } + + public long longValue(String name) { + Long value = optLong(name); + if (value == null) { + throw new ToolContractException("'" + name + "' is required"); + } + return value; + } + + public int intValue(String name) { + long value = longValue(name); + if (value < Integer.MIN_VALUE || value > Integer.MAX_VALUE) { + throw new ToolContractException("'" + name + "' is out of range"); + } + return (int) value; + } + + public Long optLong(String name) { + Object value = args.get(name); + if (value == null) { + return null; + } + if (value instanceof Number n) { + double d = n.doubleValue(); + if (d != Math.rint(d)) { + throw new ToolContractException("'" + name + "' must be an integer"); + } + return n.longValue(); + } + if (value instanceof String s) { + try { + return Long.parseLong(s.trim()); + } catch (NumberFormatException e) { + // fall through + } + } + throw new ToolContractException("'" + name + "' must be an integer"); + } + + public int optInt(String name, int defaultValue) { + return has(name) ? intValue(name) : defaultValue; + } + + public boolean bool(String name) { + Boolean value = optBool(name); + if (value == null) { + throw new ToolContractException("'" + name + "' is required (true or false)"); + } + return value; + } + + public Boolean optBool(String name) { + Object value = args.get(name); + if (value == null) { + return null; + } + if (value instanceof Boolean b) { + return b; + } + throw new ToolContractException("'" + name + "' must be true or false"); + } + + public boolean optBool(String name, boolean defaultValue) { + Boolean value = optBool(name); + return value == null ? defaultValue : value; + } + + public String[] optStringArray(String name) { + Object value = args.get(name); + if (value == null) { + return null; + } + if (value instanceof String s) { + return new String[] { s }; + } + if (value instanceof List list) { + String[] out = new String[list.size()]; + for (int i = 0; i < list.size(); i++) { + if (!(list.get(i) instanceof String s)) { + throw new ToolContractException("'" + name + "' must be an array of strings"); + } + out[i] = s; + } + return out; + } + throw new ToolContractException("'" + name + "' must be an array of strings"); + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/ToolContractException.java b/src/main/java/org/ohdsi/webapi/mcp/ToolContractException.java new file mode 100644 index 0000000000..31626e25c5 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/ToolContractException.java @@ -0,0 +1,12 @@ +package org.ohdsi.webapi.mcp; + +/** + * The call broke the tool contract (bad or missing arguments, missing rationale, a concept that + * does not match the vocabulary) and was refused before doing anything. Audited as REJECTED. + */ +public class ToolContractException extends RuntimeException { + + public ToolContractException(String message) { + super(message); + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java b/src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java new file mode 100644 index 0000000000..d5b49e3229 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java @@ -0,0 +1,156 @@ +package org.ohdsi.webapi.mcp; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * Declarative description of one MCP tool: its name, JSON input schema, the reasoning step it + * belongs to, whether it must carry a clinical rationale, and the code that runs it. + */ +public final class ToolDefinition { + + /** Whether a call must explain itself with {@code clinical_rationale} and {@code evidence_reference}. */ + public enum Rationale { + NONE, + OPTIONAL, + REQUIRED + } + + @FunctionalInterface + public interface Handler { + Object handle(ToolArguments args) throws Exception; + } + + final String name; + final String title; + final String description; + final ReasoningStep step; + final boolean readOnly; + final boolean destructive; + final Rationale rationale; + final Map properties; + final List required; + final Handler handler; + + private ToolDefinition(Builder b) { + this.name = Objects.requireNonNull(b.name); + this.title = b.title; + this.description = Objects.requireNonNull(b.description); + this.step = Objects.requireNonNull(b.step); + this.readOnly = b.readOnly; + this.destructive = b.destructive; + this.rationale = b.rationale; + this.properties = b.properties; + this.required = b.required; + this.handler = Objects.requireNonNull(b.handler); + } + + public String name() { + return name; + } + + public static Builder tool(String name) { + return new Builder(name); + } + + // ---- JSON schema helpers ------------------------------------------------------------- + + public static Map stringParam(String description) { + return param("string", description); + } + + public static Map integerParam(String description) { + return param("integer", description); + } + + public static Map booleanParam(String description) { + return param("boolean", description); + } + + public static Map objectParam(String description) { + return param("object", description); + } + + public static Map stringArrayParam(String description) { + Map p = param("array", description); + p.put("items", Map.of("type", "string")); + return p; + } + + private static Map param(String type, String description) { + Map p = new LinkedHashMap<>(); + p.put("type", type); + p.put("description", description); + return p; + } + + public static final class Builder { + private final String name; + private String title; + private String description; + private ReasoningStep step; + private boolean readOnly; + private boolean destructive; + private Rationale rationale = Rationale.NONE; + private final Map properties = new LinkedHashMap<>(); + private final List required = new ArrayList<>(); + private Handler handler; + + private Builder(String name) { + this.name = name; + } + + public Builder title(String title) { + this.title = title; + return this; + } + + public Builder description(String description) { + this.description = description; + return this; + } + + public Builder step(ReasoningStep step) { + this.step = step; + return this; + } + + public Builder readOnly() { + this.readOnly = true; + return this; + } + + public Builder destructive() { + this.destructive = true; + return this; + } + + public Builder rationale(Rationale rationale) { + this.rationale = rationale; + return this; + } + + public Builder param(String name, Map schema) { + properties.put(name, schema); + return this; + } + + public Builder requiredParam(String name, Map schema) { + properties.put(name, schema); + required.add(name); + return this; + } + + public Builder handler(Handler handler) { + this.handler = handler; + return this; + } + + public ToolDefinition build() { + return new ToolDefinition(this); + } + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/TraceContext.java b/src/main/java/org/ohdsi/webapi/mcp/TraceContext.java new file mode 100644 index 0000000000..a387ca3bd5 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/TraceContext.java @@ -0,0 +1,66 @@ +package org.ohdsi.webapi.mcp; + +import java.security.SecureRandom; +import java.util.HexFormat; +import java.util.Locale; +import java.util.regex.Pattern; + +/** + * W3C Trace Context for one MCP tool call. + * + *

The client sends its span as {@code _meta.traceparent}; WebAPI keeps the client's trace id, + * records the client's span as the parent, and allocates a fresh span id for the tool call itself. + * The returned {@link #traceparent()} is what the client should use as parent for anything the tool + * result causes it to do next. + * + * @see traceparent + */ +public record TraceContext(String traceId, String parentSpanId, String spanId, String flags, boolean generated) { + + private static final Pattern TRACEPARENT = + Pattern.compile("^([0-9a-f]{2})-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$"); + private static final String INVALID_TRACE_ID = "0".repeat(32); + private static final String INVALID_SPAN_ID = "0".repeat(16); + private static final SecureRandom RANDOM = new SecureRandom(); + + /** + * Continues the trace in {@code traceparent}. + * + * @throws IllegalArgumentException if it is not a valid version-00 traceparent + */ + public static TraceContext continueFrom(String traceparent) { + String value = traceparent.trim().toLowerCase(Locale.ROOT); + var m = TRACEPARENT.matcher(value); + if (!m.matches()) { + throw new IllegalArgumentException( + "traceparent must look like 00-<32 hex trace-id>-<16 hex parent-id>-<2 hex flags>, got: " + traceparent); + } + if ("ff".equals(m.group(1))) { + throw new IllegalArgumentException("traceparent version ff is invalid"); + } + if (INVALID_TRACE_ID.equals(m.group(2)) || INVALID_SPAN_ID.equals(m.group(3))) { + throw new IllegalArgumentException("traceparent trace-id and parent-id must not be all zeros"); + } + return new TraceContext(m.group(2), m.group(3), randomHex(8), m.group(4), false); + } + + /** Starts a new trace for a client that sent no traceparent. */ + public static TraceContext newRoot() { + return new TraceContext(randomHex(16), null, randomHex(8), "01", true); + } + + /** The traceparent for this tool call's span, to hand back to the client. */ + public String traceparent() { + return "00-" + traceId + "-" + spanId + "-" + flags; + } + + private static String randomHex(int bytes) { + byte[] buf = new byte[bytes]; + String hex; + do { + RANDOM.nextBytes(buf); + hex = HexFormat.of().formatHex(buf); + } while (hex.chars().allMatch(c -> c == '0')); + return hex; + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/VocabularyMcpTools.java b/src/main/java/org/ohdsi/webapi/mcp/VocabularyMcpTools.java new file mode 100644 index 0000000000..07a01a5d15 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/VocabularyMcpTools.java @@ -0,0 +1,210 @@ +package org.ohdsi.webapi.mcp; + +import static org.ohdsi.webapi.mcp.ToolDefinition.booleanParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.integerParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.stringArrayParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.stringParam; +import static org.ohdsi.webapi.mcp.ToolDefinition.tool; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.ohdsi.vocabulary.Concept; +import org.ohdsi.webapi.service.VocabularyService; +import org.ohdsi.webapi.source.SourceInfo; +import org.ohdsi.webapi.source.SourceService; +import org.ohdsi.webapi.vocabulary.ConceptRelationship; +import org.ohdsi.webapi.vocabulary.ConceptSearch; +import org.ohdsi.webapi.vocabulary.RelatedConcept; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Component; + +/** STEP 1 tools: find sources and explore the standardized vocabulary. All read-only. */ +@Component +@ConditionalOnProperty(prefix = "spring.ai.mcp.server", name = "enabled", havingValue = "true") +public class VocabularyMcpTools implements McpToolProvider { + + static final int DEFAULT_LIMIT = 50; + static final int MAX_LIMIT = 500; + + private final VocabularyService vocabularyService; + private final SourceService sourceService; + + public VocabularyMcpTools(VocabularyService vocabularyService, SourceService sourceService) { + this.vocabularyService = vocabularyService; + this.sourceService = sourceService; + } + + @Override + public List tools() { + return List.of( + tool("list_sources") + .title("List CDM sources") + .description("Lists the CDM sources the current user can read. Use a source_key from here in " + + "every vocabulary, validation and generation tool.") + .step(ReasoningStep.STEP_1_CONCEPT_EXPLORATION) + .readOnly() + .handler(args -> listSources()) + .build(), + tool("search_concepts") + .title("Search concepts") + .description("Searches the source's standardized vocabulary by name or code. Returns at most " + + "'limit' concepts and the total number of matches. Every concept used in a concept set " + + "or cohort must first be found here or with get_concept_details.") + .step(ReasoningStep.STEP_1_CONCEPT_EXPLORATION) + .readOnly() + .requiredParam("source_key", stringParam("Source whose vocabulary to search (see list_sources).")) + .requiredParam("query", stringParam("Text or code to search for.")) + .param("domain_id", stringArrayParam("Restrict to these domains, e.g. [\"Condition\"].")) + .param("vocabulary_id", stringArrayParam("Restrict to these vocabularies, e.g. [\"SNOMED\"].")) + .param("concept_class_id", stringArrayParam("Restrict to these concept classes.")) + .param("standard_concept", stringParam("'S' standard only, 'C' classification only, 'N' non-standard only.")) + .param("include_invalid", booleanParam("Include deprecated/updated concepts. Default false.")) + .param("limit", integerParam("Maximum concepts to return (default " + DEFAULT_LIMIT + ", max " + MAX_LIMIT + ").")) + .handler(this::searchConcepts) + .build(), + tool("get_concept_details") + .title("Get concept details") + .description("Returns one concept by id from the source's vocabulary, or an error if the id " + + "does not exist there. Use it to verify any concept id before using it.") + .step(ReasoningStep.STEP_1_CONCEPT_EXPLORATION) + .readOnly() + .requiredParam("source_key", stringParam("Source whose vocabulary to use.")) + .requiredParam("concept_id", integerParam("The concept id to look up.")) + .handler(args -> ConceptViews.view( + vocabularyService.getConcept(requireSource(args), args.longValue("concept_id")))) + .build(), + tool("get_concept_hierarchy") + .title("Get concept hierarchy") + .description("Returns the ancestors and descendants of a concept (from CONCEPT_ANCESTOR) with " + + "their levels of separation. Use it to decide whether include_descendants is " + + "clinically appropriate.") + .step(ReasoningStep.STEP_1_CONCEPT_EXPLORATION) + .readOnly() + .requiredParam("source_key", stringParam("Source whose vocabulary to use.")) + .requiredParam("concept_id", integerParam("The concept whose hierarchy to return.")) + .param("max_levels", integerParam("Only return concepts at most this many levels away.")) + .param("limit", integerParam("Maximum ancestors and maximum descendants to return (default " + + DEFAULT_LIMIT + ", max " + MAX_LIMIT + ").")) + .handler(this::hierarchy) + .build()); + } + + private List> listSources() { + Collection sources = sourceService.getSourcesEndpoint().getBody(); + List> out = new ArrayList<>(); + if (sources != null) { + for (SourceInfo s : sources) { + Map m = new LinkedHashMap<>(); + m.put("source_key", s.sourceKey); + m.put("source_name", s.sourceName); + m.put("source_dialect", s.sourceDialect); + m.put("daimons", s.daimons == null ? List.of() + : s.daimons.stream().map(d -> String.valueOf(d.getDaimonType())).distinct().toList()); + out.add(m); + } + } + return out; + } + + private Map searchConcepts(ToolArguments args) { + String sourceKey = requireSource(args); + ConceptSearch search = new ConceptSearch(); + search.query = args.string("query"); + search.domainId = args.optStringArray("domain_id"); + search.vocabularyId = args.optStringArray("vocabulary_id"); + search.conceptClassId = args.optStringArray("concept_class_id"); + String standard = args.optString("standard_concept"); + if (standard != null && !standard.isBlank()) { + if (!List.of("S", "C", "N").contains(standard)) { + throw new ToolContractException("'standard_concept' must be 'S', 'C' or 'N'"); + } + search.standardConcept = standard; + } + if (!args.optBool("include_invalid", false)) { + search.invalidReason = "V"; + } + int limit = limit(args); + + Collection found = vocabularyService.executeSearch(sourceKey, search); + List> concepts = new ArrayList<>(); + for (Concept c : found) { + if (concepts.size() >= limit) { + break; + } + concepts.add(ConceptViews.view(c)); + } + Map out = new LinkedHashMap<>(); + out.put("source_key", sourceKey); + out.put("total_matches", found.size()); + out.put("returned", concepts.size()); + out.put("concepts", concepts); + return out; + } + + private Map hierarchy(ToolArguments args) { + String sourceKey = requireSource(args); + long conceptId = args.longValue("concept_id"); + Integer maxLevels = args.has("max_levels") ? args.intValue("max_levels") : null; + int limit = limit(args); + + Concept root = vocabularyService.getConcept(sourceKey, conceptId); + Collection related = vocabularyService.getConceptAncestorAndDescendant(sourceKey, conceptId); + + List> ancestors = new ArrayList<>(); + List> descendants = new ArrayList<>(); + int ancestorTotal = 0; + int descendantTotal = 0; + List sorted = new ArrayList<>(related); + sorted.sort(java.util.Comparator.comparingInt(VocabularyMcpTools::minDistance)); + for (RelatedConcept rc : sorted) { + for (ConceptRelationship rel : rc.relationships) { + if (maxLevels != null && rel.relationshipDistance > maxLevels) { + continue; + } + boolean isAncestor = "Has ancestor of".equals(rel.relationshipName); + List> bucket = isAncestor ? ancestors : descendants; + if (isAncestor) { + ancestorTotal++; + } else { + descendantTotal++; + } + if (bucket.size() < limit) { + Map m = ConceptViews.view(rc); + m.put("levels_of_separation", rel.relationshipDistance); + bucket.add(m); + } + } + } + Map out = new LinkedHashMap<>(); + out.put("concept", ConceptViews.view(root)); + out.put("ancestor_count", ancestorTotal); + out.put("descendant_count", descendantTotal); + out.put("ancestors", ancestors); + out.put("descendants", descendants); + return out; + } + + private static int minDistance(RelatedConcept rc) { + return rc.relationships.stream().mapToInt(r -> r.relationshipDistance).min().orElse(Integer.MAX_VALUE); + } + + /** Returns the 'source_key' argument after checking that such a source exists. */ + String requireSource(ToolArguments args) { + String sourceKey = args.string("source_key"); + if (sourceService.findBySourceKey(sourceKey) == null) { + throw new ToolContractException("Unknown source_key '" + sourceKey + "'; call list_sources."); + } + return sourceKey; + } + + static int limit(ToolArguments args) { + int limit = args.optInt("limit", DEFAULT_LIMIT); + if (limit < 1) { + throw new ToolContractException("'limit' must be at least 1"); + } + return Math.min(limit, MAX_LIMIT); + } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java new file mode 100644 index 0000000000..20fd7fe11b --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java @@ -0,0 +1,138 @@ +package org.ohdsi.webapi.mcp.audit; + +import jakarta.persistence.Column; +import jakarta.persistence.Entity; +import jakarta.persistence.EnumType; +import jakarta.persistence.Enumerated; +import jakarta.persistence.GeneratedValue; +import jakarta.persistence.GenerationType; +import jakarta.persistence.Id; +import jakarta.persistence.SequenceGenerator; +import jakarta.persistence.Table; +import java.time.Instant; + +/** + * One MCP tool call. Written as {@link Outcome#STARTED} before the tool runs and completed + * exactly once afterwards; the database trigger rejects any other change. + */ +@Entity(name = "McpAuditEntry") +@Table(name = "mcp_audit_log") +public class McpAuditEntry { + + public enum Outcome { + STARTED, + SUCCESS, + /** The tool ran and failed. */ + ERROR, + /** The call broke the tool contract (missing rationale, unverified concept, bad arguments) and did not run. */ + REJECTED, + /** WebAPI authorization refused the call. */ + DENIED + } + + @Id + @Column(name = "id") + @SequenceGenerator(name = "mcp_audit_log_seq", sequenceName = "mcp_audit_log_sequence", allocationSize = 1) + @GeneratedValue(strategy = GenerationType.SEQUENCE, generator = "mcp_audit_log_seq") + private Long id; + + @Column(name = "trace_id", nullable = false, length = 32, updatable = false) + private String traceId; + + @Column(name = "parent_span_id", length = 16, updatable = false) + private String parentSpanId; + + @Column(name = "span_id", nullable = false, length = 16, updatable = false) + private String spanId; + + @Column(name = "trace_generated", nullable = false, updatable = false) + private boolean traceGenerated; + + @Column(name = "reasoning_step", length = 64, updatable = false) + private String reasoningStep; + + @Column(name = "tool_name", nullable = false, length = 128, updatable = false) + private String toolName; + + @Column(name = "arguments_json", updatable = false) + private String argumentsJson; + + @Column(name = "clinical_rationale", updatable = false) + private String clinicalRationale; + + @Column(name = "evidence_reference", updatable = false) + private String evidenceReference; + + @Enumerated(EnumType.STRING) + @Column(name = "outcome", nullable = false, length = 16) + private Outcome outcome; + + @Column(name = "result_json") + private String resultJson; + + @Column(name = "error_message") + private String errorMessage; + + @Column(name = "user_id", updatable = false) + private Long userId; + + @Column(name = "user_login", length = 255, updatable = false) + private String userLogin; + + @Column(name = "started_at", nullable = false, updatable = false) + private Instant startedAt; + + @Column(name = "finished_at") + private Instant finishedAt; + + public Long getId() { return id; } + public void setId(Long id) { this.id = id; } + + public String getTraceId() { return traceId; } + public void setTraceId(String traceId) { this.traceId = traceId; } + + public String getParentSpanId() { return parentSpanId; } + public void setParentSpanId(String parentSpanId) { this.parentSpanId = parentSpanId; } + + public String getSpanId() { return spanId; } + public void setSpanId(String spanId) { this.spanId = spanId; } + + public boolean isTraceGenerated() { return traceGenerated; } + public void setTraceGenerated(boolean traceGenerated) { this.traceGenerated = traceGenerated; } + + public String getReasoningStep() { return reasoningStep; } + public void setReasoningStep(String reasoningStep) { this.reasoningStep = reasoningStep; } + + public String getToolName() { return toolName; } + public void setToolName(String toolName) { this.toolName = toolName; } + + public String getArgumentsJson() { return argumentsJson; } + public void setArgumentsJson(String argumentsJson) { this.argumentsJson = argumentsJson; } + + public String getClinicalRationale() { return clinicalRationale; } + public void setClinicalRationale(String clinicalRationale) { this.clinicalRationale = clinicalRationale; } + + public String getEvidenceReference() { return evidenceReference; } + public void setEvidenceReference(String evidenceReference) { this.evidenceReference = evidenceReference; } + + public Outcome getOutcome() { return outcome; } + public void setOutcome(Outcome outcome) { this.outcome = outcome; } + + public String getResultJson() { return resultJson; } + public void setResultJson(String resultJson) { this.resultJson = resultJson; } + + public String getErrorMessage() { return errorMessage; } + public void setErrorMessage(String errorMessage) { this.errorMessage = errorMessage; } + + public Long getUserId() { return userId; } + public void setUserId(Long userId) { this.userId = userId; } + + public String getUserLogin() { return userLogin; } + public void setUserLogin(String userLogin) { this.userLogin = userLogin; } + + public Instant getStartedAt() { return startedAt; } + public void setStartedAt(Instant startedAt) { this.startedAt = startedAt; } + + public Instant getFinishedAt() { return finishedAt; } + public void setFinishedAt(Instant finishedAt) { this.finishedAt = finishedAt; } +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java new file mode 100644 index 0000000000..5423a0ad25 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java @@ -0,0 +1,11 @@ +package org.ohdsi.webapi.mcp.audit; + +import java.util.List; +import org.springframework.data.jpa.repository.JpaRepository; +import org.springframework.stereotype.Repository; + +@Repository +public interface McpAuditEntryRepository extends JpaRepository { + + List findByTraceIdOrderByStartedAtAscIdAsc(String traceId); +} diff --git a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java new file mode 100644 index 0000000000..bcaa6aa773 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java @@ -0,0 +1,69 @@ +package org.ohdsi.webapi.mcp.audit; + +import java.time.Instant; +import java.util.List; +import org.ohdsi.webapi.mcp.TraceContext; +import org.ohdsi.webapi.security.authz.AuthorizationService; +import org.ohdsi.webapi.security.identity.WebApiPrincipal; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; +import org.springframework.stereotype.Service; + +/** + * Writes the MCP audit trail. {@link #start} must succeed before a tool runs, so a tool never + * changes anything that is not already on record. + */ +@Service +@ConditionalOnProperty(prefix = "spring.ai.mcp.server", name = "enabled", havingValue = "true") +public class McpAuditService { + + /** Results larger than this are cut, so one big SQL rendering cannot bloat the table. */ + static final int MAX_RESULT_CHARS = 200_000; + + private final McpAuditEntryRepository repository; + private final AuthorizationService authorizationService; + + public McpAuditService(McpAuditEntryRepository repository, AuthorizationService authorizationService) { + this.repository = repository; + this.authorizationService = authorizationService; + } + + public McpAuditEntry start(TraceContext trace, String toolName, String reasoningStep, String argumentsJson, + String clinicalRationale, String evidenceReference) { + WebApiPrincipal principal = authorizationService.getAuthenticatedPrincipal(); + McpAuditEntry entry = new McpAuditEntry(); + entry.setTraceId(trace.traceId()); + entry.setParentSpanId(trace.parentSpanId()); + entry.setSpanId(trace.spanId()); + entry.setTraceGenerated(trace.generated()); + entry.setReasoningStep(reasoningStep); + entry.setToolName(toolName); + entry.setArgumentsJson(argumentsJson); + entry.setClinicalRationale(clinicalRationale); + entry.setEvidenceReference(evidenceReference); + entry.setOutcome(McpAuditEntry.Outcome.STARTED); + entry.setUserId(principal.isAnonymous() ? null : principal.getUserId()); + entry.setUserLogin(principal.getName()); + entry.setStartedAt(Instant.now()); + return repository.saveAndFlush(entry); + } + + public McpAuditEntry finish(McpAuditEntry entry, McpAuditEntry.Outcome outcome, String resultJson, + String errorMessage) { + entry.setOutcome(outcome); + entry.setResultJson(truncate(resultJson)); + entry.setErrorMessage(errorMessage); + entry.setFinishedAt(Instant.now()); + return repository.saveAndFlush(entry); + } + + public List findByTrace(String traceId) { + return repository.findByTraceIdOrderByStartedAtAscIdAsc(traceId); + } + + static String truncate(String value) { + if (value == null || value.length() <= MAX_RESULT_CHARS) { + return value; + } + return value.substring(0, MAX_RESULT_CHARS) + "...[truncated " + (value.length() - MAX_RESULT_CHARS) + " chars]"; + } +} diff --git a/src/main/resources/application.yaml b/src/main/resources/application.yaml index 4129e720ba..701f4c03a9 100644 --- a/src/main/resources/application.yaml +++ b/src/main/resources/application.yaml @@ -104,6 +104,11 @@ logging: orm: info web: info root: info +mcp: + # Refuse MCP tool calls that do not carry _meta.traceparent (otherwise a new trace is started). + requireTraceparent: false + # Minimum length of clinical_rationale on tools that require one. + minRationaleLength: 20 organization: name: OHDSI person: @@ -336,6 +341,24 @@ sparql: endpoint: http://virtuoso.ohdsi.org:8890/sparql?default-graph-uri=&query= spring: + ai: + mcp: + server: + # MCP server for AI agents (tools for vocabulary search, concept sets, cohort definitions, + # SQL validation and generation; every call is recorded in mcp_audit_log). + # Served at the streamable-http endpoint below, behind the same authentication as the REST API. + enabled: false + protocol: STATELESS + name: ohdsi-webapi + # WebAPI's tools are plain tool specifications; the annotation scanner is not needed. + annotation-scanner: + enabled: false + capabilities: + resource: false + prompt: false + completion: false + streamable-http: + mcp-endpoint: /mcp main: lazy-initialization: true autoconfigure: diff --git a/src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql b/src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql new file mode 100644 index 0000000000..3c8cfe27c1 --- /dev/null +++ b/src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql @@ -0,0 +1,69 @@ +-- Audit trail for tool calls made through the MCP endpoint (/mcp). +-- +-- Every tool call is written here before it runs (outcome STARTED) and completed once it +-- finishes, so no tool can change a concept set or cohort without leaving a row behind. +-- Rows are append-only: the trigger below allows exactly one update per row (STARTED -> +-- final outcome, touching only the result columns) and rejects every delete. +CREATE SEQUENCE IF NOT EXISTS ${ohdsiSchema}.mcp_audit_log_sequence; + +CREATE TABLE IF NOT EXISTS ${ohdsiSchema}.mcp_audit_log ( + id bigint DEFAULT nextval('${ohdsiSchema}.mcp_audit_log_sequence'::regclass) NOT NULL, + trace_id character varying(32) NOT NULL, + -- span id the client sent in _meta.traceparent (null when the client sent none) + parent_span_id character varying(16), + -- span id WebAPI allocated for this tool call; returned to the client in _meta.traceparent + span_id character varying(16) NOT NULL, + -- true when the client sent no traceparent and WebAPI started a new trace + trace_generated boolean NOT NULL, + reasoning_step character varying(64), + tool_name character varying(128) NOT NULL, + arguments_json text, + clinical_rationale text, + evidence_reference text, + outcome character varying(16) NOT NULL, + result_json text, + error_message text, + user_id bigint, + user_login character varying(255), + started_at timestamp with time zone NOT NULL, + finished_at timestamp with time zone, + CONSTRAINT pk_mcp_audit_log PRIMARY KEY (id), + CONSTRAINT chk_mcp_audit_log_outcome + CHECK (outcome IN ('STARTED', 'SUCCESS', 'ERROR', 'REJECTED', 'DENIED')) +); + +CREATE INDEX IF NOT EXISTS mcp_audit_log_trace_idx + ON ${ohdsiSchema}.mcp_audit_log USING btree (trace_id, started_at); + +CREATE OR REPLACE FUNCTION ${ohdsiSchema}.mcp_audit_log_append_only() +RETURNS trigger AS $$ +BEGIN + IF TG_OP = 'DELETE' THEN + RAISE EXCEPTION 'mcp_audit_log is append-only: rows cannot be deleted'; + END IF; + IF OLD.outcome <> 'STARTED' THEN + RAISE EXCEPTION 'mcp_audit_log is append-only: entry % is already completed', OLD.id; + END IF; + IF NEW.id IS DISTINCT FROM OLD.id + OR NEW.trace_id IS DISTINCT FROM OLD.trace_id + OR NEW.parent_span_id IS DISTINCT FROM OLD.parent_span_id + OR NEW.span_id IS DISTINCT FROM OLD.span_id + OR NEW.trace_generated IS DISTINCT FROM OLD.trace_generated + OR NEW.reasoning_step IS DISTINCT FROM OLD.reasoning_step + OR NEW.tool_name IS DISTINCT FROM OLD.tool_name + OR NEW.arguments_json IS DISTINCT FROM OLD.arguments_json + OR NEW.clinical_rationale IS DISTINCT FROM OLD.clinical_rationale + OR NEW.evidence_reference IS DISTINCT FROM OLD.evidence_reference + OR NEW.user_id IS DISTINCT FROM OLD.user_id + OR NEW.user_login IS DISTINCT FROM OLD.user_login + OR NEW.started_at IS DISTINCT FROM OLD.started_at THEN + RAISE EXCEPTION 'mcp_audit_log is append-only: only the outcome of entry % may be recorded', OLD.id; + END IF; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +DROP TRIGGER IF EXISTS mcp_audit_log_append_only ON ${ohdsiSchema}.mcp_audit_log; +CREATE TRIGGER mcp_audit_log_append_only + BEFORE UPDATE OR DELETE ON ${ohdsiSchema}.mcp_audit_log + FOR EACH ROW EXECUTE FUNCTION ${ohdsiSchema}.mcp_audit_log_append_only(); diff --git a/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java b/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java new file mode 100644 index 0000000000..df738cc130 --- /dev/null +++ b/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java @@ -0,0 +1,153 @@ +package org.ohdsi.webapi.mcp; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.fasterxml.jackson.databind.ObjectMapper; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpSchema.CallToolRequest; +import io.modelcontextprotocol.spec.McpSchema.CallToolResult; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.Before; +import org.junit.Test; +import org.ohdsi.webapi.mcp.ToolDefinition.Rationale; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry.Outcome; +import org.ohdsi.webapi.mcp.audit.McpAuditService; +import org.springframework.security.access.AccessDeniedException; + +public class McpToolFactoryTest { + + private static final String TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736"; + private static final String TRACEPARENT = "00-" + TRACE_ID + "-00f067aa0ba902b7-01"; + + private McpAuditService audit; + private McpAuditEntry entry; + + @Before + public void setUp() { + audit = mock(McpAuditService.class); + entry = new McpAuditEntry(); + entry.setId(42L); + when(audit.start(any(), anyString(), anyString(), any(), any(), any())).thenReturn(entry); + } + + private McpToolFactory factory(boolean requireTraceparent) { + return new McpToolFactory(audit, new ObjectMapper(), requireTraceparent, 20); + } + + private static ToolDefinition guarded(AtomicBoolean ran) { + return ToolDefinition.tool("add_concept_to_set") + .description("test") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .rationale(Rationale.REQUIRED) + .handler(args -> { + ran.set(true); + return Map.of("ok", true); + }) + .build(); + } + + private static CallToolRequest request(Map args, Map meta) { + return new CallToolRequest("add_concept_to_set", args, meta); + } + + @Test + public void runsTheToolAndReturnsTheToolSpanInMeta() { + AtomicBoolean ran = new AtomicBoolean(); + CallToolResult result = factory(false).call(guarded(ran), request( + Map.of("clinical_rationale", "Core diagnosis concept for the phenotype.", + "evidence_reference", "SNOMED-CT hierarchy"), + Map.of("traceparent", TRACEPARENT, "reasoning_step", "STEP_2_CONCEPT_SELECTION"))); + + assertTrue(ran.get()); + assertFalse(result.isError()); + String traceparent = (String) result.meta().get("traceparent"); + assertTrue(traceparent, traceparent.startsWith("00-" + TRACE_ID + "-")); + assertEquals(42L, result.meta().get("audit_entry_id")); + verify(audit).start(any(), eq("add_concept_to_set"), eq("STEP_2_CONCEPT_SELECTION"), anyString(), + eq("Core diagnosis concept for the phenotype."), eq("SNOMED-CT hierarchy")); + verify(audit).finish(eq(entry), eq(Outcome.SUCCESS), eq("{\"ok\":true}"), isNull()); + } + + @Test + public void refusesACallWithoutRationaleButStillAuditsIt() { + AtomicBoolean ran = new AtomicBoolean(); + CallToolResult result = factory(false).call(guarded(ran), request(Map.of(), Map.of("traceparent", TRACEPARENT))); + + assertFalse(ran.get()); + assertTrue(result.isError()); + verify(audit).finish(eq(entry), eq(Outcome.REJECTED), isNull(), any()); + } + + @Test + public void refusesATooShortRationale() { + AtomicBoolean ran = new AtomicBoolean(); + CallToolResult result = factory(false).call(guarded(ran), request( + Map.of("clinical_rationale", "fits", "evidence_reference", "SNOMED"), Map.of())); + + assertFalse(ran.get()); + assertTrue(result.isError()); + } + + @Test + public void refusesAMissingTraceparentWhenRequired() { + AtomicBoolean ran = new AtomicBoolean(); + CallToolResult result = factory(true).call(guarded(ran), request( + Map.of("clinical_rationale", "Core diagnosis concept for the phenotype.", + "evidence_reference", "SNOMED-CT hierarchy"), + Map.of())); + + assertFalse(ran.get()); + assertTrue(result.isError()); + verify(audit).finish(eq(entry), eq(Outcome.REJECTED), isNull(), any()); + } + + @Test + public void doesNotRunTheToolWhenTheAuditTrailCannotBeWritten() { + when(audit.start(any(), anyString(), anyString(), any(), any(), any())) + .thenThrow(new IllegalStateException("database down")); + AtomicBoolean ran = new AtomicBoolean(); + CallToolResult result = factory(false).call(guarded(ran), request( + Map.of("clinical_rationale", "Core diagnosis concept for the phenotype.", + "evidence_reference", "SNOMED-CT hierarchy"), + Map.of())); + + assertFalse(ran.get()); + assertTrue(result.isError()); + verify(audit, never()).finish(any(), any(), any(), any()); + } + + @Test + public void reportsAuthorizationFailuresAsDenied() { + ToolDefinition def = ToolDefinition.tool("get_concept_set") + .description("test") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .handler(args -> { + throw new AccessDeniedException("nope"); + }) + .build(); + CallToolResult result = factory(false).call(def, new CallToolRequest("get_concept_set", Map.of(), null)); + + assertTrue(result.isError()); + verify(audit).finish(eq(entry), eq(Outcome.DENIED), isNull(), any()); + } + + @Test + public void advertisesRationaleAsRequiredInTheInputSchema() { + McpSchema.Tool tool = factory(false).describe(guarded(new AtomicBoolean())); + assertTrue(tool.inputSchema().required().contains("clinical_rationale")); + assertTrue(tool.inputSchema().required().contains("evidence_reference")); + assertFalse(tool.annotations().readOnlyHint()); + } +} diff --git a/src/test/java/org/ohdsi/webapi/mcp/TraceContextTest.java b/src/test/java/org/ohdsi/webapi/mcp/TraceContextTest.java new file mode 100644 index 0000000000..c39a40ac5f --- /dev/null +++ b/src/test/java/org/ohdsi/webapi/mcp/TraceContextTest.java @@ -0,0 +1,57 @@ +package org.ohdsi.webapi.mcp; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; + +public class TraceContextTest { + + private static final String TRACE_ID = "4bf92f3577b34da6a3ce929d0e0e4736"; + private static final String PARENT = "00f067aa0ba902b7"; + + @Test + public void continuesTheCallersTraceWithANewSpan() { + TraceContext t = TraceContext.continueFrom("00-" + TRACE_ID + "-" + PARENT + "-01"); + + assertEquals(TRACE_ID, t.traceId()); + assertEquals(PARENT, t.parentSpanId()); + assertNotEquals(PARENT, t.spanId()); + assertTrue(t.spanId().matches("[0-9a-f]{16}")); + assertFalse(t.generated()); + assertEquals("00-" + TRACE_ID + "-" + t.spanId() + "-01", t.traceparent()); + } + + @Test + public void keepsTheSampledFlagAndAcceptsUpperCase() { + TraceContext t = TraceContext.continueFrom(("00-" + TRACE_ID + "-" + PARENT + "-00").toUpperCase()); + assertEquals(TRACE_ID, t.traceId()); + assertTrue(t.traceparent().endsWith("-00")); + } + + @Test + public void newRootHasNoParent() { + TraceContext t = TraceContext.newRoot(); + assertTrue(t.traceId().matches("[0-9a-f]{32}")); + assertNull(t.parentSpanId()); + assertTrue(t.generated()); + } + + @Test(expected = IllegalArgumentException.class) + public void rejectsMalformedTraceparent() { + TraceContext.continueFrom("00-" + TRACE_ID + "-" + PARENT); + } + + @Test(expected = IllegalArgumentException.class) + public void rejectsAllZeroTraceId() { + TraceContext.continueFrom("00-" + "0".repeat(32) + "-" + PARENT + "-01"); + } + + @Test(expected = IllegalArgumentException.class) + public void rejectsVersionFf() { + TraceContext.continueFrom("ff-" + TRACE_ID + "-" + PARENT + "-01"); + } +} diff --git a/src/test/java/org/ohdsi/webapi/test/McpServerIT.java b/src/test/java/org/ohdsi/webapi/test/McpServerIT.java new file mode 100644 index 0000000000..65b4a5bb64 --- /dev/null +++ b/src/test/java/org/ohdsi/webapi/test/McpServerIT.java @@ -0,0 +1,347 @@ +package org.ohdsi.webapi.test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializationFeature; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import org.junit.Before; +import org.junit.Test; +import org.ohdsi.webapi.source.SourceRepository; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.web.client.TestRestTemplate; +import org.springframework.http.HttpEntity; +import org.springframework.http.HttpHeaders; +import org.springframework.http.HttpMethod; +import org.springframework.http.HttpStatus; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.test.context.TestPropertySource; + +/** + * Drives the MCP endpoint over HTTP the way an agent would: JSON-RPC tools/list and tools/call on + * POST /mcp, authenticated with a WebAPI JWT, against the embedded CDM. + */ +@TestPropertySource(properties = "spring.ai.mcp.server.enabled=true") +public class McpServerIT extends WebApiIT { + + private static final long T2DM = 201826L; + private static final long NON_STANDARD = 44054006L; + + private final ObjectMapper mapper = new ObjectMapper(); + + @Autowired + private SourceRepository sourceRepository; + + @Autowired + private ObjectMapper appObjectMapper; + + private int rpcId; + + @Before + public void init() throws Exception { + truncateTable(String.format("%s.%s", getOhdsiSchema(), "source")); + resetSequence(String.format("%s.%s", getOhdsiSchema(), "source_sequence")); + sourceRepository.saveAndFlush(getCdmSource()); + prepareCdmSchema(); + prepareResultSchema(); + jdbcTemplate.execute("insert into " + CDM_SCHEMA_NAME + ".concept (concept_id, concept_name, domain_id, " + + "vocabulary_id, concept_class_id, standard_concept, concept_code, valid_start_date, valid_end_date, " + + "invalid_reason) values " + + "(" + T2DM + ", 'Type 2 diabetes mellitus', 'Condition', 'SNOMED', 'Clinical Finding', 'S', " + + "'44054006', '1970-01-01', '2099-12-31', null), " + + "(" + NON_STANDARD + ", 'Diabetes mellitus type 2', 'Condition', 'ICD10CM', '3-char nonbill code', " + + "null, 'E11', '1970-01-01', '2099-12-31', null)"); + } + + @Test + public void listsToolsAndRequiresRationaleOnConceptSetChanges() throws Exception { + JsonNode result = rpc("tools/list", mapper.createObjectNode()).path("result"); + List names = new ArrayList<>(); + result.path("tools").forEach(t -> names.add(t.path("name").asText())); + assertTrue(names.toString(), names.containsAll(List.of("list_sources", "search_concepts", "get_concept_details", + "get_concept_hierarchy", "create_concept_set", "add_concept_to_set", "validate_cohort_definition", + "generate_cohort", "get_cohort_generation_status", "get_audit_trail"))); + + JsonNode add = null; + for (JsonNode t : result.path("tools")) { + if ("add_concept_to_set".equals(t.path("name").asText())) { + add = t; + } + } + List required = new ArrayList<>(); + add.path("inputSchema").path("required").forEach(r -> required.add(r.asText())); + assertTrue(required.toString(), required.containsAll( + List.of("concept_id", "concept_name", "include_descendants", "clinical_rationale", "evidence_reference"))); + } + + @Test + public void searchContinuesTheCallersTraceAndIsAudited() throws Exception { + String traceId = newTraceId(); + String clientSpan = "00f067aa0ba902b7"; + JsonNode result = callTool("search_concepts", + Map.of("source_key", SOURCE_KEY, "query", "diabetes", "standard_concept", "S"), + traceparent(traceId, clientSpan), "STEP_1_CONCEPT_EXPLORATION"); + + assertFalse(result.toString(), result.path("isError").asBoolean()); + JsonNode body = mapper.readTree(result.path("content").get(0).path("text").asText()); + assertEquals(1, body.path("total_matches").asInt()); + assertEquals(T2DM, body.path("concepts").get(0).path("concept_id").asLong()); + + String returned = result.path("_meta").path("traceparent").asText(); + assertTrue(returned, returned.startsWith("00-" + traceId + "-")); + assertNotEquals("the tool call gets its own span", clientSpan, returned.split("-")[2]); + + Map row = jdbcTemplate.queryForMap("select * from " + getOhdsiSchema() + + ".mcp_audit_log where trace_id = ?", traceId); + assertEquals("SUCCESS", row.get("outcome")); + assertEquals(clientSpan, row.get("parent_span_id")); + assertEquals(returned.split("-")[2], row.get("span_id")); + assertEquals("STEP_1_CONCEPT_EXPLORATION", row.get("reasoning_step")); + assertEquals("search_concepts", row.get("tool_name")); + } + + @Test + public void addConceptRefusesAConceptWhoseNameDoesNotMatchTheVocabulary() throws Exception { + String traceId = newTraceId(); + int conceptSetId = createConceptSet(traceId); + + JsonNode result = callTool("add_concept_to_set", Map.of( + "concept_set_id", conceptSetId, + "source_key", SOURCE_KEY, + "concept_id", T2DM, + "concept_name", "Type 1 diabetes mellitus", + "include_descendants", true, + "clinical_rationale", "Core diagnosis concept for the target diabetes phenotype.", + "evidence_reference", "SNOMED-CT clinical finding hierarchy"), + traceparent(traceId, "1111111111111111"), null); + + assertTrue(result.isEmpty() ? "no result" : result.toString(), result.path("isError").asBoolean()); + assertTrue(result.toString(), text(result).contains("Type 2 diabetes mellitus")); + assertEquals(0, itemCount(conceptSetId)); + assertEquals("REJECTED", lastOutcome(traceId, "add_concept_to_set")); + } + + @Test + public void addConceptRefusesACallWithoutClinicalRationale() throws Exception { + String traceId = newTraceId(); + int conceptSetId = createConceptSet(traceId); + + JsonNode result = callTool("add_concept_to_set", Map.of( + "concept_set_id", conceptSetId, + "source_key", SOURCE_KEY, + "concept_id", T2DM, + "concept_name", "Type 2 diabetes mellitus", + "include_descendants", true), + traceparent(traceId, "2222222222222222"), null); + + assertTrue(result.path("isError").asBoolean()); + assertTrue(text(result), text(result).contains("clinical_rationale")); + assertEquals(0, itemCount(conceptSetId)); + assertEquals("REJECTED", lastOutcome(traceId, "add_concept_to_set")); + } + + @Test + public void addConceptSavesAVerifiedConceptAndTheTrailShowsTheRationale() throws Exception { + String traceId = newTraceId(); + int conceptSetId = createConceptSet(traceId); + String rationale = "Core primary diagnosis concept for adult-onset T2DM; descendants capture complications."; + + JsonNode result = callTool("add_concept_to_set", Map.of( + "concept_set_id", conceptSetId, + "source_key", SOURCE_KEY, + "concept_id", T2DM, + "concept_name", "type 2 diabetes mellitus", + "include_descendants", true, + "clinical_rationale", rationale, + "evidence_reference", "ADA Standards of Care; SNOMED-CT clinical finding hierarchy"), + traceparent(traceId, "3333333333333333"), "STEP_2_CONCEPT_SELECTION"); + + assertFalse(result.toString(), result.path("isError").asBoolean()); + assertEquals(1, itemCount(conceptSetId)); + + JsonNode trail = mapper.readTree(text(callTool("get_audit_trail", Map.of("trace_id", traceId), + traceparent(traceId, "4444444444444444"), "AUDIT"))); + JsonNode log = trail.path("decision_log"); + assertEquals("create_concept_set", log.get(0).path("tool").asText()); + assertEquals("add_concept_to_set", log.get(1).path("tool").asText()); + assertEquals(rationale, log.get(1).path("clinical_rationale").asText()); + assertEquals("STEP_2_CONCEPT_SELECTION", log.get(1).path("reasoning_step").asText()); + assertEquals("SUCCESS", log.get(1).path("outcome").asText()); + } + + @Test + public void auditRowsCannotBeChangedOrDeleted() throws Exception { + String traceId = newTraceId(); + callTool("list_sources", Map.of(), traceparent(traceId, "5555555555555555"), null); + String table = getOhdsiSchema() + ".mcp_audit_log"; + + try { + jdbcTemplate.update("update " + table + " set clinical_rationale = 'rewritten' where trace_id = ?", traceId); + fail("a completed audit row must not be updatable"); + } catch (Exception expected) { + assertTrue(expected.getMessage(), expected.getMessage().contains("append-only")); + } + try { + jdbcTemplate.update("delete from " + table + " where trace_id = ?", traceId); + fail("audit rows must not be deletable"); + } catch (Exception expected) { + assertTrue(expected.getMessage(), expected.getMessage().contains("append-only")); + } + } + + @Test + public void validationTranslatesSqlAndFlagsUnverifiedConcepts() throws Exception { + String traceId = newTraceId(); + JsonNode ok = mapper.readTree(text(callTool("validate_cohort_definition", Map.of( + "source_key", SOURCE_KEY, + "expression", cohortExpression("Type 2 diabetes mellitus"), + "target_dialect", "sql server", + "include_sql", true), + traceparent(traceId, "6666666666666666"), null))); + assertTrue(ok.toString(), ok.path("valid").asBoolean()); + assertTrue(ok.path("sql").path("translated").asBoolean()); + assertTrue(ok.path("sql").path("sql").asText().contains("#Codesets")); + + JsonNode bad = mapper.readTree(text(callTool("validate_cohort_definition", Map.of( + "source_key", SOURCE_KEY, + "expression", cohortExpression("Gestational diabetes")), + traceparent(traceId, "7777777777777777"), null))); + assertFalse(bad.path("valid").asBoolean()); + assertEquals(1, bad.path("concept_issues").size()); + } + + @Test + public void createCohortDefinitionRefusesUnverifiedConcepts() throws Exception { + String traceId = newTraceId(); + JsonNode result = callTool("create_cohort_definition", Map.of( + "name", "MCP IT cohort " + UUID.randomUUID(), + "source_key", SOURCE_KEY, + "expression", cohortExpression("Gestational diabetes"), + "clinical_rationale", "Adults with a first recorded T2DM diagnosis as index event.", + "evidence_reference", "ADA Standards of Care"), + traceparent(traceId, "8888888888888888"), null); + assertTrue(result.path("isError").asBoolean()); + assertTrue(text(result), text(result).contains("nothing was saved")); + assertEquals("REJECTED", lastOutcome(traceId, "create_cohort_definition")); + } + + @Test + public void callWithoutTraceparentStartsANewTrace() throws Exception { + JsonNode result = callTool("list_sources", Map.of(), null, null); + assertFalse(result.path("isError").asBoolean()); + String traceId = result.path("_meta").path("traceparent").asText().split("-")[1]; + Boolean generated = jdbcTemplate.queryForObject("select trace_generated from " + getOhdsiSchema() + + ".mcp_audit_log where trace_id = ?", Boolean.class, traceId); + assertTrue(generated); + } + + @Test + public void enablingMcpDoesNotReplaceTheApplicationObjectMapper() { + // Spring AI registers its own non-default-candidate "mcpServerObjectMapper" (dates as ISO strings); + // the REST API must keep the Boot mapper configured by spring.jackson (dates as timestamps). + assertTrue(appObjectMapper.isEnabled(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS)); + } + + @Test + public void endpointRequiresTheMcpAcceptHeaders() { + HttpHeaders headers = new HttpHeaders(); + headers.setContentType(MediaType.APPLICATION_JSON); + headers.setAccept(List.of(MediaType.APPLICATION_JSON)); + ResponseEntity response = getRestTemplate().exchange(getBaseUri() + "/mcp", HttpMethod.POST, + new HttpEntity<>("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"tools/list\"}", headers), String.class); + assertEquals(HttpStatus.BAD_REQUEST, response.getStatusCode()); + } + + // ---- helpers ------------------------------------------------------------------------ + + private int createConceptSet(String traceId) throws Exception { + JsonNode result = callTool("create_concept_set", Map.of( + "name", "MCP IT concept set " + UUID.randomUUID(), + "clinical_rationale", "Concept set for the adult-onset type 2 diabetes phenotype.", + "evidence_reference", "ADA Standards of Care"), + traceparent(traceId, "aaaaaaaaaaaaaaaa"), null); + assertFalse(result.toString(), result.path("isError").asBoolean()); + return mapper.readTree(text(result)).path("concept_set_id").asInt(); + } + + private Map cohortExpression(String conceptName) throws Exception { + String json = "{\"ConceptSets\":[{\"id\":0,\"name\":\"T2DM\",\"expression\":{\"items\":[{\"concept\":" + + "{\"CONCEPT_ID\":" + T2DM + ",\"CONCEPT_NAME\":\"" + conceptName + "\"},\"includeDescendants\":true}]}}]," + + "\"PrimaryCriteria\":{\"CriteriaList\":[{\"ConditionOccurrence\":{\"CodesetId\":0}}]," + + "\"ObservationWindow\":{\"PriorDays\":0,\"PostDays\":0},\"PrimaryCriteriaLimit\":{\"Type\":\"First\"}}," + + "\"QualifiedLimit\":{\"Type\":\"First\"},\"ExpressionLimit\":{\"Type\":\"First\"}," + + "\"InclusionRules\":[],\"CensoringCriteria\":[],\"CollapseSettings\":{\"CollapseType\":\"ERA\",\"EraPad\":0}}"; + @SuppressWarnings("unchecked") + Map m = mapper.readValue(json, Map.class); + return m; + } + + private int itemCount(int conceptSetId) { + return jdbcTemplate.queryForObject("select count(*) from " + getOhdsiSchema() + + ".concept_set_item where concept_set_id = ?", Integer.class, conceptSetId); + } + + private String lastOutcome(String traceId, String tool) { + return jdbcTemplate.queryForObject("select outcome from " + getOhdsiSchema() + + ".mcp_audit_log where trace_id = ? and tool_name = ? order by id desc limit 1", + String.class, traceId, tool); + } + + private static String newTraceId() { + return UUID.randomUUID().toString().replace("-", ""); + } + + private static String traceparent(String traceId, String spanId) { + return "00-" + traceId + "-" + spanId + "-01"; + } + + private static String text(JsonNode toolResult) { + return toolResult.path("content").get(0).path("text").asText(); + } + + private JsonNode callTool(String name, Map arguments, String traceparent, String step) + throws Exception { + ObjectNode params = mapper.createObjectNode(); + params.put("name", name); + params.set("arguments", mapper.valueToTree(arguments)); + if (traceparent != null || step != null) { + ObjectNode meta = params.putObject("_meta"); + if (traceparent != null) { + meta.put("traceparent", traceparent); + } + if (step != null) { + meta.put("reasoning_step", step); + } + } + JsonNode response = rpc("tools/call", params); + assertTrue(response.toString(), response.has("result")); + return response.path("result"); + } + + private JsonNode rpc(String method, ObjectNode params) throws Exception { + ObjectNode request = mapper.createObjectNode(); + request.put("jsonrpc", "2.0"); + request.put("id", ++rpcId); + request.put("method", method); + request.set("params", params); + + HttpHeaders headers = new HttpHeaders(); + headers.setContentType(MediaType.APPLICATION_JSON); + headers.setAccept(List.of(MediaType.APPLICATION_JSON, MediaType.TEXT_EVENT_STREAM)); + TestRestTemplate rest = getRestTemplate(); + ResponseEntity response = rest.exchange(getBaseUri() + "/mcp", HttpMethod.POST, + new HttpEntity<>(mapper.writeValueAsString(request), headers), String.class); + assertEquals(response.getBody(), HttpStatus.OK, response.getStatusCode()); + return mapper.readTree(response.getBody()); + } +} From b3dee5f9fe2ee67bcfe23ac08a3398250a171e98 Mon Sep 17 00:00:00 2001 From: Daijiro Wachi Date: Wed, 4 Nov 2026 21:00:00 +0900 Subject: [PATCH 2/4] Test generate_cohort end to end on the bundled Eunomia CDM McpServerIT streams person, observation_period, condition_occurrence, concept and concept_ancestor from src/test/resources/eunomia.csv.zip into the embedded CDM with COPY. Through MCP it creates a cohort for the first gastrointestinal hemorrhage with an "Age 40+" inclusion rule, generates it, and polls get_cohort_generation_status until COMPLETE. The counts and attrition it reports (479 at base, 204 kept by the rule) are checked against the same cohort counted with plain SQL on the CDM, and against the subjects written to results.cohort. --- .../org/ohdsi/webapi/test/McpServerIT.java | 112 +++++++++++++++++- 1 file changed, 110 insertions(+), 2 deletions(-) diff --git a/src/test/java/org/ohdsi/webapi/test/McpServerIT.java b/src/test/java/org/ohdsi/webapi/test/McpServerIT.java index 65b4a5bb64..261654ae58 100644 --- a/src/test/java/org/ohdsi/webapi/test/McpServerIT.java +++ b/src/test/java/org/ohdsi/webapi/test/McpServerIT.java @@ -10,13 +10,24 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializationFeature; import com.fasterxml.jackson.databind.node.ObjectNode; +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.sql.Connection; import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.UUID; +import java.util.zip.ZipEntry; +import java.util.zip.ZipInputStream; import org.junit.Before; import org.junit.Test; import org.ohdsi.webapi.source.SourceRepository; +import org.postgresql.PGConnection; +import org.postgresql.copy.CopyManager; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.test.web.client.TestRestTemplate; import org.springframework.http.HttpEntity; @@ -36,6 +47,10 @@ public class McpServerIT extends WebApiIT { private static final long T2DM = 201826L; private static final long NON_STANDARD = 44054006L; + private static final long GI_BLEED = 192671L; + /** People in Eunomia with a gastrointestinal hemorrhage inside an observation period, all ages / aged 40+. */ + private static final long EUNOMIA_GI_BLEED_PEOPLE = 479; + private static final long EUNOMIA_GI_BLEED_PEOPLE_40_PLUS = 204; private final ObjectMapper mapper = new ObjectMapper(); @@ -235,6 +250,74 @@ public void createCohortDefinitionRefusesUnverifiedConcepts() throws Exception { assertEquals("REJECTED", lastOutcome(traceId, "create_cohort_definition")); } + @Test + public void generatesACohortOnEunomiaAndReportsCountsAndAttrition() throws Exception { + loadEunomia("person", "observation_period", "condition_occurrence", "concept", "concept_ancestor"); + + // First gastrointestinal hemorrhage in an observation period; the inclusion rule keeps people aged 40+. + String traceId = newTraceId(); + Map expression = cohortExpression(GI_BLEED, "Gastrointestinal hemorrhage"); + expression.put("InclusionRules", List.of(Map.of("name", "Age 40+", "expression", Map.of( + "Type", "ALL", "CriteriaList", List.of(), "Groups", List.of(), + "DemographicCriteriaList", List.of(Map.of("Age", Map.of("Value", 40, "Op", "gte"))))))); + JsonNode created = callTool("create_cohort_definition", Map.of( + "name", "MCP IT GI bleed cohort " + UUID.randomUUID(), + "source_key", SOURCE_KEY, + "expression", expression, + "clinical_rationale", "Adults aged 40+ at their first recorded gastrointestinal hemorrhage.", + "evidence_reference", "OHDSI Eunomia GiBleed example study"), + traceparent(traceId, "9999999999999999"), "STEP_3_COHORT_CRITERIA"); + assertFalse(created.toString(), created.path("isError").asBoolean()); + int cohortId = mapper.readTree(text(created)).path("cohort_definition_id").asInt(); + + JsonNode queued = callTool("generate_cohort", Map.of("cohort_definition_id", cohortId, "source_key", SOURCE_KEY, + "clinical_rationale", "Generate to check the cohort size is plausible.", + "evidence_reference", "Phenotype review plan"), + traceparent(traceId, "bbbbbbbbbbbbbbbb"), "STEP_5_GENERATION_VERIFICATION"); + assertFalse(queued.toString(), queued.path("isError").asBoolean()); + + JsonNode status = null; + long deadline = System.currentTimeMillis() + 120_000; + while (System.currentTimeMillis() < deadline) { + status = mapper.readTree(text(callTool("get_cohort_generation_status", + Map.of("cohort_definition_id", cohortId, "source_key", SOURCE_KEY), + traceparent(traceId, "cccccccccccccccc"), "STEP_5_GENERATION_VERIFICATION"))); + if ("COMPLETE".equals(status.path("status").asText())) { + break; + } + Thread.sleep(500); + } + assertEquals(status.toString(), "COMPLETE", status.path("status").asText()); + assertTrue(status.toString(), status.path("is_valid").asBoolean()); + + // The same cohort counted independently with plain SQL on the CDM. + String firstEvents = "select co.person_id, min(co.condition_start_date) as index_date from " + + CDM_SCHEMA_NAME + ".condition_occurrence co join " + CDM_SCHEMA_NAME + ".observation_period op " + + "on op.person_id = co.person_id and co.condition_start_date between " + + "op.observation_period_start_date and op.observation_period_end_date " + + "where co.condition_concept_id = " + GI_BLEED + " group by co.person_id"; + long base = jdbcTemplate.queryForObject("select count(*) from (" + firstEvents + ") e", Long.class); + long aged40 = jdbcTemplate.queryForObject("select count(*) from (" + firstEvents + ") e join " + + CDM_SCHEMA_NAME + ".person p on p.person_id = e.person_id " + + "where extract(year from e.index_date) - p.year_of_birth >= 40", Long.class); + assertEquals(EUNOMIA_GI_BLEED_PEOPLE, base); + assertEquals(EUNOMIA_GI_BLEED_PEOPLE_40_PLUS, aged40); + + assertEquals(status.toString(), aged40, status.path("person_count").asLong()); + assertEquals(status.toString(), aged40, status.path("record_count").asLong()); + JsonNode attrition = status.path("attrition_all_events"); + assertEquals(status.toString(), base, attrition.path("base_count").asLong()); + assertEquals(status.toString(), aged40, attrition.path("final_count").asLong()); + JsonNode rule = attrition.path("inclusion_rules").get(0); + assertEquals("Age 40+", rule.path("name").asText()); + assertEquals(aged40, rule.path("count_satisfying").asLong()); + + long written = jdbcTemplate.queryForObject("select count(distinct subject_id) from " + RESULT_SCHEMA_NAME + + ".cohort where cohort_definition_id = ?", Long.class, cohortId); + assertEquals(aged40, written); + assertEquals("SUCCESS", lastOutcome(traceId, "generate_cohort")); + } + @Test public void callWithoutTraceparentStartsANewTrace() throws Exception { JsonNode result = callTool("list_sources", Map.of(), null, null); @@ -275,8 +358,12 @@ private int createConceptSet(String traceId) throws Exception { } private Map cohortExpression(String conceptName) throws Exception { - String json = "{\"ConceptSets\":[{\"id\":0,\"name\":\"T2DM\",\"expression\":{\"items\":[{\"concept\":" - + "{\"CONCEPT_ID\":" + T2DM + ",\"CONCEPT_NAME\":\"" + conceptName + "\"},\"includeDescendants\":true}]}}]," + return cohortExpression(T2DM, conceptName); + } + + private Map cohortExpression(long conceptId, String conceptName) throws Exception { + String json = "{\"ConceptSets\":[{\"id\":0,\"name\":\"Codes\",\"expression\":{\"items\":[{\"concept\":" + + "{\"CONCEPT_ID\":" + conceptId + ",\"CONCEPT_NAME\":\"" + conceptName + "\"},\"includeDescendants\":true}]}}]," + "\"PrimaryCriteria\":{\"CriteriaList\":[{\"ConditionOccurrence\":{\"CodesetId\":0}}]," + "\"ObservationWindow\":{\"PriorDays\":0,\"PostDays\":0},\"PrimaryCriteriaLimit\":{\"Type\":\"First\"}}," + "\"QualifiedLimit\":{\"Type\":\"First\"},\"ExpressionLimit\":{\"Type\":\"First\"}," @@ -286,6 +373,27 @@ private Map cohortExpression(String conceptName) throws Exceptio return m; } + /** Streams the given tables of the bundled Eunomia CDM (src/test/resources/eunomia.csv.zip) into the cdm schema. */ + private static void loadEunomia(String... tables) throws Exception { + Set wanted = new HashSet<>(Arrays.asList(tables)); + try (Connection connection = ITStarter.getDataSource().getConnection(); + ZipInputStream zip = new ZipInputStream(McpServerIT.class.getResourceAsStream("/eunomia.csv.zip"))) { + CopyManager copy = connection.unwrap(PGConnection.class).getCopyAPI(); + ZipEntry entry; + while ((entry = zip.getNextEntry()) != null) { + String table = entry.getName().replace(".csv", ""); + if (!wanted.remove(table)) { + continue; + } + BufferedReader csv = new BufferedReader(new InputStreamReader(zip, StandardCharsets.UTF_8)); + String columns = csv.readLine(); + copy.copyIn("copy " + CDM_SCHEMA_NAME + "." + table + " (" + columns + ") from stdin with (format csv)", + csv); + } + } + assertTrue("missing from eunomia.csv.zip: " + wanted, wanted.isEmpty()); + } + private int itemCount(int conceptSetId) { return jdbcTemplate.queryForObject("select count(*) from " + getOhdsiSchema() + ".concept_set_item where concept_set_id = ?", Integer.class, conceptSetId); From 7df7e7a961cc88d279bfff75db5f6a4e7b058c2b Mon Sep 17 00:00:00 2001 From: Daijiro Wachi Date: Wed, 4 Nov 2026 21:11:13 +0900 Subject: [PATCH 3/4] Record which artifact each MCP call was about, and serve its provenance Atlas needs to show, on a cohort definition or concept set, which agent calls built it and why. The audit log only had the trace, so finding the calls about one artifact meant searching result_json. - mcp_audit_log gets entity_type and entity_id (with an index). Tools that create, change, validate or generate a cohort definition or concept set declare their entity; the id is taken from the arguments before the call runs, or for a create from the result once it has. The append-only trigger lets the entity be filled in once, never changed. Read-only tools (get_*, counts, status polls) record none. - GET /cohortdefinition/{id}/mcp-provenance and GET /conceptset/{id}/mcp-provenance return those calls in order: tool, reasoning step, rationale, evidence, outcome, user, times and trace. They use the artifact's own read permission, leave out the argument and result payloads, and stay available when MCP is turned off. --- README.md | 1 + .../org/ohdsi/webapi/mcp/CohortMcpTools.java | 5 ++ .../ohdsi/webapi/mcp/ConceptSetMcpTools.java | 4 + .../org/ohdsi/webapi/mcp/McpToolFactory.java | 18 ++++- .../org/ohdsi/webapi/mcp/ToolDefinition.java | 19 +++++ .../ohdsi/webapi/mcp/audit/McpAuditEntry.java | 19 +++++ .../mcp/audit/McpAuditEntryRepository.java | 3 + .../webapi/mcp/audit/McpAuditService.java | 17 +++- .../mcp/audit/McpProvenanceController.java | 77 +++++++++++++++++++ .../V3.0.1.20261104000000__mcp_audit_log.sql | 18 ++++- .../ohdsi/webapi/mcp/McpToolFactoryTest.java | 45 +++++++++-- .../org/ohdsi/webapi/test/McpServerIT.java | 33 ++++++++ 12 files changed, 244 insertions(+), 15 deletions(-) create mode 100644 src/main/java/org/ohdsi/webapi/mcp/audit/McpProvenanceController.java diff --git a/README.md b/README.md index 157a1c6298..cde87553bd 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,7 @@ Traceability rules the server enforces: - **Verified concepts.** `add_concept_to_set`, `create_cohort_definition` and `update_cohort_definition` look up every concept id in the source vocabulary. If an id does not exist, or the name the caller gave does not match the vocabulary's name, the call is refused and nothing is saved. - **Tracing.** Send a W3C `traceparent` in the call's `_meta`, plus an optional `reasoning_step`. WebAPI keeps the trace id, records your span as the parent, and returns the tool call's own span in the result's `_meta.traceparent`, along with `_meta.audit_entry_id`. A call without a traceparent starts a new trace, unless `mcp.requireTraceparent: true`. - **Audit.** Every call is written to `mcp_audit_log` before it runs, then completed with its outcome: `SUCCESS`, `ERROR`, `REJECTED` or `DENIED`. A database trigger makes the table append-only. `get_audit_trail` returns a trace's entries in order, as the decision log for an audit manifest. +- **Provenance.** Calls that create, change, validate or generate a cohort definition or concept set are also recorded against that artifact. `GET /cohortdefinition/{id}/mcp-provenance` and `GET /conceptset/{id}/mcp-provenance` return that history (tool, reasoning step, rationale, evidence, outcome, user, time and trace id; not the argument or result payloads) to anyone who can read the artifact, so Atlas can show why an agent built it that way. These endpoints stay available when the MCP server is turned off. ## Geospatial support diff --git a/src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java b/src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java index cae98f5b90..61f5b839f7 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java +++ b/src/main/java/org/ohdsi/webapi/mcp/CohortMcpTools.java @@ -33,6 +33,7 @@ import org.ohdsi.webapi.cohortdefinition.InclusionRuleReport; import org.ohdsi.webapi.job.JobExecutionResource; import org.ohdsi.webapi.mcp.ToolDefinition.Rationale; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry.EntityType; import org.ohdsi.webapi.service.VocabularyService; import org.ohdsi.webapi.source.Source; import org.ohdsi.webapi.source.SourceDaimon.DaimonType; @@ -78,6 +79,7 @@ public List tools() { .requiredParam("expression", objectParam("The Circe cohort expression (ConceptSets, " + "PrimaryCriteria, InclusionRules, EndStrategy, ...), as JSON.")) .requiredParam("source_key", stringParam("Vocabulary source used to verify the concepts.")) + .entity(EntityType.COHORT_DEFINITION, "cohort_definition_id") .handler(this::create) .build(), tool("update_cohort_definition") @@ -92,6 +94,7 @@ public List tools() { .requiredParam("source_key", stringParam("Vocabulary source used to verify the concepts.")) .param("name", stringParam("New name; unchanged if omitted.")) .param("description", stringParam("New description; unchanged if omitted.")) + .entity(EntityType.COHORT_DEFINITION, "cohort_definition_id") .handler(this::update) .build(), tool("get_cohort_definition") @@ -118,6 +121,7 @@ public List tools() { + "'sql server', 'oracle', 'redshift', 'spark', 'snowflake', 'bigquery'. Defaults to " + "the source's dialect.")) .param("include_sql", booleanParam("Return the translated SQL. Default false.")) + .entity(EntityType.COHORT_DEFINITION, "cohort_definition_id") .handler(this::validate) .build(), tool("generate_cohort") @@ -129,6 +133,7 @@ public List tools() { .rationale(Rationale.OPTIONAL) .requiredParam("cohort_definition_id", integerParam("The cohort definition to generate.")) .requiredParam("source_key", stringParam("CDM source to generate on.")) + .entity(EntityType.COHORT_DEFINITION, "cohort_definition_id") .handler(this::generate) .build(), tool("get_cohort_generation_status") diff --git a/src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java b/src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java index 0ef54c2d6c..52be547f0e 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java +++ b/src/main/java/org/ohdsi/webapi/mcp/ConceptSetMcpTools.java @@ -14,6 +14,7 @@ import org.ohdsi.webapi.conceptset.ConceptSetItem; import org.ohdsi.webapi.conceptset.ConceptSetService; import org.ohdsi.webapi.mcp.ToolDefinition.Rationale; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry.EntityType; import org.ohdsi.webapi.service.VocabularyService; import org.ohdsi.webapi.service.dto.ConceptSetDTO; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; @@ -50,6 +51,7 @@ public List tools() { .rationale(Rationale.REQUIRED) .requiredParam("name", stringParam("Concept set name; must be unique in WebAPI.")) .param("description", stringParam("What the concept set is meant to capture.")) + .entity(EntityType.CONCEPT_SET, "concept_set_id") .handler(this::create) .build(), tool("get_concept_set") @@ -78,6 +80,7 @@ public List tools() { .requiredParam("include_descendants", booleanParam("Include all descendants of the concept.")) .param("include_mapped", booleanParam("Include source concepts mapped to it. Default false.")) .param("is_excluded", booleanParam("Exclude rather than include the concept. Default false.")) + .entity(EntityType.CONCEPT_SET, "concept_set_id") .handler(this::addConcept) .build(), tool("remove_concept_from_set") @@ -88,6 +91,7 @@ public List tools() { .rationale(Rationale.REQUIRED) .requiredParam("concept_set_id", integerParam("The concept set to change.")) .requiredParam("concept_id", integerParam("The concept id to remove.")) + .entity(EntityType.CONCEPT_SET, "concept_set_id") .handler(this::removeConcept) .build(), tool("count_concept_set_concepts") diff --git a/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java b/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java index 836404eb9a..47a7719a08 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java +++ b/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java @@ -118,7 +118,9 @@ CallToolResult call(ToolDefinition def, CallToolRequest request) { McpAuditEntry entry; try { - entry = audit.start(trace, def.name, step, toJson(args.asMap()), rationale, evidence); + Integer entityId = def.entityType == null ? null : entityId(args.raw(def.entityIdField)); + entry = audit.start(trace, def.name, step, toJson(args.asMap()), rationale, evidence, + entityId == null ? null : def.entityType, entityId); } catch (Exception e) { log.error("MCP audit trail unavailable; refusing to run tool {}", def.name, e); return error(trace, null, Outcome.ERROR, "The audit trail could not be written, so the tool was not run."); @@ -151,7 +153,9 @@ CallToolResult call(ToolDefinition def, CallToolRequest request) { String resultJson = outcome == Outcome.SUCCESS ? toJson(result) : null; try { - audit.finish(entry, outcome, resultJson, message); + Integer createdId = def.entityType != null && result instanceof Map m + ? entityId(m.get(def.entityIdField)) : null; + audit.finish(entry, outcome, resultJson, message, createdId == null ? null : def.entityType, createdId); } catch (Exception e) { // The action already happened and its STARTED row is on record; do not hide the result. log.error("Could not record outcome {} of MCP audit entry {}", outcome, entry.getId(), e); @@ -167,6 +171,16 @@ CallToolResult call(ToolDefinition def, CallToolRequest request) { .build(); } + private static Integer entityId(Object value) { + if (value instanceof Number n) { + return n.intValue(); + } + if (value instanceof String s && s.matches("\\d{1,9}")) { + return Integer.valueOf(s); + } + return null; + } + private void checkRationale(ToolDefinition def, String rationale, String evidence) { if (def.rationale != Rationale.REQUIRED) { return; diff --git a/src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java b/src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java index d5b49e3229..ef3be9d783 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java +++ b/src/main/java/org/ohdsi/webapi/mcp/ToolDefinition.java @@ -5,6 +5,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry.EntityType; /** * Declarative description of one MCP tool: its name, JSON input schema, the reasoning step it @@ -34,6 +35,10 @@ public interface Handler { final Map properties; final List required; final Handler handler; + /** The artifact a call is about, recorded in the audit log so Atlas can show its provenance. */ + final EntityType entityType; + /** Argument, or for a create the result field, that holds the artifact's id. */ + final String entityIdField; private ToolDefinition(Builder b) { this.name = Objects.requireNonNull(b.name); @@ -46,6 +51,8 @@ private ToolDefinition(Builder b) { this.properties = b.properties; this.required = b.required; this.handler = Objects.requireNonNull(b.handler); + this.entityType = b.entityType; + this.entityIdField = b.entityIdField; } public String name() { @@ -98,6 +105,8 @@ public static final class Builder { private final Map properties = new LinkedHashMap<>(); private final List required = new ArrayList<>(); private Handler handler; + private EntityType entityType; + private String entityIdField; private Builder(String name) { this.name = name; @@ -144,6 +153,16 @@ public Builder requiredParam(String name, Map schema) { return this; } + /** + * Records calls to this tool against the cohort definition or concept set whose id is in + * {@code idField}: the argument when the caller passes it, otherwise the result field of that name. + */ + public Builder entity(EntityType entityType, String idField) { + this.entityType = Objects.requireNonNull(entityType); + this.entityIdField = Objects.requireNonNull(idField); + return this; + } + public Builder handler(Handler handler) { this.handler = handler; return this; diff --git a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java index 20fd7fe11b..12f41f4af3 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java +++ b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntry.java @@ -30,6 +30,12 @@ public enum Outcome { DENIED } + /** What kind of WebAPI artifact a call was about; see {@code entity_type} in {@code mcp_audit_log}. */ + public enum EntityType { + COHORT_DEFINITION, + CONCEPT_SET + } + @Id @Column(name = "id") @SequenceGenerator(name = "mcp_audit_log_seq", sequenceName = "mcp_audit_log_sequence", allocationSize = 1) @@ -79,6 +85,13 @@ public enum Outcome { @Column(name = "user_login", length = 255, updatable = false) private String userLogin; + @Enumerated(EnumType.STRING) + @Column(name = "entity_type", length = 32) + private EntityType entityType; + + @Column(name = "entity_id") + private Integer entityId; + @Column(name = "started_at", nullable = false, updatable = false) private Instant startedAt; @@ -135,4 +148,10 @@ public enum Outcome { public Instant getFinishedAt() { return finishedAt; } public void setFinishedAt(Instant finishedAt) { this.finishedAt = finishedAt; } + + public EntityType getEntityType() { return entityType; } + public void setEntityType(EntityType entityType) { this.entityType = entityType; } + + public Integer getEntityId() { return entityId; } + public void setEntityId(Integer entityId) { this.entityId = entityId; } } diff --git a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java index 5423a0ad25..b03ef42dac 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java +++ b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditEntryRepository.java @@ -8,4 +8,7 @@ public interface McpAuditEntryRepository extends JpaRepository { List findByTraceIdOrderByStartedAtAscIdAsc(String traceId); + + List findByEntityTypeAndEntityIdOrderByStartedAtAscIdAsc(McpAuditEntry.EntityType entityType, + Integer entityId); } diff --git a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java index bcaa6aa773..bc7b692442 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java +++ b/src/main/java/org/ohdsi/webapi/mcp/audit/McpAuditService.java @@ -28,7 +28,8 @@ public McpAuditService(McpAuditEntryRepository repository, AuthorizationService } public McpAuditEntry start(TraceContext trace, String toolName, String reasoningStep, String argumentsJson, - String clinicalRationale, String evidenceReference) { + String clinicalRationale, String evidenceReference, McpAuditEntry.EntityType entityType, + Integer entityId) { WebApiPrincipal principal = authorizationService.getAuthenticatedPrincipal(); McpAuditEntry entry = new McpAuditEntry(); entry.setTraceId(trace.traceId()); @@ -44,11 +45,23 @@ public McpAuditEntry start(TraceContext trace, String toolName, String reasoning entry.setUserId(principal.isAnonymous() ? null : principal.getUserId()); entry.setUserLogin(principal.getName()); entry.setStartedAt(Instant.now()); + if (entityType != null && entityId != null) { + entry.setEntityType(entityType); + entry.setEntityId(entityId); + } return repository.saveAndFlush(entry); } + /** + * Completes an entry. {@code entityType}/{@code entityId} are recorded only if the entry has no entity + * yet, which is the case for a create: the new artifact's id is known only once the tool has run. + */ public McpAuditEntry finish(McpAuditEntry entry, McpAuditEntry.Outcome outcome, String resultJson, - String errorMessage) { + String errorMessage, McpAuditEntry.EntityType entityType, Integer entityId) { + if (entry.getEntityId() == null && entityType != null && entityId != null) { + entry.setEntityType(entityType); + entry.setEntityId(entityId); + } entry.setOutcome(outcome); entry.setResultJson(truncate(resultJson)); entry.setErrorMessage(errorMessage); diff --git a/src/main/java/org/ohdsi/webapi/mcp/audit/McpProvenanceController.java b/src/main/java/org/ohdsi/webapi/mcp/audit/McpProvenanceController.java new file mode 100644 index 0000000000..a32b8acc6a --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/mcp/audit/McpProvenanceController.java @@ -0,0 +1,77 @@ +package org.ohdsi.webapi.mcp.audit; + +import java.time.Instant; +import java.util.Date; +import java.util.List; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry.EntityType; +import org.springframework.http.MediaType; +import org.springframework.security.access.prepost.PreAuthorize; +import org.springframework.transaction.annotation.Transactional; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.PathVariable; +import org.springframework.web.bind.annotation.RestController; + +/** + * Read-only provenance of cohort definitions and concept sets changed through the MCP server: every + * audited tool call about the artifact, in order, with the agent's stated rationale and the outcome. + * Anyone who can read the artifact can read its provenance. + * + *

Unlike the MCP endpoint this controller is not switched off with {@code spring.ai.mcp.server.enabled}, + * so the history stays visible after MCP is disabled. + */ +@RestController +@Transactional(readOnly = true) +public class McpProvenanceController { + + /** One audited tool call, without its arguments and result payloads. */ + public record ProvenanceEntry( + Long id, + String traceId, + String spanId, + String parentSpanId, + String reasoningStep, + String tool, + String outcome, + String clinicalRationale, + String evidenceReference, + String errorMessage, + String userLogin, + Date startedAt, + Date finishedAt) { + + static ProvenanceEntry of(McpAuditEntry e) { + return new ProvenanceEntry(e.getId(), e.getTraceId(), e.getSpanId(), e.getParentSpanId(), + e.getReasoningStep(), e.getToolName(), e.getOutcome().name(), e.getClinicalRationale(), + e.getEvidenceReference(), e.getErrorMessage(), e.getUserLogin(), date(e.getStartedAt()), + date(e.getFinishedAt())); + } + + private static Date date(Instant instant) { + return instant == null ? null : Date.from(instant); + } + } + + private final McpAuditEntryRepository repository; + + public McpProvenanceController(McpAuditEntryRepository repository) { + this.repository = repository; + } + + @GetMapping(value = "/cohortdefinition/{id}/mcp-provenance", produces = MediaType.APPLICATION_JSON_VALUE) + @PreAuthorize("isOwner(#id, COHORT_DEFINITION) or isAnyPermitted(anyOf('read:cohort-definition','write:cohort-definition')) or hasEntityAccess(#id, COHORT_DEFINITION, READ)") + public List getCohortDefinitionProvenance(@PathVariable("id") final int id) { + return find(EntityType.COHORT_DEFINITION, id); + } + + @GetMapping(value = "/conceptset/{id}/mcp-provenance", produces = MediaType.APPLICATION_JSON_VALUE) + @PreAuthorize("isOwner(#id, CONCEPT_SET) or isPermitted('read:conceptset') or isPermitted('write:conceptset') or hasEntityAccess(#id, CONCEPT_SET, READ)") + public List getConceptSetProvenance(@PathVariable("id") final int id) { + return find(EntityType.CONCEPT_SET, id); + } + + private List find(EntityType type, int id) { + return repository.findByEntityTypeAndEntityIdOrderByStartedAtAscIdAsc(type, id).stream() + .map(ProvenanceEntry::of) + .toList(); + } +} diff --git a/src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql b/src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql index 3c8cfe27c1..213860bbfc 100644 --- a/src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql +++ b/src/main/resources/db/migration/postgresql/V3.0.1.20261104000000__mcp_audit_log.sql @@ -3,7 +3,8 @@ -- Every tool call is written here before it runs (outcome STARTED) and completed once it -- finishes, so no tool can change a concept set or cohort without leaving a row behind. -- Rows are append-only: the trigger below allows exactly one update per row (STARTED -> --- final outcome, touching only the result columns) and rejects every delete. +-- final outcome, touching only the result columns, plus the entity when it was not yet known) +-- and rejects every delete. CREATE SEQUENCE IF NOT EXISTS ${ohdsiSchema}.mcp_audit_log_sequence; CREATE TABLE IF NOT EXISTS ${ohdsiSchema}.mcp_audit_log ( @@ -25,16 +26,25 @@ CREATE TABLE IF NOT EXISTS ${ohdsiSchema}.mcp_audit_log ( error_message text, user_id bigint, user_login character varying(255), + -- the cohort definition or concept set the call created, changed, validated or generated + -- (null for calls not about one); set when known before the call runs, otherwise when it finishes + entity_type character varying(32), + entity_id integer, started_at timestamp with time zone NOT NULL, finished_at timestamp with time zone, CONSTRAINT pk_mcp_audit_log PRIMARY KEY (id), CONSTRAINT chk_mcp_audit_log_outcome - CHECK (outcome IN ('STARTED', 'SUCCESS', 'ERROR', 'REJECTED', 'DENIED')) + CHECK (outcome IN ('STARTED', 'SUCCESS', 'ERROR', 'REJECTED', 'DENIED')), + CONSTRAINT chk_mcp_audit_log_entity_type + CHECK (entity_type IS NULL OR entity_type IN ('COHORT_DEFINITION', 'CONCEPT_SET')) ); CREATE INDEX IF NOT EXISTS mcp_audit_log_trace_idx ON ${ohdsiSchema}.mcp_audit_log USING btree (trace_id, started_at); +CREATE INDEX IF NOT EXISTS mcp_audit_log_entity_idx + ON ${ohdsiSchema}.mcp_audit_log USING btree (entity_type, entity_id, started_at); + CREATE OR REPLACE FUNCTION ${ohdsiSchema}.mcp_audit_log_append_only() RETURNS trigger AS $$ BEGIN @@ -56,7 +66,9 @@ BEGIN OR NEW.evidence_reference IS DISTINCT FROM OLD.evidence_reference OR NEW.user_id IS DISTINCT FROM OLD.user_id OR NEW.user_login IS DISTINCT FROM OLD.user_login - OR NEW.started_at IS DISTINCT FROM OLD.started_at THEN + OR NEW.started_at IS DISTINCT FROM OLD.started_at + OR (OLD.entity_type IS NOT NULL AND NEW.entity_type IS DISTINCT FROM OLD.entity_type) + OR (OLD.entity_id IS NOT NULL AND NEW.entity_id IS DISTINCT FROM OLD.entity_id) THEN RAISE EXCEPTION 'mcp_audit_log is append-only: only the outcome of entry % may be recorded', OLD.id; END IF; RETURN NEW; diff --git a/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java b/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java index df738cc130..39c1e3a84a 100644 --- a/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java +++ b/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java @@ -22,6 +22,7 @@ import org.junit.Test; import org.ohdsi.webapi.mcp.ToolDefinition.Rationale; import org.ohdsi.webapi.mcp.audit.McpAuditEntry; +import org.ohdsi.webapi.mcp.audit.McpAuditEntry.EntityType; import org.ohdsi.webapi.mcp.audit.McpAuditEntry.Outcome; import org.ohdsi.webapi.mcp.audit.McpAuditService; import org.springframework.security.access.AccessDeniedException; @@ -39,7 +40,7 @@ public void setUp() { audit = mock(McpAuditService.class); entry = new McpAuditEntry(); entry.setId(42L); - when(audit.start(any(), anyString(), anyString(), any(), any(), any())).thenReturn(entry); + when(audit.start(any(), anyString(), anyString(), any(), any(), any(), any(), any())).thenReturn(entry); } private McpToolFactory factory(boolean requireTraceparent) { @@ -76,8 +77,8 @@ public void runsTheToolAndReturnsTheToolSpanInMeta() { assertTrue(traceparent, traceparent.startsWith("00-" + TRACE_ID + "-")); assertEquals(42L, result.meta().get("audit_entry_id")); verify(audit).start(any(), eq("add_concept_to_set"), eq("STEP_2_CONCEPT_SELECTION"), anyString(), - eq("Core diagnosis concept for the phenotype."), eq("SNOMED-CT hierarchy")); - verify(audit).finish(eq(entry), eq(Outcome.SUCCESS), eq("{\"ok\":true}"), isNull()); + eq("Core diagnosis concept for the phenotype."), eq("SNOMED-CT hierarchy"), isNull(), isNull()); + verify(audit).finish(eq(entry), eq(Outcome.SUCCESS), eq("{\"ok\":true}"), isNull(), any(), any()); } @Test @@ -87,7 +88,7 @@ public void refusesACallWithoutRationaleButStillAuditsIt() { assertFalse(ran.get()); assertTrue(result.isError()); - verify(audit).finish(eq(entry), eq(Outcome.REJECTED), isNull(), any()); + verify(audit).finish(eq(entry), eq(Outcome.REJECTED), isNull(), any(), any(), any()); } @Test @@ -110,12 +111,12 @@ public void refusesAMissingTraceparentWhenRequired() { assertFalse(ran.get()); assertTrue(result.isError()); - verify(audit).finish(eq(entry), eq(Outcome.REJECTED), isNull(), any()); + verify(audit).finish(eq(entry), eq(Outcome.REJECTED), isNull(), any(), any(), any()); } @Test public void doesNotRunTheToolWhenTheAuditTrailCannotBeWritten() { - when(audit.start(any(), anyString(), anyString(), any(), any(), any())) + when(audit.start(any(), anyString(), anyString(), any(), any(), any(), any(), any())) .thenThrow(new IllegalStateException("database down")); AtomicBoolean ran = new AtomicBoolean(); CallToolResult result = factory(false).call(guarded(ran), request( @@ -125,7 +126,7 @@ public void doesNotRunTheToolWhenTheAuditTrailCannotBeWritten() { assertFalse(ran.get()); assertTrue(result.isError()); - verify(audit, never()).finish(any(), any(), any(), any()); + verify(audit, never()).finish(any(), any(), any(), any(), any(), any()); } @Test @@ -140,7 +141,7 @@ public void reportsAuthorizationFailuresAsDenied() { CallToolResult result = factory(false).call(def, new CallToolRequest("get_concept_set", Map.of(), null)); assertTrue(result.isError()); - verify(audit).finish(eq(entry), eq(Outcome.DENIED), isNull(), any()); + verify(audit).finish(eq(entry), eq(Outcome.DENIED), isNull(), any(), any(), any()); } @Test @@ -150,4 +151,32 @@ public void advertisesRationaleAsRequiredInTheInputSchema() { assertTrue(tool.inputSchema().required().contains("evidence_reference")); assertFalse(tool.annotations().readOnlyHint()); } + + @Test + public void recordsTheEntityFromTheArgumentsBeforeTheToolRuns() { + ToolDefinition def = ToolDefinition.tool("generate_cohort") + .description("test") + .step(ReasoningStep.STEP_5_GENERATION_VERIFICATION) + .entity(EntityType.COHORT_DEFINITION, "cohort_definition_id") + .handler(args -> Map.of("status", "STARTING")) + .build(); + factory(false).call(def, new CallToolRequest("generate_cohort", Map.of("cohort_definition_id", 7), Map.of())); + verify(audit).start(any(), eq("generate_cohort"), anyString(), anyString(), isNull(), isNull(), + eq(EntityType.COHORT_DEFINITION), eq(7)); + } + + @Test + public void recordsTheEntityACreateReturns() { + ToolDefinition def = ToolDefinition.tool("create_concept_set") + .description("test") + .step(ReasoningStep.STEP_2_CONCEPT_SET_CONSOLIDATION) + .entity(EntityType.CONCEPT_SET, "concept_set_id") + .handler(args -> Map.of("concept_set_id", 12)) + .build(); + factory(false).call(def, new CallToolRequest("create_concept_set", Map.of("name", "x"), Map.of())); + verify(audit).start(any(), eq("create_concept_set"), anyString(), anyString(), isNull(), isNull(), isNull(), + isNull()); + verify(audit).finish(eq(entry), eq(Outcome.SUCCESS), anyString(), isNull(), eq(EntityType.CONCEPT_SET), + eq(12)); + } } diff --git a/src/test/java/org/ohdsi/webapi/test/McpServerIT.java b/src/test/java/org/ohdsi/webapi/test/McpServerIT.java index 261654ae58..6e1f760900 100644 --- a/src/test/java/org/ohdsi/webapi/test/McpServerIT.java +++ b/src/test/java/org/ohdsi/webapi/test/McpServerIT.java @@ -144,6 +144,12 @@ public void addConceptRefusesAConceptWhoseNameDoesNotMatchTheVocabulary() throws assertTrue(result.toString(), text(result).contains("Type 2 diabetes mellitus")); assertEquals(0, itemCount(conceptSetId)); assertEquals("REJECTED", lastOutcome(traceId, "add_concept_to_set")); + + // The refused attempt is part of the concept set's provenance. + JsonNode provenance = provenance("conceptset", conceptSetId); + assertEquals(provenance.toString(), 2, provenance.size()); + assertEquals("add_concept_to_set", provenance.get(1).path("tool").asText()); + assertEquals("REJECTED", provenance.get(1).path("outcome").asText()); } @Test @@ -192,6 +198,19 @@ public void addConceptSavesAVerifiedConceptAndTheTrailShowsTheRationale() throws assertEquals(rationale, log.get(1).path("clinical_rationale").asText()); assertEquals("STEP_2_CONCEPT_SELECTION", log.get(1).path("reasoning_step").asText()); assertEquals("SUCCESS", log.get(1).path("outcome").asText()); + + // Atlas reads the same history from the concept set's provenance endpoint. + JsonNode provenance = provenance("conceptset", conceptSetId); + assertEquals(provenance.toString(), 2, provenance.size()); + assertEquals("create_concept_set", provenance.get(0).path("tool").asText()); + assertEquals("add_concept_to_set", provenance.get(1).path("tool").asText()); + assertEquals(rationale, provenance.get(1).path("clinicalRationale").asText()); + assertEquals(traceId, provenance.get(1).path("traceId").asText()); + assertEquals("SUCCESS", provenance.get(1).path("outcome").asText()); + assertEquals("anonymous", provenance.get(1).path("userLogin").asText()); + assertTrue(provenance.get(1).path("startedAt").isNumber()); + assertFalse("payloads are not exposed", provenance.get(1).has("argumentsJson")); + assertEquals(0, provenance("conceptset", conceptSetId + 1000).size()); } @Test @@ -316,6 +335,13 @@ public void generatesACohortOnEunomiaAndReportsCountsAndAttrition() throws Excep + ".cohort where cohort_definition_id = ?", Long.class, cohortId); assertEquals(aged40, written); assertEquals("SUCCESS", lastOutcome(traceId, "generate_cohort")); + + // Status polls are reads and stay out of the cohort's provenance. + JsonNode provenance = provenance("cohortdefinition", cohortId); + List tools = new ArrayList<>(); + provenance.forEach(e -> tools.add(e.path("tool").asText())); + assertEquals(List.of("create_cohort_definition", "generate_cohort"), tools); + assertEquals("STEP_3_COHORT_CRITERIA", provenance.get(0).path("reasoningStep").asText()); } @Test @@ -394,6 +420,13 @@ private static void loadEunomia(String... tables) throws Exception { assertTrue("missing from eunomia.csv.zip: " + wanted, wanted.isEmpty()); } + private JsonNode provenance(String resource, int id) throws Exception { + ResponseEntity response = getRestTemplate().getForEntity( + getBaseUri() + "/" + resource + "/" + id + "/mcp-provenance", String.class); + assertEquals(response.getBody(), HttpStatus.OK, response.getStatusCode()); + return mapper.readTree(response.getBody()); + } + private int itemCount(int conceptSetId) { return jdbcTemplate.queryForObject("select count(*) from " + getOhdsiSchema() + ".concept_set_item where concept_set_id = ?", Integer.class, conceptSetId); From 42ea43e991e9971b756de2558523a412e55323d5 Mon Sep 17 00:00:00 2001 From: Daijiro Wachi Date: Wed, 4 Nov 2026 21:26:58 +0900 Subject: [PATCH 4/4] Trace MCP tool calls as OpenTelemetry spans The tool-call span ids were WebAPI's own and never exported, so an audit row's trace id could not be found in a tracing backend. Also, the OpenTelemetry Java agent traced the HTTP request separately, under the traceparent header rather than the one MCP carries in _meta. Follow the OpenTelemetry semantic conventions for MCP and SEP-414 instead: each tools/call is a SERVER span "tools/call ". Its parent is the W3C context in params._meta, it links the ambient (HTTP request) span, and it is current while the tool runs, so the agent's JDBC spans nest under it. The audit row and the result's _meta.traceparent carry that span's ids, so the audit log and the tracing backend share trace ids. Only opentelemetry-api is added (version from Spring Boot). With the agent it is the agent's implementation; without it, it is a no-op and WebAPI allocates the ids as before. Span attributes are mcp.method.name, gen_ai.tool.name, gen_ai.operation.name, error.type=tool_error and WebAPI's step, outcome, audit entry id and entity. The arguments, rationale and error messages stay in the audit log. --- README.md | 2 +- pom.xml | 11 ++ .../org/ohdsi/webapi/mcp/McpToolFactory.java | 146 ++++++++++++++++-- .../ohdsi/webapi/mcp/McpToolFactoryTest.java | 140 ++++++++++++++++- 4 files changed, 284 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index cde87553bd..84d30fa503 100644 --- a/README.md +++ b/README.md @@ -178,7 +178,7 @@ Traceability rules the server enforces: - **Rationale.** Tools that change concept sets or cohort definitions require `clinical_rationale` (at least `mcp.minRationaleLength` characters, default 20) and `evidence_reference`. A call that leaves them out is refused. - **Verified concepts.** `add_concept_to_set`, `create_cohort_definition` and `update_cohort_definition` look up every concept id in the source vocabulary. If an id does not exist, or the name the caller gave does not match the vocabulary's name, the call is refused and nothing is saved. -- **Tracing.** Send a W3C `traceparent` in the call's `_meta`, plus an optional `reasoning_step`. WebAPI keeps the trace id, records your span as the parent, and returns the tool call's own span in the result's `_meta.traceparent`, along with `_meta.audit_entry_id`. A call without a traceparent starts a new trace, unless `mcp.requireTraceparent: true`. +- **Tracing.** Send a W3C `traceparent` in the call's `_meta`, plus an optional `reasoning_step`. Each tool call is a span that follows the [OpenTelemetry semantic conventions for MCP](https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/): its parent is the `_meta` traceparent, and it links the HTTP request's span. WebAPI returns that span in the result's `_meta.traceparent`, along with `_meta.audit_entry_id`. A call without a traceparent starts a new trace, unless `mcp.requireTraceparent: true`. When WebAPI runs with the OpenTelemetry Java agent, the span is exported together with the SQL it ran, and the audit log records the exported span's ids. An audit row can then be found in the tracing backend by its trace id, and a span by its `ohdsi.mcp.audit_entry_id`. Spans carry the tool, step and outcome, but never the arguments, rationale or error messages. Those stay in the audit log. Without the agent, WebAPI allocates the ids itself. - **Audit.** Every call is written to `mcp_audit_log` before it runs, then completed with its outcome: `SUCCESS`, `ERROR`, `REJECTED` or `DENIED`. A database trigger makes the table append-only. `get_audit_trail` returns a trace's entries in order, as the decision log for an audit manifest. - **Provenance.** Calls that create, change, validate or generate a cohort definition or concept set are also recorded against that artifact. `GET /cohortdefinition/{id}/mcp-provenance` and `GET /conceptset/{id}/mcp-provenance` return that history (tool, reasoning step, rationale, evidence, outcome, user, time and trace id; not the argument or result payloads) to anyone who can read the artifact, so Atlas can show why an agent built it that way. These endpoints stay available when the MCP server is turned off. diff --git a/pom.xml b/pom.xml index 473447e02c..6c13eed3e2 100644 --- a/pom.xml +++ b/pom.xml @@ -463,6 +463,17 @@ org.springframework.ai spring-ai-starter-mcp-server-webmvc + + + io.opentelemetry + opentelemetry-api + + + io.opentelemetry + opentelemetry-sdk-testing + test + diff --git a/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java b/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java index 47a7719a08..086e32db6c 100644 --- a/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java +++ b/src/main/java/org/ohdsi/webapi/mcp/McpToolFactory.java @@ -2,6 +2,19 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.context.propagation.TextMapGetter; import io.modelcontextprotocol.server.McpStatelessServerFeatures.SyncToolSpecification; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.CallToolRequest; @@ -16,6 +29,8 @@ import org.ohdsi.webapi.mcp.audit.McpAuditService; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.springframework.beans.factory.ObjectProvider; +import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Value; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.security.access.AccessDeniedException; @@ -25,7 +40,8 @@ /** * Turns a {@link ToolDefinition} into an MCP tool whose every call is traced and audited. * - *

For each call: the {@code _meta.traceparent} is continued (or a trace started), the audit row is + *

For each call: an OpenTelemetry server span is opened as a child of {@code _meta.traceparent} + * (MCP semantic conventions; a no-op without OpenTelemetry), the audit row is * written before the tool runs, the rationale contract is enforced, the tool runs under the * caller's WebAPI identity (so the services' {@code @PreAuthorize} rules apply unchanged), and the * audit row is completed with the outcome. The result carries the tool span's {@code traceparent} @@ -43,18 +59,57 @@ public class McpToolFactory { private static final Logger log = LoggerFactory.getLogger(McpToolFactory.class); + // OpenTelemetry semantic conventions for MCP (gen-ai/mcp) plus WebAPI's own attributes. + static final String MCP_METHOD = "tools/call"; + static final AttributeKey ATTR_MCP_METHOD = AttributeKey.stringKey("mcp.method.name"); + static final AttributeKey ATTR_OPERATION = AttributeKey.stringKey("gen_ai.operation.name"); + static final AttributeKey ATTR_TOOL_NAME = AttributeKey.stringKey("gen_ai.tool.name"); + static final AttributeKey ATTR_ERROR_TYPE = AttributeKey.stringKey("error.type"); + static final AttributeKey ATTR_REASONING_STEP = AttributeKey.stringKey("ohdsi.mcp.reasoning_step"); + static final AttributeKey ATTR_OUTCOME = AttributeKey.stringKey("ohdsi.mcp.outcome"); + static final AttributeKey ATTR_AUDIT_ENTRY_ID = AttributeKey.longKey("ohdsi.mcp.audit_entry_id"); + static final AttributeKey ATTR_ENTITY_TYPE = AttributeKey.stringKey("ohdsi.mcp.entity.type"); + static final AttributeKey ATTR_ENTITY_ID = AttributeKey.longKey("ohdsi.mcp.entity.id"); + + /** Reads W3C trace context from {@code params._meta}, where MCP carries it instead of HTTP headers. */ + private static final TextMapGetter> META_GETTER = new TextMapGetter<>() { + @Override + public Iterable keys(Map carrier) { + return carrier.keySet(); + } + + @Override + public String get(Map carrier, String key) { + return carrier != null && carrier.get(key) instanceof String value ? value : null; + } + }; + private final McpAuditService audit; private final ObjectMapper mapper; private final boolean requireTraceparent; private final int minRationaleLength; + private final Tracer tracer; + /** + * Uses the application's {@link OpenTelemetry} bean if there is one, otherwise the global instance: the + * OpenTelemetry Java agent's when WebAPI runs with it, a no-op otherwise. + */ + @Autowired public McpToolFactory(McpAuditService audit, ObjectMapper mapper, @Value("${mcp.requireTraceparent:false}") boolean requireTraceparent, - @Value("${mcp.minRationaleLength:20}") int minRationaleLength) { + @Value("${mcp.minRationaleLength:20}") int minRationaleLength, + ObjectProvider openTelemetry) { + this(audit, mapper, requireTraceparent, minRationaleLength, + openTelemetry.getIfAvailable(GlobalOpenTelemetry::get)); + } + + McpToolFactory(McpAuditService audit, ObjectMapper mapper, boolean requireTraceparent, int minRationaleLength, + OpenTelemetry openTelemetry) { this.audit = audit; this.mapper = mapper; this.requireTraceparent = requireTraceparent; this.minRationaleLength = minRationaleLength; + this.tracer = openTelemetry.getTracer("org.ohdsi.webapi.mcp"); } public SyncToolSpecification build(ToolDefinition def) { @@ -94,27 +149,73 @@ McpSchema.Tool describe(ToolDefinition def) { CallToolResult call(ToolDefinition def, CallToolRequest request) { Map meta = request.meta() == null ? Map.of() : request.meta(); - ToolArguments args = new ToolArguments(request.arguments()); - TraceContext trace; + // The client's traceparent is checked against the W3C rules whether or not OpenTelemetry is active, + // so a malformed one is refused the same way in both cases. + TraceContext clientTrace = null; String traceProblem = null; Object traceparent = meta.get(META_TRACEPARENT); if (traceparent instanceof String tp && !tp.isBlank()) { try { - trace = TraceContext.continueFrom(tp); + clientTrace = TraceContext.continueFrom(tp); } catch (IllegalArgumentException e) { - trace = TraceContext.newRoot(); traceProblem = e.getMessage(); } - } else { - trace = TraceContext.newRoot(); - if (requireTraceparent) { - traceProblem = "_meta.traceparent is required by this server"; - } + } else if (requireTraceparent) { + traceProblem = "_meta.traceparent is required by this server"; + } + + // MCP semantic conventions: the server span's parent is the context in params._meta, and the + // ambient context (the HTTP request's span, when the OpenTelemetry agent runs) is linked. + Span ambient = Span.current(); + Context parent = clientTrace != null + ? W3CTraceContextPropagator.getInstance().extract(Context.current(), meta, META_GETTER) + : Context.current(); + SpanBuilder spanBuilder = tracer.spanBuilder(MCP_METHOD + " " + def.name) + .setSpanKind(SpanKind.SERVER) + .setParent(parent) + .setAttribute(ATTR_MCP_METHOD, MCP_METHOD) + .setAttribute(ATTR_OPERATION, "execute_tool") + .setAttribute(ATTR_TOOL_NAME, def.name); + if (clientTrace != null && ambient.getSpanContext().isValid()) { + spanBuilder.addLink(ambient.getSpanContext()); + } + Span span = spanBuilder.startSpan(); + try (Scope ignored = span.makeCurrent()) { + return run(def, request, meta, traceFor(span, clientTrace, ambient), traceProblem, span); + } finally { + span.end(); } + } + + /** + * The ids recorded in the audit log and returned to the client. With OpenTelemetry active they are the + * exported span's own, so an audit row and its trace in the tracing backend share one trace id. Without + * it (no agent, no SDK) WebAPI allocates them itself. + */ + static TraceContext traceFor(Span span, TraceContext clientTrace, Span ambient) { + SpanContext context = span.getSpanContext(); + if (!context.isValid()) { + return clientTrace != null ? clientTrace : TraceContext.newRoot(); + } + String parentSpanId = null; + if (clientTrace != null) { + parentSpanId = clientTrace.parentSpanId(); + } else if (ambient.getSpanContext().isValid() + && ambient.getSpanContext().getTraceId().equals(context.getTraceId())) { + parentSpanId = ambient.getSpanContext().getSpanId(); + } + return new TraceContext(context.getTraceId(), parentSpanId, context.getSpanId(), + context.getTraceFlags().asHex(), clientTrace == null); + } + + private CallToolResult run(ToolDefinition def, CallToolRequest request, Map meta, + TraceContext trace, String traceProblem, Span span) { + ToolArguments args = new ToolArguments(request.arguments()); String step = meta.get(META_REASONING_STEP) instanceof String s && !s.isBlank() ? s : def.step.name(); String rationale = args.raw(CLINICAL_RATIONALE) instanceof String s ? s : null; String evidence = args.raw(EVIDENCE_REFERENCE) instanceof String s ? s : null; + span.setAttribute(ATTR_REASONING_STEP, step); McpAuditEntry entry; try { @@ -123,8 +224,10 @@ CallToolResult call(ToolDefinition def, CallToolRequest request) { entityId == null ? null : def.entityType, entityId); } catch (Exception e) { log.error("MCP audit trail unavailable; refusing to run tool {}", def.name, e); + markOutcome(span, Outcome.ERROR); return error(trace, null, Outcome.ERROR, "The audit trail could not be written, so the tool was not run."); } + span.setAttribute(ATTR_AUDIT_ENTRY_ID, entry.getId()); Outcome outcome; Object result = null; @@ -152,14 +255,19 @@ CallToolResult call(ToolDefinition def, CallToolRequest request) { } String resultJson = outcome == Outcome.SUCCESS ? toJson(result) : null; + Integer createdId = def.entityType != null && result instanceof Map m + ? entityId(m.get(def.entityIdField)) : null; try { - Integer createdId = def.entityType != null && result instanceof Map m - ? entityId(m.get(def.entityIdField)) : null; audit.finish(entry, outcome, resultJson, message, createdId == null ? null : def.entityType, createdId); } catch (Exception e) { // The action already happened and its STARTED row is on record; do not hide the result. log.error("Could not record outcome {} of MCP audit entry {}", outcome, entry.getId(), e); } + markOutcome(span, outcome); + if (entry.getEntityType() != null) { + span.setAttribute(ATTR_ENTITY_TYPE, entry.getEntityType().name()); + span.setAttribute(ATTR_ENTITY_ID, entry.getEntityId().longValue()); + } if (outcome != Outcome.SUCCESS) { return error(trace, entry.getId(), outcome, message); @@ -171,6 +279,18 @@ CallToolResult call(ToolDefinition def, CallToolRequest request) { .build(); } + /** + * Outcome on the span. Only the outcome is recorded: arguments, rationale and error messages stay in the + * audit log, which has per-artifact access control, and out of the tracing backend, which does not. + */ + private static void markOutcome(Span span, Outcome outcome) { + span.setAttribute(ATTR_OUTCOME, outcome.name()); + if (outcome != Outcome.SUCCESS) { + span.setAttribute(ATTR_ERROR_TYPE, "tool_error"); + span.setStatus(StatusCode.ERROR, outcome.name()); + } + } + private static Integer entityId(Object value) { if (value instanceof Number n) { return n.intValue(); diff --git a/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java b/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java index 39c1e3a84a..5f4561ee03 100644 --- a/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java +++ b/src/test/java/org/ohdsi/webapi/mcp/McpToolFactoryTest.java @@ -3,6 +3,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertTrue; +import static org.junit.Assert.assertNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; @@ -12,7 +13,21 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import org.mockito.ArgumentCaptor; + import com.fasterxml.jackson.databind.ObjectMapper; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import java.util.List; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.CallToolRequest; import io.modelcontextprotocol.spec.McpSchema.CallToolResult; @@ -44,7 +59,7 @@ public void setUp() { } private McpToolFactory factory(boolean requireTraceparent) { - return new McpToolFactory(audit, new ObjectMapper(), requireTraceparent, 20); + return new McpToolFactory(audit, new ObjectMapper(), requireTraceparent, 20, OpenTelemetry.noop()); } private static ToolDefinition guarded(AtomicBoolean ran) { @@ -179,4 +194,127 @@ public void recordsTheEntityACreateReturns() { verify(audit).finish(eq(entry), eq(Outcome.SUCCESS), anyString(), isNull(), eq(EntityType.CONCEPT_SET), eq(12)); } + + // ---- OpenTelemetry --------------------------------------------------------------------- + + private static final class Otel { + final InMemorySpanExporter exporter = InMemorySpanExporter.create(); + final OpenTelemetrySdk sdk = OpenTelemetrySdk.builder() + .setTracerProvider(SdkTracerProvider.builder() + .addSpanProcessor(SimpleSpanProcessor.create(exporter)) + .build()) + .build(); + + SpanData onlySpan() { + List spans = exporter.getFinishedSpanItems(); + assertEquals(spans.toString(), 1, spans.size()); + return spans.get(0); + } + } + + private TraceContext auditedTrace() { + ArgumentCaptor trace = ArgumentCaptor.forClass(TraceContext.class); + verify(audit).start(trace.capture(), anyString(), anyString(), any(), any(), any(), any(), any()); + return trace.getValue(); + } + + @Test + public void opensAServerSpanUnderTheMetaTraceparentAndAuditsItsIds() { + Otel otel = new Otel(); + McpToolFactory factory = new McpToolFactory(audit, new ObjectMapper(), false, 20, otel.sdk); + + CallToolResult result = factory.call(guarded(new AtomicBoolean()), request( + Map.of("clinical_rationale", "Core diagnosis concept for the phenotype.", + "evidence_reference", "SNOMED-CT hierarchy"), + Map.of("traceparent", TRACEPARENT, "reasoning_step", "STEP_2_CONCEPT_SELECTION"))); + + SpanData span = otel.onlySpan(); + assertEquals("tools/call add_concept_to_set", span.getName()); + assertEquals(SpanKind.SERVER, span.getKind()); + assertEquals(TRACE_ID, span.getTraceId()); + assertEquals("00f067aa0ba902b7", span.getParentSpanId()); + assertEquals("tools/call", span.getAttributes().get(AttributeKey.stringKey("mcp.method.name"))); + assertEquals("add_concept_to_set", span.getAttributes().get(AttributeKey.stringKey("gen_ai.tool.name"))); + assertEquals("SUCCESS", span.getAttributes().get(AttributeKey.stringKey("ohdsi.mcp.outcome"))); + assertEquals(Long.valueOf(42), span.getAttributes().get(AttributeKey.longKey("ohdsi.mcp.audit_entry_id"))); + assertEquals("STEP_2_CONCEPT_SELECTION", + span.getAttributes().get(AttributeKey.stringKey("ohdsi.mcp.reasoning_step"))); + assertNull("the rationale stays out of the tracing backend", + span.getAttributes().get(AttributeKey.stringKey("clinical_rationale"))); + + // The audit row and the returned traceparent carry the exported span's own ids. + TraceContext audited = auditedTrace(); + assertEquals(TRACE_ID, audited.traceId()); + assertEquals("00f067aa0ba902b7", audited.parentSpanId()); + assertEquals(span.getSpanId(), audited.spanId()); + assertEquals("00-" + TRACE_ID + "-" + span.getSpanId() + "-01", result.meta().get("traceparent")); + } + + @Test + public void linksTheHttpRequestSpanWhenTheClientSentItsOwnTrace() { + Otel otel = new Otel(); + McpToolFactory factory = new McpToolFactory(audit, new ObjectMapper(), false, 20, otel.sdk); + Span http = otel.sdk.getTracer("test").spanBuilder("POST /mcp").setSpanKind(SpanKind.SERVER).startSpan(); + try (Scope ignored = http.makeCurrent()) { + factory.call(guarded(new AtomicBoolean()), request( + Map.of("clinical_rationale", "Core diagnosis concept for the phenotype.", + "evidence_reference", "SNOMED-CT hierarchy"), + Map.of("traceparent", TRACEPARENT))); + } finally { + http.end(); + } + + SpanData tool = otel.exporter.getFinishedSpanItems().stream() + .filter(s -> s.getName().startsWith("tools/call")).findFirst().orElseThrow(); + assertEquals(TRACE_ID, tool.getTraceId()); + assertEquals(1, tool.getLinks().size()); + assertEquals(http.getSpanContext().getSpanId(), tool.getLinks().get(0).getSpanContext().getSpanId()); + } + + @Test + public void joinsTheHttpRequestTraceWhenTheClientSentNone() { + Otel otel = new Otel(); + McpToolFactory factory = new McpToolFactory(audit, new ObjectMapper(), false, 20, otel.sdk); + Span http = otel.sdk.getTracer("test").spanBuilder("POST /mcp").setSpanKind(SpanKind.SERVER).startSpan(); + try (Scope ignored = http.makeCurrent()) { + factory.call(guarded(new AtomicBoolean()), request( + Map.of("clinical_rationale", "Core diagnosis concept for the phenotype.", + "evidence_reference", "SNOMED-CT hierarchy"), + Map.of())); + } finally { + http.end(); + } + + TraceContext audited = auditedTrace(); + assertEquals(http.getSpanContext().getTraceId(), audited.traceId()); + assertEquals(http.getSpanContext().getSpanId(), audited.parentSpanId()); + assertTrue("no client traceparent", audited.generated()); + } + + @Test + public void marksARefusedCallAsAToolError() { + Otel otel = new Otel(); + McpToolFactory factory = new McpToolFactory(audit, new ObjectMapper(), false, 20, otel.sdk); + + factory.call(guarded(new AtomicBoolean()), request(Map.of(), Map.of("traceparent", TRACEPARENT))); + + SpanData span = otel.onlySpan(); + assertEquals("REJECTED", span.getAttributes().get(AttributeKey.stringKey("ohdsi.mcp.outcome"))); + assertEquals("tool_error", span.getAttributes().get(AttributeKey.stringKey("error.type"))); + assertEquals(StatusCode.ERROR, span.getStatus().getStatusCode()); + assertEquals("REJECTED", span.getStatus().getDescription()); + } + + @Test + public void withoutOpenTelemetryWebApiAllocatesTheIdsItself() { + factory(false).call(guarded(new AtomicBoolean()), request( + Map.of("clinical_rationale", "Core diagnosis concept for the phenotype.", + "evidence_reference", "SNOMED-CT hierarchy"), + Map.of("traceparent", TRACEPARENT))); + + TraceContext audited = auditedTrace(); + assertEquals(TRACE_ID, audited.traceId()); + assertEquals("00f067aa0ba902b7", audited.parentSpanId()); + assertEquals(16, audited.spanId().length()); + } }