popagent

Latest change 7f6c1d0ed24ffc264e50e9966eaf45024245ba71 - feat: add per-task agent communication by AkurAI Build

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

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",
  selfUpdateWorkspacePath: "popagent.selfUpdateWorkspacePath",
  runtimeSettings: "popagent.runtimeSettings",
  hookRuntime: "popagent.hookRuntime",
  iterationObserver: "popagent.iterationObserver",
  repositoryBriefReady: "popagent.repositoryBriefReady",
  communicationSession: "popagent.communicationSession",
} as const;

export type AgentContextValues = {
  "popagent.resourceId": string;
  "popagent.sessionId": string;
  "popagent.turnId": string;
  "popagent.model": string;
  "popagent.workspaceId": string;
  "popagent.executionSource": AgentExecutionSource;
  "popagent.taskId"?: string;
  "popagent.scheduleId"?: string;
  "popagent.selfUpdateWorkspacePath"?: string;
  "popagent.runtimeSettings"?: AgentRuntimeSettings;
  "popagent.hookRuntime"?: HookRuntime;
  "popagent.iterationObserver"?: (context: IterationCompleteContext) => Promise<void>;
  "popagent.repositoryBriefReady"?: boolean;
  "popagent.communicationSession"?: TaskAgentCommunicationSession;
};

export type AgentRequestContext = RequestContext<AgentContextValues>;

export function createAgentRequestContext(input: {
  resourceId: string;
  sessionId: string;
  turnId: string;
  model: string;
  workspaceId?: string;
  executionSource?: AgentExecutionSource;
  taskId?: string;
  scheduleId?: string;
  selfUpdateWorkspacePath?: string;
  runtimeSettings?: AgentRuntimeSettings;
  hookRuntime?: HookRuntime;
  iterationObserver?: (context: IterationCompleteContext) => Promise<void>;
  repositoryBriefReady?: boolean;
  communicationSession?: TaskAgentCommunicationSession;
}): 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.selfUpdateWorkspacePath) {
    context.set(AGENT_CONTEXT_KEYS.selfUpdateWorkspacePath, input.selfUpdateWorkspacePath);
  }
  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);
  if (input.communicationSession) {
    context.set(AGENT_CONTEXT_KEYS.communicationSession, input.communicationSession);
  }
  if (input.repositoryBriefReady) {
    context.set(AGENT_CONTEXT_KEYS.repositoryBriefReady, true);
  }
  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;
}