Add generate-image agent skill

frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/index.jsx

@@ -33,6 +33,7 @@ export default function AgentSkillsTab({
   // All skill state management
   const {
     fileSystemAgentAvailable,
+    imageGenerationAvailable,
     importedSkills,
     flows,
     mcpServers,
@@ -51,6 +52,7 @@ export default function AgentSkillsTab({
 
   const configurableSkills = getConfigurableSkills(t, {
     fileSystemAgentAvailable,
+    imageGenerationAvailable,
   });
 
   // UI state

frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/useAgentSkillsState.js

@@ -16,6 +16,8 @@ export default function useAgentSkillsState(defaultSkills) {
   // Core skill state
   const [fileSystemAgentAvailable, setFileSystemAgentAvailable] =
     useState(false);
+  const [imageGenerationAvailable, setImageGenerationAvailable] =
+    useState(false);
   const [isMultiUser, setIsMultiUser] = useState(false);
   const [disabledDefaults, setDisabledDefaults] = useState([]);
   const [enabledConfigurable, setEnabledConfigurable] = useState([]);
@@ -37,7 +39,7 @@ export default function useAgentSkillsState(defaultSkills) {
   async function fetchSkillSettings() {
     try {
       const subSkillPrefKeys = getSubSkillPreferenceKeys();
-      const [prefs, flowsRes, fsAgentAvailable, multiUserMode] =
+      const [prefs, flowsRes, fsAgentAvailable, multiUserMode, keys] =
         await Promise.all([
           Admin.systemPreferencesByFields([
             "disabled_agent_skills",
@@ -48,6 +50,7 @@ export default function useAgentSkillsState(defaultSkills) {
           AgentFlows.listFlows(),
           System.isFileSystemAgentAvailable(),
           System.isMultiUserMode(),
+          System.keys(),
         ]);
 
       if (prefs?.settings) {
@@ -58,6 +61,7 @@ export default function useAgentSkillsState(defaultSkills) {
       }
       if (flowsRes?.flows) setFlows(flowsRes.flows);
       setFileSystemAgentAvailable(fsAgentAvailable);
+      setImageGenerationAvailable(!!keys?.ImageGenerationProvider);
       setIsMultiUser(!!multiUserMode);
     } catch (e) {
       console.error(e);
@@ -172,6 +176,7 @@ export default function useAgentSkillsState(defaultSkills) {
   return {
     // State
     fileSystemAgentAvailable,
+    imageGenerationAvailable,
     isMultiUser,
     disabledDefaults,
     enabledConfigurable,

frontend/src/locales/en/common.js

@@ -320,6 +320,11 @@ const TRANSLATIONS = {
         description:
           "Enable the default agent to generate various types of charts from data provided or given in chat.",
       },
+      generateImage: {
+        title: "Generate images",
+        description:
+          "Allow the agent to generate images from chat, or edit images attached to the conversation, using your configured image generation provider.",
+      },
       web: {
         title: "Web Search",
         description:

frontend/src/pages/Admin/Agents/skills.jsx

@@ -15,11 +15,13 @@ import {
   FolderOpen,
   FilePlus,
   CalendarCheck,
+  ImageSquare,
 } from "@phosphor-icons/react";
 import RAGImage from "@/media/agents/rag-memory.png";
 import SummarizeImage from "@/media/agents/view-summarize.png";
 import ScrapeWebsitesImage from "@/media/agents/scrape-websites.png";
 import GenerateChartsImage from "@/media/agents/generate-charts.png";
+import GenerateImageImage from "@/media/agents/generate-image.png";
 import GenerateSaveImages from "@/media/agents/generate-save-files.png";
 import FileSystemImage from "@/media/agents/file-system.png";
 import GMailIcon from "./GMailSkillPanel/gmail.png";
@@ -60,11 +62,18 @@ export const getDefaultSkills = (t) => ({
  * @param {object} options - The options for the configurable skills.
  * @param {boolean} options.fileSystemAgentAvailable - Whether the file system agent is available.
  * @param {boolean} options.createFilesAgentAvailable - Whether the create files agent is available.
+ * @param {boolean} options.imageGenerationAvailable - Whether an image generation provider is
+ * configured. Only the in-chat skills menu passes this - the admin page always lists the skill so
+ * it can be enabled before a provider is set up.
  * @returns {object} The configurable skills.
  */
 export const getConfigurableSkills = (
   t,
-  { fileSystemAgentAvailable = true, createFilesAgentAvailable = true } = {}
+  {
+    fileSystemAgentAvailable = true,
+    createFilesAgentAvailable = true,
+    imageGenerationAvailable = true,
+  } = {}
 ) => ({
   ...(fileSystemAgentAvailable && {
     "filesystem-agent": {
@@ -94,6 +103,16 @@ export const getConfigurableSkills = (
     icon: ChartBar,
     image: GenerateChartsImage,
   },
+  ...(imageGenerationAvailable && {
+    "generate-image": {
+      title: t("agent.skill.generateImage.title"),
+      description: t("agent.skill.generateImage.description"),
+      component: GenericSkillPanel,
+      skill: "generate-image",
+      icon: ImageSquare,
+      image: GenerateImageImage,
+    },
+  }),
   "web-browsing": {
     title: t("agent.skill.web.title"),
     description: t("agent.skill.web.description"),

frontend/src/utils/chat/agent.js

@@ -65,6 +65,7 @@ const handledEvents = [
   "statusResponse",
   "fileDownloadCard",
   "imageGenerationCard",
+  "imageGenerationPending",
   "scheduledJobCreated",
   "awaitingFeedback",
   "wssFailure",
@@ -314,16 +315,41 @@ export default function handleSocketResponse(socket, event, setChatHistory) {
     });
   }
 
-  if (data.type === "imageGenerationCard") {
+  if (data.type === "imageGenerationPending") {
     return setChatHistory((prev) => {
       return [
         ...prev.filter((msg) => !!msg.content),
+        {
+          uuid: data.content.pendingId,
+          type: "imageGenerationPending",
+          content: data.content.prompt,
+          role: "assistant",
+          sources: [],
+          closed: false,
+          error: null,
+          animate: false,
+          pending: true,
+          metrics: {},
+        },
+      ];
+    });
+  }
+
+  if (data.type === "imageGenerationCard") {
+    return setChatHistory((prev) => {
+      // Drops the placeholder card this result belongs to, if there was one.
+      const history = prev.filter(
+        (msg) => !!msg.content && msg.uuid !== data.content.pendingId
+      );
+      if (data.content.failed) return history;
+      return [
+        ...history,
         {
           uuid: v4(),
           type: "textResponse",
-          content: data.content,
-          outputs: data.outputs || [],
-          chatId: data.chatId || null,
+          content: data.content.text,
+          outputs: data.content.outputs || [],
+          chatId: data.content.chatId || null,
           role: "assistant",
           sources: [],
           closed: true,

server/__tests__/utils/agents/aibitat/plugins/generate-image.test.js

@@ -0,0 +1,188 @@
+/* eslint-env jest */
+const fs = require("fs");
+const os = require("os");
+const path = require("path");
+
+// `utils/files` resolves its storage paths from STORAGE_DIR at require time.
+process.env.STORAGE_DIR = fs.mkdtempSync(
+  path.join(os.tmpdir(), "generate-image-test-")
+);
+jest.mock("../../../../../utils/ImageGenerators", () => ({
+  generateImageForWorkspace: jest.fn(),
+  editImageForWorkspace: jest.fn(),
+}));
+jest.mock("../../../../../utils/helpers", () => ({
+  getImageGeneratorProvider: jest.fn(),
+}));
+jest.mock("../../../../../models/workspaceChats", () => ({
+  WorkspaceChats: { _update: jest.fn() },
+}));
+
+const {
+  generateImageForWorkspace,
+  editImageForWorkspace,
+} = require("../../../../../utils/ImageGenerators");
+const { WorkspaceChats } = require("../../../../../models/workspaceChats");
+const {
+  generateImage,
+} = require("../../../../../utils/agents/aibitat/plugins/generate-image.js");
+
+const SAVED_IMAGE = {
+  storageFilename: "img-11111111-2222-3333-4444-555555555555.png",
+  filename: "a-fox.png",
+  fileSize: 100,
+  buffer: Buffer.from("image-bytes"),
+};
+
+beforeAll(() => {
+  const dir = path.join(process.env.STORAGE_DIR, "generated-images");
+  fs.mkdirSync(dir, { recursive: true });
+  fs.writeFileSync(
+    path.join(dir, SAVED_IMAGE.storageFilename),
+    SAVED_IMAGE.buffer
+  );
+});
+
+afterAll(() =>
+  fs.rmSync(process.env.STORAGE_DIR, { recursive: true, force: true })
+);
+
+function setupPlugin(chats = []) {
+  const aibitat = {
+    chats,
+    introspect: jest.fn(),
+    handlerProps: { log: jest.fn() },
+    socket: { send: jest.fn() },
+    function: (config) => (aibitat._fn = config),
+  };
+  generateImage.plugin.call(generateImage).setup(aibitat);
+  return { aibitat, handler: aibitat._fn.handler.bind(aibitat._fn) };
+}
+
+beforeEach(() => jest.clearAllMocks());
+
+describe("generate-image agent skill", () => {
+  test("generates an image and registers the card for live + historical render", async () => {
+    generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE);
+    const { aibitat, handler } = setupPlugin();
+
+    await handler({ prompt: "a red fox in the snow", size: "512x512" });
+
+    expect(generateImageForWorkspace).toHaveBeenCalledWith(
+      expect.objectContaining({
+        prompt: "a red fox in the snow",
+        size: "512x512",
+      })
+    );
+    const expectedOutput = {
+      type: "imageGenerationCard",
+      payload: {
+        storageFilename: SAVED_IMAGE.storageFilename,
+        filename: SAVED_IMAGE.filename,
+        fileSize: SAVED_IMAGE.fileSize,
+        prompt: "a red fox in the snow",
+      },
+    };
+    // The card must only reference the stored file - no image bytes travel over
+    // the socket or into the chat record.
+    expect(aibitat._pendingOutputs).toEqual([expectedOutput]);
+    const [, card] = aibitat.socket.send.mock.calls.find(
+      ([type]) => type === "imageGenerationCard"
+    );
+    expect(card.outputs).toEqual([expectedOutput]);
+  });
+
+  test("writes the image reference to the reserved chat before showing the card", async () => {
+    generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE);
+    const { aibitat, handler } = setupPlugin();
+    aibitat.trackedChatId = 42;
+
+    await handler({ prompt: "a fox" });
+
+    // The serve endpoint only authorizes files referenced by a chat record, so
+    // the reference has to land before the frontend requests the image.
+    expect(WorkspaceChats._update).toHaveBeenCalledWith(42, {
+      response: JSON.stringify({ outputs: aibitat._pendingOutputs }),
+    });
+    expect(WorkspaceChats._update.mock.invocationCallOrder[0]).toBeLessThan(
+      aibitat.socket.send.mock.invocationCallOrder[1]
+    );
+  });
+
+  test("shows a placeholder card and swaps it for the result", async () => {
+    generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE);
+    const { aibitat, handler } = setupPlugin();
+
+    await handler({ prompt: "a fox" });
+
+    const [pending, card] = aibitat.socket.send.mock.calls;
+    expect(pending[0]).toBe("imageGenerationPending");
+    expect(card[0]).toBe("imageGenerationCard");
+    expect(card[1].pendingId).toBe(pending[1].pendingId);
+  });
+
+  test("clears the placeholder card when generation fails", async () => {
+    generateImageForWorkspace.mockRejectedValue(new Error("provider is down"));
+    const { aibitat, handler } = setupPlugin();
+
+    const reply = await handler({ prompt: "a fox" });
+
+    const [pending, failure] = aibitat.socket.send.mock.calls;
+    expect(failure[0]).toBe("imageGenerationCard");
+    expect(failure[1]).toEqual({
+      pendingId: pending[1].pendingId,
+      text: "provider is down",
+      failed: true,
+    });
+    expect(reply).toContain("provider is down");
+  });
+
+  test("drops a size the model invented instead of failing the call", async () => {
+    generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE);
+    const { handler } = setupPlugin();
+
+    await handler({ prompt: "a fox", size: "large" });
+
+    expect(generateImageForWorkspace).toHaveBeenCalledWith(
+      expect.objectContaining({ size: null })
+    );
+  });
+
+  test("edits using the images attached to the last user message", async () => {
+    editImageForWorkspace.mockResolvedValue(SAVED_IMAGE);
+    const { handler } = setupPlugin([
+      { from: "USER", content: "make it blue" },
+      {
+        from: "USER",
+        content: "here is my photo",
+        attachments: [
+          {
+            mime: "image/png",
+            contentString: `data:image/png;base64,${Buffer.from("source").toString("base64")}`,
+          },
+        ],
+      },
+    ]);
+
+    await handler({ prompt: "make it blue", edit: true });
+
+    expect(generateImageForWorkspace).not.toHaveBeenCalled();
+    expect(editImageForWorkspace).toHaveBeenCalledWith(
+      expect.objectContaining({ images: [Buffer.from("source")] })
+    );
+  });
+
+  test("falls back to the image generated earlier in the session when editing", async () => {
+    generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE);
+    editImageForWorkspace.mockResolvedValue(SAVED_IMAGE);
+    const { aibitat, handler } = setupPlugin();
+
+    await handler({ prompt: "a fox" });
+    await handler({ prompt: "now make it blue", edit: true });
+
+    expect(aibitat._lastGeneratedImage).toBe(SAVED_IMAGE.storageFilename);
+    expect(editImageForWorkspace).toHaveBeenCalledWith(
+      expect.objectContaining({ images: [SAVED_IMAGE.buffer] })
+    );
+  });
+});

server/endpoints/agentFileServer.js

@@ -182,12 +182,16 @@ async function findInWorkspaceChats(storageFilename, { user, isMultiUser }) {
   // DB-level filter so we don't load every chat into memory.
   const chats = await WorkspaceChats.where({
     workspaceId: { in: workspaceIds },
-    include: true,
     response: { contains: storageFilename },
   });
 
   for (const chat of chats) {
-    const { outputs = [] } = safeJsonParse(chat.response, { outputs: [] });
+    const { outputs = [], text } = safeJsonParse(chat.response, {
+      outputs: [],
+    });
+    // A hidden chat is either an agent reply still being written (no text yet)
+    // or history the user cleared - only the former can still serve its files.
+    if (!chat.include && !!text) continue;
     const output = outputs.find(
       (o) => o?.payload?.storageFilename === storageFilename
     );