diff --git a/CHANGELOG.md b/CHANGELOG.md index d860a2b803..ac83f96d95 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -41,6 +41,12 @@ and this project adheres to ### Changed +- Global chat can now change multiple workflow steps in a single response. It + receives a full workflow YAML from the Apollo AI server with each step's job + code embedded, and applies the changes together. When a step is open in the + editor, its diff is previewed before applying; previewing several step diffs + at once is a follow-up. + [#4890](https://github.com/OpenFn/lightning/issues/4890) - Redesigned the trigger inspector in the collaborative editor: selecting a trigger now opens a read-only resting panel with an **Edit** button that leads into a guided wizard (Choose → Configure → Finish), replacing the previous diff --git a/assets/js/collaborative-editor/components/AIAssistantPanelWrapper.tsx b/assets/js/collaborative-editor/components/AIAssistantPanelWrapper.tsx index b2c158545a..6858cbf25b 100644 --- a/assets/js/collaborative-editor/components/AIAssistantPanelWrapper.tsx +++ b/assets/js/collaborative-editor/components/AIAssistantPanelWrapper.tsx @@ -560,36 +560,55 @@ export function AIAssistantPanelWrapper({ ); // Hook to handle workflow/job code application logic - const { handleApplyWorkflow, handlePreviewJobCode, handleApplyJobCode } = - useAIWorkflowApplications({ - sessionId, - page: aiMode?.page || 'workflow_template', - currentSession: - sessionId && messages.length > 0 - ? { - messages, - workflowTemplateContext, - } - : null, - currentUserId: user?.id, - aiMode, - workflowActions: { - importWorkflow, - startApplyingWorkflow, - doneApplyingWorkflow, - startApplyingJobCode, - doneApplyingJobCode, - updateJob, - }, - monacoRef, - jobs, - canApplyChanges, - connectionState, - setPreviewingMessageId, - previewingMessageId, - setApplyingMessageId, - appliedMessageIdsRef, - }); + const { + handleApplyWorkflow, + handlePreviewJobCode, + handlePreviewGlobalStep, + handleApplyJobCode, + } = useAIWorkflowApplications({ + sessionId, + page: aiMode?.page || 'workflow_template', + currentSession: + sessionId && messages.length > 0 + ? { + messages, + workflowTemplateContext, + } + : null, + currentUserId: user?.id, + aiMode, + workflowActions: { + importWorkflow, + startApplyingWorkflow, + doneApplyingWorkflow, + startApplyingJobCode, + doneApplyingJobCode, + updateJob, + }, + monacoRef, + jobs, + canApplyChanges, + connectionState, + setPreviewingMessageId, + previewingMessageId, + setApplyingMessageId, + appliedMessageIdsRef, + }); + + // Route auto-preview to the right handler: global messages carry a full + // workflow YAML (the open step's diff is extracted from it), job-code + // messages carry the job body directly. + const handleAutoPreview = useCallback( + (code: string, messageId: string) => { + const message = messages.find(m => m.id === messageId); + if (message?.from_global) { + handlePreviewGlobalStep(code, messageId); + } else { + handlePreviewJobCode(code, messageId); + } + }, + [messages, handlePreviewGlobalStep, handlePreviewJobCode] + ); // Auto-preview job code when AI responds with code // Only for the user who authored the triggering message @@ -599,7 +618,7 @@ export function AIAssistantPanelWrapper({ ? { id: sessionId, session_type: 'workflow_template', messages } : null, currentUserId: user?.id, - onPreview: handlePreviewJobCode, + onPreview: handleAutoPreview, }); // Auto-apply streaming changes as soon as they arrive (before text finishes) @@ -707,7 +726,12 @@ export function AIAssistantPanelWrapper({ messages={messages} isLoading={isLoading} onApplyWorkflow={ - aiMode?.page === 'workflow_template' && !isApplyingWorkflow + (aiMode?.page === 'workflow_template' || + // Global messages apply the full workflow even while a + // job is open; handleApplyWorkflow still no-ops for + // non-global messages outside workflow_template mode. + messages.some(m => m.from_global && m.code)) && + !isApplyingWorkflow ? (yaml, messageId) => { void handleApplyWorkflow(yaml, messageId); } @@ -723,6 +747,8 @@ export function AIAssistantPanelWrapper({ onPreviewJobCode={ aiMode?.page === 'job_code' ? handlePreviewJobCode : undefined } + onPreviewGlobalStep={handlePreviewGlobalStep} + canPreviewGlobalStep={aiMode?.page === 'job_code'} applyingMessageId={ // If anyone is applying (including other users), pass the message ID // to show "APPLYING..." state. Prioritize stored message ID from store, diff --git a/assets/js/collaborative-editor/components/MessageList.tsx b/assets/js/collaborative-editor/components/MessageList.tsx index 48e75a9fca..3845b90976 100644 --- a/assets/js/collaborative-editor/components/MessageList.tsx +++ b/assets/js/collaborative-editor/components/MessageList.tsx @@ -396,6 +396,10 @@ interface MessageListProps { onApplyWorkflow?: ((yaml: string, messageId: string) => void) | undefined; onApplyJobCode?: ((code: string, messageId: string) => void) | undefined; onPreviewJobCode?: ((code: string, messageId: string) => void) | undefined; + /** Per-step diff preview for global messages (extracts the open job's body from the YAML) */ + onPreviewGlobalStep?: ((yaml: string, messageId: string) => void) | undefined; + /** Whether a job is open in the IDE, enabling preview for global messages */ + canPreviewGlobalStep?: boolean; applyingMessageId?: string | null | undefined; previewingMessageId?: string | null | undefined; showAddButtons?: boolean; @@ -412,6 +416,8 @@ export function MessageList({ onApplyWorkflow, onApplyJobCode, onPreviewJobCode, + onPreviewGlobalStep, + canPreviewGlobalStep = false, applyingMessageId, previewingMessageId, showAddButtons = false, @@ -618,7 +624,10 @@ export function MessageList({ code={message.code} showAdd={showAddButtons} showApply={showApplyButton} - showPreview={!!message.job_id} + showPreview={ + !!message.job_id || + (!!message.from_global && canPreviewGlobalStep) + } onApply={() => { if (message.job_id) { onApplyJobCode?.(message.code!, message.id); @@ -627,7 +636,12 @@ export function MessageList({ } }} onPreview={() => { - onPreviewJobCode?.(message.code!, message.id); + if (message.from_global) { + // Per-step diff from the full workflow YAML + onPreviewGlobalStep?.(message.code!, message.id); + } else { + onPreviewJobCode?.(message.code!, message.id); + } }} isApplying={!!applyingMessageId} isPreviewActive={previewingMessageId === message.id} diff --git a/assets/js/collaborative-editor/hooks/useAIWorkflowApplications.ts b/assets/js/collaborative-editor/hooks/useAIWorkflowApplications.ts index 729775704e..bd2f552836 100644 --- a/assets/js/collaborative-editor/hooks/useAIWorkflowApplications.ts +++ b/assets/js/collaborative-editor/hooks/useAIWorkflowApplications.ts @@ -170,7 +170,14 @@ export function useAIWorkflowApplications({ */ const handleApplyWorkflow = useCallback( async (yaml: string, messageId: string) => { - if (!aiMode || aiMode.page !== 'workflow_template') { + if (!aiMode) return; + // Global messages carry a full workflow YAML and may be applied even + // while a job is open (job_code mode). Non-global workflow chat keeps + // the workflow_template-only guard so its Apply stays a no-op when a + // job is open. + const isGlobal = !!currentSession?.messages.find(m => m.id === messageId) + ?.from_global; + if (aiMode.page !== 'workflow_template' && !isGlobal) { console.error( '[AI Assistant] Cannot apply workflow - not in workflow mode', { @@ -179,6 +186,15 @@ export function useAIWorkflowApplications({ ); return; } + + // A global message applied while a step is open leaves an active diff in + // the open step. Clear it so the editor returns to an editable state. + const monaco = monacoRef?.current; + if (previewingMessageId && monaco) { + monaco.clearDiff(); + setPreviewingMessageId(null); + } + setApplyingMessageId(messageId); // Signal to all collaborators that we're starting to apply @@ -218,12 +234,16 @@ export function useAIWorkflowApplications({ } }, [ - aiMode?.page, + aiMode, + currentSession, importWorkflow, startApplyingWorkflow, doneApplyingWorkflow, jobs, setApplyingMessageId, + monacoRef, + previewingMessageId, + setPreviewingMessageId, ] ); @@ -295,6 +315,83 @@ export function useAIWorkflowApplications({ [aiMode, jobs, previewingMessageId, monacoRef, setPreviewingMessageId] ); + /** + * Preview the open job's diff from a global full-workflow YAML message + * + * Mirrors handlePreviewJobCode, but extracts the open job's body from the + * workflow YAML (global messages carry the whole workflow in `code`). + * Shows a diff only when the open step's body actually changed; clears any + * stale diff otherwise. + */ + const handlePreviewGlobalStep = useCallback( + (yaml: string, messageId: string) => { + if (!aiMode || aiMode.page !== 'job_code') return; // only when a step is open + const jobId = (aiMode.context as JobCodeContext).job_id; + if (!jobId) return; + + // Same dedup guards as handlePreviewJobCode + if (previewingMessageId === messageId) return; + if (previewingMessageId === '__streaming__') { + setPreviewingMessageId(messageId); + return; + } + + const currentJob = jobs.find(j => j.id === jobId); + const currentBody = currentJob?.body ?? ''; + + let newBody: string | undefined; + try { + const spec = parseWorkflowYAML(yaml); + // ids from the YAML are preserved, so we match the open step by id + const state = convertWorkflowSpecToState(spec); + newBody = state.jobs.find(j => j.id === jobId)?.body; + } catch (error) { + console.error( + '[AI Assistant] Failed to parse global workflow YAML:', + error + ); + notifications.alert({ + title: 'Could not preview step', + description: + error instanceof Error + ? error.message + : 'The AI server returned invalid workflow YAML.', + }); + return; + } + + if (newBody === undefined) { + // Open step's id wasn't in the YAML, so the server likely didn't preserve it + console.warn( + '[AI Assistant] Open step not found in global workflow YAML', + { jobId } + ); + notifications.warning({ + title: 'Could not preview this step', + description: `Step "${ + currentJob?.name ?? jobId + }" was not found in the AI response (id: ${jobId}). Its ID may not have been preserved by the server.`, + }); + if (previewingMessageId) monacoRef?.current?.clearDiff(); + return; + } + + if (newBody === currentBody) { + // open step genuinely unchanged -> ensure no stale diff is shown + if (previewingMessageId) monacoRef?.current?.clearDiff(); + return; + } + + const monaco = monacoRef?.current; + if (previewingMessageId && monaco) monaco.clearDiff(); + if (monaco) { + monaco.showDiff(currentBody, newBody); + setPreviewingMessageId(messageId); + } + }, + [aiMode, jobs, previewingMessageId, monacoRef, setPreviewingMessageId] + ); + /** * Apply job code to Y.Doc * @@ -453,6 +550,7 @@ export function useAIWorkflowApplications({ return { handleApplyWorkflow, handlePreviewJobCode, + handlePreviewGlobalStep, handleApplyJobCode, }; } diff --git a/assets/js/collaborative-editor/hooks/useAutoPreview.ts b/assets/js/collaborative-editor/hooks/useAutoPreview.ts index 41d45ac074..ac68678000 100644 --- a/assets/js/collaborative-editor/hooks/useAutoPreview.ts +++ b/assets/js/collaborative-editor/hooks/useAutoPreview.ts @@ -94,9 +94,10 @@ export function useAutoPreview({ new Date(b.inserted_at).getTime() - new Date(a.inserted_at).getTime() )[0]; - if (!latestCodeMessage || !latestCodeMessage.job_id) { - return; - } + if (!latestCodeMessage) return; + // Global messages have no job_id but carry a full workflow YAML; the + // onPreview callback handles extracting the open step's body from it. + if (!latestCodeMessage.job_id && !latestCodeMessage.from_global) return; // Skip if we've already auto-previewed this message if (stateRef.current.lastAutoPreviewedMessageId === latestCodeMessage.id) { diff --git a/assets/js/collaborative-editor/lib/AIChannelRegistry.ts b/assets/js/collaborative-editor/lib/AIChannelRegistry.ts index 3fd0484046..01a06deec5 100644 --- a/assets/js/collaborative-editor/lib/AIChannelRegistry.ts +++ b/assets/js/collaborative-editor/lib/AIChannelRegistry.ts @@ -915,14 +915,23 @@ export class AIChannelRegistry { if (context.workflow_id) { params['workflow_id'] = context.workflow_id; } - if (context.code) { - params['code'] = context.code; - } if (context.content) { params['content'] = context.content; } } + // Workflow YAML (applicable to both session types). For global chat the + // `code` slot carries the FULL serialized workflow YAML (every step body + // embedded), not a single job's code. When a step is open the context is + // JobCodeContext-shaped and took the branch above, which does not forward + // `code`. Forwarding it here (outside the branch) ensures the YAML reaches + // Apollo on the *first* turn — the only message sent via the channel join. + // Later turns go through `new_message` and were never affected. Plain job + // chat never sets `context.code`, so this is a no-op there. + if ('code' in context && context.code) { + params['code'] = context.code; + } + // Global assistant flags (applicable to both session types) if ('use_global_assistant' in context && context.use_global_assistant) { params['use_global_assistant'] = true; diff --git a/assets/js/collaborative-editor/types/ai-assistant.ts b/assets/js/collaborative-editor/types/ai-assistant.ts index 0afb9deae0..90bf9d8234 100644 --- a/assets/js/collaborative-editor/types/ai-assistant.ts +++ b/assets/js/collaborative-editor/types/ai-assistant.ts @@ -54,6 +54,11 @@ export interface Message { user_id?: string; user?: MessageUser | null; job_id?: string; + /** + * True when this message came from the global AI assistant. Global + * messages carry a full workflow YAML in `code` and never a `job_id`. + */ + from_global?: boolean; } /** @@ -72,6 +77,10 @@ export interface JobCodeContext { job_body?: string; job_adaptor?: string; workflow_id?: string; + + // Full serialized workflow YAML attached by global chat (sent as the message + // `code`). Present even with a step open, so it lives on the job context too. + code?: string; } /** @@ -101,6 +110,10 @@ export type WorkflowTemplateContext = workflow_id?: string; project_id: string; + + // Full serialized workflow YAML attached by global chat (sent as the + // message `code`), present even when a step is open. + code?: string; }; /** diff --git a/assets/test/collaborative-editor/components/MessageList.test.tsx b/assets/test/collaborative-editor/components/MessageList.test.tsx index 2536d468fb..10a6ffdcf8 100644 --- a/assets/test/collaborative-editor/components/MessageList.test.tsx +++ b/assets/test/collaborative-editor/components/MessageList.test.tsx @@ -823,6 +823,115 @@ describe('MessageList', () => { }); }); + describe('Global Messages (from_global)', () => { + it('renders global message as Generated Workflow with Apply and no Preview on canvas', async () => { + const onApplyWorkflow = vi.fn(); + const onApplyJobCode = vi.fn(); + const messages = [ + createMockAIMessage({ + id: 'msg-global', + role: 'assistant', + code: 'name: Test\njobs: {}', + from_global: true, + }), + ]; + + render( + + ); + + // No job_id -> existing "Generated Workflow" branch, workflow apply + expect(screen.getByText('Generated Workflow')).toBeInTheDocument(); + expect(screen.queryByText('Preview')).not.toBeInTheDocument(); + + await userEvent.click(screen.getByText('Apply')); + expect(onApplyWorkflow).toHaveBeenCalledWith( + 'name: Test\njobs: {}', + 'msg-global' + ); + expect(onApplyJobCode).not.toHaveBeenCalled(); + }); + + it('shows Preview routed to onPreviewGlobalStep when a step is open', async () => { + const onPreviewGlobalStep = vi.fn(); + const onPreviewJobCode = vi.fn(); + const messages = [ + createMockAIMessage({ + id: 'msg-global', + role: 'assistant', + code: 'name: Test\njobs: {}', + from_global: true, + }), + ]; + + render( + + ); + + // Label stays "Generated Workflow" — it is one + expect(screen.getByText('Generated Workflow')).toBeInTheDocument(); + + await userEvent.click(screen.getByText('Preview')); + expect(onPreviewGlobalStep).toHaveBeenCalledWith( + 'name: Test\njobs: {}', + 'msg-global' + ); + expect(onPreviewJobCode).not.toHaveBeenCalled(); + }); + + it('keeps job-code messages unchanged: Preview routes to onPreviewJobCode', async () => { + const onPreviewGlobalStep = vi.fn(); + const onPreviewJobCode = vi.fn(); + const onApplyWorkflow = vi.fn(); + const onApplyJobCode = vi.fn(); + const messages = [ + createMockAIMessage({ + id: 'msg-job', + role: 'assistant', + code: 'fn(state => state)', + job_id: 'job-1', + }), + ]; + + render( + + ); + + expect(screen.getByText('Generated Job Code')).toBeInTheDocument(); + + await userEvent.click(screen.getByText('Preview')); + expect(onPreviewJobCode).toHaveBeenCalledWith( + 'fn(state => state)', + 'msg-job' + ); + expect(onPreviewGlobalStep).not.toHaveBeenCalled(); + + await userEvent.click(screen.getByText('Apply')); + expect(onApplyJobCode).toHaveBeenCalledWith( + 'fn(state => state)', + 'msg-job' + ); + expect(onApplyWorkflow).not.toHaveBeenCalled(); + }); + }); + describe('Props Handling', () => { it('should handle undefined messages prop', () => { render(); diff --git a/assets/test/collaborative-editor/hooks/useAIWorkflowApplications.globalStep.test.ts b/assets/test/collaborative-editor/hooks/useAIWorkflowApplications.globalStep.test.ts new file mode 100644 index 0000000000..ff054d7a8f --- /dev/null +++ b/assets/test/collaborative-editor/hooks/useAIWorkflowApplications.globalStep.test.ts @@ -0,0 +1,324 @@ +/** + * useAIWorkflowApplications - Global Message Tests + * + * Tests behavior specific to global AI assistant messages (full workflow + * YAML, `from_global: true`, no job_id): + * - handlePreviewGlobalStep: per-step diff extracted from the workflow YAML + * - handleApplyWorkflow: relaxed mode guard for global messages only + */ + +import { renderHook, waitFor } from '@testing-library/react'; +import { describe, it, expect, beforeEach, vi } from 'vitest'; + +import type { MonacoHandle } from '../../../js/collaborative-editor/components/CollaborativeMonaco'; +import type { AIModeResult } from '../../../js/collaborative-editor/hooks/useAIMode'; +import { useAIWorkflowApplications } from '../../../js/collaborative-editor/hooks/useAIWorkflowApplications'; +import { notifications } from '../../../js/collaborative-editor/lib/notifications'; +import type { Job } from '../../../js/collaborative-editor/types'; +import type { Message } from '../../../js/collaborative-editor/types/ai-assistant'; +import { + convertWorkflowSpecToState, + parseWorkflowYAML, +} from '../../../js/yaml/util'; + +vi.mock('../../../js/yaml/util', () => ({ + parseWorkflowYAML: vi.fn(), + convertWorkflowSpecToState: vi.fn(), + applyJobCredsToWorkflowState: vi.fn((state: unknown) => state), + extractJobCredentials: vi.fn(() => ({})), +})); + +vi.mock('../../../js/collaborative-editor/lib/notifications', () => ({ + notifications: { + alert: vi.fn(), + warning: vi.fn(), + success: vi.fn(), + }, +})); + +describe('useAIWorkflowApplications - global messages', () => { + const mockImportWorkflow = vi.fn(() => Promise.resolve()); + const mockStartApplyingWorkflow = vi.fn(() => Promise.resolve(true)); + const mockDoneApplyingWorkflow = vi.fn(() => Promise.resolve()); + const mockStartApplyingJobCode = vi.fn(() => Promise.resolve(true)); + const mockDoneApplyingJobCode = vi.fn(() => Promise.resolve()); + const mockUpdateJob = vi.fn(); + const mockSetPreviewingMessageId = vi.fn(); + const mockSetApplyingMessageId = vi.fn(); + const mockClearDiff = vi.fn(); + const mockShowDiff = vi.fn(); + + const mockWorkflowActions = { + importWorkflow: mockImportWorkflow, + startApplyingWorkflow: mockStartApplyingWorkflow, + doneApplyingWorkflow: mockDoneApplyingWorkflow, + startApplyingJobCode: mockStartApplyingJobCode, + doneApplyingJobCode: mockDoneApplyingJobCode, + updateJob: mockUpdateJob, + }; + + const createMockMonacoRef = () => ({ + current: { + clearDiff: mockClearDiff, + showDiff: mockShowDiff, + } as MonacoHandle, + }); + + const createMockAIMode = ( + page: 'workflow_template' | 'job_code', + context: Record = {} + ): AIModeResult => ({ + mode: 'workflow_template', + page, + context, + storageKey: `ai-${page}`, + }); + + const createMockJob = (overrides: Partial = {}): Job => + ({ + id: 'job-1', + name: 'Test Job', + body: 'old body', + adaptor: '@openfn/language-http@latest', + enabled: true, + ...overrides, + }) as Job; + + const createGlobalMessage = (overrides: Partial = {}): Message => ({ + id: 'msg-global', + role: 'assistant', + content: 'Updated your workflow', + code: 'name: Test Workflow', + status: 'success', + inserted_at: new Date().toISOString(), + user_id: 'user-123', + from_global: true, + ...overrides, + }); + + /** Sets the body that the parsed YAML reports for job-1 */ + const mockYamlJobBody = (body: string) => { + vi.mocked(convertWorkflowSpecToState).mockReturnValue({ + id: 'wf-1', + name: 'Test Workflow', + jobs: [ + { + id: 'job-1', + name: 'Test Job', + adaptor: '@openfn/language-http@latest', + body, + keychain_credential_id: null, + project_credential_id: null, + }, + ], + triggers: [], + edges: [], + positions: null, + }); + }; + + const renderApplications = ( + overrides: Partial[0]> = {} + ) => + renderHook(() => + useAIWorkflowApplications({ + sessionId: 'session-1', + page: 'job_code', + currentSession: null, + currentUserId: 'user-123', + aiMode: createMockAIMode('job_code', { job_id: 'job-1' }), + workflowActions: mockWorkflowActions, + monacoRef: createMockMonacoRef(), + jobs: [createMockJob()], + canApplyChanges: true, + connectionState: 'connected', + setPreviewingMessageId: mockSetPreviewingMessageId, + previewingMessageId: null, + setApplyingMessageId: mockSetApplyingMessageId, + appliedMessageIdsRef: { current: new Set() }, + ...overrides, + }) + ); + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(parseWorkflowYAML).mockImplementation((yaml: string) => { + if (yaml.includes('invalid')) { + throw new Error('Invalid YAML syntax'); + } + return { name: 'Test Workflow', jobs: {}, triggers: {}, edges: {} }; + }); + mockYamlJobBody('new body'); + }); + + describe('handlePreviewGlobalStep', () => { + it('shows a diff when the open step body changed in the YAML', () => { + const { result } = renderApplications(); + + result.current.handlePreviewGlobalStep('name: Test Workflow', 'msg-1'); + + expect(mockShowDiff).toHaveBeenCalledWith('old body', 'new body'); + expect(mockSetPreviewingMessageId).toHaveBeenCalledWith('msg-1'); + }); + + it('shows no diff and clears a stale one when the open step is unchanged', () => { + mockYamlJobBody('old body'); // same as current job body + + const { result } = renderApplications({ + previewingMessageId: 'previous-msg', + }); + + result.current.handlePreviewGlobalStep('name: Test Workflow', 'msg-1'); + + expect(mockClearDiff).toHaveBeenCalled(); + expect(mockShowDiff).not.toHaveBeenCalled(); + expect(mockSetPreviewingMessageId).not.toHaveBeenCalled(); + }); + + it('warns when the open step is missing from the YAML (id not preserved)', () => { + vi.mocked(convertWorkflowSpecToState).mockReturnValue({ + id: 'wf-1', + name: 'Test Workflow', + jobs: [], + triggers: [], + edges: [], + positions: null, + }); + + const { result } = renderApplications(); + + result.current.handlePreviewGlobalStep('name: Test Workflow', 'msg-1'); + + expect(mockShowDiff).not.toHaveBeenCalled(); + expect(mockSetPreviewingMessageId).not.toHaveBeenCalled(); + expect(notifications.warning).toHaveBeenCalled(); + }); + + it('alerts for invalid YAML', () => { + const { result } = renderApplications(); + + result.current.handlePreviewGlobalStep('invalid yaml', 'msg-1'); + + expect(mockShowDiff).not.toHaveBeenCalled(); + expect(mockClearDiff).not.toHaveBeenCalled(); + expect(mockSetPreviewingMessageId).not.toHaveBeenCalled(); + expect(notifications.alert).toHaveBeenCalled(); + }); + + it('does nothing when no job is open (workflow_template mode)', () => { + const { result } = renderApplications({ + aiMode: createMockAIMode('workflow_template'), + page: 'workflow_template', + }); + + result.current.handlePreviewGlobalStep('name: Test Workflow', 'msg-1'); + + expect(parseWorkflowYAML).not.toHaveBeenCalled(); + expect(mockShowDiff).not.toHaveBeenCalled(); + }); + + it('deduplicates previews like handlePreviewJobCode', () => { + // Already previewing this message -> no-op + const { result } = renderApplications({ + previewingMessageId: 'msg-1', + }); + result.current.handlePreviewGlobalStep('name: Test Workflow', 'msg-1'); + expect(mockShowDiff).not.toHaveBeenCalled(); + expect(mockSetPreviewingMessageId).not.toHaveBeenCalled(); + + // Streaming preview already shown -> only swap the message id + const { result: streaming } = renderApplications({ + previewingMessageId: '__streaming__', + }); + streaming.current.handlePreviewGlobalStep('name: Test Workflow', 'msg-1'); + expect(mockShowDiff).not.toHaveBeenCalled(); + expect(mockSetPreviewingMessageId).toHaveBeenCalledWith('msg-1'); + }); + }); + + describe('handleApplyWorkflow with global messages', () => { + it('applies a global message while a job is open (job_code mode)', async () => { + const globalMessage = createGlobalMessage(); + const { result } = renderApplications({ + currentSession: { messages: [globalMessage] }, + }); + + await result.current.handleApplyWorkflow( + globalMessage.code!, + globalMessage.id + ); + + await waitFor(() => { + expect(mockStartApplyingWorkflow).toHaveBeenCalledWith('msg-global'); + expect(mockImportWorkflow).toHaveBeenCalled(); + expect(mockDoneApplyingWorkflow).toHaveBeenCalledWith('msg-global'); + }); + }); + + it('clears an active step diff when applying a global message', async () => { + const globalMessage = createGlobalMessage(); + const { result } = renderApplications({ + currentSession: { messages: [globalMessage] }, + previewingMessageId: globalMessage.id, + }); + + await result.current.handleApplyWorkflow( + globalMessage.code!, + globalMessage.id + ); + + await waitFor(() => { + expect(mockClearDiff).toHaveBeenCalled(); + expect(mockSetPreviewingMessageId).toHaveBeenCalledWith(null); + expect(mockImportWorkflow).toHaveBeenCalled(); + }); + }); + + it('still no-ops for a non-global message in job_code mode', async () => { + const consoleError = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + const workflowMessage = createGlobalMessage({ + id: 'msg-workflow', + from_global: false, + }); + + const { result } = renderApplications({ + currentSession: { messages: [workflowMessage] }, + }); + + await result.current.handleApplyWorkflow( + workflowMessage.code!, + workflowMessage.id + ); + + expect(mockStartApplyingWorkflow).not.toHaveBeenCalled(); + expect(mockImportWorkflow).not.toHaveBeenCalled(); + expect(mockSetApplyingMessageId).not.toHaveBeenCalled(); + consoleError.mockRestore(); + }); + + it('still applies non-global messages in workflow_template mode', async () => { + const workflowMessage = createGlobalMessage({ + id: 'msg-workflow', + from_global: false, + }); + + const { result } = renderApplications({ + aiMode: createMockAIMode('workflow_template'), + page: 'workflow_template', + currentSession: { messages: [workflowMessage] }, + }); + + await result.current.handleApplyWorkflow( + workflowMessage.code!, + workflowMessage.id + ); + + await waitFor(() => { + expect(mockImportWorkflow).toHaveBeenCalled(); + expect(mockDoneApplyingWorkflow).toHaveBeenCalledWith('msg-workflow'); + }); + }); + }); +}); diff --git a/assets/test/collaborative-editor/hooks/useAutoPreview.test.ts b/assets/test/collaborative-editor/hooks/useAutoPreview.test.ts index 5c505bc63f..4bb35d696f 100644 --- a/assets/test/collaborative-editor/hooks/useAutoPreview.test.ts +++ b/assets/test/collaborative-editor/hooks/useAutoPreview.test.ts @@ -209,6 +209,102 @@ describe('useAutoPreview', () => { }); }); + describe('Global messages', () => { + const olderUserMessage: Message = { + id: 'user-msg-1', + role: 'user', + content: 'Earlier question', + status: 'success', + inserted_at: new Date(Date.now() - 3000).toISOString(), + user_id: mockCurrentUserId, + }; + const triggeringUserMessage: Message = { + id: 'user-msg-2', + role: 'user', + content: 'Update my workflow', + status: 'success', + inserted_at: new Date(Date.now() - 1000).toISOString(), + user_id: mockCurrentUserId, + }; + const globalYaml = 'name: My Workflow\njobs:\n job-1:\n name: Step 1'; + const createGlobalMessage = (): Message => + createMockMessage({ + id: 'global-msg', + code: globalYaml, + job_id: undefined, // global messages never carry a job_id + from_global: true, + inserted_at: new Date().toISOString(), + }); + + it('auto-previews a new global full-YAML message when a job is open', () => { + const firstAssistant = createMockMessage({ + id: 'first-msg', + inserted_at: new Date(Date.now() - 2000).toISOString(), + }); + const initialSession = createMockSession([ + olderUserMessage, + firstAssistant, + ]); + + const { rerender } = renderHook( + ({ session }) => + useAutoPreview({ + aiMode: createMockJobCodeMode(), + session, + currentUserId: mockCurrentUserId, + onPreview: mockOnPreview, + }), + { initialProps: { session: initialSession } } + ); + + // Initial load - no preview + expect(mockOnPreview).not.toHaveBeenCalled(); + + // New global message arrives - preview receives the full YAML + rerender({ + session: createMockSession([ + olderUserMessage, + firstAssistant, + triggeringUserMessage, + createGlobalMessage(), + ]), + }); + + expect(mockOnPreview).toHaveBeenCalledWith(globalYaml, 'global-msg'); + }); + + it('does not auto-preview a global message when no job is open', () => { + const workflowMode: AIModeResult = { + mode: 'workflow_template', + page: 'workflow_template', + context: { project_id: 'proj-1', workflow_id: 'wf-1' }, + storageKey: 'ai-workflow-wf-1', + }; + const initialSession = createMockSession([olderUserMessage]); + + const { rerender } = renderHook( + ({ session }) => + useAutoPreview({ + aiMode: workflowMode, + session, + currentUserId: mockCurrentUserId, + onPreview: mockOnPreview, + }), + { initialProps: { session: initialSession } } + ); + + rerender({ + session: createMockSession([ + olderUserMessage, + triggeringUserMessage, + createGlobalMessage(), + ]), + }); + + expect(mockOnPreview).not.toHaveBeenCalled(); + }); + }); + describe('Session mount behavior', () => { it('should not auto-preview on initial session load', () => { const session = createMockSession([createMockMessage()]); diff --git a/lib/lightning/ai_assistant/ai_assistant.ex b/lib/lightning/ai_assistant/ai_assistant.ex index b95886e1a4..1e31a955ee 100644 --- a/lib/lightning/ai_assistant/ai_assistant.ex +++ b/lib/lightning/ai_assistant/ai_assistant.ex @@ -780,14 +780,26 @@ defmodule Lightning.AiAssistant do end defp prepare_message_attrs(message_attrs, session, code) do - message_attrs - |> Enum.into(%{}, fn {k, v} -> {to_string(k), v} end) - |> Map.put("chat_session_id", session.id) - |> Map.put("code", code) - |> maybe_put_job_id_from_session(session) - |> maybe_put_unsaved_job_meta(session) + attrs = + message_attrs + |> Enum.into(%{}, fn {k, v} -> {to_string(k), v} end) + |> Map.put("chat_session_id", session.id) + |> Map.put("code", code) + + # Global messages carry a full workflow YAML, never a job — + # even when the session was started from a job step. + if global_message?(attrs) do + attrs + else + attrs + |> maybe_put_job_id_from_session(session) + |> maybe_put_unsaved_job_meta(session) + end end + defp global_message?(%{"meta" => %{"from_global" => true}}), do: true + defp global_message?(_), do: false + defp maybe_put_unsaved_job_meta(attrs, session) do is_assistant = to_string(Map.get(attrs, "role")) == "assistant" unsaved_job = get_in(session.meta || %{}, ["unsaved_job"]) @@ -1104,11 +1116,7 @@ defmodule Lightning.AiAssistant do ) do {:ok, %Tesla.Env{status: status, body: body}} when status in @success_status_range -> - process_stream( - session, - body, - &build_global_message(&1, session) - ) + process_stream(session, body, &build_global_message/1) error -> handle_error_response(error, session) @@ -1479,93 +1487,29 @@ defmodule Lightning.AiAssistant do {message_attrs, opts} end - defp build_global_message(body, session) do - {code, job, job_key} = - extract_global_code_and_job(body["attachments"], session) + defp build_global_message(body) do + code = extract_global_workflow_yaml(body["attachments"]) message_attrs = %{ role: :assistant, - content: body["response"] + content: body["response"], + meta: %{"from_global" => true} } - # Set job on message for "Generated Job Code" rendering. - # For saved jobs, set the job association directly. - # For unsaved jobs, set from_unsaved_job in meta so - # format_message can use it as a fallback job_id. - message_attrs = - cond do - job -> - Map.put(message_attrs, :job, job) - - job_key -> - Map.put(message_attrs, :meta, %{"from_global_job_code" => job_key}) - - true -> - message_attrs - end - - opts = [ - usage: body["usage"] || %{}, - meta: body["meta"], - code: code - ] - + opts = [usage: body["usage"] || %{}, meta: body["meta"], code: code] {message_attrs, opts} end - # Extracts the appropriate code artifact and optional job from global chat - # attachments. On a job step, prefers job_code (renders as code diff). - # On the workflow overview, prefers workflow_yaml (renders as YAML card). - defp extract_global_code_and_job(attachments, session) - when is_list(attachments) do - page = get_in(session.meta || %{}, ["message_options", "page"]) - on_job_step = page && length(String.split(page, "/")) >= 3 - - job_code_attachment = - Enum.find(attachments, &match?(%{"type" => "job_code"}, &1)) - - if on_job_step && job_code_attachment do - job_key = job_code_attachment["job_key"] - - job = - resolve_job_from_key(session.workflow_id, job_key) - - {job_code_attachment["content"], job, job_key} - else - workflow_yaml = - Enum.find_value(attachments, fn - %{"type" => "workflow_yaml", "content" => content} -> content - _ -> nil - end) - - {workflow_yaml, nil, nil} - end - end - - defp extract_global_code_and_job(_, _), do: {nil, nil, nil} - - defp resolve_job_from_key(nil, _), do: nil - defp resolve_job_from_key(_, nil), do: nil - - defp resolve_job_from_key(workflow_id, job_key) do - import Ecto.Query - - Lightning.Workflows.Job - |> where([j], j.workflow_id == ^workflow_id) - |> Repo.all() - |> Enum.find(fn job -> - normalize_job_name(job.name) == normalize_job_name(job_key) + # Global chat always returns a full workflow YAML (job bodies embedded). + # The frontend handles per-step diffing and full-workflow apply. + defp extract_global_workflow_yaml(attachments) when is_list(attachments) do + Enum.find_value(attachments, fn + %{"type" => "workflow_yaml", "content" => content} -> content + _ -> nil end) end - defp normalize_job_name(name) when is_binary(name) do - name - |> String.downcase() - |> String.replace(~r/[^a-z0-9]+/, "-") - |> String.trim("-") - end - - defp normalize_job_name(_), do: "" + defp extract_global_workflow_yaml(_), do: nil defp build_history(session) do messages = session.messages || [] diff --git a/lib/lightning/ai_assistant/message_processor.ex b/lib/lightning/ai_assistant/message_processor.ex index 735c280d17..6b0cde0659 100644 --- a/lib/lightning/ai_assistant/message_processor.ex +++ b/lib/lightning/ai_assistant/message_processor.ex @@ -153,6 +153,13 @@ defmodule Lightning.AiAssistant.MessageProcessor do workflow_yaml = message.code page = get_in(session.meta, ["message_options", "page"]) + if workflow_yaml in [nil, ""] do + Logger.warning( + "[AI Assistant] Global chat message #{message.id} has no workflow YAML; " <> + "Apollo will receive no workflow context (session #{session.id}, page #{inspect(page)})" + ) + end + AiAssistant.query_global_stream(session, message.content, workflow_yaml: workflow_yaml, page: page diff --git a/lib/lightning_web/channels/ai_assistant_channel.ex b/lib/lightning_web/channels/ai_assistant_channel.ex index 1e7e920d16..3837c0e822 100644 --- a/lib/lightning_web/channels/ai_assistant_channel.ex +++ b/lib/lightning_web/channels/ai_assistant_channel.ex @@ -395,6 +395,20 @@ defmodule LightningWeb.AiAssistantChannel do end end + # Global chat is always a workflow_template session, even when launched with a + # step open (which puts a job_id in params). Route it to the workflow_template + # path so the full workflow YAML (`code`) and global meta are persisted and no + # job is attached to the message — matching the global-chat receive design and + # the turn-2 (new_message) path. Without this, the job_id clause below would + # drop `code` and Apollo would receive no workflow context on the first turn. + defp load_or_create_session( + session_id, + %{"use_global_assistant" => true} = params, + user + ) do + load_or_create_workflow_template_session(session_id, params, user) + end + defp load_or_create_session(session_id, %{"job_id" => job_id} = params, user) when not is_nil(job_id) do case session_id do @@ -404,6 +418,10 @@ defmodule LightningWeb.AiAssistantChannel do end defp load_or_create_session(session_id, params, user) do + load_or_create_workflow_template_session(session_id, params, user) + end + + defp load_or_create_workflow_template_session(session_id, params, user) do case session_id do "new" -> with {:project_id, project_id} when not is_nil(project_id) <- @@ -725,14 +743,12 @@ defmodule LightningWeb.AiAssistantChannel do when not is_nil(unsaved_job_id) -> unsaved_job_id - %{"from_global_job_code" => job_key} - when not is_nil(job_key) -> - job_key - _ -> message.job_id end + from_global = match?(%{"from_global" => true}, message.meta) + %{ id: message.id, content: message.content, @@ -742,7 +758,8 @@ defmodule LightningWeb.AiAssistantChannel do inserted_at: message.inserted_at, user_id: message.user_id, user: format_user(message.user), - job_id: job_id + job_id: job_id, + from_global: from_global } end diff --git a/test/lightning/ai_assistant/ai_assistant_test.exs b/test/lightning/ai_assistant/ai_assistant_test.exs index d57b1e3f15..405046073f 100644 --- a/test/lightning/ai_assistant/ai_assistant_test.exs +++ b/test/lightning/ai_assistant/ai_assistant_test.exs @@ -986,6 +986,60 @@ defmodule Lightning.AiAssistantTest do assert saved_message.role == :user end + test "global assistant message in a job session gets no job and no from_unsaved_job" do + user = insert(:user) + job = insert(:job, workflow: build(:workflow)) + unsaved_job_id = Ecto.UUID.generate() + + session = + insert(:chat_session, + job: job, + user: user, + meta: %{"unsaved_job" => %{"id" => unsaved_job_id}} + ) + + {:ok, updated_session} = + AiAssistant.save_message(session, %{ + role: :assistant, + content: "Global response", + meta: %{"from_global" => true} + }) + + saved_message = List.last(updated_session.messages) + + assert %{ + role: :assistant, + job_id: nil, + meta: %{"from_global" => true} + } = saved_message + + refute Map.has_key?(saved_message.meta, "from_unsaved_job") + end + + test "non-global assistant message in a job session still gets job and from_unsaved_job" do + user = insert(:user) + %{id: job_id} = job = insert(:job, workflow: build(:workflow)) + unsaved_job_id = Ecto.UUID.generate() + + session = + insert(:chat_session, + job: job, + user: user, + meta: %{"unsaved_job" => %{"id" => unsaved_job_id}} + ) + + {:ok, updated_session} = + AiAssistant.save_message(session, %{ + role: :assistant, + content: "Job chat response" + }) + + saved_message = List.last(updated_session.messages) + + assert %{role: :assistant, job_id: ^job_id} = saved_message + assert saved_message.meta["from_unsaved_job"] == unsaved_job_id + end + test "enqueues user message for processing when pending", %{user: user} do job = insert(:job, workflow: build(:workflow)) session = insert(:chat_session, job: job, user: user) @@ -3593,13 +3647,16 @@ defmodule Lightning.AiAssistantTest do assert assistant_msg.content == "Here is your updated workflow" assert assistant_msg.code == "workflow:\n name: updated" assert assistant_msg.role == :assistant + assert assistant_msg.meta == %{"from_global" => true} + assert is_nil(assistant_msg.job_id) end - test "extracts job_code on job step pages and resolves job", %{ - user: user, - project: project, - workflow: workflow - } do + test "uses workflow_yaml even on job step pages with a job_code attachment", + %{ + user: user, + project: project, + workflow: workflow + } do %{jobs: [job | _]} = workflow session = @@ -3660,12 +3717,13 @@ defmodule Lightning.AiAssistantTest do assistant_msg = List.last(updated_session.messages) assert assistant_msg.content == "Fixed the job" - # On a job step page, prefers job_code over workflow_yaml - assert assistant_msg.code == "fn(state => state.data)" - assert assistant_msg.job_id == job.id + # Global responses always use workflow_yaml; job_code is ignored + assert assistant_msg.code == "workflow:\n name: full" + assert assistant_msg.meta == %{"from_global" => true} + assert is_nil(assistant_msg.job_id) end - test "falls back to workflow_yaml when no job_code attachment", %{ + test "uses workflow_yaml when no job_code attachment", %{ user: user, project: project, workflow: workflow @@ -3715,7 +3773,6 @@ defmodule Lightning.AiAssistantTest do AiAssistant.query_global_stream(session, "overview") assistant_msg = List.last(updated_session.messages) - # When no job_code attachment, falls back to workflow_yaml assert assistant_msg.code == "workflow:\n name: overview" assert is_nil(assistant_msg.job_id) end @@ -3858,7 +3915,7 @@ defmodule Lightning.AiAssistantTest do assert is_nil(assistant_msg.job_id) end - test "handles resolve_job_from_key with nil inputs", %{ + test "stores no code when only job_code attachments are present", %{ user: user, project: project, workflow: workflow @@ -3884,64 +3941,6 @@ defmodule Lightning.AiAssistantTest do user: user }) - complete_payload = - Jason.encode!(%{ - "response" => "Fixed", - "attachments" => [ - %{ - "type" => "job_code", - "content" => "fn(state => state);", - "job_key" => nil - } - ], - "usage" => %{}, - "meta" => %{} - }) - - sse_stream = [%{event: "complete", data: complete_payload}] - - Mox.expect(Lightning.Tesla.Mock, :call, fn _env, _opts -> - {:ok, %Tesla.Env{status: 200, body: sse_stream}} - end) - - {:ok, updated_session} = - AiAssistant.query_global_stream(session, "fix code", - workflow_yaml: "name: test", - page: "workflows/test/Some-job" - ) - - assistant_msg = List.last(updated_session.messages) - assert assistant_msg.code == "fn(state => state);" - # job_key is nil, so job shouldn't be resolved - assert is_nil(assistant_msg.job_id) - end - - test "handles resolve_job_from_key with nil workflow_id", %{ - user: user, - project: project - } do - # Session without a workflow (workflow_id is nil) - session = - insert(:chat_session, - user: user, - project: project, - workflow: nil, - session_type: "workflow_template", - meta: %{ - "message_options" => %{ - "use_global_assistant" => true, - "page" => "workflows/test/Some-job" - } - } - ) - - {:ok, session} = - AiAssistant.save_message(session, %{ - role: :user, - content: "fix code", - user: user - }) - complete_payload = Jason.encode!(%{ "response" => "Fixed", @@ -3969,67 +3968,9 @@ defmodule Lightning.AiAssistantTest do ) assistant_msg = List.last(updated_session.messages) - assert assistant_msg.code == "fn(state => state);" - # workflow_id is nil, so job can't be resolved - assert is_nil(assistant_msg.job_id) - end - - test "handles non-string job_key in attachment", %{ - user: user, - project: project, - workflow: workflow - } do - session = - insert(:chat_session, - user: user, - project: project, - workflow: workflow, - session_type: "workflow_template", - meta: %{ - "message_options" => %{ - "use_global_assistant" => true, - "page" => "workflows/test/Some-job" - } - } - ) - - {:ok, session} = - AiAssistant.save_message(session, %{ - role: :user, - content: "fix code", - user: user - }) - - # job_key is an integer instead of a string - complete_payload = - Jason.encode!(%{ - "response" => "Fixed", - "attachments" => [ - %{ - "type" => "job_code", - "content" => "fn(state => state);", - "job_key" => 12345 - } - ], - "usage" => %{}, - "meta" => %{} - }) - - sse_stream = [%{event: "complete", data: complete_payload}] - - Mox.expect(Lightning.Tesla.Mock, :call, fn _env, _opts -> - {:ok, %Tesla.Env{status: 200, body: sse_stream}} - end) - - {:ok, updated_session} = - AiAssistant.query_global_stream(session, "fix code", - workflow_yaml: "name: test", - page: "workflows/test/Some-job" - ) - - assistant_msg = List.last(updated_session.messages) - assert assistant_msg.code == "fn(state => state);" - # Non-string job_key won't match any job name + # job_code attachments are ignored; only workflow_yaml is stored + assert is_nil(assistant_msg.code) + assert assistant_msg.meta == %{"from_global" => true} assert is_nil(assistant_msg.job_id) end diff --git a/test/lightning_web/channels/ai_assistant_channel_test.exs b/test/lightning_web/channels/ai_assistant_channel_test.exs index a0637d988e..32ab28d4f0 100644 --- a/test/lightning_web/channels/ai_assistant_channel_test.exs +++ b/test/lightning_web/channels/ai_assistant_channel_test.exs @@ -208,6 +208,53 @@ defmodule LightningWeb.AiAssistantChannelTest do end end + describe "message serialization" do + test "serializes from_global marker with nil job_id", %{ + socket: socket, + job: job, + user: user + } do + session = + insert(:chat_session, + job: job, + user: user, + session_type: "job_code", + messages: [ + %{ + role: :assistant, + content: "Global workflow response", + status: :success, + meta: %{"from_global" => true}, + code: "workflow:\n name: updated" + }, + %{ + role: :assistant, + content: "Job chat response", + status: :success, + inserted_at: DateTime.utc_now() |> DateTime.add(1) + } + ] + ) + + assert {:ok, %{messages: messages}, _socket} = + subscribe_and_join( + socket, + AiAssistantChannel, + "ai_assistant:job_code:#{session.id}", + %{} + ) + + assert [ + %{ + from_global: true, + job_id: nil, + code: "workflow:\n name: updated" + }, + %{from_global: false} + ] = messages + end + end + describe "workflow_template sessions" do test "successfully creates workflow template session", %{ socket: socket, @@ -3199,5 +3246,50 @@ defmodule LightningWeb.AiAssistantChannelTest do assert message_options["use_global_assistant"] == true assert message_options["page"] == "/projects/p1/workflows/w1/jobs/j1" end + + test "first-turn global session opened with a step persists code and no job_id", + %{ + socket: socket, + project: project, + workflow: workflow, + job: job + } do + # Global chat launched with a step open sends job_id in the join params. + # The session must still be created on the workflow_template path so the + # full workflow YAML is stored on the message, and no job is attached. + yaml = "workflow:\n name: opened-from-step" + + params = %{ + "job_id" => job.id, + "workflow_id" => workflow.id, + "project_id" => project.id, + "content" => "what does this do", + "use_global_assistant" => true, + "page" => "workflows/Test Workflow/Test Job", + "code" => yaml + } + + assert {:ok, response, _socket} = + subscribe_and_join( + socket, + AiAssistantChannel, + "ai_assistant:workflow_template:new", + params + ) + + assert response.session_type == "workflow_template" + + session = AiAssistant.get_session!(response.session_id) + + user_msg = + Enum.find(session.messages, fn m -> + m.role == :user && m.content == "what does this do" + end) + + # The full workflow YAML reaches Apollo via the message code... + assert user_msg.code == yaml + # ...and the message carries no job_id, even though a step was open. + assert user_msg.job_id == nil + end end end