Calculated Column Expression Assistant

api/src/org/labkey/api/mcp/AbstractAgentAction.java

@@ -4,32 +4,48 @@
 import com.google.genai.errors.ServerException;
 import jakarta.servlet.http.HttpSession;
 import org.apache.commons.lang3.StringUtils;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.json.JSONObject;
 import org.labkey.api.action.ReadOnlyApiAction;
+import org.labkey.api.util.GUID;
 import org.labkey.api.util.HtmlString;
+import org.labkey.api.util.SessionHelper;
 import org.springframework.ai.chat.client.ChatClient;
 import org.springframework.validation.BindException;
+import org.springframework.validation.Errors;
 
+import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.Map;
+import java.util.Set;
 
 import static org.apache.commons.lang3.StringUtils.isNotBlank;
+import static org.labkey.api.action.SpringActionController.ERROR_GENERIC;
 
 /**
- * "agent" it is too strong a word, but if you want to create a tools specific chat endpoint then
+ * "Agent" it is too strong a word, but if you want to create a tools-specific chat endpoint, then
  * start here.
  * First implement getServicePrompt() to tell your "agent its mission.  You can also listen in on the
- * conversation to help you user get the right results.
+ * conversation to help your user get the right results.
  */
 public abstract class AbstractAgentAction<F extends PromptForm> extends ReadOnlyApiAction<F>
 {
+    private static final int MAX_ISSUED_CONVERSATION_IDS = 16;
+
+    protected GUID conversationId;
+
     protected abstract String getAgentName();
 
     protected abstract String getServicePrompt();
 
     protected ChatClient getChat(boolean create)
     {
-        HttpSession session = getViewContext().getRequest().getSession(true);
-        ChatClient chatSession = McpService.get().getChat(session, getAgentName(), this::getServicePrompt, create);
+        String conversationName = getAgentName() + ":" + getConversationId();
+
+        HttpSession session = getViewContext().getSession();
+        ChatClient chatSession = McpService.get().getChat(session, conversationName, this::getServicePrompt, create);
+
         return chatSession;
     }
 
@@ -43,38 +59,68 @@ protected String handleEscape(String prompt)
                 ChatClient chatSession = getChat(false); // CONSIDER: getChat(boolean ifStarted)
                 if (null != chatSession)
                     McpService.get().close(getViewContext().getSession(), chatSession);
-                 return "OK, let's start over.";
+                getIssuedConversationIds(getViewContext().getSession()).remove(conversationId);
+                return "OK, let's start over.";
             }
         }
         return null;
     }
 
+    @Override
+    public void validateForm(F form, Errors errors)
+    {
+        // Only honor a client-supplied conversationId if this agent previously issued this session that
+        // id. Otherwise, generate a fresh one. This prevents a same-session caller from splicing into
+        // another conversation by guessing or replaying a GUID.
+        Set<GUID> issued = getIssuedConversationIds(getViewContext().getSession());
+        String supplied = form.getConversationId();
+        if (supplied != null)
+        {
+            GUID candidate;
+
+            try
+            {
+                candidate = new GUID(supplied);
+            }
+            catch (IllegalArgumentException e)
+            {
+                errors.rejectValue("conversationId", ERROR_GENERIC, "Invalid conversationId");
+                return;
+            }
+
+            if (issued.contains(candidate))
+            {
+                conversationId = candidate;
+                return;
+            }
+        }
+        conversationId = new GUID();
+        issued.add(conversationId);
+    }
+
     @Override
     public Object execute(PromptForm form, BindException errors) throws Exception
     {
-        try (var mcpPush = McpContext.withContext(getViewContext()))
+        try (var _ = McpContext.withContext(getViewContext()))
         {
             String prompt = form.getPrompt();
 
-            String escapeResponse = handleEscape(prompt);
+            JSONObject escapeResponse = escapeResponse(prompt);
             if (null != escapeResponse)
-            {
-                return new JSONObject(Map.of(
-                        "contentType", "text/plain",
-                        "response", escapeResponse,
-                        "success", Boolean.TRUE));
-            }
+                return escapeResponse;
 
             // call getChat() after handleEscape()
             ChatClient chatSession = getChat(true);
             if (null == chatSession)
+            {
                 return new JSONObject(Map.of(
                         "contentType", "text/plain",
                         "response", "Service is not ready yet",
                         "success", Boolean.FALSE));
+            }
 
             McpService.MessageResponse response = McpService.get().sendMessage(chatSession, prompt);
-            var ret = new JSONObject(Map.of("success", Boolean.TRUE));
+            var ret = new JSONObject(Map.of("success", Boolean.TRUE, "conversationId", conversationId));
             if (!HtmlString.isBlank(response.html()))
             {
                 ret.put("contentType", "text/html");
@@ -101,11 +147,52 @@ else if (isNotBlank(response.text()))
         }
         catch (ClientException ex)
         {
-            var ret = new JSONObject(Map.of(
-                    "text", ex.getMessage(),
-                    "user", getViewContext().getUser().getName(),
-                    "success", Boolean.FALSE));
-            return ret;
+            return errorResponse(ex);
         }
     }
+
+    protected @NotNull JSONObject errorResponse(Exception ex)
+    {
+        return new JSONObject(Map.of(
+                "text", ex.getMessage(),
+                "user", getViewContext().getUser().getName(),
+                "success", Boolean.FALSE));
+    }
+
+    protected @Nullable JSONObject escapeResponse(String prompt)
+    {
+        String escapeResponse = handleEscape(prompt);
+        if (null == escapeResponse)
+            return null;
+
+        return new JSONObject(Map.of(
+                "contentType", "text/plain",
+                "text", escapeResponse,
+                "conversationId", conversationId,
+                "success", Boolean.TRUE));
+    }
+
+    protected GUID getConversationId()
+    {
+        return conversationId;
+    }
+
+    private String getIssuedConversationIdsKey()
+    {
+        return getAgentName() + "#issuedConversationIds";
+    }
+
+    private Set<GUID> getIssuedConversationIds(HttpSession session)
+    {
+        return SessionHelper.getAttribute(session, getIssuedConversationIdsKey(), () ->
+                Collections.synchronizedSet(Collections.newSetFromMap(
+                        new LinkedHashMap<>(16, 0.75f, false)
+                        {
+                            @Override
+                            protected boolean removeEldestEntry(Map.Entry<GUID, Boolean> eldest)
+                            {
+                                return size() > MAX_ISSUED_CONVERSATION_IDS;
+                            }
+                        })));
+    }
 }

api/src/org/labkey/api/mcp/McpContext.java

@@ -8,6 +8,8 @@
 import org.labkey.api.view.UnauthorizedException;
 import org.labkey.api.writer.ContainerUser;
 import org.springframework.ai.chat.model.ToolContext;
+
+import java.util.HashMap;
 import java.util.Map;
 
 /**
@@ -19,11 +21,11 @@ public class McpContext implements ContainerUser
 {
     final User user;
     final Container container;
+    final Map<String, Object> attributes = new HashMap<>();
 
     public McpContext(ContainerUser ctx)
     {
-        this.container = ctx.getContainer();
-        this.user = ctx.getUser();
+        this(ctx.getContainer(), ctx.getUser());
     }
 
     public McpContext(Container container, User user)
@@ -36,9 +38,24 @@ public McpContext(Container container, User user)
 
     public ToolContext getToolContext()
     {
-        return new ToolContext(Map.of("container", getContainer(), "user", getUser()));
+        Map<String, Object> map = new HashMap<>(attributes);
+        map.put("container", getContainer());
+        map.put("user", getUser());
+        return new ToolContext(map);
     }
 
+    public McpContext put(String key, Object value)
+    {
+        if ("container".equals(key) || "user".equals(key))
+            throw new IllegalArgumentException("Reserved key: " + key);
+        attributes.put(key, value);
+        return this;
+    }
+
+    public Object get(String key)
+    {
+        return attributes.get(key);
+    }
 
     @Override
     public Container getContainer()
@@ -52,10 +69,9 @@ public User getUser()
         return user;
     }
 
-
     //
     // I'd like to get away from using ThreadLocal, but I haven't
-    // researched if there are other ways to pass context around to Tools registerd by McpService
+    // researched if there are other ways to pass context around to Tools registered by McpService
     //
 
     private static final ThreadLocal<McpContext> contexts = new ThreadLocal();

api/src/org/labkey/api/mcp/McpInternal.java

@@ -0,0 +1,22 @@
+package org.labkey.api.mcp;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Marks an MCP tool as internal-only. Internal tools are available to in-process callers (e.g.,
+ * {@link AbstractAgentAction}-backed agents that supply tool-specific context via {@link McpContext})
+ * but are filtered out of the externally exposed MCP server: they do not appear in {@code listTools}
+ * and direct {@code callTool} requests for them are rejected.
+ * <p>
+ * Use this for tools that require request-scoped context the LLM cannot reasonably supply on its own.
+ */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+public @interface McpInternal
+{
+    /** Optional reason surfaced in startup logs. */
+    String value() default "";
+}

api/src/org/labkey/api/mcp/McpService.java

@@ -1,6 +1,5 @@
 package org.labkey.api.mcp;
 
-
 import io.modelcontextprotocol.server.McpServerFeatures;
 import jakarta.servlet.http.HttpSession;
 import org.apache.logging.log4j.Logger;
@@ -43,20 +42,20 @@
 ///     permission annotation is required, otherwise your tool will not be registered.**
 /// 4.  Add `ToolContext` as the first parameter to the method
 /// 5.  Add additional required or optional parameters to the method signature, as needed. Note that "required" is the
-///     default. Again here, the parameter descriptions are very important. Provide examples of parameter values.
+///     default. Again, here, the parameter descriptions are very important. Provide examples of parameter values.
 /// 6.  Use the helper method `getContext(ToolContext)` to retrieve the current `Container` and `User`
 /// 7.  Use the helper method `getUser(ToolContext)` in the rare cases where you need just a `User`
 /// 8.  Perform additional permissions checking (beyond what the annotations offer), where appropriate
 /// 9.  Filter all results to the current container, of course
 /// 10. For any error conditions, throw exceptions with detailed information. These will get translated into appropriate
-///     failure responses and the LLM client will attempt to correct any problems (hopefully).
+///     failure responses, and the LLM client will attempt to correct any problems (hopefully).
 /// 11. For success cases, return a String with a message or JSON content, for example, `JSONObject.toString()`. Spring
 ///     has some limited ability to convert other objects into JSON strings, but we haven't experimented with that. See
 ///     `DefaultToolCallResultConverter` and the ability to provide a custom result converter via the `@Tool` annotation.
 ///
 /// At registration time, the framework will:
 /// - Ensure all tools are annotated for permissions
-/// - Ensure there aren't multiple tools with the same name
+/// - Ensure there are not multiple tools with the same name
 ///
 /// On every tool request, before invoking any tool code, the framework will:
 /// - Authenticate the user or provide a guest user
@@ -88,7 +87,7 @@ interface McpImpl
     {
         default ContainerUser getContext(ToolContext toolContext)
         {
-            User user = (User)toolContext.getContext().get("user");
+            User user = getUser(toolContext);
             Container container = (Container)toolContext.getContext().get("container");
             if (container == null)
                 throw new McpException("No container path is set. Ask the user which container/folder they want to use (you can call listContainers to show available options), then call setContainer before retrying.");
@@ -103,7 +102,7 @@ default User getUser(ToolContext toolContext)
         // Every MCP resource should call this on every invocation
         default void incrementResourceRequestCount(String resource)
         {
-            if (!OptionalFeatureService.get().isFeatureEnabled(ENABLE_MCP_SERVER_FLAG))
+            if (!get().isEnabled())
                 throw new RuntimeException("The MCP server is not enabled for external requests. Consider toggling the experimental feature flag.");
 
             get().incrementResourceRequestCount(resource);
@@ -123,6 +122,11 @@ static void setInstance(McpService service)
         ServiceRegistry.get().registerService(McpService.class, service);
     }
 
+    default boolean isEnabled()
+    {
+        return OptionalFeatureService.get().isFeatureEnabled(ENABLE_MCP_SERVER_FLAG);
+    }
+
     boolean isReady();
 
     // Register MCPs in Module.startup()
@@ -153,22 +157,22 @@ default void register(McpImpl mcp)
     @Override
     ToolCallback @NotNull [] getToolCallbacks();
 
-    default ChatClient getChat(HttpSession session, String agentName, Supplier<String> systemPromptSupplier)
+    default ChatClient getChat(HttpSession session, String conversationName, Supplier<String> systemPromptSupplier)
     {
-        return getChat(session, agentName, systemPromptSupplier, true);
+        return getChat(session, conversationName, systemPromptSupplier, true);
     }
 
     void saveSessionContainer(ToolContext context, Container container);
 
     void incrementResourceRequestCount(String resource);
 
-    ChatClient getChat(HttpSession session, String agentName, Supplier<String> systemPromptSupplier, boolean createIfNotExists);
+    ChatClient getChat(HttpSession session, String conversationName, Supplier<String> systemPromptSupplier, boolean createIfNotExists);
 
     void close(HttpSession session, ChatClient chat);
 
     record MessageResponse(String contentType, String text, HtmlString html) {}
 
-    /** get consolidated response (good for many text oriented agents/use-cases) */
+    /** get a consolidated response (good for many text-oriented agents/use-cases) */
     MessageResponse sendMessage(ChatClient chat, String message);
 
     /** get individual response parts, useful for agents that generate SQL or programmatic responses */