AkurAI Build
Menu

popagent

public

Latest change da13a7bebe63bf4b2693180d2d4850aabeaa0807 - Add autonomous evolution and self-healing by AkurAI Build

import type { DelegationConfig } from "@mastra/core/agent";
import type { Processor, ProcessAPIErrorArgs, ProcessOutputResultArgs } from "@mastra/core/processors";
import type { ToolHooks } from "@mastra/core/tools";
import type { AgentRuntimeSettings } from "./api-types";
import { createHookEvent, HookBlockedError, type HookRuntime } from "./hooks";
import type { AgentTaskLifecycle } from "./tasks";
import { isPlaintextSecretToolName } from "./chat-messages";

type DelegationRuntimeSettings = Pick<
  AgentRuntimeSettings,
  | "delegationContextMessages"
  | "delegationResultCharacters"
  | "delegationFailureFeedback"
  | "delegationResultTruncationMarker"
>;

function boundDelegationResult(
  result: string,
  settings: DelegationRuntimeSettings,
): string {
  if (result.length <= settings.delegationResultCharacters) return result;
  return result.slice(
    0,
    settings.delegationResultCharacters - settings.delegationResultTruncationMarker.length,
  ) + settings.delegationResultTruncationMarker;
}

function delegationFailureFeedback(
  agentId: string,
  error: unknown,
  settings: DelegationRuntimeSettings,
): string {
  const message = error == null ? "delegation failed without a reported error" : errorDetail(error).message.slice(0, 2_000);
  return settings.delegationFailureFeedback
    .replaceAll("{{agentId}}", agentId)
    .replaceAll("{{error}}", message);
}

type HookExecutionContext = {
  sessionId: string;
  turnId: string;
  model: string;
};

export class HookLifecycleProcessor implements Processor<"popagent-hooks"> {
  readonly id = "popagent-hooks" as const;

  constructor(
    private readonly runtime: HookRuntime,
    private readonly context: HookExecutionContext,
  ) {}

  async processOutputResult(args: ProcessOutputResultArgs) {
    const result = await this.runtime.dispatch(createHookEvent("Stop", {
      ...this.context,
      detail: {
        response: args.result.text,
        finishReason: args.result.finishReason,
        continuationCount: args.retryCount,
      },
    }));
    if (result.continue) args.abort(result.continue.additionalContext, { retry: true });
    return args.messageList;
  }

  async processAPIError(args: ProcessAPIErrorArgs) {
    await this.runtime.dispatch(createHookEvent("StopFailure", {
      ...this.context,
      detail: { error: errorDetail(args.error) },
    }));
    return { retry: false };
  }
}

export function createToolHooks(
  runtime: HookRuntime,
  context: HookExecutionContext,
  fixedPolicy?: (toolName: string, input: unknown) => void,
): ToolHooks {
  return {
    beforeToolCall: async ({ toolName, input, context: toolContext }) => {
      fixedPolicy?.(toolName, input);
      const toolCallId = toolCallIdFrom(toolContext);
      const sensitive = isPlaintextSecretToolName(toolName);
      const result = await runtime.dispatch(createHookEvent("PreToolUse", {
        ...context,
        toolCallId,
        detail: { tool: { name: toolName, input: sensitive ? "[redacted]" : input } },
      }));
      if (
        !sensitive
        && result.replacement !== undefined
        && isMutableObject(input)
        && isMutableObject(result.replacement)
      ) {
        for (const key of Object.keys(input)) delete input[key];
        Object.assign(input, result.replacement);
      }
      fixedPolicy?.(toolName, input);
    },
    afterToolCall: async ({ toolName, input, output, error, context: toolContext }) => {
      const toolCallId = toolCallIdFrom(toolContext);
      const sensitive = isPlaintextSecretToolName(toolName);
      await runtime.dispatch(createHookEvent(error ? "PostToolUseFailure" : "PostToolUse", {
        ...context,
        toolCallId,
        detail: {
          tool: {
            name: toolName,
            input: sensitive ? "[redacted]" : input,
            ...(error
              ? { error: sensitive ? "Sensitive tool failed" : errorDetail(error) }
              : { output: sensitive ? "[redacted]" : output }),
          },
        },
      }));
    },
  };
}

export function createDelegationHooks(
  runtime: HookRuntime,
  context: HookExecutionContext,
  settings: DelegationRuntimeSettings,
  observe?: (agentId: string, success: boolean) => void,
): DelegationConfig {
  return {
    includeSubAgentToolResultsInModelContext: false,
    messageFilter: ({ messages }) => messages.slice(-settings.delegationContextMessages),
    onDelegationStart: async (delegation) => {
      try {
        const result = await runtime.dispatch(createHookEvent("SubagentStart", {
          ...context,
          toolCallId: delegation.toolCallId,
          detail: {
            agent: {
              id: delegation.primitiveId,
              type: delegation.primitiveType,
              parentId: delegation.parentAgentId,
            },
            prompt: delegation.prompt,
            iteration: delegation.iteration,
            requestedMaxSteps: delegation.params.maxSteps ?? null,
          },
        }));
        return result.replacement === undefined
          ? { proceed: true }
          : { proceed: true, modifiedPrompt: result.replacement as string };
      } catch (error) {
        return {
          proceed: false,
          rejectionReason: error instanceof HookBlockedError
            ? error.reason
            : "Managed delegation hook failed",
        };
      }
    },
    onDelegationComplete: async (delegation) => {
      const response = boundDelegationResult(delegation.result.text, settings);
      observe?.(delegation.primitiveId, delegation.success);
      await runtime.dispatch(createHookEvent("SubagentStop", {
        ...context,
        toolCallId: delegation.toolCallId,
        detail: {
          agent: {
            id: delegation.primitiveId,
            type: delegation.primitiveType,
            parentId: delegation.parentAgentId,
          },
          success: delegation.success,
          durationMs: delegation.duration,
          response,
          ...(delegation.error ? { error: errorDetail(delegation.error) } : {}),
        },
      }));
      if (!delegation.success) {
        return { feedback: delegationFailureFeedback(delegation.primitiveId, delegation.error, settings) };
      }
      if (response !== delegation.result.text) return { resultText: response };
    },
  };
}

export function createTaskLifecycleHooks(runtime: HookRuntime): AgentTaskLifecycle {
  return async ({ task, turnId, status, error }) => {
    await runtime.dispatch(createHookEvent(status === "running" ? "TaskStart" : "TaskStop", {
      sessionId: task.sessionId ?? `task:${task.id}`,
      turnId,
      model: task.model,
      detail: {
        task: {
          id: task.id,
          scheduleId: task.scheduleId,
          source: task.scheduleId ? "schedule" : "manual",
        },
        status,
        ...(error ? { error } : {}),
      },
    }));
  };
}

function toolCallIdFrom(context: unknown): string | undefined {
  return (context as { agent?: { toolCallId?: string } } | undefined)?.agent?.toolCallId;
}

function isMutableObject(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

function errorDetail(error: unknown) {
  return { name: error instanceof Error ? error.name : "Error", message: error instanceof Error ? error.message : String(error) };
}