refactor: remove dead flow rawscript helper

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

ai_evals/adapters/frontend/core/flow/fileHelpers.ts

@@ -89,13 +89,14 @@ export async function createFlowFileHelpers(
 		preprocessorModule,
 		failureModule
 	}) => {
-		applyFlowJsonUpdate(flow, inlineScriptSession, {
+		const result = applyFlowJsonUpdate(flow, inlineScriptSession, {
 			modules,
 			schema,
 			preprocessorModule,
 			failureModule
 		})
 		await persistFlow()
+		return result
 	}
 
 	const helpers: FlowAIChatHelpers = {

ai_evals/adapters/frontend/core/shared/baseEvalRunner.ts

@@ -145,6 +145,20 @@ export async function runEval<THelpers, TOutput>(
 			skipResponsesApi: modelProvider.provider !== 'openai'
 		})
 
+		if (result.hitMaxIterations) {
+			return {
+				success: false,
+				output: getOutput(),
+				error: `Reached max turns (${maxIterations})`,
+				tokenUsage: result.tokenUsage,
+				toolCallsCount,
+				toolsCalled,
+				toolCallDetails,
+				iterations: Math.max(1, result.addedMessages.filter((m) => m.role === 'assistant').length),
+				messages
+			}
+		}
+
 		return {
 			success: true,
 			output: getOutput(),

ai_evals/cases/flow.yaml

@@ -189,6 +189,8 @@
     If validation passes, let the save continue normally.
     Update `save_results` so it uses the validation outcome instead of bypassing it.
   initial: ai_evals/fixtures/frontend/flow/initial/test5_initial.json
+  runtime:
+    maxTurns: 8
   validate:
     topLevelStepIds:
       - fetch_data

ai_evals/core/runSuite.ts

@@ -100,6 +100,7 @@ async function runCaseAttempts<TInitial, TExpected, TActual>(input: {
       const initial = await input.modeRunner.loadInitial(input.evalCase.initialPath);
       const expected = await input.modeRunner.loadExpected(input.evalCase.expectedPath);
       const run = await input.modeRunner.run(input.evalCase.prompt, initial, {
+        evalCase: input.evalCase,
         caseId: input.evalCase.id,
         caseNumber: input.caseIndex + 1,
         totalCases: input.totalCases,
@@ -181,6 +182,7 @@ async function runCaseAttempts<TInitial, TExpected, TActual>(input: {
             actual: run.actual,
             run,
             context: {
+              evalCase: input.evalCase,
               caseId: input.evalCase.id,
               caseNumber: input.caseIndex + 1,
               totalCases: input.totalCases,

ai_evals/core/types.ts

@@ -8,6 +8,7 @@ export interface EvalCaseRuntimeBackendPreview {
 }
 
 export interface EvalCaseRuntimeSpec {
+  maxTurns?: number;
   backendPreview?: EvalCaseRuntimeBackendPreview;
 }
 
@@ -109,6 +110,7 @@ export interface ModeRunOutput<TActual> {
 }
 
 export interface ModeRunContext {
+  evalCase?: EvalCase;
   caseId: string;
   caseNumber: number;
   totalCases: number;

ai_evals/modes/app.ts

@@ -23,6 +23,7 @@ export function createAppModeRunner(
       const result = await runAppEval(prompt, getFrontendApiKey(modelConfig.provider), {
         initialFrontend: initial?.frontend,
         initialBackend: initial?.backend as AppFiles["backend"] | undefined,
+        maxIterations: context.evalCase?.runtime?.maxTurns,
         provider: modelConfig.provider,
         model: modelConfig.model,
         runContext: context,

ai_evals/modes/flow.ts

@@ -37,6 +37,7 @@ export function createFlowModeRunner(
       const result = await runFlowEval(prompt, getFrontendApiKey(modelConfig.provider), {
         initialFlow: initial?.flowFixture,
         workspaceFixtures: initial?.workspace,
+        maxIterations: context.evalCase?.runtime?.maxTurns,
         provider: modelConfig.provider,
         model: modelConfig.model,
         runContext: context,

ai_evals/modes/script.ts

@@ -29,6 +29,7 @@ export function createScriptModeRunner(
 
       const result = await runScriptEval(prompt, getFrontendApiKey(modelConfig.provider), {
         initialScript: initial,
+        maxIterations: context.evalCase?.runtime?.maxTurns,
         provider: modelConfig.provider,
         model: modelConfig.model,
         runContext: context,

frontend/src/lib/components/copilot/chat/chatLoop.ts

@@ -55,6 +55,7 @@ export interface ChatLoopConfig {
 export interface ChatLoopResult {
 	addedMessages: ChatCompletionMessageParam[]
 	tokenUsage: ChatTokenUsage
+	hitMaxIterations: boolean
 }
 
 export async function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResult> {
@@ -74,9 +75,11 @@ export async function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResul
 	const addedMessages: ChatCompletionMessageParam[] = []
 	let tokenUsage = emptyChatTokenUsage()
 	let iterations = 0
+	let hitMaxIterations = false
 
 	while (true) {
 		if (maxIterations !== undefined && iterations >= maxIterations) {
+			hitMaxIterations = true
 			break
 		}
 		iterations++
@@ -218,5 +221,5 @@ export async function runChatLoop(config: ChatLoopConfig): Promise<ChatLoopResul
 		}
 	}
 
-	return { addedMessages, tokenUsage }
+	return { addedMessages, tokenUsage, hitMaxIterations }
 }

frontend/src/lib/components/copilot/chat/flow/FlowAIChat.svelte

@@ -14,7 +14,7 @@
 	import {
 		applyFlowJsonUpdate,
 		getFlowModuleById,
-		getRawScriptModuleById
+		updateRawScriptModuleContent
 	} from './helperUtils'
 
 	let {
@@ -83,11 +83,6 @@
 
 		// ai chat tools
 		setCode: async (id: string, code: string) => {
-			const module = getRawScriptModuleById(flowStore.val, id)
-			if (!module) {
-				throw new Error('Module not found or is not a rawscript')
-			}
-
 			// 1. Take snapshot only if none exists (preserves baseline for cumulative changes)
 			if (!diffManager?.beforeFlow) {
 				const snapshot = $state.snapshot(flowStore).val
@@ -96,7 +91,11 @@
 			}
 
 			// 2. Apply the code change
-			module.value.content = code
+			const module = updateRawScriptModuleContent(flowStore.val, id, code)
+			if (!module) {
+				throw new Error('Module not found or is not a rawscript')
+			}
+
 			inlineScriptSession.set(id, code)
 			const { input_transforms, schema } = await loadSchemaFromModule(module)
 			module.value.input_transforms = input_transforms
@@ -207,7 +206,7 @@
 					}
 				}
 
-				applyFlowJsonUpdate(flowStore.val, inlineScriptSession, {
+				const result = applyFlowJsonUpdate(flowStore.val, inlineScriptSession, {
 					modules,
 					schema,
 					preprocessorModule,
@@ -216,6 +215,7 @@
 
 				// Refresh the state store to update UI
 				refreshStateStore(flowStore)
+				return result
 			} catch (error) {
 				throw new Error(
 					`Failed to parse or apply JSON: ${error instanceof Error ? error.message : String(error)}`

frontend/src/lib/components/copilot/chat/flow/core.ts

@@ -35,6 +35,7 @@ import {
 import type { ContextElement } from '../context'
 import type { ExtendedOpenFlow } from '$lib/components/flows/types'
 import { createInlineScriptSession, type InlineScriptSession } from './inlineScriptsUtils'
+import type { FlowJsonUpdateResult } from './helperUtils'
 import { flowModuleSchema, flowModulesSchema } from './openFlowZod'
 import { collectAllModuleIdsFromArray } from './utils'
 import { FLOW_CHAT_SPECIAL_MODULES, getFlowPrompt } from '$system_prompts'
@@ -276,6 +277,17 @@ type EditableFlowJson = {
 	failure_module: FlowModule | null
 }
 
+function formatEmptyInlineScriptWarning({
+	emptyInlineScriptModuleIds
+}: FlowJsonUpdateResult): string {
+	if (emptyInlineScriptModuleIds.length === 0) {
+		return ''
+	}
+
+	const moduleList = emptyInlineScriptModuleIds.map((id) => `'${id}'`).join(', ')
+	return ` Warning: inline scripts ${moduleList} are empty for now. Use set_module_code to fill them in.`
+}
+
 function validateFlowModules(rawModules: unknown): FlowModule[] {
 	if (!Array.isArray(rawModules)) {
 		throw new Error('Flow modules must be an array')
@@ -492,7 +504,7 @@ export interface FlowAIChatHelpers {
 
 	// ai chat tools
 	setCode: (id: string, code: string) => Promise<void>
-	setFlowJson: (update: FlowJsonUpdate) => Promise<void>
+	setFlowJson: (update: FlowJsonUpdate) => Promise<FlowJsonUpdateResult>
 	getFlowInputsSchema: () => Promise<Record<string, any>>
 	/** Update exprsToSet store for InputTransformForm components (only if module is selected) */
 	updateExprsToSet: (id: string, inputTransforms: Record<string, InputTransform>) => void
@@ -1018,12 +1030,13 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 				helpers.inlineScriptSession.set(moduleId, content)
 			}
 
-			await helpers.setFlowJson({
+			const updateResult = await helpers.setFlowJson({
 				modules: parsedFlow.modules,
 				schema: parsedFlow.schema,
 				preprocessorModule: parsedFlow.preprocessor_module,
 				failureModule: parsedFlow.failure_module
 			})
+			const warning = formatEmptyInlineScriptWarning(updateResult)
 
 			const selectedModule = findModuleInEditableFlow(parsedFlow, selectedId)
 			if (
@@ -1039,7 +1052,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 				result: 'Success'
 			})
 
-			return `Flow JSON updated`
+			return `Flow JSON updated.${warning}`
 		}
 	},
 	{
@@ -1058,7 +1071,8 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 				content:
 					parsedModule === null ? 'Removing preprocessor module...' : 'Setting preprocessor module...'
 			})
-			await helpers.setFlowJson({ preprocessorModule: parsedModule })
+			const updateResult = await helpers.setFlowJson({ preprocessorModule: parsedModule })
+			const warning = formatEmptyInlineScriptWarning(updateResult)
 
 			if (
 				parsedModule &&
@@ -1070,13 +1084,12 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 			}
 
 			toolCallbacks.setToolStatus(toolId, {
-				content:
-					parsedModule === null ? 'Preprocessor module removed' : 'Preprocessor module updated',
+				content: parsedModule === null ? 'Preprocessor module removed' : 'Preprocessor module updated',
 				result: 'Success'
 			})
 			return parsedModule === null
 				? 'Preprocessor module removed'
-				: 'Preprocessor module updated successfully.'
+				: `Preprocessor module updated successfully.${warning}`
 		}
 	},
 	{
@@ -1094,7 +1107,8 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 			toolCallbacks.setToolStatus(toolId, {
 				content: parsedModule === null ? 'Removing failure module...' : 'Setting failure module...'
 			})
-			await helpers.setFlowJson({ failureModule: parsedModule })
+			const updateResult = await helpers.setFlowJson({ failureModule: parsedModule })
+			const warning = formatEmptyInlineScriptWarning(updateResult)
 
 			if (
 				parsedModule &&
@@ -1111,7 +1125,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 			})
 			return parsedModule === null
 				? 'Failure module removed'
-				: 'Failure module updated successfully.'
+				: `Failure module updated successfully.${warning}`
 		}
 	},
 	{
@@ -1181,7 +1195,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 			toolCallbacks.setToolStatus(toolId, {
 				content: `Setting flow...`
 			})
-			await helpers.setFlowJson({
+			const updateResult = await helpers.setFlowJson({
 				...(parsedModules !== undefined ? { modules: parsedModules } : {}),
 				...(parsedSchema !== undefined ? { schema: parsedSchema } : {}),
 				...(parsedPreprocessorModule !== undefined
@@ -1189,6 +1203,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 					: {}),
 				...(parsedFailureModule !== undefined ? { failureModule: parsedFailureModule } : {})
 			})
+			const warning = formatEmptyInlineScriptWarning(updateResult)
 
 			// Update exprsToSet if the selected module has input_transforms
 			if (
@@ -1218,7 +1233,7 @@ export const flowTools: Tool<FlowAIChatHelpers>[] = [
 				content: `Flow updated`,
 				result: 'Success'
 			})
-			return `Flow updated`
+			return `Flow updated.${warning}`
 		}
 	},
 	{
@@ -1490,6 +1505,10 @@ To reduce token usage, rawscript content in the flow you receive is replaced wit
 **To inspect existing code:**
 - Use \`inspect_inline_script\` tool to view the current code: \`inspect_inline_script({ moduleId: "step_a" })\`
 
+**If a flow update tool warns that inline scripts are empty:**
+- The module structure was created successfully, but the code is still empty
+- Immediately call \`set_module_code\` for each warned module ID
+
 ### Writing Code for Modules
 
 **IMPORTANT: Before writing any code for a rawscript module, you MUST call the \`get_instructions_for_code_generation\` tool with the target language.** This tool provides essential language-specific instructions.

frontend/src/lib/components/copilot/chat/flow/helperUtils.test.ts

@@ -0,0 +1,151 @@
+import { describe, expect, it, vi } from 'vitest'
+import type { FlowModule } from '$lib/gen'
+import { applyFlowJsonUpdate, updateRawScriptModuleContent } from './helperUtils'
+import { createInlineScriptSession } from './inlineScriptsUtils'
+
+vi.mock('../shared', () => ({
+	SPECIAL_MODULE_IDS: {
+		PREPROCESSOR: 'preprocessor',
+		FAILURE: 'failure'
+	}
+}))
+
+vi.mock('$lib/components/flows/previousResults', () => ({
+	dfs: () => []
+}))
+
+function makeRawScriptModule(id: string, content: string): FlowModule {
+	return {
+		id,
+		summary: id,
+		value: {
+			type: 'rawscript',
+			language: 'bun',
+			content,
+			input_transforms: {}
+		}
+	} as FlowModule
+}
+
+function makeAiAgentWithTool(agentId: string, toolId: string, toolContent: string): FlowModule {
+	return {
+		id: agentId,
+		summary: agentId,
+		value: {
+			type: 'aiagent',
+			tools: [
+				{
+					id: toolId,
+					summary: toolId,
+					value: {
+						tool_type: 'flowmodule',
+						type: 'rawscript',
+						language: 'bun',
+						content: toolContent,
+						input_transforms: {}
+					}
+				}
+			],
+			input_transforms: {}
+		}
+	} as FlowModule
+}
+
+describe('applyFlowJsonUpdate', () => {
+	it('accepts new self-referenced inline scripts and initializes them as empty', () => {
+		const flow = {
+			value: {
+				modules: [makeRawScriptModule('process_data', 'existing code')]
+			}
+		}
+		const inlineScriptSession = createInlineScriptSession()
+		inlineScriptSession.set('process_data', 'existing code')
+
+		const result = applyFlowJsonUpdate(flow as any, inlineScriptSession, {
+			modules: [
+				makeRawScriptModule('process_data', 'inline_script.process_data'),
+				makeRawScriptModule('validate_data', 'inline_script.validate_data')
+			]
+		})
+		const [processDataModule, validateDataModule] = flow.value.modules as Array<FlowModule & { value: any }>
+
+		expect(result.emptyInlineScriptModuleIds).toEqual(['validate_data'])
+		expect(inlineScriptSession.has('validate_data')).toBe(false)
+		expect(processDataModule?.value.type).toBe('rawscript')
+		expect(processDataModule?.value.content).toBe('existing code')
+		expect(validateDataModule?.value.type).toBe('rawscript')
+		expect(validateDataModule?.value.content).toBe('')
+	})
+
+	it('still rejects unresolved inline script references that do not match the module id', () => {
+		const flow = {
+			value: {
+				modules: []
+			}
+		}
+		const inlineScriptSession = createInlineScriptSession()
+
+		expect(() =>
+			applyFlowJsonUpdate(flow as any, inlineScriptSession, {
+				modules: [makeRawScriptModule('validate_data', 'inline_script.other_module')]
+			})
+			).toThrow('Unresolved inline script references: other_module')
+	})
+
+	it('keeps the inline script session unchanged after a failed update so retries still warn', () => {
+		const flow = {
+			value: {
+				modules: [makeRawScriptModule('process_data', 'existing code')]
+			}
+		}
+		const inlineScriptSession = createInlineScriptSession()
+		inlineScriptSession.set('process_data', 'existing code')
+
+		expect(() =>
+			applyFlowJsonUpdate(flow as any, inlineScriptSession, {
+				modules: [
+					makeRawScriptModule('validate_data', 'inline_script.validate_data'),
+					makeRawScriptModule('save_results', 'inline_script.other_module')
+				]
+			})
+		).toThrow('Unresolved inline script references: other_module')
+
+		expect(inlineScriptSession.getAll()).toEqual({
+			process_data: 'existing code'
+		})
+		expect((flow.value.modules as Array<FlowModule & { value: any }>)[0]?.value.content).toBe(
+			'existing code'
+		)
+
+		const result = applyFlowJsonUpdate(flow as any, inlineScriptSession, {
+			modules: [
+				makeRawScriptModule('process_data', 'inline_script.process_data'),
+				makeRawScriptModule('validate_data', 'inline_script.validate_data')
+			]
+		})
+
+		expect(result.emptyInlineScriptModuleIds).toEqual(['validate_data'])
+		expect(inlineScriptSession.has('validate_data')).toBe(false)
+	})
+
+	it('updates ai agent rawscript tools in place when changing module code', () => {
+		const flow = {
+			value: {
+				modules: [makeAiAgentWithTool('agent', 'sum', '')]
+			}
+		}
+
+		const updatedModule = updateRawScriptModuleContent(
+			flow as any,
+			'sum',
+			'export async function main(numbers: number[]) { return 0 }'
+		)
+
+		expect(updatedModule?.value.content).toBe(
+			'export async function main(numbers: number[]) { return 0 }'
+		)
+		expect((flow.value.modules[0] as any).value.tools[0].value.content).toBe(
+			'export async function main(numbers: number[]) { return 0 }'
+		)
+	})
+})

frontend/src/lib/components/copilot/chat/flow/helperUtils.ts

@@ -1,4 +1,5 @@
 import type { FlowModule, OpenFlow, RawScript } from '$lib/gen'
+import { forEachFlowModule } from '$lib/components/flows/dfs'
 import { dfs } from '$lib/components/flows/previousResults'
 import { SPECIAL_MODULE_IDS } from '../shared'
 import type { InlineScriptSession } from './inlineScriptsUtils'
@@ -14,6 +15,10 @@ export interface FlowJsonUpdate {
 	failureModule?: FlowModule | null
 }
 
+export interface FlowJsonUpdateResult {
+	emptyInlineScriptModuleIds: string[]
+}
+
 export function getFlowModuleById(flow: FlowLike | undefined, id: string): FlowModule | undefined {
 	if (!flow) {
 		return undefined
@@ -30,16 +35,30 @@ export function getFlowModuleById(flow: FlowLike | undefined, id: string): FlowM
 	return dfs(id, flow as OpenFlow, false)[0]
 }
 
-export function getRawScriptModuleById(
+export function getMutableRawScriptModuleById(
 	flow: FlowLike | undefined,
 	id: string
 ): (FlowModule & { value: RawScript }) | undefined {
-	const module = getFlowModuleById(flow, id)
-	if (!module || module.value.type !== 'rawscript') {
+	if (!flow) {
 		return undefined
 	}
 
-	return module as FlowModule & { value: RawScript }
+	if (flow.value.preprocessor_module?.id === id && flow.value.preprocessor_module.value.type === 'rawscript') {
+		return flow.value.preprocessor_module as FlowModule & { value: RawScript }
+	}
+
+	if (flow.value.failure_module?.id === id && flow.value.failure_module.value.type === 'rawscript') {
+		return flow.value.failure_module as FlowModule & { value: RawScript }
+	}
+
+	let matchedModule: (FlowModule & { value: RawScript }) | undefined
+	forEachFlowModule(flow.value.modules, (module) => {
+		if (!matchedModule && module.id === id && module.value.type === 'rawscript') {
+			matchedModule = module as FlowModule & { value: RawScript }
+		}
+	})
+
+	return matchedModule
 }
 
 export function updateRawScriptModuleContent(
@@ -47,7 +66,7 @@ export function updateRawScriptModuleContent(
 	id: string,
 	code: string
 ): (FlowModule & { value: RawScript }) | undefined {
-	const rawScriptModule = getRawScriptModuleById(flow, id)
+	const rawScriptModule = getMutableRawScriptModuleById(flow, id)
 	if (!rawScriptModule) {
 		return undefined
 	}
@@ -60,9 +79,11 @@ export function applyFlowJsonUpdate(
 	flow: FlowLike,
 	inlineScriptSession: InlineScriptSession,
 	{ modules, schema, preprocessorModule, failureModule }: FlowJsonUpdate
-): void {
+): FlowJsonUpdateResult {
+	const emptyInlineScriptModuleIds = new Set<string>()
+
 	if (modules !== undefined) {
-		flow.value.modules = restoreFlowModules(modules, inlineScriptSession)
+		flow.value.modules = restoreFlowModules(modules, inlineScriptSession, emptyInlineScriptModuleIds)
 	}
 
 	if (schema !== undefined) {
@@ -73,29 +94,39 @@ export function applyFlowJsonUpdate(
 		flow.value.preprocessor_module =
 			preprocessorModule === null
 				? undefined
-				: restoreFlowModule(preprocessorModule, inlineScriptSession)
+				: restoreFlowModule(preprocessorModule, inlineScriptSession, emptyInlineScriptModuleIds)
 	}
 
 	if (failureModule !== undefined) {
 		flow.value.failure_module =
-			failureModule === null ? undefined : restoreFlowModule(failureModule, inlineScriptSession)
+			failureModule === null
+				? undefined
+				: restoreFlowModule(failureModule, inlineScriptSession, emptyInlineScriptModuleIds)
+	}
+
+	return {
+		emptyInlineScriptModuleIds: Array.from(emptyInlineScriptModuleIds)
 	}
 }
 
 function restoreFlowModules(
 	modules: FlowModule[],
-	inlineScriptSession: InlineScriptSession
+	inlineScriptSession: InlineScriptSession,
+	emptyInlineScriptModuleIds: Set<string>
 ): FlowModule[] {
 	const restoredModules = inlineScriptSession.restoreInlineScriptReferences(modules)
+	replaceNewInlineScriptRefsWithEmptyCode(restoredModules, emptyInlineScriptModuleIds)
 	assertResolvedInlineScripts(restoredModules, inlineScriptSession)
 	return restoredModules
 }
 
 function restoreFlowModule(
 	module: FlowModule,
-	inlineScriptSession: InlineScriptSession
+	inlineScriptSession: InlineScriptSession,
+	emptyInlineScriptModuleIds: Set<string>
 ): FlowModule {
 	const [restoredModule] = inlineScriptSession.restoreInlineScriptReferences([module])
+	replaceNewInlineScriptRefsWithEmptyCode([restoredModule], emptyInlineScriptModuleIds)
 	assertResolvedInlineScripts([restoredModule], inlineScriptSession)
 	return restoredModule
 }
@@ -109,3 +140,24 @@ function assertResolvedInlineScripts(
 		throw new Error(`Unresolved inline script references: ${unresolvedRefs.join(', ')}`)
 	}
 }
+
+function replaceNewInlineScriptRefsWithEmptyCode(
+	modules: FlowModule[],
+	emptyInlineScriptModuleIds: Set<string>
+): void {
+	function replaceInlineScriptRefWithEmptyCode(ownerId: string, content: string): string {
+		const match = content.match(/^inline_script\.(.+)$/)
+		if (!match || match[1] !== ownerId) {
+			return content
+		}
+
+		emptyInlineScriptModuleIds.add(ownerId)
+		return ''
+	}
+
+	forEachFlowModule(modules, (module) => {
+		if (module.value.type === 'rawscript' && module.value.content) {
+			module.value.content = replaceInlineScriptRefWithEmptyCode(module.id, module.value.content)
+		}
+	})
+}

frontend/src/lib/components/copilot/chat/flow/utils.test.ts

@@ -0,0 +1,70 @@
+import { describe, expect, it } from 'vitest'
+import type { FlowModule, OpenFlow } from '$lib/gen'
+import { getIndexInNestedModules, getModuleById } from './utils'
+
+function createAiAgentTool(id: string, content: string) {
+	return {
+		id,
+		summary: id,
+		value: {
+			tool_type: 'flowmodule',
+			type: 'rawscript',
+			content,
+			language: 'bun',
+			input_transforms: {}
+		}
+	}
+}
+
+function createAiAgentModule(id: string, tools: ReturnType<typeof createAiAgentTool>[]): FlowModule {
+	return {
+		id,
+		value: {
+			type: 'aiagent',
+			tools,
+			input_transforms: {}
+		} as any
+	}
+}
+
+describe('chat flow utils', () => {
+	it('resolves ai agent tools by id against the stored flow node', () => {
+		const flow = {
+			value: {
+				modules: [createAiAgentModule('agent', [createAiAgentTool('sum', 'before-tool')])]
+			}
+		} as OpenFlow
+
+		const tool = getModuleById(flow, 'sum')
+
+		expect(tool?.id).toBe('sum')
+		expect((tool?.value as any).content).toBe('before-tool')
+
+		;(tool?.value as any).content = 'after-tool'
+		expect(((flow.value.modules[0].value as any).tools[0].value as any).content).toBe('after-tool')
+	})
+
+	it('returns the ai agent tools array for nested index operations', () => {
+		const flow = {
+			value: {
+				modules: [
+					createAiAgentModule('agent', [
+						createAiAgentTool('lookup', 'lookup-tool'),
+						createAiAgentTool('sum', 'sum-tool')
+					])
+				]
+			}
+		} as OpenFlow
+
+		const result = getIndexInNestedModules(flow, 'sum')
+
+		expect(result).not.toBeNull()
+		expect(result?.index).toBe(1)
+		expect(result?.modules.map((module) => module.id)).toEqual(['lookup', 'sum'])
+
+		result?.modules.splice(result.index, 1)
+		expect(((flow.value.modules[0].value as any).tools ?? []).map((tool: any) => tool.id)).toEqual([
+			'lookup'
+		])
+	})
+})

frontend/src/lib/components/copilot/chat/flow/utils.ts

@@ -1,69 +1,79 @@
-import { dfs } from '$lib/components/flows/previousResults'
+import { findModuleInFlow, findModuleParent } from '$lib/components/flows/flowDiff'
 import type { FlowModule, OpenFlow } from '$lib/gen'
 
 // Helper to find module by ID in a flow
 export function getModuleById(flow: OpenFlow, moduleId: string): FlowModule | undefined {
-	const allModules = dfs(moduleId, flow, false)
-	return allModules[0]
+	return findModuleInFlow(flow.value, moduleId) ?? undefined
 }
 
 export function getIndexInNestedModules(
 	flow: OpenFlow,
 	id: string
 ): { index: number; modules: FlowModule[] } | null {
-	const accessingModules = dfs(id, flow, true).reverse()
-
-	if (accessingModules.length === 0) {
-		// Module not found in flow
+	const parentLocation = findModuleParent(flow.value, id)
+	if (!parentLocation) {
+		// Module not found in flow.
 		return null
 	}
 
-	let parent = flow.value.modules
-	let lastIndex = -1
-	for (const [ai, am] of accessingModules.entries()) {
-		const index = parent.findIndex((m) => m.id === am.id)
+	if (parentLocation.type === 'failure' || parentLocation.type === 'preprocessor') {
+		return null
+	}
 
-		if (index === -1) {
-			// Module no longer exists in expected location (may have been deleted with parent)
-			return null
+	if (parentLocation.type === 'root') {
+		return {
+			index: parentLocation.index,
+			modules: flow.value.modules
 		}
+	}
 
-		lastIndex = index
+	const parent = findModuleInFlow(flow.value, parentLocation.parentId)
+	if (!parent) {
+		return null
+	}
 
-		if (ai === accessingModules.length - 1) {
-			break
-		}
-
-		if (
-			parent[index].value.type === 'forloopflow' ||
-			parent[index].value.type === 'whileloopflow'
-		) {
-			parent = parent[index].value.modules
-		} else if (
-			parent[index].value.type === 'branchall' ||
-			parent[index].value.type === 'branchone'
-		) {
-			const branchIdx = parent[index].value.branches.findIndex((b) =>
-				b.modules.some((m) => m.id === accessingModules[ai + 1].id)
-			)
-			if (branchIdx === -1) {
-				// Module no longer exists in branch (may have been deleted)
+	switch (parentLocation.type) {
+		case 'forloop':
+			if (parent.value.type !== 'forloopflow') {
 				return null
 			}
-			parent = parent[index].value.branches[branchIdx].modules
-		} else {
-			// Unexpected module type in path
+			return { index: parentLocation.index, modules: parent.value.modules }
+		case 'whileloop':
+			if (parent.value.type !== 'whileloopflow') {
+				return null
+			}
+			return { index: parentLocation.index, modules: parent.value.modules }
+		case 'branchone-default':
+			if (parent.value.type !== 'branchone') {
+				return null
+			}
+			return { index: parentLocation.index, modules: parent.value.default }
+		case 'branchone-branch':
+			if (parent.value.type !== 'branchone') {
+				return null
+			}
+			return {
+				index: parentLocation.index,
+				modules: parent.value.branches[parentLocation.branchIndex]?.modules ?? []
+			}
+		case 'branchall-branch':
+			if (parent.value.type !== 'branchall') {
+				return null
+			}
+			return {
+				index: parentLocation.index,
+				modules: parent.value.branches[parentLocation.branchIndex]?.modules ?? []
+			}
+		case 'aiagent':
+			if (parent.value.type !== 'aiagent') {
+				return null
+			}
+			return {
+				index: parentLocation.index,
+				modules: (parent.value.tools as FlowModule[]) ?? []
+			}
+		default:
 			return null
-		}
-	}
-
-	if (lastIndex === -1) {
-		return null
-	}
-
-	return {
-		index: lastIndex,
-		modules: parent
 	}
 }
 

frontend/src/lib/components/flows/dfs.ts

@@ -1,34 +1,60 @@
 import type { FlowModule } from '$lib/gen'
 
-export function dfs<T>(
+type FlowDfsOptions = { skipToolNodes?: boolean }
+
+type FlowModuleVisitor<T> = (
+	x: FlowModule,
 	modules: FlowModule[],
-	f: (x: FlowModule, modules: FlowModule[], branches: FlowModule[][]) => T,
-	opts: { skipToolNodes?: boolean } = {}
-): T[] {
-	let result: T[] = []
+	branches: FlowModule[][]
+) => T
+
+function traverseFlowModules(
+	modules: FlowModule[],
+	visit: FlowModuleVisitor<void>,
+	opts: FlowDfsOptions = {}
+): void {
 	for (const module of modules) {
 		if (module.value.type == 'forloopflow' || module.value.type == 'whileloopflow') {
-			result = result.concat(f(module, modules, [module.value.modules]))
-			result = result.concat(dfs(module.value.modules, f, opts))
+			visit(module, modules, [module.value.modules])
+			traverseFlowModules(module.value.modules, visit, opts)
 		} else if (module.value.type == 'branchone') {
 			const allBranches = [module.value.default, ...module.value.branches.map((b) => b.modules)]
-			result = result.concat(f(module, modules, allBranches))
+			visit(module, modules, allBranches)
 
 			for (const branch of allBranches) {
-				result = result.concat(dfs(branch, f, opts))
+				traverseFlowModules(branch, visit, opts)
 			}
 		} else if (module.value.type == 'branchall') {
 			const allBranches = module.value.branches.map((b) => b.modules)
-			result = result.concat(f(module, modules, allBranches))
+			visit(module, modules, allBranches)
 			for (const branch of allBranches) {
-				result = result.concat(dfs(branch, f, opts))
+				traverseFlowModules(branch, visit, opts)
 			}
 		} else if (module.value.type == 'aiagent' && !opts.skipToolNodes) {
-			result = result.concat(f(module, modules, [module.value.tools as FlowModule[]]))
-			result = result.concat(dfs(module.value.tools as FlowModule[], f, opts))
+			visit(module, modules, [module.value.tools as FlowModule[]])
+			traverseFlowModules(module.value.tools as FlowModule[], visit, opts)
 		} else {
-			result.push(f(module, modules, []))
+			visit(module, modules, [])
 		}
 	}
+}
+
+export function dfs<T>(
+	modules: FlowModule[],
+	f: FlowModuleVisitor<T>,
+	opts: FlowDfsOptions = {}
+): T[] {
+	let result: T[] = []
+	traverseFlowModules(modules, (module, parentModules, branches) => {
+		result.push(f(module, parentModules, branches))
+	}, opts)
 	return result
 }
+
+export function forEachFlowModule(
+	modules: FlowModule[],
+	f: FlowModuleVisitor<void>,
+	opts: FlowDfsOptions = {}
+): void {
+	traverseFlowModules(modules, f, opts)
+}

frontend/src/lib/components/flows/flowDiffManager.svelte.ts

@@ -11,6 +11,7 @@ import type { FlowModule, FlowValue } from '$lib/gen'
 import type { ModuleActionInfo } from './flowDiff'
 import {
 	buildFlowTimeline,
+	findModuleInFlow,
 	insertModuleIntoFlow,
 	findModuleParent,
 	locationsEqual,
@@ -20,7 +21,6 @@ import {
 import { refreshStateStore } from '$lib/svelte5Utils.svelte'
 import type { StateStore } from '$lib/utils'
 import { getIndexInNestedModules } from '../copilot/chat/flow/utils'
-import { dfs } from './previousResults'
 import type DiffDrawer from '../DiffDrawer.svelte'
 import { SPECIAL_MODULE_IDS } from '../copilot/chat/shared'
 
@@ -185,13 +185,7 @@ export function createFlowDiffManager({ testMode = false } = {}) {
 	 * Helper to get a module from a flow by ID
 	 */
 	function getModuleFromFlow(id: string, flow: ExtendedOpenFlow): FlowModule | undefined {
-		if (flow.value.preprocessor_module?.id === id) {
-			return flow.value.preprocessor_module
-		} else if (flow.value.failure_module?.id === id) {
-			return flow.value.failure_module
-		} else {
-			return dfs(id, flow, false)[0]
-		}
+		return findModuleInFlow(flow.value, id) ?? undefined
 	}
 
 	/**
@@ -478,7 +472,7 @@ export function createFlowDiffManager({ testMode = false } = {}) {
 				} else if (currentFlow.failure_module?.id === moduleId) {
 					afterModule = currentFlow.failure_module
 				} else {
-					afterModule = dfs(moduleId, { value: currentFlow, summary: '' }, false)[0]
+					afterModule = findModuleInFlow(currentFlow, moduleId) ?? undefined
 				}
 			}