popagent

Latest change 48f38756b02cd3f67b173ba9e5d7a5e122c03126 - Route read-only agents through fast model by Ólafur Búi Ólafsson

import { Agent } from "@mastra/core/agent";
import type { RequestContext } from "@mastra/core/request-context";
import { Mastra } from "@mastra/core/mastra";
import {
  createFinalResponseGuard,
  specialistExecutionOptions,
  toolCallConcurrency,
} from "./agent-autonomy";
import { AGENT_CONTEXT_KEYS, agentContextValue } from "./agent-context";
import { agentRuntimeSettings } from "./agent-runtime-settings";
import { agentSettings } from "./agent-settings";
import { agentSkills } from "./agent-skills";
import {
  DEFAULT_WORKSPACE_ID,
  type AgentRuntimeSettings,
  type AgentSettings,
} from "./api-types";
import { evolutionStore } from "./evolution-store";
import { resolveModel } from "./models";
import { logger, observability, setAppLogger } from "./observability";
import { memory, storage } from "./storage";
import { getTime } from "./tools/get-time";
import { editMemory, recallMemory, reflectMemory, retainMemory } from "./tools/memory";
import { useBrowserSecret } from "./tools/secrets";
import { webSearch } from "./tools/web-search";
import { searchDocumentation } from "./tools/search-documentation";
import { agentWorkspaceRuntime } from "./workspace";

async function requiredAgentProfile(id: string): Promise<AgentSettings> {
  const settings = await agentSettings.get(id);
  if (!settings) throw new Error(`Agent settings not found: ${id}`);
  return settings;
}

async function renderedInstructions(settings: AgentSettings): Promise<string> {
  const instructions = settings.instructions.trim();
  if (!instructions) throw new Error(`Agent instructions are empty: ${settings.id}`);
  const overlay = await evolutionStore.learnedOverlay(settings.id);
  const base = `Display identity: ${settings.name.trim()}.\n\n${instructions}`;
  return overlay ? `${base}\n\n## Learned behavior overlay\n${overlay}` : base;
}

async function configuredRuntime(
  requestContext: RequestContext | undefined,
): Promise<AgentRuntimeSettings> {
  return agentContextValue(requestContext, AGENT_CONTEXT_KEYS.runtimeSettings)
    ?? agentRuntimeSettings.get();
}

const modelForContext = async (
  { requestContext }: { requestContext: RequestContext },
) => {
  const selected = agentContextValue(requestContext, AGENT_CONTEXT_KEYS.model);
  const settings = selected ? undefined : await configuredRuntime(requestContext);
  return resolveModel(selected ?? settings!.defaultModel);
};
const modelForRole = (id: string) => async (
  { requestContext }: { requestContext: RequestContext },
) => {
  const profile = await requiredAgentProfile(id);
  return profile.model ? resolveModel(profile.model) : modelForContext({ requestContext });
};

const sharedSubagentOptions = {
  defaultOptions: specialistExecutionOptions,
};
const toolRegistry = {
  getTime,
  retainMemory,
  recallMemory,
  reflectMemory,
  editMemory,
  useBrowserSecret,
  webSearch,
  searchDocumentation,
};
async function configuredTools(id: string, requestContext?: RequestContext) {
  const profile = await requiredAgentProfile(id);
  const selfUpdate = agentContextValue(
    requestContext,
    AGENT_CONTEXT_KEYS.executionSource,
  ) === "self-update";
  return Object.fromEntries(
    profile.tools
      .filter((tool) => !selfUpdate || !["retainMemory", "editMemory", "useBrowserSecret"].includes(tool))
      .map((tool) => [tool, toolRegistry[tool]]),
  ) as Partial<typeof toolRegistry>;
}
const workspaceFor = (id: string) =>
  async ({ requestContext }: { requestContext: RequestContext }) => {
    const profile = await requiredAgentProfile(id);
    if (profile.workspaceAccess === "none") return undefined;
    const selectedWorkspace = agentContextValue(requestContext, AGENT_CONTEXT_KEYS.workspaceId);
    const selfUpdate = agentContextValue(
      requestContext,
      AGENT_CONTEXT_KEYS.executionSource,
    ) === "self-update";
    const readOnly = profile.workspaceAccess === "read-only"
      || (selfUpdate && id !== "popagent" && id !== "implementer");
    if (!selfUpdate) return agentWorkspaceRuntime.resolve(selectedWorkspace, readOnly);
    const path = agentContextValue(requestContext, AGENT_CONTEXT_KEYS.selfUpdateWorkspacePath);
    if (!path) throw new Error("Self-update workspace context is unavailable");
    return agentWorkspaceRuntime.resolveSelfUpdate(
      selectedWorkspace ?? DEFAULT_WORKSPACE_ID,
      path,
      readOnly,
    );
  };

let [supervisorProfile, researcherProfile, implementerProfile, reviewerProfile] = await Promise.all([
  requiredAgentProfile("popagent"),
  requiredAgentProfile("researcher"),
  requiredAgentProfile("implementer"),
  requiredAgentProfile("reviewer"),
]);

let researcherAgent: Agent;
researcherAgent = new Agent({
  id: "researcher",
  name: researcherProfile.name,
  description: researcherProfile.description,
  instructions: async () => {
    researcherProfile = await requiredAgentProfile("researcher");
    researcherAgent.name = researcherProfile.name;
    return renderedInstructions(researcherProfile);
  },
  tools: ({ requestContext }) => configuredTools("researcher", requestContext),
  skills: () => agentSkills.resolve("researcher"),
  workspace: workspaceFor("researcher"),
  ...sharedSubagentOptions,
  model: modelForRole("researcher"),
});
researcherAgent.getDescription = () => researcherProfile.description;

let implementerAgent: Agent;
implementerAgent = new Agent({
  id: "implementer",
  name: implementerProfile.name,
  description: implementerProfile.description,
  instructions: async () => {
    implementerProfile = await requiredAgentProfile("implementer");
    implementerAgent.name = implementerProfile.name;
    return renderedInstructions(implementerProfile);
  },
  tools: ({ requestContext }) => configuredTools("implementer", requestContext),
  skills: () => agentSkills.resolve("implementer"),
  workspace: workspaceFor("implementer"),
  ...sharedSubagentOptions,
  model: modelForRole("implementer"),
});
implementerAgent.getDescription = () => implementerProfile.description;

let reviewerAgent: Agent;
reviewerAgent = new Agent({
  id: "reviewer",
  name: reviewerProfile.name,
  description: reviewerProfile.description,
  instructions: async () => {
    reviewerProfile = await requiredAgentProfile("reviewer");
    reviewerAgent.name = reviewerProfile.name;
    return renderedInstructions(reviewerProfile);
  },
  tools: ({ requestContext }) => configuredTools("reviewer", requestContext),
  skills: () => agentSkills.resolve("reviewer"),
  workspace: workspaceFor("reviewer"),
  ...sharedSubagentOptions,
  model: modelForRole("reviewer"),
});
reviewerAgent.getDescription = () => reviewerProfile.description;

const specialistAgents = {
  researcher: researcherAgent,
  implementer: implementerAgent,
  reviewer: reviewerAgent,
};

async function refreshSpecialistProfiles(
  { requestContext }: { requestContext: RequestContext },
) {
  [researcherProfile, implementerProfile, reviewerProfile] = await Promise.all([
    requiredAgentProfile("researcher"),
    requiredAgentProfile("implementer"),
    requiredAgentProfile("reviewer"),
  ]);
  researcherAgent.name = researcherProfile.name;
  implementerAgent.name = implementerProfile.name;
  reviewerAgent.name = reviewerProfile.name;
  const enabled = Object.entries(specialistAgents).filter(([id]) => {
    if (id === "researcher") return researcherProfile.delegationEnabled;
    if (id === "implementer") return implementerProfile.delegationEnabled;
    return reviewerProfile.delegationEnabled;
  });
  if (agentContextValue(requestContext, AGENT_CONTEXT_KEYS.executionSource) !== "self-update") {
    return Object.fromEntries(enabled) as typeof specialistAgents;
  }
  return Object.fromEntries(enabled.map(([id, specialist]) => {
    const fork = specialist.__fork();
    fork.setBrowser(undefined);
    return [id, fork];
  })) as typeof specialistAgents;
}

let supervisorAgent: Agent;
supervisorAgent = new Agent({
  id: "popagent",
  name: supervisorProfile.name,
  description: supervisorProfile.description,
  instructions: async () => {
    supervisorProfile = await requiredAgentProfile("popagent");
    supervisorAgent.name = supervisorProfile.name;
    return renderedInstructions(supervisorProfile);
  },
  model: modelForRole("popagent"),
  defaultOptions: async ({ requestContext }) => {
    const settings = await configuredRuntime(requestContext);
    return {
      maxSteps: settings.supervisorMaxSteps,
      toolCallConcurrency: toolCallConcurrency(settings),
      onIterationComplete: createFinalResponseGuard(settings),
    };
  },
  tools: ({ requestContext }) => configuredTools("popagent", requestContext),
  memory,
  workspace: workspaceFor("popagent"),
  skills: () => agentSkills.resolve("popagent"),
  agents: refreshSpecialistProfiles,
});
supervisorAgent.getDescription = () => supervisorProfile.description;

export const agent = supervisorAgent;
export const mastra = new Mastra({
  agents: { popagent: agent },
  storage,
  logger,
  observability,
  environment: process.env.POPAGENT_ENVIRONMENT ?? process.env.NODE_ENV,
});
setAppLogger(mastra.getLogger());