popagent

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

import { RequestContext } from "@mastra/core/request-context";
import type { IterationCompleteContext } from "@mastra/core/agent";
import { DEFAULT_WORKSPACE_ID, type AgentRuntimeSettings } from "./api-types";
import type { HookRuntime } from "./hooks";

export const AGENT_CONTEXT_KEYS = {
  resourceId: "popagent.resourceId",
  sessionId: "popagent.sessionId",
  turnId: "popagent.turnId",
  model: "popagent.model",
  workspaceId: "popagent.workspaceId",
  executionSource: "popagent.executionSource",
  taskId: "popagent.taskId",
  scheduleId: "popagent.scheduleId",
  runtimeSettings: "popagent.runtimeSettings",
  hookRuntime: "popagent.hookRuntime",
  iterationObserver: "popagent.iterationObserver",
} as const;

export type AgentContextValues = {
  "popagent.resourceId": string;
  "popagent.sessionId": string;
  "popagent.turnId": string;
  "popagent.model": string;
  "popagent.workspaceId": string;
  "popagent.executionSource": "chat" | "task" | "schedule";
  "popagent.taskId"?: string;
  "popagent.scheduleId"?: string;
  "popagent.runtimeSettings"?: AgentRuntimeSettings;
  "popagent.hookRuntime"?: HookRuntime;
  "popagent.iterationObserver"?: (context: IterationCompleteContext) => Promise<void>;
};

export type AgentRequestContext = RequestContext<AgentContextValues>;

export function createAgentRequestContext(input: {
  resourceId: string;
  sessionId: string;
  turnId: string;
  model: string;
  workspaceId?: string;
  executionSource?: "chat" | "task" | "schedule";
  taskId?: string;
  scheduleId?: string;
  runtimeSettings?: AgentRuntimeSettings;
  hookRuntime?: HookRuntime;
  iterationObserver?: (context: IterationCompleteContext) => Promise<void>;
}): AgentRequestContext {
  const context = new RequestContext<AgentContextValues>();
  context.set(AGENT_CONTEXT_KEYS.resourceId, input.resourceId);
  context.set(AGENT_CONTEXT_KEYS.sessionId, input.sessionId);
  context.set(AGENT_CONTEXT_KEYS.turnId, input.turnId);
  context.set(AGENT_CONTEXT_KEYS.model, input.model);
  context.set(AGENT_CONTEXT_KEYS.workspaceId, input.workspaceId ?? DEFAULT_WORKSPACE_ID);
  context.set(AGENT_CONTEXT_KEYS.executionSource, input.executionSource ?? "chat");
  if (input.taskId) context.set(AGENT_CONTEXT_KEYS.taskId, input.taskId);
  if (input.scheduleId) context.set(AGENT_CONTEXT_KEYS.scheduleId, input.scheduleId);
  if (input.runtimeSettings) context.set(AGENT_CONTEXT_KEYS.runtimeSettings, input.runtimeSettings);
  if (input.hookRuntime) context.set(AGENT_CONTEXT_KEYS.hookRuntime, input.hookRuntime);
  if (input.iterationObserver) context.set(AGENT_CONTEXT_KEYS.iterationObserver, input.iterationObserver);
  return context;
}

export function agentContextValue<Key extends keyof AgentContextValues>(
  requestContext: AgentRequestContext | RequestContext | undefined,
  key: Key,
): AgentContextValues[Key] | undefined {
  return (requestContext as RequestContext | undefined)?.get(key) as AgentContextValues[Key] | undefined;
}