popagent

Latest change 7cecf6a89b6f39dae8c5678f5b369014199aeb3c - Add self-hosted Mastra observability by AkurAI Build

import { describe, expect, test } from "bun:test";
import type { Agent, IterationCompleteContext } from "@mastra/core/agent";
import {
  type AgentRuntimeSettings,
  type AgentTask,
  type AgentTaskProgress,
  type MemorySettings,
} from "./api-types";
import { agentActivity, AgentExecutionRuntime, createAgentTaskExecutor } from "./agent-execution";
import { HookRuntime, type HookEvent } from "./hooks";

const memoryConfiguration: MemorySettings = {
  autoCompact: false,
  observationTokens: 30_000,
  reflectionTokens: 40_000,
  recentMessagePercent: 20,
  bufferOnIdle: true,
  updatedAt: "",
};
const runtimeConfiguration: AgentRuntimeSettings = {
  defaultModel: "test/default",
  supervisorMaxSteps: 9,
  specialistMaxSteps: 7,
  toolConcurrency: 2,
  delegationContextMessages: 5,
  delegationResultCharacters: 4_000,
  maxProcessorRetries: 4,
  finalResponseFeedback: "Return the configured final response",
  delegationFailureFeedback: "Delegate {{agentId}} failed: {{error}}",
  delegationResultTruncationMarker: "\n[truncated]",
  taskConcurrency: 2,
  taskPollIntervalMs: 1_000,
  taskTimeoutMs: 30_000,
  taskStaleAfterMs: 60_000,
  updatedAt: "",
};

describe("AgentExecutionRuntime", () => {
  test("gives background tasks the chat execution context", async () => {
    const events: HookEvent[] = [];
    const episodes: Array<{ sessionId: string; userText: string; assistantText: string }> = [];
    const progressUpdates: AgentTaskProgress[] = [];
    let appliedAgent: Agent | undefined;
    let streamedPrompt: string | undefined;
    let streamedOptions: Record<string, unknown> | undefined;
    const hooks = new HookRuntime({
      config: {
        version: 1,
        hooks: {
          UserPromptSubmit: [{ id: "normalize", type: "http", url: "https://hooks.example/normalize" }],
        },
      },
      transport: async (_handler, event) => {
        events.push(event);
        return {
          schemaVersion: 1,
          outcome: "modify",
          reason: "normalize",
          replacement: "normalized task prompt",
          additionalContext: "managed hook context",
        };
      },
    });
    const agent = {
      stream: async (prompt: string, options: Record<string, unknown>) => {
        streamedPrompt = prompt;
        streamedOptions = options;
        await (options.onIterationComplete as (
          context: IterationCompleteContext,
        ) => Promise<unknown>)({
          iteration: 2,
          maxIterations: runtimeConfiguration.supervisorMaxSteps,
          text: "",
          toolCalls: [{ id: "tool-1", name: "webSearch", args: {} }],
          toolResults: [],
          isFinal: false,
          finishReason: "tool-calls",
          runId: "run-1",
          agentId: "popagent",
          agentName: "popagent",
          messages: [],
        });
        const observeSpecialist = (
          options.requestContext as { get: (key: string) => unknown }
        ).get("popagent.iterationObserver") as (context: IterationCompleteContext) => Promise<void>;
        await observeSpecialist({
          iteration: 4,
          maxIterations: 12,
          text: "",
          toolCalls: [{ id: "tool-2", name: "webSearch", args: {} }],
          toolResults: [],
          isFinal: false,
          finishReason: "tool-calls",
          runId: "run-2",
          agentId: "researcher",
          agentName: "Researcher",
          messages: [],
        });
        await (options.onFinish as (result: { text: string }) => Promise<void>)({ text: "task result" });
        return { text: Promise.resolve("task result") };
      },
    } as unknown as Agent;
    const execution = new AgentExecutionRuntime(
      agent,
      hooks,
      {
        formatRecall: async ({ resourceId, query }) => {
          expect(resourceId).toBe("popagent-user");
          expect(query).toBe("normalized task prompt");
          return "recalled durable fact";
        },
        retainEpisode: async ({ sessionId, userText, assistantText }) => {
          episodes.push({ sessionId, userText, assistantText });
        },
      },
      {
        apply: async (candidate) => { appliedAgent = candidate; },
      },
      {
        get: async () => memoryConfiguration,
      },
      {
        get: async () => runtimeConfiguration,
      },
    );
    const task: AgentTask = {
      id: "task-1",
      sessionId: "session-1",
      scheduleId: null,
      workspaceId: "default",
      prompt: "original task prompt",
      model: runtimeConfiguration.defaultModel,
      status: "queued",
      output: null,
      error: null,
      stepsCompleted: 0,
      progress: null,
      recoveryCount: 0,
      createdAt: "",
      startedAt: null,
      completedAt: null,
    };
    const controller = new AbortController();

    const output = await createAgentTaskExecutor(agent, execution)(
      task,
      controller.signal,
      "turn-1",
      async (progress) => { progressUpdates.push(progress); },
    );

    expect(output).toBe("task result");
    expect(progressUpdates).toEqual([
      { stepsCompleted: 2, progress: "Using webSearch" },
      { stepsCompleted: 4, progress: "Researcher: Using webSearch" },
    ]);
    expect(streamedPrompt).toBe("normalized task prompt");
    expect(appliedAgent).toBe(agent);
    expect(events).toEqual([expect.objectContaining({
      eventName: "UserPromptSubmit",
      sessionId: "session-1",
      detail: { prompt: "original task prompt" },
    })]);
    expect(streamedOptions).toEqual(expect.objectContaining({
      abortSignal: controller.signal,
      memory: { thread: "session-1", resource: "popagent-user" },
      hooks: expect.objectContaining({
        beforeToolCall: expect.any(Function),
        afterToolCall: expect.any(Function),
      }),
      delegation: expect.objectContaining({
        onDelegationStart: expect.any(Function),
        onDelegationComplete: expect.any(Function),
      }),
      outputProcessors: [expect.anything()],
      errorProcessors: [expect.anything()],
      memoryConfig: { observationalMemory: { enabled: false } },
      maxSteps: runtimeConfiguration.supervisorMaxSteps,
      toolCallConcurrency: { limit: runtimeConfiguration.toolConcurrency, strategy: "called" },
      tracingOptions: expect.objectContaining({
        traceId: expect.stringMatching(/^[a-f0-9]{32}$/),
        tags: ["task"],
        metadata: expect.objectContaining({
          sessionId: "session-1",
          turnId: "turn-1",
          workspaceId: "default",
          executionSource: "task",
          taskId: "task-1",
        }),
      }),
    }));
    expect((streamedOptions?.requestContext as { get: (key: string) => unknown }).get("popagent.sessionId"))
      .toBe("session-1");
    expect((streamedOptions?.requestContext as { get: (key: string) => unknown }).get("popagent.model"))
      .toBe(runtimeConfiguration.defaultModel);
    expect((streamedOptions?.requestContext as { get: (key: string) => unknown }).get("popagent.executionSource"))
      .toBe("task");
    expect((streamedOptions?.requestContext as { get: (key: string) => unknown }).get("popagent.taskId"))
      .toBe("task-1");
    expect(streamedOptions?.context).toEqual([
      { role: "user", content: "<memory-context>\nmanaged hook context\n</memory-context>" },
      { role: "user", content: "<memory-context>\nrecalled durable fact\n</memory-context>" },
    ]);
    expect(episodes).toEqual([{
      sessionId: "session-1",
      userText: "normalized task prompt",
      assistantText: "task result",
    }]);
  });
  test("serializes specialist iterations for live chat activity", () => {
    const activity = agentActivity("turn-7", {
      iteration: 3,
      maxIterations: 12,
      text: "Verified the repository contract.",
      toolCalls: [
        { id: "tool-1", name: "read", args: { path: "src" } },
        { id: "tool-2", name: "read", args: { path: "test" } },
      ],
      toolResults: [],
      isFinal: true,
      finishReason: "stop",
      runId: "run-7",
      agentId: "reviewer",
      agentName: "Reviewer",
      messages: [],
    });

    expect(activity).toEqual({
      turnId: "turn-7",
      runId: "run-7",
      agentId: "reviewer",
      agentName: "Reviewer",
      iteration: 3,
      maxIterations: 12,
      isFinal: true,
      finishReason: "stop",
      text: "Verified the repository contract.",
      tools: ["read"],
    });
  });

});