@swizzy_ai/kit 1.0.4 → 1.0.6

package/dist/index.js

@@ -47,6 +47,7 @@ class Step {
         this.beforeRun = config.beforeRun;
         this.afterRun = config.afterRun;
         this.model = config.model;
+        this.stream = config.stream ?? true; // Default to streaming enabled
     }
     validate(data) {
         const result = this.schema.safeParse(data);
@@ -55,8 +56,8 @@ class Step {
         }
         return result.data;
     }
-    getContext(workflowContext) {
-        return this.context ? this.context(workflowContext) : workflowContext;
+    async getContext(workflowContext) {
+        return this.context ? await this.context(workflowContext) : workflowContext;
     }
 }
 
@@ -380,18 +381,27 @@ class BungeeBuilder {
         /**
          * Add a single step execution.
          */
-        this.add = (stepId) => {
-            this._plan.destinations.push({ type: 'step', targetId: stepId });
+        this.add = (stepId, config) => {
+            this._plan.destinations.push({ type: 'step', targetId: stepId, config });
             return this;
         };
         /**
          * Add multiple executions based on count with config function.
          */
-        this.batch = (stepId, count, configFn) => {
+        this.batch = (stepId, count, configFn, options) => {
             for (let i = 0; i < count; i++) {
-                this._plan.destinations.push({ type: 'step', targetId: stepId });
+                this._plan.destinations.push({ type: 'step', targetId: stepId, config: configFn(i) });
+            }
+            // Apply batch-specific options
+            if (options?.optimistic !== undefined) {
+                this._plan.optimistic = options.optimistic;
+            }
+            if (options?.returnToAnchor !== undefined) {
+                this._plan.returnToAnchor = options.returnToAnchor;
+            }
+            if (options?.failWizardOnFailure !== undefined) {
+                this._plan.failWizardOnFailure = options.failWizardOnFailure;
             }
-            this._plan.configFn = configFn;
             return this;
         };
         /**
@@ -401,6 +411,22 @@ class BungeeBuilder {
             if (options.concurrency !== undefined) {
                 this._plan.concurrency = options.concurrency;
             }
+            if (options.optimistic !== undefined) {
+                this._plan.optimistic = options.optimistic;
+            }
+            if (options.returnToAnchor !== undefined) {
+                this._plan.returnToAnchor = options.returnToAnchor;
+            }
+            if (options.failWizardOnFailure !== undefined) {
+                this._plan.failWizardOnFailure = options.failWizardOnFailure;
+            }
+            return this;
+        };
+        /**
+         * Set completion callback.
+         */
+        this.onComplete = (callback) => {
+            this._plan.onComplete = callback;
             return this;
         };
         /**
@@ -428,8 +454,8 @@ class SchemaUtils {
             const shape = schema._def.shape();
             const fields = Object.entries(shape).map(([key, fieldSchema]) => {
                 const type = this.getSchemaType(fieldSchema);
-                const xmlExample = this.getXmlExample(key, type);
-                return `${key}: ${type} - ${xmlExample}`;
+                this.getXmlExample(key, type);
+                return `${key}: ${type}`;
             });
             description = `Object with fields:\n${fields.join('\n')}`;
         }
@@ -479,8 +505,8 @@ class SchemaUtils {
     }
     static getXmlExample(key, type) {
         switch (type) {
-            case 'string': return `<${key} tag-category="wizard" type="string">example`;
-            case 'number': return `<${key} tag-category="wizard" type="number">123`;
+            case 'string': return `<${key} tag-category="wizard" type="string">[your text should be here]`;
+            case 'number': return `<${key} tag-category="wizard" type="number">[number should be here]`;
             case 'boolean': return `<${key} tag-category="wizard" type="boolean">true`;
             case 'array': return `<${key} tag-category="wizard" type="array">["item1", "item2"]`;
             default:
@@ -488,7 +514,7 @@ class SchemaUtils {
                     const values = type.split(': ')[1].split(', ');
                     return `<${key} tag-category="wizard" type="string">${values[0]}`;
                 }
-                return `<${key} tag-category="wizard" type="object"><subfield type="string">value</subfield>`;
+                return `<${key} tag-category="wizard" type="object"><subfield type="string">[text value should be here]</subfield>`;
         }
     }
     static objectToXml(obj, rootName = 'context') {
@@ -533,6 +559,20 @@ class VisualizationManager {
         this.connectedClients = new Set();
         this.maxWebSocketConnections = 10;
         this.wsIntervals = new WeakMap();
+        // Listen for state update events
+        this.wizard.on('state:update', (data) => {
+            this.sendToClients({
+                type: 'state_update',
+                ...data
+            });
+        });
+        // Listen for wizard stop events
+        this.wizard.on('wizard:stop', (data) => {
+            this.sendToClients({
+                type: 'wizard_stop',
+                ...data
+            });
+        });
     } // Wizard instance for callbacks
     getStepsInfo() {
         return this.wizard.steps.map((item) => {
@@ -696,6 +736,12 @@ class VisualizationManager {
                         this.wizard.isPaused = false;
                         this.wizard.isStepMode = false;
                         this.sendToClients({ type: 'status_update', status: { isRunning: false, isPaused: false, isStepMode: false } });
+                        // Emit wizard:stop event for manual stops
+                        this.wizard.events.emit('wizard:stop', {
+                            reason: 'manual_stop',
+                            finalState: this.wizard.stateManager.getState(),
+                            timestamp: Date.now()
+                        });
                         break;
                     case 'replay':
                         console.log('🔄 Replaying wizard - resetting state');
@@ -803,24 +849,33 @@ class BungeeExecutor {
             // Launch worker
             const workerPromise = this.launchBungeeWorker(plan, i);
             activeWorkers.add(workerPromise);
-            // Respect concurrency limit
-            if (activeWorkers.size >= plan.concurrency) {
-                await Promise.race(activeWorkers);
-                // Clean up completed workers
-                for (const promise of activeWorkers) {
-                    if (promise !== workerPromise) {
-                        activeWorkers.delete(promise);
-                    }
+        }
+        // Wait for all workers to complete unless optimistic
+        if (!plan.optimistic) {
+            try {
+                await Promise.all(activeWorkers);
+            }
+            catch (error) {
+                if (plan.failWizardOnFailure !== false) { // Default true
+                    throw error; // Re-throw to stop wizard
                 }
+                // If failWizardOnFailure is false, just log and continue
+                console.error(`Bungee plan ${plan.id} had failures but continuing:`, error);
             }
         }
-        // Wait for all workers to complete
-        await Promise.all(activeWorkers);
-        console.log(`✅ Bungee plan ${plan.id} completed, returning to anchor ${plan.anchorId}`);
+        if (plan.onComplete) {
+            return plan.onComplete(this.wizard);
+        }
+        if (plan.returnToAnchor !== false) { // Default true
+            console.log(`✅ Bungee plan ${plan.id} completed, returning to anchor ${plan.anchorId}`);
+        }
+        else {
+            console.log(`✅ Bungee plan ${plan.id} completed, proceeding to next step`);
+        }
     }
     async launchBungeeWorker(plan, index) {
         const destination = plan.destinations[index];
-        const telescope = plan.configFn ? plan.configFn(index) : {};
+        const telescope = destination.config || {};
         const workerId = `${plan.id}_${destination.targetId}_${index}_${Date.now()}`;
         const telescopeContext = this.createTelescopeContext(this.wizard.workflowContext, telescope);
         const promise = this.executeWorkerStep(destination.targetId, telescopeContext);
@@ -837,10 +892,6 @@ class BungeeExecutor {
         try {
             await promise;
         }
-        catch (error) {
-            console.error(`Bungee worker ${workerId} failed:`, error);
-            this.wizard.workflowContext[`${workerId}_error`] = error.message;
-        }
         finally {
             // Clean up
             const planWorkers = this.bungeeWorkers.get(plan.id);
@@ -848,8 +899,10 @@ class BungeeExecutor {
                 planWorkers.delete(workerId);
                 if (planWorkers.size === 0) {
                     this.bungeeWorkers.delete(plan.id);
-                    // Trigger reentry to anchor
-                    this.pendingReentry.add(plan.anchorId);
+                    // Trigger reentry to anchor if configured
+                    if (plan.returnToAnchor !== false) {
+                        this.pendingReentry.add(plan.anchorId);
+                    }
                 }
             }
         }
@@ -872,9 +925,7 @@ class BungeeExecutor {
         return await step.update(stepData, telescopeContext, actions);
     }
     mergeWorkerResults(updates, telescope) {
-        Object.entries(updates).forEach(([key, value]) => {
-            this.wizard.workflowContext[key] = value;
-        });
+        this.wizard.updateContext(updates);
     }
     async retriggerAnchor(anchorId) {
         const anchorStep = this.wizard.findStep(anchorId);
@@ -993,24 +1044,49 @@ class UsageTracker {
     }
 }
 
-class ContextManager {
-    constructor() {
-        this.workflowContext = {};
+class StateManager {
+    constructor(events) {
+        this.events = events;
+        this.state = {};
+    }
+    setState(updates) {
+        let newUpdates;
+        if (typeof updates === 'function') {
+            // Higher-order function pattern like React setState
+            newUpdates = updates(this.state);
+        }
+        else {
+            // Object pattern
+            newUpdates = updates;
+        }
+        const previousState = { ...this.state };
+        this.state = { ...this.state, ...newUpdates };
+        // Emit state update event
+        this.events.emit('state:update', {
+            previousState,
+            newState: this.state,
+            updates: newUpdates,
+            timestamp: Date.now()
+        });
+    }
+    getState() {
+        return this.state;
     }
+    // Legacy methods for backward compatibility
     setContext(context) {
-        this.workflowContext = { ...this.workflowContext, ...context };
+        this.setState(context);
     }
     getContext() {
-        return this.workflowContext;
+        return this.getState();
     }
     updateContext(updates) {
-        this.workflowContext = { ...this.workflowContext, ...updates };
+        this.setState(updates);
     }
     getWorkflowContext() {
-        return this.workflowContext;
+        return this.getState();
     }
-    setWorkflowContext(context) {
-        this.workflowContext = context;
+    setWorkflowContext(state) {
+        this.state = state;
     }
 }
 
@@ -1042,10 +1118,10 @@ class Wizard {
     }
     // Getters for manager methods
     get workflowContext() {
-        return this.contextManager.getContext();
+        return this.stateManager.getContext();
     }
     set workflowContext(value) {
-        this.contextManager.setWorkflowContext(value);
+        this.stateManager.setWorkflowContext(value);
     }
     get log() {
         return this.isLoggingEnabled ? this.logger.log.bind(this.logger) : () => { };
@@ -1065,7 +1141,7 @@ class Wizard {
         this.isStepMode = false;
         this.skipStartWait = false;
         this.debouncedSendContextUpdate = this.debounce(() => {
-            this.visualizationManager.sendContextUpdate(this.contextManager.getContext());
+            this.visualizationManager.sendContextUpdate(this.stateManager.getContext());
         }, 100);
         this.id = config.id;
         const registry = new ProviderRegistry();
@@ -1076,13 +1152,13 @@ class Wizard {
         this.isLoggingEnabled = config.logging ?? (process.env.NODE_ENV === 'development' || !process.env.NODE_ENV);
         // Initialize managers
         this.logger = new Logger(this.id);
-        this.contextManager = new ContextManager();
+        this.events = new EventEmitter();
+        this.stateManager = new StateManager(this.events);
         this.visualizationManager = new VisualizationManager(this);
         this.usageTracker = new UsageTracker(config.onUsage, (totalTokens, rate) => {
             this.visualizationManager.sendTokenUpdate(totalTokens, rate);
         });
         this.bungeeExecutor = new BungeeExecutor(this);
-        this.events = new EventEmitter();
     }
     addStep(config) {
         const step = new Step(config);
@@ -1126,10 +1202,12 @@ class Wizard {
         return this;
     }
     clearStepError(stepId) {
-        const context = this.contextManager.getContext();
-        delete context[`${stepId}_error`];
-        delete context[`${stepId}_retryCount`];
-        this.contextManager.setWorkflowContext(context);
+        this.stateManager.setState((prevState) => {
+            const newState = { ...prevState };
+            delete newState[`${stepId}_error`];
+            delete newState[`${stepId}_retryCount`];
+            return newState;
+        });
     }
     isStringSignal(signal) {
         return typeof signal === 'string';
@@ -1152,11 +1230,7 @@ class Wizard {
                 this.currentStepIndex++;
                 return true;
             default:
-                if (this.isBungeeJumpSignal(signal)) {
-                    await this.bungeeExecutor.executeBungeePlan(signal.plan);
-                    return true;
-                }
-                else if (this.isStringSignal(signal) && signal.startsWith('GOTO ')) {
+                if (this.isStringSignal(signal) && signal.startsWith('GOTO ')) {
                     const targetStepId = signal.substring(5);
                     const targetIndex = this.findStepIndex(targetStepId);
                     if (targetIndex !== -1) {
@@ -1167,6 +1241,28 @@ class Wizard {
                     }
                     return true;
                 }
+                else if (this.isBungeeJumpSignal(signal)) {
+                    try {
+                        const result = await this.bungeeExecutor.executeBungeePlan(signal.plan);
+                        if (result) {
+                            return await this.handleFlowControlSignal(result);
+                        }
+                        if (signal.plan.returnToAnchor === false) {
+                            this.currentStepIndex++; // Proceed to next step when not returning to anchor
+                        }
+                        return true;
+                    }
+                    catch (error) {
+                        console.error('Bungee plan failed:', error);
+                        this.workflowContext[`bungee_error`] = error.message;
+                        if (signal.plan.failWizardOnFailure !== false) { // Default true
+                            this.isRunning = false; // Stop the wizard on bungee failure
+                            return false;
+                        }
+                        // If failWizardOnFailure is false, continue
+                        return true;
+                    }
+                }
         }
         return true;
     }
@@ -1261,10 +1357,17 @@ class Wizard {
             totalSteps: this.steps.length,
             timestamp: endTime
         });
+        // Emit wizard:stop event
+        this.events.emit('wizard:stop', {
+            reason: 'completed',
+            finalState: this.stateManager.getState(),
+            timestamp: endTime
+        });
     }
     createBaseActions() {
         return {
             updateContext: (updates) => this.updateContext(updates),
+            setState: (updates) => this.stateManager.setState(updates),
             llmClient: this.llmClient,
             goto: (stepId) => this.goto(stepId),
             next: () => this.next(),
@@ -1276,6 +1379,7 @@ class Wizard {
     createWizardActions(anchorStepId = '') {
         return {
             ...this.createBaseActions(),
+            setState: (updates) => this.stateManager.setState(updates),
             bungee: {
                 init: () => new BungeeBuilder(anchorStepId)
             }
@@ -1286,6 +1390,9 @@ class Wizard {
             updateContext: (updates) => {
                 this.bungeeExecutor.mergeWorkerResults(updates, telescope);
             },
+            setState: (updates) => {
+                this.bungeeExecutor.mergeWorkerResults(updates, telescope);
+            },
             llmClient: this.llmClient,
             goto: () => Wizard.STOP,
             next: () => Wizard.STOP,
@@ -1324,7 +1431,7 @@ class Wizard {
             instruction: step.instruction,
             timestamp: stepStartTime
         });
-        const stepContext = step.getContext(this.contextManager.getContext());
+        const stepContext = await step.getContext(this.stateManager.getContext());
         let processedInstruction = step.instruction;
         if (step.contextType === 'template' || step.contextType === 'both') {
             processedInstruction = this.applyTemplate(step.instruction, stepContext);
@@ -1519,124 +1626,98 @@ class Wizard {
      * Uses streaming for regular steps to provide real-time parsing and UI updates.
      * TextStep and ComputeStep use different approaches (non-streaming).
      */
+    /**
+      * Generates data for a wizard step by calling the LLM.
+      * CLEANED VERSION: Separates formatting rules from logic to prevent hallucination.
+      */
     async generateStepData(step, stepContext) {
         const systemContext = this.systemPrompt ? `${this.systemPrompt}\n\n` : '';
+        // Build context strings
         const errorContext = this.workflowContext[`${step.id}_error`] ?
-            `\n\nPREVIOUS ERROR (attempt ${this.workflowContext[`${step.id}_retryCount`] || 1}):\n${this.workflowContext[`${step.id}_error`]}\nPlease fix this.` : '';
+            `\n\n!!! PREVIOUS ERROR (Attempt ${this.workflowContext[`${step.id}_retryCount`] || 1}) !!!\nThe previous output caused this error: ${this.workflowContext[`${step.id}_error`]}\nYOU MUST FIX THIS.` : '';
         let processedInstruction = step.instruction;
         if (step.contextType === 'template' || step.contextType === 'both') {
             processedInstruction = this.applyTemplate(step.instruction, stepContext);
         }
-        this.log(() => `Processed instruction for step ${step.id}: ${processedInstruction}`);
         let contextSection = '';
         if (step.contextType === 'xml' || step.contextType === 'both' || !step.contextType) {
-            contextSection = `\n\nSTEP CONTEXT:\n${this.objectToXml(stepContext)}`;
+            contextSection = `\n\n### CURRENT CONTEXT ###\n${this.objectToXml(stepContext)}`;
         }
-        this.log(() => `Context section for step ${step.id}: ${contextSection}`);
+        // --- Text Step Handling ---
         if (step instanceof TextStep) {
-            const prompt = `${systemContext}You are executing a wizard step. Generate text for this step.
+            const prompt = `${systemContext}
+TASK: Generate content for step "${step.id}".
 
-STEP: ${step.id}
-INSTRUCTION: ${processedInstruction}${errorContext}${contextSection}
+INSTRUCTION:
+${processedInstruction}
+${contextSection}
+${errorContext}
 
+OUTPUT:
 Generate the text response now.`;
-            this.log(() => `Full prompt for step ${step.id}: ${prompt}`);
             let fullText = '';
-            await this.llmClient.complete({
+            console.log("full prompt", prompt);
+            const result = await this.llmClient.complete({
                 prompt,
                 model: step.model,
                 maxTokens: 1000,
                 temperature: 0.3,
-                stream: true,
+                stream: step.stream,
                 onChunk: (chunk) => {
                     fullText += chunk;
-                    // Emit raw chunk event for streaming text
-                    this.events.emit('step:chunk', {
-                        stepId: step.id,
-                        chunk: chunk,
-                        timestamp: Date.now()
-                    });
+                    this.events.emit('step:chunk', { stepId: step.id, chunk, timestamp: Date.now() });
                 },
                 onUsage: this.usageTracker.updateUsage.bind(this.usageTracker)
             });
-            this.log(() => `LLM response for step ${step.id}: ${fullText}`);
-            console.log(`LLM response for step ${step.id}:`, fullText);
+            if (result.text) {
+                fullText = result.text;
+            }
             return fullText;
         }
-        // For regular steps, use streaming XML parsing
-        // This allows real-time processing of LLM responses as they arrive
+        // --- Regular XML Step Handling ---
         const parser = this.createStreamingXmlParser();
         let latestResult = {};
         const schemaDescription = SchemaUtils.describeSchema(step.schema, step.id);
-        const prompt = `${systemContext}You are executing a wizard step. Generate data for this step.
-
-STEP: ${step.id}
-INSTRUCTION: ${processedInstruction}${errorContext}${contextSection}
-
-SCHEMA REQUIREMENTS:
-${schemaDescription}
-
-REQUIRED OUTPUT FORMAT:
-Return a plain XML response with a root <response> tag.
-CRITICAL: Every field MUST include tag-category="wizard" attribute. This is MANDATORY.
-Every field MUST also include a type attribute (e.g., type="string", type="number", type="boolean", type="array").
+        // CLEANER PROMPT STRUCTURE
+        const prompt = `${systemContext}
+=== GOAL ===
+You are an intelligent agent executing step: "${step.id}".
+Your task is to generate data that satisfies the INSTRUCTION below based on the CONTEXT.
 
-ARRAY FORMATTING RULES (CRITICAL):
-- Arrays MUST be valid JSON on a SINGLE line
-- Use double quotes, not single quotes: ["a", "b"] NOT ['a', 'b']
-- NO trailing commas: ["a", "b"] NOT ["a", "b",]
-- NO line breaks inside arrays
-- Example: <items tag-category="wizard" type="array">["apple", "banana", "orange"]
+=== INSTRUCTION ===
+${processedInstruction}
+${contextSection}
+${errorContext}
 
-IMPORTANT PARSING RULES:
-- Fields with tag-category="wizard" do NOT need closing tags
-- Content ends when the next tag with tag-category="wizard" begins, OR when </response> is reached
-- This means you can include ANY content (including code with <>, XML snippets, etc.) without worrying about breaking the parser
-- Only fields marked with tag-category="wizard" will be parsed
+=== RESPONSE FORMAT ===
+You must output a VALID XML object inside a <response> tag.
+1. Every field must have: tag-category="wizard" and a type attribute (string, number, boolean, array).
+2. Arrays must be single-line JSON: <tags tag-category="wizard" type="array">["a", "b"]
 
-Example:
-<response>
-  <name tag-category="wizard" type="string">John Smith
-  <age tag-category="wizard" type="number">25
-  <tags tag-category="wizard" type="array">["a", "b", "c"]
-</response>
+=== SCHEMA DEFINITION ===
+${schemaDescription}
 
-Notice: Arrays are compact JSON on one line! No closing tags needed for wizard fields.
+*** CRITICAL RULES ***
+1. Do NOT copy values from the schema definition or examples above.
+2. Generate NEW values based strictly on the "INSTRUCTION" and "CURRENT CONTEXT".
+3. If the instruction implies a selection (like an ID), ensure the ID exists in the Context.
 
-Generate the XML response now.`;
-        this.log(() => `Full prompt for step ${step.id}: ${prompt}`);
-        // Initiate streaming LLM call with chunk processing
-        // Each chunk is pushed to the parser for incremental XML parsing
-        await this.llmClient.complete({
+Generate the XML <response> now.`;
+        // console.log(prompt); // Uncomment to debug the cleaner prompt
+        step.stream !== false;
+        const result = await this.llmClient.complete({
             prompt,
             model: step.model,
             maxTokens: 1000,
-            temperature: 0.3,
-            stream: true,
+            temperature: 0.3, // Lower temp for precision
+            stream: step.stream,
             onChunk: (chunk) => {
-                // Emit raw chunk event
-                this.events.emit('step:chunk', {
-                    stepId: step.id,
-                    chunk: chunk,
-                    timestamp: Date.now()
-                });
-                // Process each incoming chunk through the streaming parser
+                this.events.emit('step:chunk', { stepId: step.id, chunk, timestamp: Date.now() });
                 const parseResult = parser.push(chunk);
                 if (parseResult && !parseResult.done) {
                     latestResult = parseResult.result;
-                    // Emit parsed streaming event
-                    this.events.emit('step:streaming', {
-                        stepId: step.id,
-                        data: latestResult,
-                        timestamp: Date.now()
-                    });
-                    // Stream partial results out to connected clients via WebSocket
-                    // This broadcasts real-time updates to UI and external consumers
-                    this.visualizationManager.sendStepUpdate({
-                        stepId: step.id,
-                        status: 'streaming',
-                        data: latestResult
-                    });
+                    this.events.emit('step:streaming', { stepId: step.id, data: latestResult, timestamp: Date.now() });
+                    this.visualizationManager.sendStepUpdate({ stepId: step.id, status: 'streaming', data: latestResult });
                 }
                 else if (parseResult?.done) {
                     latestResult = parseResult.result;
@@ -1644,20 +1725,19 @@ Generate the XML response now.`;
             },
             onUsage: this.usageTracker.updateUsage.bind(this.usageTracker)
         });
-        this.log(() => `Final parsed data: ${JSON.stringify(latestResult)}`);
-        this.log(() => `Parsed JSON data for step ${step.id}: ${JSON.stringify(latestResult)}`);
+        if (result.text) {
+            latestResult = result.text;
+        }
         try {
             return step.validate(latestResult);
         }
         catch (validationError) {
-            this.log(() => `Validation failed for step ${step.id}: ${validationError.message}`);
+            // Logic for repair remains the same...
             try {
                 const repairedData = await this.repairSchemaData(latestResult, step.schema, validationError.message, step.id);
-                this.log(() => `Repaired data for step ${step.id}: ${JSON.stringify(repairedData)}`);
                 return step.validate(repairedData);
             }
             catch (repairError) {
-                this.log(() => `Repair failed for step ${step.id}: ${repairError.message}`);
                 return { __validationFailed: true, error: validationError.message };
             }
         }
@@ -1700,16 +1780,17 @@ Fix the data to match the schema and generate the XML response now.`;
         return repairedJsonData;
     }
     /**
-     * Creates a streaming XML parser for incremental processing of LLM responses.
+     * Creates an improved streaming XML parser for incremental processing of LLM responses.
      *
-     * This parser processes XML chunks as they arrive from the LLM, extracting
-     * fields marked with tag-category="wizard" and parsing them based on their type attributes.
+     * This parser is designed to handle partial chunks robustly and provides better error recovery.
+     * It processes XML chunks as they arrive, extracting fields marked with tag-category="wizard".
      *
-     * Key features:
-     * - Incremental parsing: processes data as it streams in
-     * - Buffer management: accumulates chunks until complete fields are available
-     * - Type-aware parsing: handles strings, numbers, booleans, arrays, objects
-     * - Real-time updates: returns partial results as fields complete
+     * Key improvements:
+     * - Better partial tag handling
+     * - More robust regex matching
+     * - Improved buffer management
+     * - Better error recovery for malformed chunks
+     * - State machine approach for parsing
      *
      * @returns An object with a push method that accepts text chunks and returns parse results
      */
@@ -1718,57 +1799,131 @@ Fix the data to match the schema and generate the XML response now.`;
         let inResponse = false; // Tracks if we've entered the <response> tag
         const result = {}; // The final parsed JSON object
         let currentField = null; // Currently parsing field
+        let parseErrors = 0; // Track consecutive parse errors
         return {
             push: (chunk) => {
-                buffer += chunk;
-                // Wait for <response> tag to start parsing
-                if (!inResponse && buffer.includes('<response>')) {
-                    inResponse = true;
-                    buffer = buffer.slice(buffer.indexOf('<response>') + 10); // Remove <response> from buffer
-                }
-                if (!inResponse)
-                    return null; // Nothing to parse yet
-                // Look for wizard-tagged fields using regex pattern
-                const tagMatch = buffer.match(Wizard.WIZARD_TAG_PATTERN);
-                if (tagMatch) {
-                    // If we were parsing a field, finalize it before starting new one
-                    if (currentField) {
-                        result[currentField.name] = this.parseValueByType(currentField.content.trim(), currentField.type);
+                try {
+                    buffer += chunk;
+                    // Wait for <response> tag to start parsing
+                    if (!inResponse) {
+                        const responseStart = buffer.indexOf('<response>');
+                        if (responseStart !== -1) {
+                            inResponse = true;
+                            buffer = buffer.slice(responseStart + 10); // Remove <response> from buffer
+                        }
+                        else {
+                            return null; // Still waiting for response start
+                        }
                     }
-                    // Extract field name and type from the matched tag
-                    const typeMatch = tagMatch[2].match(/type=["']([^"']+)["']/);
-                    currentField = {
-                        name: tagMatch[1], // Field name from tag
-                        type: typeMatch?.[1]?.toLowerCase() || 'string', // Type attribute, default to string
-                        content: '' // Will accumulate field content
-                    };
-                    // Remove the processed tag from buffer
-                    buffer = buffer.slice(tagMatch.index + tagMatch[0].length);
-                }
-                // Accumulate content for the current field
-                if (currentField) {
-                    // Find the next wizard tag or end of response
-                    const nextTagIndex = buffer.search(/<\w+\s+[^>]*tag-category=["']wizard["']/);
-                    if (nextTagIndex !== -1) {
-                        // Found next tag, accumulate content up to it
-                        currentField.content += buffer.slice(0, nextTagIndex);
-                        buffer = buffer.slice(nextTagIndex);
+                    // Process buffer for wizard tags
+                    let processedSomething = false;
+                    // Continue processing while we have data
+                    while (buffer.length > 0) {
+                        // If we have a current field, try to accumulate content
+                        if (currentField) {
+                            // Look for the next wizard tag or end of response
+                            const nextWizardTag = buffer.match(/<\w+\s+[^>]*tag-category=["']wizard["'][^>]*>/);
+                            const responseEnd = buffer.indexOf('</response>');
+                            if (nextWizardTag && nextWizardTag.index !== undefined) {
+                                // Found next tag, finalize current field
+                                const contentEnd = nextWizardTag.index;
+                                currentField.content += buffer.slice(0, contentEnd);
+                                result[currentField.name] = this.parseValueByType(currentField.content.trim(), currentField.type);
+                                currentField = null;
+                                buffer = buffer.slice(contentEnd);
+                                processedSomething = true;
+                            }
+                            else if (responseEnd !== -1) {
+                                // End of response, finalize current field
+                                currentField.content += buffer.slice(0, responseEnd);
+                                result[currentField.name] = this.parseValueByType(currentField.content.trim(), currentField.type);
+                                return { done: true, result }; // Parsing complete
+                            }
+                            else {
+                                // No complete field yet, but we might have a partial tag at the end
+                                // Check if buffer ends with a partial wizard tag
+                                const partialTagMatch = buffer.match(/<\w+\s+[^>]*tag-category=["']wizard["'][^>]*$/);
+                                if (partialTagMatch) {
+                                    // Buffer ends with partial tag, keep it for next chunk
+                                    break;
+                                }
+                                else {
+                                    // Safe to accumulate entire buffer
+                                    currentField.content += buffer;
+                                    buffer = '';
+                                    processedSomething = true;
+                                }
+                            }
+                        }
+                        else {
+                            // No current field, look for a new wizard tag
+                            const tagMatch = buffer.match(Wizard.WIZARD_TAG_PATTERN);
+                            if (tagMatch && tagMatch.index === 0) {
+                                // Tag starts at beginning of buffer
+                                const typeMatch = tagMatch[2].match(/type=["']([^"']+)["']/);
+                                currentField = {
+                                    name: tagMatch[1],
+                                    type: typeMatch?.[1]?.toLowerCase() || 'string',
+                                    content: ''
+                                };
+                                buffer = buffer.slice(tagMatch[0].length);
+                                processedSomething = true;
+                            }
+                            else if (tagMatch && tagMatch.index !== undefined && tagMatch.index > 0) {
+                                // Tag exists but not at start - might be partial
+                                const partialTagMatch = buffer.match(/<\w+\s+[^>]*tag-category=["']wizard["'][^>]*$/);
+                                if (partialTagMatch) {
+                                    // Buffer ends with partial tag, wait for more data
+                                    break;
+                                }
+                                else {
+                                    // Tag is in middle, process up to it
+                                    // This shouldn't happen in well-formed XML, but handle gracefully
+                                    buffer = buffer.slice(tagMatch.index);
+                                    continue;
+                                }
+                            }
+                            else {
+                                // No wizard tag found
+                                if (buffer.includes('</response>')) {
+                                    // End of response without finalizing a field
+                                    return { done: true, result };
+                                }
+                                // Check for partial tag at end
+                                const partialTagMatch = buffer.match(/<\w+\s+[^>]*tag-category=["']wizard["'][^>]*$/);
+                                if (partialTagMatch) {
+                                    break; // Wait for more data
+                                }
+                                // No partial tag, buffer might contain non-wizard content
+                                // This is unusual but we'll keep it for now
+                                break;
+                            }
+                        }
                     }
-                    else if (buffer.includes('</response>')) {
-                        // End of response reached, finalize current field
-                        const endIndex = buffer.indexOf('</response>');
-                        currentField.content += buffer.slice(0, endIndex);
-                        result[currentField.name] = this.parseValueByType(currentField.content.trim(), currentField.type);
-                        return { done: true, result }; // Parsing complete
+                    // Reset error counter on successful processing
+                    if (processedSomething) {
+                        parseErrors = 0;
                     }
-                    else {
-                        // No complete field yet, accumulate entire buffer
-                        currentField.content += buffer;
-                        buffer = '';
+                    // Return partial result for UI updates
+                    return { done: false, result: { ...result } };
+                }
+                catch (error) {
+                    parseErrors++;
+                    console.warn(`Streaming XML parse error (attempt ${parseErrors}):`, error instanceof Error ? error.message : String(error));
+                    // If we have too many consecutive errors, try to recover
+                    if (parseErrors > 3) {
+                        console.error('Too many parse errors, attempting recovery');
+                        // Try to find next valid wizard tag and restart from there
+                        const recoveryMatch = buffer.match(/<\w+\s+[^>]*tag-category=["']wizard["'][^>]*>/);
+                        if (recoveryMatch && recoveryMatch.index !== undefined && recoveryMatch.index > 0) {
+                            buffer = buffer.slice(recoveryMatch.index);
+                            parseErrors = 0; // Reset error counter
+                            return { done: false, result: { ...result } };
+                        }
                     }
+                    // Return current result even with errors
+                    return { done: false, result: { ...result } };
                 }
-                // Return partial result for UI updates
-                return { done: false, result: { ...result } };
             }
         };
     }