Skip to content

State Management

State in Helix Agents consists of two parts: built-in state (conversation history, status, etc.) and custom state (your application-specific data). This guide focuses on custom state.

Why Custom State?

Custom state solves several problems:

  1. Tool Coordination - Tools can share data without passing through the LLM
  2. Progress Tracking - Track metrics, counts, and phases of execution
  3. Dynamic Behavior - System prompts and LLM config can adapt based on state
  4. Persistence - State survives across steps and can be resumed after crashes

Defining State Schemas

Use Zod to define your state schema:

typescript
import { z } from 'zod';
import { defineAgent } from '@helix-agents/sdk';

const StateSchema = z.object({
  // Counters and metrics
  searchCount: z.number().default(0),
  tokensUsed: z.number().default(0),

  // Phase tracking
  phase: z.enum(['gathering', 'analyzing', 'summarizing']).default('gathering'),

  // Collections
  findings: z.array(z.string()).default([]),
  sourcesVisited: z.array(z.string()).default([]),

  // Flags
  hasFoundAnswer: z.boolean().default(false),

  // Optional values
  currentQuery: z.string().optional(),
  errorCount: z.number().default(0),
});

const agent = defineAgent({
  name: 'researcher',
  stateSchema: StateSchema,
  // ...
});

Default Values

Always provide defaults using .default():

typescript
// Good: All fields have defaults
const StateSchema = z.object({
  count: z.number().default(0),
  items: z.array(z.string()).default([]),
  config: z
    .object({
      maxRetries: z.number().default(3),
    })
    .default({}),
});

// Bad: No defaults - will fail at runtime
const StateSchema = z.object({
  count: z.number(), // Error: no default
  items: z.array(z.string()), // Error: no default
});

Nested Objects

For nested objects, provide defaults at both levels:

typescript
const StateSchema = z.object({
  metadata: z
    .object({
      createdAt: z.number().optional(),
      updatedAt: z.number().optional(),
      author: z.string().optional(),
    })
    .default({}), // Default empty object

  settings: z
    .object({
      maxResults: z.number().default(10),
      includeImages: z.boolean().default(false),
    })
    .default({}),
});

Reading State

In Tools

Use context.getState<T>() to read state:

typescript
const searchTool = defineTool({
  name: 'search',
  inputSchema: z.object({ query: z.string() }),
  execute: async (input, context) => {
    // Get typed state
    const state = context.getState<z.infer<typeof StateSchema>>();

    // Read values
    console.log(`Search count: ${state.searchCount}`);
    console.log(`Phase: ${state.phase}`);
    console.log(`Previous findings: ${state.findings.length}`);

    // Check conditions
    if (state.searchCount >= 10) {
      throw new Error('Search limit reached');
    }

    if (state.hasFoundAnswer) {
      return { results: [], note: 'Already found answer' };
    }

    // ...perform search
  },
});

Read-Only Snapshot

getState() returns a read-only snapshot. Do not modify the returned object - use updateState() instead.

In System Prompts

System prompt functions receive custom state:

typescript
const agent = defineAgent({
  name: 'adaptive-assistant',
  stateSchema: StateSchema,
  systemPrompt: (state) => {
    let prompt = 'You are a research assistant.\n\n';

    // Adapt based on phase
    if (state.phase === 'gathering') {
      prompt += 'Focus on finding relevant information. Use search liberally.\n';
    } else if (state.phase === 'analyzing') {
      prompt += 'Analyze the findings. No more searching needed.\n';
    } else {
      prompt += 'Provide a final summary of your findings.\n';
    }

    // Include context from state
    if (state.findings.length > 0) {
      prompt += `\nFindings so far:\n${state.findings.map((f, i) => `${i + 1}. ${f}`).join('\n')}`;
    }

    return prompt;
  },
  // ...
});

In LLM Config Override

Access state when adjusting LLM config:

typescript
const agent = defineAgent({
  name: 'smart-agent',
  stateSchema: StateSchema,
  llmConfig: {
    model: openai('gpt-4o-mini'),
    temperature: 0.7,
  },
  llmConfigOverride: (customState, stepCount) => {
    // Use more capable model for complex phases
    if (customState.phase === 'analyzing') {
      return {
        model: openai('gpt-4o'),
        temperature: 0.3, // More focused
      };
    }

    // Reduce creativity as we progress
    if (stepCount > 10) {
      return { temperature: 0.2 };
    }

    return {};
  },
  // ...
});

Updating State

Basic Updates

Use context.updateState<T>(updater) with Immer's draft pattern:

typescript
execute: async (input, context) => {
  // Update state by mutating the draft
  context.updateState<z.infer<typeof StateSchema>>((draft) => {
    draft.searchCount++;
    draft.currentQuery = input.query;
  });

  // ... rest of tool execution
};

Immer Draft Pattern

Immer lets you write "mutating" code that produces immutable updates:

typescript
context.updateState<State>((draft) => {
  // Direct property assignment
  draft.count = 5;

  // Increment/decrement
  draft.count++;
  draft.count--;

  // Array push (most common pattern)
  draft.items.push('new item');

  // Array operations
  draft.items.pop();
  draft.items.shift();
  draft.items.unshift('first');
  draft.items.splice(1, 0, 'inserted');

  // Nested object modification
  draft.config.maxRetries = 5;

  // Set optional value
  draft.currentQuery = 'new query';

  // Clear optional value
  draft.currentQuery = undefined;
});

Multiple Updates

You can call updateState multiple times in one execution:

typescript
execute: async (input, context) => {
  // First update: mark as started
  context.updateState<State>((draft) => {
    draft.phase = 'processing';
    draft.startedAt = Date.now();
  });

  // Do some work...
  const results = await performWork();

  // Second update: record results
  context.updateState<State>((draft) => {
    draft.results.push(...results);
    draft.processedCount += results.length;
  });

  // Third update: mark as done
  context.updateState<State>((draft) => {
    draft.phase = 'done';
    draft.completedAt = Date.now();
  });

  return { success: true };
};

State Persistence

How It Works

  1. After each step, state is saved to the StateStore
  2. State stores (Memory, Redis, Cloudflare) handle persistence
  3. On resume, state is loaded from the store to continue execution
Step 1 → LLM Call → Tool Execution → State Saved

Step 2 → LLM Call → Tool Execution → State Saved

                                    [Crash/Restart]

Resume  → State Loaded → Step 3 continues...

JSON Serialization

State must be JSON-serializable:

typescript
// Good: JSON-serializable types
const StateSchema = z.object({
  count: z.number(),
  name: z.string(),
  tags: z.array(z.string()),
  metadata: z.record(z.string()),
  timestamp: z.number(), // Store dates as numbers
  data: z.unknown(), // Any JSON value
});

// Bad: Non-serializable types
const StateSchema = z.object({
  callback: z.function(), // Functions can't serialize
  date: z.date(), // Dates become strings
  map: z.map(), // Maps become objects
  set: z.set(), // Sets become arrays
  regex: z.instanceof(RegExp), // Becomes empty object
});

Working with dates:

typescript
const StateSchema = z.object({
  createdAt: z.number().default(() => Date.now()),
  updatedAt: z.number().optional(),
});

// In tool
context.updateState<State>((draft) => {
  draft.updatedAt = Date.now();
});

// Reading
const state = context.getState<State>();
const date = new Date(state.createdAt);

Advanced Patterns

Phase-Based Execution

Use state to track execution phases:

typescript
const StateSchema = z.object({
  phase: z.enum(['init', 'search', 'analyze', 'summarize', 'done']).default('init'),
  phaseData: z.record(z.unknown()).default({}),
});

const agent = defineAgent({
  name: 'phased-agent',
  stateSchema: StateSchema,
  systemPrompt: (state) => {
    const instructions = {
      init: 'Start by understanding the user request.',
      search: 'Search for relevant information.',
      analyze: 'Analyze what you found.',
      summarize: 'Provide a final summary.',
      done: 'Task complete.',
    };
    return instructions[state.phase];
  },
  // ...
});

// Tool to advance phase
const advancePhaseTool = defineTool({
  name: 'advance_phase',
  description: 'Move to the next execution phase',
  inputSchema: z.object({}),
  execute: async (_, context) => {
    const state = context.getState<z.infer<typeof StateSchema>>();
    const phases = ['init', 'search', 'analyze', 'summarize', 'done'];
    const currentIndex = phases.indexOf(state.phase);

    if (currentIndex < phases.length - 1) {
      context.updateState<typeof state>((draft) => {
        draft.phase = phases[currentIndex + 1] as typeof state.phase;
      });
    }

    return { newPhase: phases[currentIndex + 1] };
  },
});

Accumulating Results

Common pattern for gathering results:

typescript
const StateSchema = z.object({
  results: z
    .array(
      z.object({
        source: z.string(),
        data: z.unknown(),
        quality: z.number(),
      })
    )
    .default([]),
  bestResult: z
    .object({
      source: z.string(),
      data: z.unknown(),
      quality: z.number(),
    })
    .optional(),
});

const gatherTool = defineTool({
  name: 'gather',
  inputSchema: z.object({ source: z.string() }),
  execute: async (input, context) => {
    const data = await fetchFromSource(input.source);
    const quality = assessQuality(data);

    context.updateState<z.infer<typeof StateSchema>>((draft) => {
      // Add to results
      draft.results.push({ source: input.source, data, quality });

      // Track best result
      if (!draft.bestResult || quality > draft.bestResult.quality) {
        draft.bestResult = { source: input.source, data, quality };
      }
    });

    return { quality };
  },
});

Rate Limiting

Track usage and enforce limits:

typescript
const StateSchema = z.object({
  apiCalls: z.number().default(0),
  lastCallTime: z.number().default(0),
  rateLimitErrors: z.number().default(0),
});

const apiTool = defineTool({
  name: 'api_call',
  inputSchema: z.object({ endpoint: z.string() }),
  execute: async (input, context) => {
    const state = context.getState<z.infer<typeof StateSchema>>();

    // Check rate limit
    const timeSinceLastCall = Date.now() - state.lastCallTime;
    if (timeSinceLastCall < 1000) {
      // 1 second minimum between calls
      await new Promise((resolve) => setTimeout(resolve, 1000 - timeSinceLastCall));
    }

    // Check total calls
    if (state.apiCalls >= 100) {
      throw new Error('API call limit reached (100 calls)');
    }

    // Track the call
    context.updateState<typeof state>((draft) => {
      draft.apiCalls++;
      draft.lastCallTime = Date.now();
    });

    try {
      return await makeApiCall(input.endpoint);
    } catch (error) {
      if (error.message.includes('rate limit')) {
        context.updateState<typeof state>((draft) => {
          draft.rateLimitErrors++;
        });
      }
      throw error;
    }
  },
});

Gotchas and Limitations

1. No Direct Array Index Modification in Parallel

When tools run in parallel, avoid modifying arrays by index:

typescript
// Risky with parallel execution
context.updateState<State>((draft) => {
  draft.items[0] = 'modified'; // May affect wrong item
  draft.items.splice(1, 1); // Index may have shifted
});

// Safe: push to arrays, or use maps with keys
context.updateState<State>((draft) => {
  draft.items.push(newItem); // Safe: appends to end
  draft.itemsById[id] = item; // Safe: keyed access
});

2. Type Casting Required

TypeScript can't infer the state type from context alone:

typescript
// Required: explicit type parameter
context.updateState<z.infer<typeof StateSchema>>((draft) => {
  draft.count++; // TypeScript knows about 'count'
});

// Common pattern: define type alias
type State = z.infer<typeof StateSchema>;

context.updateState<State>((draft) => {
  draft.count++;
});

3. Default Values Must Be Provided

State is initialized from schema defaults on first run:

typescript
// Schema with defaults
const StateSchema = z.object({
  count: z.number().default(0), // ✓ Default provided
  items: z.array(z.string()).default([]), // ✓ Default provided
  name: z.string().optional(), // ✓ Optional fields OK
});

// First run: state is { count: 0, items: [], name: undefined }

4. State Changes Stream to Clients

Updates are sent as RFC 6902 patches. Large state updates mean more streaming data:

typescript
// Efficient: small incremental updates
context.updateState<State>((draft) => {
  draft.items.push(newItem); // Small patch: add operation
});

// Less efficient: replacing large arrays
context.updateState<State>((draft) => {
  draft.items = [...draft.items, newItem]; // Replace entire array
});

Next Steps

Released under the MIT License.