From d5032b7285e1cc07e4d7d132dec6b848e2766585 Mon Sep 17 00:00:00 2001 From: Adam Rauch Date: Wed, 18 Mar 2026 18:16:20 -0700 Subject: [PATCH 1/5] MCP server production MVP --- core/src/org/labkey/core/CoreMcp.java | 13 +- .../org/labkey/core/mcp/McpServiceImpl.java | 113 +++++++++++------- 2 files changed, 75 insertions(+), 51 deletions(-) diff --git a/core/src/org/labkey/core/CoreMcp.java b/core/src/org/labkey/core/CoreMcp.java index 06482fe5627..de0c5f9a7c1 100644 --- a/core/src/org/labkey/core/CoreMcp.java +++ b/core/src/org/labkey/core/CoreMcp.java @@ -1,8 +1,8 @@ package org.labkey.core; +import io.modelcontextprotocol.server.McpSyncServerExchange; import org.json.JSONObject; import org.labkey.api.data.Container; -import org.labkey.api.mcp.McpContext; import org.labkey.api.mcp.McpService; import org.labkey.api.security.User; import org.labkey.api.settings.AppProps; @@ -10,6 +10,7 @@ import org.labkey.api.study.Study; import org.labkey.api.study.StudyService; import org.labkey.api.util.HtmlString; +import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.tool.annotation.Tool; import java.util.Map; @@ -19,13 +20,11 @@ public class CoreMcp implements McpService.McpImpl { - // TODO ChatSessions are currently per session. The McpService should detect change of folder. - @Tool(description = "Call this tool before answering any prompts! This tool provides useful context information about the current user (name, userid), webserver (name, url, description), and current folder (name, path, url, description).") - String whereAmIWhoAmITalkingTo() + @Tool(description = "Call this tool before answering any prompts! This tool provides useful context information about the current user (name, userid), webserver (name, url, description), and current folder (name, path, url, description).") + String whereAmIWhoAmITalkingTo(ToolContext context) { - McpContext context = McpContext.get(); - User user = context.getUser(); - Container folder = context.getContainer(); + User user = (User)context.getContext().get("user"); + Container folder = (Container)context.getContext().get("container"); AppProps appProps = AppProps.getInstance(); Study study = null != StudyService.get() ? Objects.requireNonNull(StudyService.get()).getStudy(folder) : null; LookAndFeelProperties laf = LookAndFeelProperties.getInstance(folder); diff --git a/core/src/org/labkey/core/mcp/McpServiceImpl.java b/core/src/org/labkey/core/mcp/McpServiceImpl.java index 0937f308a20..eca9b1a8ae6 100644 --- a/core/src/org/labkey/core/mcp/McpServiceImpl.java +++ b/core/src/org/labkey/core/mcp/McpServiceImpl.java @@ -1,8 +1,10 @@ package org.labkey.core.mcp; +import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.genai.Client; import com.google.genai.types.ClientOptions; +import io.modelcontextprotocol.common.McpTransportContext; import io.modelcontextprotocol.json.McpJsonMapper; import io.modelcontextprotocol.server.McpServer; import io.modelcontextprotocol.server.McpServerFeatures; @@ -23,9 +25,12 @@ import org.jetbrains.annotations.NotNull; import org.jspecify.annotations.NonNull; import org.labkey.api.collections.CopyOnWriteHashMap; +import org.labkey.api.data.Container; +import org.labkey.api.data.ContainerManager; import org.labkey.api.markdown.MarkdownService; import org.labkey.api.mcp.McpContext; import org.labkey.api.mcp.McpService; +import org.labkey.api.security.User; import org.labkey.api.util.ContextListener; import org.labkey.api.util.FileUtil; import org.labkey.api.util.HtmlString; @@ -35,13 +40,6 @@ import org.labkey.api.util.logging.LogHelper; import org.springframework.ai.anthropic.AnthropicChatModel; import org.springframework.ai.anthropic.AnthropicChatOptions; -import org.springframework.ai.google.genai.common.GoogleGenAiThinkingLevel; -import org.springframework.ai.openai.OpenAiChatModel; -import org.springframework.ai.openai.OpenAiChatOptions; -import org.springframework.ai.openai.OpenAiEmbeddingModel; -import org.springframework.ai.openai.OpenAiEmbeddingOptions; -import org.springframework.ai.openai.api.OpenAiApi; -import org.springframework.ai.document.MetadataMode; import org.springframework.ai.anthropic.api.AnthropicApi; import org.springframework.ai.chat.client.ChatClient; import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor; @@ -53,8 +51,11 @@ import org.springframework.ai.chat.memory.MessageWindowChatMemory; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.document.Document; +import org.springframework.ai.document.MetadataMode; import org.springframework.ai.embedding.EmbeddingModel; import org.springframework.ai.google.genai.GoogleGenAiChatModel; import org.springframework.ai.google.genai.GoogleGenAiChatOptions; @@ -62,11 +63,14 @@ import org.springframework.ai.google.genai.text.GoogleGenAiTextEmbeddingModel; import org.springframework.ai.google.genai.text.GoogleGenAiTextEmbeddingOptions; import org.springframework.ai.mcp.McpToolUtils; -import org.springframework.ai.chat.model.ToolContext; +import org.springframework.ai.openai.OpenAiChatModel; +import org.springframework.ai.openai.OpenAiChatOptions; +import org.springframework.ai.openai.OpenAiEmbeddingModel; +import org.springframework.ai.openai.OpenAiEmbeddingOptions; +import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.definition.ToolDefinition; import org.springframework.ai.tool.metadata.ToolMetadata; -import org.springframework.ai.document.Document; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.SimpleVectorStore; import org.springframework.ai.vectorstore.VectorStore; @@ -82,6 +86,7 @@ import java.util.Arrays; import java.util.ConcurrentModificationException; import java.util.List; +import java.util.Map; import java.util.NoSuchElementException; import java.util.Objects; import java.util.function.Supplier; @@ -222,24 +227,72 @@ private class _McpServlet extends HttpServlet // wraps HttpServletSseServerTrans _McpServlet(ObjectMapper objectMapper, String messageEndpoint, String sseEndpoint) { transportProvider = HttpServletStreamableServerTransportProvider.builder() - .jsonMapper(McpJsonMapper.getDefault()) - .mcpEndpoint(messageEndpoint) - .build(); + .jsonMapper(McpJsonMapper.getDefault()) + .mcpEndpoint(messageEndpoint) + .contextExtractor(req -> { + User user = (User) req.getUserPrincipal(); + return McpTransportContext.create(Map.of( + "container", ContainerManager.getHomeContainer(), + "user", user + ) + ); + }) + .build(); } void startMcpServer() { - List tools = Arrays.stream(getToolCallbacks()).map(McpToolUtils::toSyncToolSpecification).toList(); + List tools = Arrays.stream(getToolCallbacks()) + .map(this::toSyncToolSpecification) + .toList(); + List resources = new ArrayList<>(resourceMap.values()); mcpServer = McpServer.sync(transportProvider) - .tools(tools) - .resources(resources) + .tools(tools) + .resources(resources) // .capabilities(new McpSchema.ServerCapabilities()) - .build(); + .build(); ContextListener.addShutdownListener(new _ShutdownListener()); } + private McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCallback toolCallback) + { + var toolDef = toolCallback.getToolDefinition(); + var schema = McpSchema.Tool.builder() + .name(toolDef.name()) + .description(toolDef.description()) + .inputSchema(McpJsonMapper.getDefault(), toolDef.inputSchema()) + .build(); + + return new McpServerFeatures.SyncToolSpecification(schema, (exchange, args) -> { + var transportCtx = exchange.transportContext(); + var container = (Container) transportCtx.get("container"); // TODO: Pull container from session instead. Or insist that LLM provides it? + var user = (User) transportCtx.get("user"); + var sessionId = exchange.sessionId(); + LOG.info("MCP sessionId: {}", sessionId); + + var toolContext = new ToolContext(Map.of("container", container, "user", user, "sessionId", sessionId)); + + String toolInput = /* serialize args to JSON */ null; + try + { + toolInput = JsonUtil.DEFAULT_MAPPER.writeValueAsString(args); + } + catch (JsonProcessingException e) + { + throw new RuntimeException(e); + } + String result = toolCallback.call(toolInput, toolContext); + return new McpSchema.CallToolResult( + List.of( + new McpSchema.TextContent(result) + ), + false + ); + }); + } + @Override public void service(ServletRequest sreq, ServletResponse sres) throws ServletException, IOException { @@ -255,34 +308,6 @@ public void service(ServletRequest sreq, ServletResponse sres) throws ServletExc return; } - if ("POST".equals(req.getMethod())) - { - if (null == req.getParameter("sessionId") && null == req.getSession(true).getAttribute("McpServiceImpl#mcpSessionId")) - { - // USE SSE endpoint to get a sessionId - MockHttpServletRequest mockRequest = new MockHttpServletRequest(req.getServletContext(), "GET", SSE_ENDPOINT); - mockRequest.setAsyncSupported(true); - MockHttpServletResponse mockResponse = new MockHttpServletResponse(); - transportProvider.service(mockRequest, mockResponse); - String body = new String(mockResponse.getContentAsByteArray(), StandardCharsets.UTF_8); - String mcpSessionId = StringUtils.substringBetween(body, "sessionId=", "\n"); - req.getSession(true).setAttribute("McpServiceImpl#mcpSessionId", mcpSessionId); - mockRequest.close(); - mockResponse.getOutputStream().close(); - } - - req = new HttpServletRequestWrapper(req) - { - @Override - public String getParameter(String name) - { - var ret = super.getParameter(name); - if (null == ret && "sessionId".equals(name)) - return String.valueOf(Objects.requireNonNull(((HttpServletRequest) getRequest()).getSession(true).getAttribute("McpServiceImpl#mcpSessionId"))); - return ret; - } - }; - } transportProvider.service(req, res); } From 5a656f36481e28fb8d3a51137771a61fc92e5809 Mon Sep 17 00:00:00 2001 From: Adam Rauch Date: Wed, 18 Mar 2026 18:36:49 -0700 Subject: [PATCH 2/5] Imports --- core/src/org/labkey/core/CoreMcp.java | 1 - core/src/org/labkey/core/mcp/McpServiceImpl.java | 7 ------- 2 files changed, 8 deletions(-) diff --git a/core/src/org/labkey/core/CoreMcp.java b/core/src/org/labkey/core/CoreMcp.java index de0c5f9a7c1..4d49e9d4c5c 100644 --- a/core/src/org/labkey/core/CoreMcp.java +++ b/core/src/org/labkey/core/CoreMcp.java @@ -1,6 +1,5 @@ package org.labkey.core; -import io.modelcontextprotocol.server.McpSyncServerExchange; import org.json.JSONObject; import org.labkey.api.data.Container; import org.labkey.api.mcp.McpService; diff --git a/core/src/org/labkey/core/mcp/McpServiceImpl.java b/core/src/org/labkey/core/mcp/McpServiceImpl.java index eca9b1a8ae6..01da6597676 100644 --- a/core/src/org/labkey/core/mcp/McpServiceImpl.java +++ b/core/src/org/labkey/core/mcp/McpServiceImpl.java @@ -17,10 +17,8 @@ import jakarta.servlet.ServletResponse; import jakarta.servlet.http.HttpServlet; import jakarta.servlet.http.HttpServletRequest; -import jakarta.servlet.http.HttpServletRequestWrapper; import jakarta.servlet.http.HttpServletResponse; import jakarta.servlet.http.HttpSession; -import org.apache.commons.lang3.StringUtils; import org.apache.logging.log4j.Logger; import org.jetbrains.annotations.NotNull; import org.jspecify.annotations.NonNull; @@ -62,7 +60,6 @@ import org.springframework.ai.google.genai.GoogleGenAiEmbeddingConnectionDetails; import org.springframework.ai.google.genai.text.GoogleGenAiTextEmbeddingModel; import org.springframework.ai.google.genai.text.GoogleGenAiTextEmbeddingOptions; -import org.springframework.ai.mcp.McpToolUtils; import org.springframework.ai.openai.OpenAiChatModel; import org.springframework.ai.openai.OpenAiChatOptions; import org.springframework.ai.openai.OpenAiEmbeddingModel; @@ -75,12 +72,9 @@ import org.springframework.ai.vectorstore.SimpleVectorStore; import org.springframework.ai.vectorstore.VectorStore; import org.springframework.ai.vectorstore.filter.Filter; -import org.springframework.mock.web.MockHttpServletRequest; -import org.springframework.mock.web.MockHttpServletResponse; import reactor.core.publisher.Mono; import java.io.IOException; -import java.nio.charset.StandardCharsets; import java.time.Duration; import java.util.ArrayList; import java.util.Arrays; @@ -88,7 +82,6 @@ import java.util.List; import java.util.Map; import java.util.NoSuchElementException; -import java.util.Objects; import java.util.function.Supplier; import static org.apache.commons.lang3.StringUtils.isBlank; From bf1cf2bf90f686d66ea722d1251f1380f664b0a4 Mon Sep 17 00:00:00 2001 From: Adam Rauch Date: Thu, 19 Mar 2026 12:18:32 -0700 Subject: [PATCH 3/5] listContainers --- core/src/org/labkey/core/CoreMcp.java | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/core/src/org/labkey/core/CoreMcp.java b/core/src/org/labkey/core/CoreMcp.java index 4d49e9d4c5c..555a782e423 100644 --- a/core/src/org/labkey/core/CoreMcp.java +++ b/core/src/org/labkey/core/CoreMcp.java @@ -1,9 +1,12 @@ package org.labkey.core; import org.json.JSONObject; +import org.labkey.api.collections.LabKeyCollectors; import org.labkey.api.data.Container; +import org.labkey.api.data.ContainerManager; import org.labkey.api.mcp.McpService; import org.labkey.api.security.User; +import org.labkey.api.security.permissions.ReadPermission; import org.labkey.api.settings.AppProps; import org.labkey.api.settings.LookAndFeelProperties; import org.labkey.api.study.Study; @@ -61,4 +64,15 @@ String whereAmIWhoAmITalkingTo(ToolContext context) "site", siteObj )).toString(); } + + @Tool(description = "List the hierarchical path for every container in the server where the user has read permissions.") + String listContainers(ToolContext context) + { + User user = (User)context.getContext().get("user"); + return ContainerManager.getAllChildren(ContainerManager.getRoot(), user, ReadPermission.class) + .stream() + .map(Container::getPath) + .collect(LabKeyCollectors.toJSONArray()) + .toString(); + } } From 8ee2e9fdd104418d9dcf287015b7a1d5e40ef30d Mon Sep 17 00:00:00 2001 From: Adam Rauch Date: Thu, 19 Mar 2026 16:11:24 -0700 Subject: [PATCH 4/5] Ask user for container path and cache it. Clean up all endpoints. --- api/src/org/labkey/api/mcp/McpService.java | 16 ++++- core/src/org/labkey/core/CoreMcp.java | 25 ++++++-- .../org/labkey/core/mcp/McpServiceImpl.java | 25 ++++++-- .../labkey/query/controllers/QueryMcp.java | 63 ++++--------------- .../src/org/labkey/search/SearchModule.java | 11 ++-- 5 files changed, 70 insertions(+), 70 deletions(-) diff --git a/api/src/org/labkey/api/mcp/McpService.java b/api/src/org/labkey/api/mcp/McpService.java index 67854a58c61..c8d7da6115a 100644 --- a/api/src/org/labkey/api/mcp/McpService.java +++ b/api/src/org/labkey/api/mcp/McpService.java @@ -5,11 +5,15 @@ import jakarta.servlet.http.HttpSession; import org.jetbrains.annotations.NotNull; import org.jspecify.annotations.NonNull; +import org.labkey.api.data.Container; import org.labkey.api.module.McpProvider; +import org.labkey.api.security.User; import org.labkey.api.services.ServiceRegistry; import org.labkey.api.util.HtmlString; +import org.labkey.api.writer.ContainerUser; import org.springaicommunity.mcp.provider.resource.SyncMcpResourceProvider; import org.springframework.ai.chat.client.ChatClient; +import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.support.ToolCallbacks; import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.ToolCallbackProvider; @@ -31,7 +35,17 @@ public interface McpService extends ToolCallbackProvider { // marker interface for classes that we will "ingest" using Spring annotations - interface McpImpl {} + interface McpImpl + { + default ContainerUser getContext(ToolContext toolContext) + { + User user = (User)toolContext.getContext().get("user"); + Container container = (Container)toolContext.getContext().get("container"); + if (container == null) + throw new IllegalArgumentException("You need to set a container path before invoking this tool"); + return ContainerUser.create(container, user); + } + } static @NotNull McpService get() { diff --git a/core/src/org/labkey/core/CoreMcp.java b/core/src/org/labkey/core/CoreMcp.java index 555a782e423..33e88b9b108 100644 --- a/core/src/org/labkey/core/CoreMcp.java +++ b/core/src/org/labkey/core/CoreMcp.java @@ -14,19 +14,22 @@ import org.labkey.api.util.HtmlString; import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.tool.annotation.Tool; +import org.springframework.ai.tool.annotation.ToolParam; import java.util.Map; import java.util.Objects; import static org.apache.commons.lang3.StringUtils.isNotBlank; +import static org.labkey.core.mcp.McpServiceImpl.PATH_CACHE; public class CoreMcp implements McpService.McpImpl { @Tool(description = "Call this tool before answering any prompts! This tool provides useful context information about the current user (name, userid), webserver (name, url, description), and current folder (name, path, url, description).") String whereAmIWhoAmITalkingTo(ToolContext context) { - User user = (User)context.getContext().get("user"); - Container folder = (Container)context.getContext().get("container"); + var cu = getContext(context); + User user = cu.getUser(); + Container folder = cu.getContainer(); AppProps appProps = AppProps.getInstance(); Study study = null != StudyService.get() ? Objects.requireNonNull(StudyService.get()).getStudy(folder) : null; LookAndFeelProperties laf = LookAndFeelProperties.getInstance(folder); @@ -66,13 +69,25 @@ String whereAmIWhoAmITalkingTo(ToolContext context) } @Tool(description = "List the hierarchical path for every container in the server where the user has read permissions.") - String listContainers(ToolContext context) + String listContainers(ToolContext toolContext) { - User user = (User)context.getContext().get("user"); - return ContainerManager.getAllChildren(ContainerManager.getRoot(), user, ReadPermission.class) + return ContainerManager.getAllChildren(ContainerManager.getRoot(), getContext(toolContext).getUser(), ReadPermission.class) .stream() .map(Container::getPath) .collect(LabKeyCollectors.toJSONArray()) .toString(); } + + @Tool(description = "Every tool in this MCP requires a container path, e.g. /MyProject/MyFolder. A container is also called a folder or project. Please prompt the user for a container path and use this tool to save the path for this session.") + String setContainer(ToolContext context, @ToolParam(description = "Container path, e.g. /MyProject/MyFolder", required = true) String containerPath) + { + Container container = ContainerManager.getForPath(containerPath); + if (container != null) + { + PATH_CACHE.put((String) context.getContext().get("sessionId"), containerPath); + return "OK!"; + } + + return "That's not a valid container path. Try using listContainers to see them."; + } } diff --git a/core/src/org/labkey/core/mcp/McpServiceImpl.java b/core/src/org/labkey/core/mcp/McpServiceImpl.java index 01da6597676..2ed41ed863e 100644 --- a/core/src/org/labkey/core/mcp/McpServiceImpl.java +++ b/core/src/org/labkey/core/mcp/McpServiceImpl.java @@ -22,6 +22,8 @@ import org.apache.logging.log4j.Logger; import org.jetbrains.annotations.NotNull; import org.jspecify.annotations.NonNull; +import org.labkey.api.cache.Cache; +import org.labkey.api.cache.CacheManager; import org.labkey.api.collections.CopyOnWriteHashMap; import org.labkey.api.data.Container; import org.labkey.api.data.ContainerManager; @@ -79,6 +81,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.ConcurrentModificationException; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.NoSuchElementException; @@ -212,6 +215,8 @@ public List tools() ).toList(); } + public final static Cache PATH_CACHE = CacheManager.getCache(1000, CacheManager.DAY, "MCP container paths"); + private class _McpServlet extends HttpServlet // wraps HttpServletSseServerTransportProvider { HttpServletStreamableServerTransportProvider transportProvider; @@ -225,10 +230,8 @@ private class _McpServlet extends HttpServlet // wraps HttpServletSseServerTrans .contextExtractor(req -> { User user = (User) req.getUserPrincipal(); return McpTransportContext.create(Map.of( - "container", ContainerManager.getHomeContainer(), "user", user - ) - ); + )); }) .build(); } @@ -260,12 +263,22 @@ private McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCall return new McpServerFeatures.SyncToolSpecification(schema, (exchange, args) -> { var transportCtx = exchange.transportContext(); - var container = (Container) transportCtx.get("container"); // TODO: Pull container from session instead. Or insist that LLM provides it? var user = (User) transportCtx.get("user"); var sessionId = exchange.sessionId(); - LOG.info("MCP sessionId: {}", sessionId); - var toolContext = new ToolContext(Map.of("container", container, "user", user, "sessionId", sessionId)); + Map map = new HashMap<>(); + map.put("user", user); + map.put("sessionId", sessionId); + + String containerPath = PATH_CACHE.get(exchange.sessionId()); + if (containerPath != null) + { + Container container = ContainerManager.getForPath(containerPath); + map.put("container", container); + map.put("containerPath", containerPath); + } + + var toolContext = new ToolContext(map); String toolInput = /* serialize args to JSON */ null; try diff --git a/query/src/org/labkey/query/controllers/QueryMcp.java b/query/src/org/labkey/query/controllers/QueryMcp.java index 99bb2df99c0..18b6364da0d 100644 --- a/query/src/org/labkey/query/controllers/QueryMcp.java +++ b/query/src/org/labkey/query/controllers/QueryMcp.java @@ -5,10 +5,8 @@ import org.apache.commons.lang3.StringUtils; import org.json.JSONArray; import org.json.JSONObject; -import org.labkey.api.action.SpringActionController; import org.labkey.api.collections.CaseInsensitiveHashSet; import org.labkey.api.data.ColumnInfo; -import org.labkey.api.data.ContainerManager; import org.labkey.api.data.PropertyManager; import org.labkey.api.data.TableDescription; import org.labkey.api.data.TableInfo; @@ -22,9 +20,10 @@ import org.labkey.api.query.SchemaKey; import org.labkey.api.query.SimpleSchemaTreeVisitor; import org.labkey.api.query.UserSchema; -import org.labkey.api.security.UserManager; +import org.labkey.api.writer.ContainerUser; import org.labkey.query.sql.SqlParser; import org.springaicommunity.mcp.annotation.McpResource; +import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.tool.annotation.Tool; import org.springframework.ai.tool.annotation.ToolParam; @@ -66,18 +65,18 @@ String listColumnMetaData(@ToolParam(description = "Fully qualified table name a } @Tool(description = "Provide list of tables within the provided schema.") - String listTablesForSchema(@ToolParam(description = "Fully qualified schema name as it would appear in SQL e.g. \"schema\"") String quotedSchemaName) + String listTablesForSchema(ToolContext toolContext, @ToolParam(description = "Fully qualified schema name as it would appear in SQL e.g. \"schema\"") String quotedSchemaName) { - var json = _listTablesForSchema(quotedSchemaName); + var json = _listTablesForSchema(quotedSchemaName, getContext(toolContext)); // can I just return a JSONObject return json.toString(); } @Tool(description = "Provide list of database schemas") - String listSchemas() + String listSchemas(ToolContext toolContext) { - McpContext context = getContext(); - var map = _listAllSchemas(DefaultSchema.get(context.getUser(), context.getContainer())); + ContainerUser cu = getContext(toolContext); + var map = _listAllSchemas(DefaultSchema.get(cu.getUser(), cu.getContainer())); var array = new JSONArray(); for (var entry : map.entrySet()) { @@ -92,54 +91,15 @@ String listSchemas() @Tool(description = "Provide the SQL source for a saved query.") - String getSourceForSavedQuery(@ToolParam(description = "Fully qualified query name as it would appear in SQL e.g. \"schema\".\"table or query\"") String fullQuotedTableName) + String getSourceForSavedQuery(ToolContext toolContext, @ToolParam(description = "Fully qualified query name as it would appear in SQL e.g. \"schema\".\"table or query\"") String fullQuotedTableName) { - var json = _listTablesForSchema(fullQuotedTableName); + var json = _listTablesForSchema(fullQuotedTableName, getContext(toolContext)); if (json.has("sql")) return "```sql\n" + json.getString("sql") + "\n```\n"; else return "I could not find the source for " + fullQuotedTableName; } - - @Tool(description = """ - Save addition information for database columns. If additional metadata is gathered via - chat, it can be saved to improve further interactions. - """) - String saveColumnDescription( - @ToolParam(description = "Fully qualified table or query name as it would appear in SQL e.g. \"schema\".\"table or query\"") - String fullQuotedTableName, - @ToolParam(description = "Quoted column name as it would appear in SQL e.g. \"column name\"") - String quotedColumnName, - @ToolParam(description = "Additional metadata to remember for future use. This will replace any currently saved value") - String columnMetadata - ) - { - McpContext context = McpContext.get(); - var map = PropertyManager.getWritableProperties(context.getContainer(), "QueryMCP.annotations", true); - String fullPath = normalizeIdentifier(fullQuotedTableName + "." + quotedColumnName); - map.put(fullPath, columnMetadata); - try (var ignore = SpringActionController.ignoreSqlUpdates()) - { - map.save(); - } - return new JSONObject(Map.of("success",Boolean.TRUE)).toString(); - } - - /* TODO McpContext setup */ - - static McpContext getContext() - { - try - { - return McpContext.get(); - } - catch (Exception x) - { - return new McpContext(ContainerManager.getHomeContainer(), UserManager.getGuestUser()); - } - } - /* For now, list all schemas. CONSIDER support incremental querying. */ public static Map _listAllSchemas(DefaultSchema root) { @@ -179,7 +139,7 @@ public Map reduce(Map r1, Map Date: Thu, 19 Mar 2026 17:41:48 -0700 Subject: [PATCH 5/5] Provide less severe guidance when container is missing --- api/src/org/labkey/api/mcp/McpException.java | 11 +++++++++++ api/src/org/labkey/api/mcp/McpService.java | 7 ++++++- core/src/org/labkey/core/CoreMcp.java | 2 +- core/src/org/labkey/core/mcp/McpServiceImpl.java | 16 +++++++++++++++- 4 files changed, 33 insertions(+), 3 deletions(-) create mode 100644 api/src/org/labkey/api/mcp/McpException.java diff --git a/api/src/org/labkey/api/mcp/McpException.java b/api/src/org/labkey/api/mcp/McpException.java new file mode 100644 index 00000000000..e9270a8f88f --- /dev/null +++ b/api/src/org/labkey/api/mcp/McpException.java @@ -0,0 +1,11 @@ +package org.labkey.api.mcp; + +// A special exception that MCP endpoints can throw when they want to provide guidance to the client without making +// it a big red error. The message will be extracted and sent as text to the client. +public class McpException extends RuntimeException +{ + public McpException(String message) + { + super(message); + } +} diff --git a/api/src/org/labkey/api/mcp/McpService.java b/api/src/org/labkey/api/mcp/McpService.java index c8d7da6115a..cc6273a7d47 100644 --- a/api/src/org/labkey/api/mcp/McpService.java +++ b/api/src/org/labkey/api/mcp/McpService.java @@ -42,9 +42,14 @@ default ContainerUser getContext(ToolContext toolContext) User user = (User)toolContext.getContext().get("user"); Container container = (Container)toolContext.getContext().get("container"); if (container == null) - throw new IllegalArgumentException("You need to set a container path before invoking this tool"); + throw new McpException("You need to set a container path before invoking this tool"); return ContainerUser.create(container, user); } + + default User getUser(ToolContext toolContext) + { + return (User)toolContext.getContext().get("user"); + } } static @NotNull McpService get() diff --git a/core/src/org/labkey/core/CoreMcp.java b/core/src/org/labkey/core/CoreMcp.java index 33e88b9b108..55ef74486da 100644 --- a/core/src/org/labkey/core/CoreMcp.java +++ b/core/src/org/labkey/core/CoreMcp.java @@ -71,7 +71,7 @@ String whereAmIWhoAmITalkingTo(ToolContext context) @Tool(description = "List the hierarchical path for every container in the server where the user has read permissions.") String listContainers(ToolContext toolContext) { - return ContainerManager.getAllChildren(ContainerManager.getRoot(), getContext(toolContext).getUser(), ReadPermission.class) + return ContainerManager.getAllChildren(ContainerManager.getRoot(), getUser(toolContext), ReadPermission.class) .stream() .map(Container::getPath) .collect(LabKeyCollectors.toJSONArray()) diff --git a/core/src/org/labkey/core/mcp/McpServiceImpl.java b/core/src/org/labkey/core/mcp/McpServiceImpl.java index 2ed41ed863e..48dce674460 100644 --- a/core/src/org/labkey/core/mcp/McpServiceImpl.java +++ b/core/src/org/labkey/core/mcp/McpServiceImpl.java @@ -29,6 +29,7 @@ import org.labkey.api.data.ContainerManager; import org.labkey.api.markdown.MarkdownService; import org.labkey.api.mcp.McpContext; +import org.labkey.api.mcp.McpException; import org.labkey.api.mcp.McpService; import org.labkey.api.security.User; import org.labkey.api.util.ContextListener; @@ -69,6 +70,7 @@ import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.definition.ToolDefinition; +import org.springframework.ai.tool.execution.ToolExecutionException; import org.springframework.ai.tool.metadata.ToolMetadata; import org.springframework.ai.vectorstore.SearchRequest; import org.springframework.ai.vectorstore.SimpleVectorStore; @@ -289,7 +291,19 @@ private McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCall { throw new RuntimeException(e); } - String result = toolCallback.call(toolInput, toolContext); + String result; + try + { + result = toolCallback.call(toolInput, toolContext); + } + catch (ToolExecutionException e) + { + // If a tool threw McpException then just send back the message without making a big fuss + if (e.getCause() instanceof McpException) + result = e.getMessage(); + else + throw e; + } return new McpSchema.CallToolResult( List.of( new McpSchema.TextContent(result)