From 90f35658c51b68741e3625c407534ab50ae4e413 Mon Sep 17 00:00:00 2001 From: Etienne <45695613+etiennejouan@users.noreply.github.com> Date: Fri, 3 Jul 2026 10:47:12 +0200 Subject: [PATCH] fix(ai) - sync AI agent step output schema when agent response format changes (#22466) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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). Review in cubic --- .../workflow-schema.workspace-service.ts | 5 +- .../tools/__tests__/update-agent.tool.spec.ts | 151 ++++++++++++++++++ .../workflow-tools/tools/update-agent.tool.ts | 75 ++++++++- 3 files changed, 229 insertions(+), 2 deletions(-) create mode 100644 packages/twenty-server/src/modules/workflow/workflow-tools/tools/__tests__/update-agent.tool.spec.ts diff --git a/packages/twenty-server/src/modules/workflow/workflow-builder/workflow-schema/workflow-schema.workspace-service.ts b/packages/twenty-server/src/modules/workflow/workflow-builder/workflow-schema/workflow-schema.workspace-service.ts index 8748fe2004..9a1071e8d0 100644 --- a/packages/twenty-server/src/modules/workflow/workflow-builder/workflow-schema/workflow-schema.workspace-service.ts +++ b/packages/twenty-server/src/modules/workflow/workflow-builder/workflow-schema/workflow-schema.workspace-service.ts @@ -175,7 +175,10 @@ export class WorkflowSchemaWorkspaceService { workspaceId: string; workflowVersionId: string; }): Promise { - const BACKEND_ENRICHED_TYPES = [WorkflowActionType.ITERATOR]; + const BACKEND_ENRICHED_TYPES = [ + WorkflowActionType.ITERATOR, + WorkflowActionType.AI_AGENT, + ]; if (!BACKEND_ENRICHED_TYPES.includes(step.type)) { return step; diff --git a/packages/twenty-server/src/modules/workflow/workflow-tools/tools/__tests__/update-agent.tool.spec.ts b/packages/twenty-server/src/modules/workflow/workflow-tools/tools/__tests__/update-agent.tool.spec.ts new file mode 100644 index 0000000000..76a8ec58ed --- /dev/null +++ b/packages/twenty-server/src/modules/workflow/workflow-tools/tools/__tests__/update-agent.tool.spec.ts @@ -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; + + 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; + + 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; + + expect(result.success).toBe(true); + expect(result.message).toContain('failed to resync'); + }); +}); diff --git a/packages/twenty-server/src/modules/workflow/workflow-tools/tools/update-agent.tool.ts b/packages/twenty-server/src/modules/workflow/workflow-tools/tools/update-agent.tool.ts index 179d3a30ae..9f1ff3b5d0 100644 --- a/packages/twenty-server/src/modules/workflow/workflow-tools/tools/update-agent.tool.ts +++ b/packages/twenty-server/src/modules/workflow/workflow-tools/tools/update-agent.tool.ts @@ -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 => { + await deps.flatEntityMapsCacheService.invalidateFlatEntityMaps({ + workspaceId, + flatMapsKeys: ['flatAgentMaps'], + }); + + const workflowVersionRepository = + await deps.globalWorkspaceOrmManager.getRepository( + 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, + 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}`,