​
Overview

Middleware in Magma allows you to intercept and modify agent behavior at key points in the execution flow. Using middleware, you can add custom logic that runs before or after completions are generated and tools are executed. You can manipulate the inputs of middleware directly, or throw an error to modify the flow of execution.

​
Middleware Types

Magma supports four types of middleware:

  • preCompletion: Runs before an LLM completion is generated
  • onCompletion: Runs after an LLM completion is generated
  • preToolExecution: Runs before a tool is executed
  • onToolExecution: Runs after a tool is executed
  • onMainFinish: Runs after the agent has generated its last completion in the main loop
  • postProcess: Runs after the agent has generated its last completion in the main loop (difference between this and onMainFinish explained in code snippet)
@middleware("preCompletion")
async preCompletionMiddleware(userContent: string) {
    // Modify the text directly by returning
    // This modifies the message in the agent message history as well
    return userContent.replace(/badword/gi, '');

    // or throw an error to modify the flow of execution
    throw new Error("Message contains bad words");
}

Throwing an error from preCompletion middleware will remove the failing user message from the message history, and return the error message as the result from main. No completion will be generated.

@middleware("onCompletion")
async onCompletionMiddleware(agentContent: string) {
    // Modify the message directly by returning
    // This modifies the message in the agent message history as well
    return agentContent.replace(/badword/gi, '');

    // or throw an error to modify the flow of execution
    throw new Error("Message contains bad words");
}

Throwing an error from onCompletion middleware will remove the failing assistant message from the message history, and trigger a regeneration of the completion with your error message as context. This middleware obeys the middleware failure rules

@middleware("preToolExecution")
async preToolExecutionMiddleware(toolCall: MagmaToolCall) {
    // Modify the call directly by returning
    // This modifies the call in the agent message history as well
    return {
        ...toolCall,
        fn_name: message.fn_name.toUpperCase();
    };

    // or throw an error to modify the flow of execution
    throw new Error("Tool call contains bad words");
}

Throwing an error from preToolExecution middleware will remove the failing tool call from the message history, and trigger a regeneration of the completion with your error message as context. This middleware obeys the middleware failure rules

@middleware("onToolExecution")
async onToolExecutionMiddleware(toolResult: MagmaToolResult) {
    // Modify the result directly
    // This modifies the result in the agent message history as well
    return {
        ...toolResult,
        result: result.replace(/badword/gi, ''),
    };

    // or throw an error to modify the flow of execution
    throw new Error("Tool response contains bad words");
}

Throwing an error from onToolExecution middleware will replace the return value of the tool with the error that was thrown.

@middleware('onMainFinish')
async onMainFinish(content: string) {
    // Modify the message directly by returning
    // This modifies the message in the agent message history as well
    return agentContent.replace(/badword/gi, '');

    // or throw an error to modify the flow of execution
    throw new Error("Message contains bad words");
}

Throwing an error from onMainFinish middleware will remove the failing assistant message from the message history, and trigger a regeneration of the completion with your error message as context. This middleware obeys the middleware failure rules

@middleware('postProcess')
async postProcess(content: string) {
    // Modify the message directly by returning
    // This DOES NOT modify the message in the agent message history
    return agentContent.replace(/badword/gi, '');

    // or throw an error to modify the flow of execution
    throw new Error("Message contains bad words");
}

PostProcess works the same as onMainFinish when throwing an error. However, if you return a value from PostProcess, the value is only changed in the return from main, not in the agent message history. This can be particularly useful for translating agent responses into different languages without affecting the agent message history.

​
Middleware Error handling

If a middleware function throws an error, the agent flow will be modified to handle it accordingly.

  • preCompletion: The agent will not generate a completion, and the error message will be returned as if it were the agent’s response.
  • onCompletion: The agent will be notified of the error, and a completion will be re-generated.
  • preToolExecution: The agent will not execute the tool, and the error message will be returned as if it were the tool’s response.
  • onToolExecution: The agent will replace the tool’s response with the error message.
  • onMainFinish: The agent will be notified of the error, and a completion will be re-generated.
  • postProcess: The agent will be notified of the error, and a completion will be re-generated.

​
Middleware Failure Rules

For middleware errors that trigger completion regenerations, they can fail a maximum of 5 times. After 5 retries, the behavior depends on whether the failing middleware is marked as critical

@middleware('onMainFinish', { critical: true })
async onMainFinish(content: string) {
    console.log('On main finish');
    if (!content.includes('FizzBuzz')) {
        console.log(this['middlewareRetries']);
        throw new Error('uh oh');
    }
}

If a critical middleware condition cannot be met, the agent will throw an error.

If a non-critical middleware condition cannot be met, the agent will proceed normally, using the last generation.

​
Implementation

To add middleware to your agent, use the @middleware decorator:

import { MagmaAgent } from "@pompeii-labs/magma";
import { middleware } from "@pompeii-labs/magma/decorators";

class MyAgent extends MagmaAgent {
    @middleware("preCompletion")
    async validateMessage(message: string) {
        if (message.length > 1000) {
            throw new Error(
                "Message is too long. Please keep responses under 1000 characters."
            );
        }
    }
}

​
Chaining Middleware

You can add as many middleware functions as you want to your agent. Middleware functions of the same type are executed in order of definition.

Or, if you would like more explicit control, you can specify the order argument in the middleware decorator. When using order, middleware functions will be executed from least to greatest order value.

class MyAgent extends MagmaAgent {

    @middleware("preCompletion", { order: -100 })
    async filterContent(message: string) {
        return message.replace(/badword/gi, '');
    }
    
    @middleware("onCompletion", { order: 10 })
    async logResponse(message: string) {
        console.log(Agent response: ${message});
        return message;
    }
}

​
Common Use Cases

Middleware is particularly useful for:

  • Input validation and sanitization
  • Response filtering and modification
  • Logging and monitoring
  • Rate limiting
  • Access control
  • Response formatting

​
Best Practices

  • Keep middleware functions focused and single-purpose
  • Handle errors appropriately within middleware
  • Use TypeScript types for better type safety
  • Add clear documentation about what each middleware does
  • Consider the order of middleware execution
  • Return meaningful error messages when validation fails