fix(ai) - sync AI agent step output schema when agent response format changes (#22466)

## Summary When an AI agent workflow step is built via the AI chat tools, its persisted `settings.outputSchema` was left empty (or stale as a text `{ response }` schema) even after the agent was given a structured JSON `responseFormat`. The workflow still executed correctly (runtime uses actual step results), but the builder UI resolves downstream variables (`{{stepId.fieldName}}`) exclusively from the persisted `outputSchema`, so those variables showed as **"Not Found"**. Root cause: the `update_agent` tool only mutated the agent entity and never re-derived the linked step's `outputSchema`, and `enrichOutputSchema` did not handle `AI_AGENT` steps at all. ## What changed - **Enrich AI_AGENT output schema on the backend**: added `AI_AGENT` to `BACKEND_ENRICHED_TYPES` in `WorkflowSchemaWorkspaceService.enrichOutputSchema`, so a step's `outputSchema` is computed from the agent's `responseFormat` on every create/update (text → `{ response }`, JSON → one field per property). - **Re-sync the step when the agent's response format changes**: after `update_agent` sets a `responseFormat`, the tool now finds the draft workflow version(s) whose `AI_AGENT` step references that agent and re-runs the step update so the persisted `outputSchema` is regenerated. - **Fix stale-cache read**: `updateOneAgent` reads `flatAgentMaps` before its migration, which can leave a memoized/local stale copy for a few seconds. The resync now invalidates `flatAgentMaps` before re-enriching, so the fresh `responseFormat` is used. - **Surface failures**: resync errors are logged (`UpdateAgentTool`) instead of failing silently; the agent update itself still succeeds. - Added unit tests for the `update_agent` resync behavior (fires on `responseFormat` change, invalidates the cache, skips unrelated agents, and reports success when the resync fails).  <a href="https://cubic.dev/pr/twentyhq/twenty/pull/22466?utm_source=github" target="_blank" rel="noopener noreferrer" data-no-image-dialog="true"><picture><source media="(prefers-color-scheme: dark)" srcset="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://www.cubic.dev/buttons/review-in-cubic-light.svg"><img alt="Review in cubic" src="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"></picture></a>
2026-07-03 10:47:12 +02:00
parent caa282f5da
commit 90f35658c5
3 changed files with 229 additions and 2 deletions
@@ -175,7 +175,10 @@ export class WorkflowSchemaWorkspaceService {
    workspaceId: string;
    workflowVersionId: string;
  }): Promise<WorkflowAction> {
-    const BACKEND_ENRICHED_TYPES = [WorkflowActionType.ITERATOR];
+    const BACKEND_ENRICHED_TYPES = [
+      WorkflowActionType.ITERATOR,
+      WorkflowActionType.AI_AGENT,
+    ];

    if (!BACKEND_ENRICHED_TYPES.includes(step.type)) {
      return step;
@@ -0,0 +1,151 @@
+import { createUpdateAgentTool } from 'src/modules/workflow/workflow-tools/tools/update-agent.tool';
+
+const AGENT_ID = 'f47ac10b-58cc-4372-a567-0e02b2c3d479';
+const WORKSPACE_ID = 'workspace-id';
+
+const buildAiAgentStep = (agentId: string, stepId = 'step-1') => ({
+  id: stepId,
+  name: 'AI Agent',
+  type: 'AI_AGENT',
+  valid: true,
+  settings: { input: { agentId }, outputSchema: {} },
+});
+
+const buildTool = ({
+  draftVersions = [],
+}: {
+  draftVersions?: { id: string; status: string; steps: unknown[] | null }[];
+} = {}) => {
+  const agentService = {
+    updateOneAgent: jest.fn().mockResolvedValue({ id: AGENT_ID }),
+  };
+  const workflowVersionStepService = {
+    updateWorkflowVersionStep: jest.fn().mockResolvedValue(undefined),
+  };
+  const workflowVersionRepository = {
+    find: jest.fn().mockResolvedValue(draftVersions),
+  };
+  const globalWorkspaceOrmManager = {
+    getRepository: jest.fn().mockResolvedValue(workflowVersionRepository),
+  };
+  const flatEntityMapsCacheService = {
+    invalidateFlatEntityMaps: jest.fn().mockResolvedValue(undefined),
+  };
+
+  const tool = createUpdateAgentTool(
+    {
+      agentService,
+      workflowVersionStepService,
+      globalWorkspaceOrmManager,
+      flatEntityMapsCacheService,
+    } as never,
+    { workspaceId: WORKSPACE_ID },
+  );
+
+  return {
+    tool,
+    agentService,
+    workflowVersionStepService,
+    globalWorkspaceOrmManager,
+    flatEntityMapsCacheService,
+  };
+};
+
+describe('createUpdateAgentTool', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('should resync the linked AI_AGENT step output schema when responseFormat changes', async () => {
+    const step = buildAiAgentStep(AGENT_ID);
+    const { tool, workflowVersionStepService, flatEntityMapsCacheService } =
+      buildTool({
+        draftVersions: [{ id: 'version-1', status: 'DRAFT', steps: [step] }],
+      });
+
+    const result = (await tool.execute({
+      agentId: AGENT_ID,
+      responseFormat: {
+        type: 'json',
+        schema: {
+          type: 'object',
+          properties: { summary: { type: 'string' } },
+        },
+      },
+    } as never)) as Record<string, unknown>;
+
+    expect(result.success).toBe(true);
+    expect(
+      flatEntityMapsCacheService.invalidateFlatEntityMaps,
+    ).toHaveBeenCalledWith({
+      workspaceId: WORKSPACE_ID,
+      flatMapsKeys: ['flatAgentMaps'],
+    });
+    expect(
+      workflowVersionStepService.updateWorkflowVersionStep,
+    ).toHaveBeenCalledWith({
+      workspaceId: WORKSPACE_ID,
+      workflowVersionId: 'version-1',
+      step,
+    });
+  });
+
+  it('should not resync when responseFormat is not provided', async () => {
+    const step = buildAiAgentStep(AGENT_ID);
+    const { tool, workflowVersionStepService, globalWorkspaceOrmManager } =
+      buildTool({
+        draftVersions: [{ id: 'version-1', status: 'DRAFT', steps: [step] }],
+      });
+
+    const result = (await tool.execute({
+      agentId: AGENT_ID,
+      prompt: 'You are a helpful assistant',
+    } as never)) as Record<string, unknown>;
+
+    expect(result.success).toBe(true);
+    expect(globalWorkspaceOrmManager.getRepository).not.toHaveBeenCalled();
+    expect(
+      workflowVersionStepService.updateWorkflowVersionStep,
+    ).not.toHaveBeenCalled();
+  });
+
+  it('should skip steps referencing a different agent', async () => {
+    const { tool, workflowVersionStepService } = buildTool({
+      draftVersions: [
+        {
+          id: 'version-1',
+          status: 'DRAFT',
+          steps: [buildAiAgentStep('another-agent-id')],
+        },
+      ],
+    });
+
+    await tool.execute({
+      agentId: AGENT_ID,
+      responseFormat: { type: 'text' },
+    } as never);
+
+    expect(
+      workflowVersionStepService.updateWorkflowVersionStep,
+    ).not.toHaveBeenCalled();
+  });
+
+  it('should still report agent update success when the resync fails', async () => {
+    const step = buildAiAgentStep(AGENT_ID);
+    const { tool, workflowVersionStepService } = buildTool({
+      draftVersions: [{ id: 'version-1', status: 'DRAFT', steps: [step] }],
+    });
+
+    workflowVersionStepService.updateWorkflowVersionStep.mockRejectedValue(
+      new Error('boom'),
+    );
+
+    const result = (await tool.execute({
+      agentId: AGENT_ID,
+      responseFormat: { type: 'text' },
+    } as never)) as Record<string, unknown>;
+
+    expect(result.success).toBe(true);
+    expect(result.message).toContain('failed to resync');
+  });
+});
@@ -1,8 +1,13 @@
 import { isDefined } from 'twenty-shared/utils';
+import { WorkflowActionType } from 'twenty-shared/workflow';
 import { z } from 'zod';

 import { type AgentResponseFormat } from 'src/engine/metadata-modules/ai/ai-agent/types/agent-response-format.type';
 import { type ModelId } from 'src/engine/metadata-modules/ai/ai-models/types/model-id.type';
+import {
+  WorkflowVersionStatus,
+  type WorkflowVersionWorkspaceEntity,
+} from 'src/modules/workflow/common/standard-objects/workflow-version.workspace-entity';
 import {
  type WorkflowToolContext,
  type WorkflowToolDependencies,
@@ -53,8 +58,64 @@ const updateAgentSchema = z.object({
    ),
 });

+const resyncAiAgentStepOutputSchemas = async (
+  deps: Pick<
+    WorkflowToolDependencies,
+    | 'workflowVersionStepService'
+    | 'globalWorkspaceOrmManager'
+    | 'flatEntityMapsCacheService'
+  >,
+  { workspaceId, agentId }: { workspaceId: string; agentId: string },
+): Promise<void> => {
+  await deps.flatEntityMapsCacheService.invalidateFlatEntityMaps({
+    workspaceId,
+    flatMapsKeys: ['flatAgentMaps'],
+  });
+
+  const workflowVersionRepository =
+    await deps.globalWorkspaceOrmManager.getRepository<WorkflowVersionWorkspaceEntity>(
+      workspaceId,
+      'workflowVersion',
+      { shouldBypassPermissionChecks: true },
+    );
+
+  const draftVersions = await workflowVersionRepository.find({
+    where: { status: WorkflowVersionStatus.DRAFT },
+  });
+
+  for (const version of draftVersions) {
+    const steps = version.steps;
+
+    if (!isDefined(steps)) {
+      continue;
+    }
+
+    const matchingStep = steps.find(
+      (step) =>
+        step.type === WorkflowActionType.AI_AGENT &&
+        step.settings?.input?.agentId === agentId,
+    );
+
+    if (!isDefined(matchingStep)) {
+      continue;
+    }
+
+    await deps.workflowVersionStepService.updateWorkflowVersionStep({
+      workspaceId,
+      workflowVersionId: version.id,
+      step: matchingStep,
+    });
+  }
+};
+
 export const createUpdateAgentTool = (
-  deps: Pick<WorkflowToolDependencies, 'agentService'>,
+  deps: Pick<
+    WorkflowToolDependencies,
+    | 'agentService'
+    | 'workflowVersionStepService'
+    | 'globalWorkspaceOrmManager'
+    | 'flatEntityMapsCacheService'
+  >,
  context: WorkflowToolContext,
 ) => ({
  name: 'update_agent' as const,
@@ -88,6 +149,18 @@ To find the agentId, look at the AI_AGENT step's settings.input.agentId field.`,
        workspaceId,
      });

+      if (isDefined(responseFormat)) {
+        try {
+          await resyncAiAgentStepOutputSchemas(deps, { workspaceId, agentId });
+        } catch (resyncError) {
+          return {
+            success: true,
+            message: `Successfully updated agent ${agentId}, but failed to resync workflow step output schema: ${resyncError.message}`,
+            agentId: updatedAgent.id,
+          };
+        }
+      }
+
      return {
        success: true,
        message: `Successfully updated agent ${agentId}`,