Skip to content

Local Environment Documentation

Overview

The Local Environment is a TypeScript implementation of the Environment interface that manages tool registration, execution, and lifecycle within the Azad system. It provides a bridge between the WebSocket client and tool implementations, handling tool approval, execution, and observer pattern integration.

Architecture

The Local Environment architecture consists of several key components:

  1. Environment Interface: Defines the contract for environment implementations
  2. LocalEnvironment: Concrete implementation of the Environment interface
  3. ToolRegistry: Manages tool registration and lifecycle
  4. Authority: Handles tool approval decisions
  5. Observers: Monitor tool execution and provide additional insights
┌─────────────┐      ┌────────────────┐
│  AzadClient │◄────►│ WebSocketClient│
└──────┬──────┘      └────────────────┘
       │                     ▲
       │                     │
       ▼                     │
┌──────────────┐      ┌─────────────┐
│ Environment  │◄────►│  Authority  │
└──────┬───────┘      └─────────────┘
       │
       ▼
┌───────────────┐     ┌───────────┐
│ Tool Registry │◄───►│   Tools   │
└───────────────┘     └───────────┘
       │
       ▼
┌───────────────┐
│   Observers   │
└───────────────┘

Key Interfaces

Environment Interface

The Environment interface defines the contract that environment implementations must fulfill:

interface Environment extends BaseEnvironment {
  approveTool(
    toolCall: ToolCallPart
  ): Promise<EnvironmentResponse<ApprovalResponse>>;
  executeTool(
    toolCall: ToolCallPart
  ): Promise<EnvironmentResponse<ToolResultPart>>;
}

interface BaseEnvironment {
  setClient(client: WebSocketClient): void;
}

LocalEnvironment

The LocalEnvironment class implements the Environment interface:

class LocalEnvironment implements Environment {
  // Tool Registry
  private _toolRegistry: ReturnType<typeof toolRegistry.createInstanceRegistry>;
  private requestToolCalls = new Map<string, AnyTool>();

  // Observers
  private observers = new Map<string, Observer>();

  // Client and Control
  private client: WebSocketClient | null = null;
  private _currentAbortController: AbortController | null = null;

  // Callbacks
  private toolApprovalCallback?: ToolApprovalCallback;
  private toolExecutionCallback?: (
    toolCall: ToolCallPart,
    toolResult?: ToolResultPart
  ) => void;

  // Authority
  public authority: CallbackAuthority;

  // Methods
  constructor();
  async approveTool(
    toolCall: ToolCallPart
  ): Promise<EnvironmentResponse<ApprovalResponse>>;
  async executeTool(
    toolCall: ToolCallPart
  ): Promise<EnvironmentResponse<ToolResultPart>>;
  registerRequestToolCall(toolCall: ToolCallPart): void;
  updateRequestToolParams(toolCall: ToolCallPart, isFinal?: boolean): void;
  registerTool(tool: AnyTool): void;
  abortStep(): void;
  setClient(client: WebSocketClient): void;
  dispose(): void;

  // Callback Registration
  registerToolApprovalCallback(callback: ToolApprovalCallback): void;
  clearToolApprovalCallback(): void;
  registerExecuteCallback(
    startCallback: (toolCall: ToolCallPart) => void,
    endCallback: (toolCall: ToolCallPart) => void
  ): void;
  clearExecuteCallback(): void;
}

Tool Registry

The Tool Registry manages tool registration and lifecycle:

class ToolRegistry {
  private _tools = new Map<string, AnyTool>();

  // Methods
  registerTool(tool: AnyTool): void;
  cloneTool(toolName: string): AnyTool | null;
  createInstanceRegistry(environment: Environment): {
    tools: Map<string, AnyTool>;
    registerTool: (tool: AnyTool) => void;
    getToolMetadata: () => ToolMetadata[];
    disposeAll: () => void;
    cloneTool: (toolName: string) => AnyTool | null;
  };
}

Authority

The Authority handles tool approval decisions:

class CallbackAuthority {
  constructor(
    private approveToolCallback: (
      toolCall: ToolCallPart
    ) => Promise<EnvironmentResponse<ApprovalResponse>>
  );

  // Methods
  async approveTool(
    toolCall: ToolCallPart
  ): Promise<EnvironmentResponse<ApprovalResponse>>;
}

Usage Examples

Creating and Using a Local Environment

import { LocalEnvironment } from './environment/environment';
import { executeCommandTool } from './environment/tools/definitions/execute-command';
import { createToolInstance } from './environment/tools/utils/tool-registry';

// Create a local environment
const environment = new LocalEnvironment();

// Register tools
const execTool = createToolInstance(executeCommandTool);
environment.registerTool(execTool);

// Register tool approval callback
environment.registerToolApprovalCallback(async (toolCall) => {
  console.log(`Tool approval requested: ${toolCall.tool_name}`);
  console.log('Arguments:', toolCall.args);

  // Automatically approve all tools
  return { approved: true, images: [], feedback: null };
});

// Execute a tool
const result = await environment.executeTool({
  type: 'toolCall',
  tool_call_id: 'tool-call-123',
  tool_name: 'execute_command',
  args: { command: 'ls -la' },
  is_approval_pending: false,
  is_loading: false,
});

console.log('Tool execution result:', result);

// Clean up
environment.dispose();

Integrated with AzadClient

import { AzadClient } from './server/azad-client';
import { executeCommandTool } from './environment/tools/definitions/execute-command';

// Create the client
const client = new AzadClient(process.cwd());

// Execute a task with tools
await client.executeTask({
  model: {
    id: 'gpt-4',
    apiKey: 'your-api-key',
  },
  tools: [executeCommandTool],
  task: 'List all files in the current directory',

  // Tool approval callback
  onToolApproval: async (toolName, toolCallId, args) => {
    return { approved: true, images: [], feedback: null };
  },
});

Tool Execution Flow

The tool execution flow in the LocalEnvironment is as follows:

  1. Tool Registration: Tools are registered with the environment using registerTool()
  2. Tool Call Received: A tool call is received and registered using registerRequestToolCall()
  3. Parameter Updates: Tool parameters are updated through updateRequestToolParams()
  4. Parameter Finalization: Once parameters are complete, updateRequestToolParams(isFinal = true) is called
  5. Tool Approval: The tool is sent for approval through the authority
  6. Tool Execution: If approved, the tool is executed
  7. Observer Execution: All registered observers are executed
  8. Result Return: The tool result is returned
  9. Tool Disposal: The tool instance is disposed
sequenceDiagram
    participant Client as AzadClient
    participant Env as LocalEnvironment
    participant Auth as Authority
    participant Tool as Tool
    participant Observers as Observers

    Client->>Env: registerRequestToolCall()
    Client->>Env: updateRequestToolParams()
    Client->>Env: updateRequestToolParams(isFinal=true)
    Env->>Auth: approveTool()
    Auth-->>Env: ApprovalResponse

    alt Tool Approved
        Env->>Tool: executeTool()
        Tool-->>Env: ToolResultPart
        Env->>Observers: execute()
        Observers-->>Env: ObserverResults
        Env-->>Client: Success Response
    else Tool Rejected
        Env-->>Client: Error Response
    end

    Env->>Tool: dispose()

Tool Approval Mechanism

The tool approval mechanism in the LocalEnvironment is handled by the Authority:

  1. Callback Registration: A tool approval callback is registered using registerToolApprovalCallback()
  2. Approval Request: When a tool needs approval, the callback is invoked with the tool call
  3. Approval Decision: The callback returns an ApprovalResponse with approved flag and optional feedback
  4. Decision Processing: If approved, the tool is executed; if rejected, an error response is returned
// Register approval callback
environment.registerToolApprovalCallback(async (toolCall) => {
  // Custom approval logic
  if (toolCall.tool_name === 'execute_command') {
    const command = toolCall.args.command;

    // Reject potentially harmful commands
    if (command.includes('rm -rf') || command.includes('sudo')) {
      return {
        approved: false,
        images: [],
        feedback: 'Potentially harmful command rejected',
      };
    }
  }

  // Approve other tools
  return { approved: true, images: [], feedback: null };
});

Observer Pattern

The LocalEnvironment implements the observer pattern to monitor tool execution:

  1. Observer Registration: Observers are registered with the environment
  2. Tool Execution: When a tool is executed, all observers are notified
  3. Observer Processing: Each observer processes the tool call and result
  4. Result Collection: Observer results are collected and attached to the tool result
// Example observer
class LoggingObserver implements Observer {
  name = 'logging_observer';
  description = 'Logs all tool calls and results';

  async execute(
    toolCall: ToolCallPart,
    toolResult: ToolResultPart
  ): Promise<ObserverResult> {
    console.log(`Tool executed: ${toolCall.tool_name}`);
    console.log('Arguments:', toolCall.args);
    console.log('Result:', toolResult);

    return {
      observer_name: this.name,
      observer_result: {
        timestamp: new Date().toISOString(),
        logged: true,
      },
    };
  }
}

// Register observer
const observer = new LoggingObserver();
environment.observers.set(observer.name, observer);

Tool Lifecycle

Tools in the LocalEnvironment go through the following lifecycle:

  1. Registration: The tool is registered with the environment
  2. Initialization: The tool is initialized with the environment
  3. Execution: The tool is executed with parameters
  4. Disposal: The tool is disposed after execution
// Tool lifecycle example
const tool = createToolInstance(myToolSchema);

// Registration
environment.registerTool(tool);

// Initialization (automatic during registration)
// tool.init(environment);

// Execution
await environment.executeTool({
  type: 'toolCall',
  tool_call_id: 'tool-call-123',
  tool_name: tool.name,
  args: {
    /* arguments */
  },
  is_approval_pending: false,
  is_loading: false,
});

// Disposal (automatic after execution)
// tool.dispose();

Error Handling and Abort Mechanism

The LocalEnvironment includes error handling and an abort mechanism:

  1. Error Handling: Tool execution errors are caught and returned as error responses
  2. Abort Mechanism: The environment maintains an AbortController that can be used to abort tool execution
  3. Abort Signal: The abort signal is passed to tools during execution
// Abort example
try {
  // Start executing a long-running tool
  const executePromise = environment.executeTool({
    type: 'toolCall',
    tool_call_id: 'tool-call-123',
    tool_name: 'long_running_tool',
    args: {
      /* arguments */
    },
    is_approval_pending: false,
    is_loading: false,
  });

  // Abort the execution after 5 seconds
  setTimeout(() => {
    environment.abortStep();
  }, 5000);

  // Wait for the result or abort
  const result = await executePromise;
  console.log('Tool execution result:', result);
} catch (error) {
  console.error('Tool execution aborted:', error);
}

Implementation Details

Tool Registration and Management

The LocalEnvironment uses the ToolRegistry to manage tools:

  1. Tool Registration: Tools are registered with the registry using registerTool()
  2. Tool Instance Creation: Tool instances are created using createInstanceRegistry().cloneTool()
  3. Tool Metadata: Tool metadata is retrieved using getToolMetadata()
  4. Tool Disposal: Tools are disposed using disposeAll()

Tool Execution

The tool execution process in detail:

  1. Tool Call Validation: The tool call is validated to ensure it contains required fields
  2. Tool Lookup: The tool is looked up in the registry using the tool name
  3. Tool Context Creation: A context is created for the tool execution with abort signal
  4. Tool Execution: The tool is executed with the parameters and context
  5. Observer Execution: All registered observers are executed with the tool call and result
  6. Result Processing: The tool result is processed and returned

Authority Implementation

The CallbackAuthority delegates approval decisions to a callback function:

  1. Callback Invocation: The callback is invoked with the tool call
  2. Response Handling: The callback response is processed and returned
  3. Auto-Approval: If no callback is registered, tools are auto-approved (for headless/testing)

Advanced Topics

Custom Observers

Custom observers can be created by implementing the Observer interface:

class CustomObserver implements Observer {
  name = 'custom_observer';
  description = 'Custom observer for tool execution';

  async execute(
    toolCall: ToolCallPart,
    toolResult: ToolResultPart
  ): Promise<ObserverResult> {
    // Custom observation logic
    const data = {
      toolName: toolCall.tool_name,
      argsCount: Object.keys(toolCall.args).length,
      resultSuccess: toolResult.success,
      timestamp: Date.now(),
    };

    // Return observation result
    return {
      observer_name: this.name,
      observer_result: data,
    };
  }
}

Tool UI Integration

The LocalEnvironment supports tool UI integration through the loadUI method:

  1. Partial Parameters: When partial parameters are received, updateRequestToolParams() is called
  2. UI Loading: The tool's loadUI method is called with the partial parameters
  3. UI Update: The UI can be updated based on the partial parameters

Parallel Tool Execution

The LocalEnvironment supports parallel tool execution:

  1. Parallel Flag: Tools can indicate support for parallel execution through the supportsParllelExecution flag
  2. Execution Context: The execution context includes the parallel flag
  3. Abort Handling: Each parallel execution gets its own abort controller

Best Practices

  1. Tool Registration: Register tools before using them
  2. Tool Approval: Implement proper tool approval callbacks to ensure security
  3. Error Handling: Handle tool execution errors gracefully
  4. Resource Management: Always call dispose() when done with the environment
  5. Observer Usage: Use observers for monitoring and logging, not for critical functionality
  6. Abort Handling: Implement proper abort handling in long-running tools
  7. Context Management: Use tool context for persistent state across invocations