refactor: streamline agent runtime options and enhance turn context handling

Author: NoOne7135

agent/runtime/AgentModels.ts

@@ -10,20 +10,21 @@ import { contextSchema, toLangchainAgentContext } from "./AgentContext.js";
 
 export type AgentRuntimeOptions = {
   name: string;
-  adminforth: IAdminForth;
-  checkpointer: BaseCheckpointSaver;
+  getAdminforth: () => IAdminForth;
+  getCheckpointer: () => BaseCheckpointSaver;
   toolProvider: AgentToolProvider;
 };
 
 export class AgentRuntime {
   constructor(private readonly options: AgentRuntimeOptions) {}
 
   async stream(input: AgentRuntimeRunInput) {
-    const tools = await this.options.toolProvider.getTools();
     const apiBasedTools = this.options.toolProvider.getApiBasedTools();
+    const tools = await this.options.toolProvider.getTools(apiBasedTools);
+    const adminforth = this.options.getAdminforth();
     const apiBasedToolsMiddleware = createApiBasedToolsMiddleware(
       apiBasedTools,
-      this.options.adminforth,
+      adminforth,
     );
     const sequenceDebugMiddleware = createSequenceDebugMiddleware(
       input.observability.sequenceDebugSink,
@@ -42,7 +43,7 @@ export class AgentRuntime {
     const agent = createAgent({
       name: this.options.name,
       model: input.models.model,
-      checkpointer: this.options.checkpointer,
+      checkpointer: this.options.getCheckpointer(),
       tools,
       contextSchema,
       middleware,
@@ -58,7 +59,7 @@ export class AgentRuntime {
       },
       context: toLangchainAgentContext({
         ...input.context,
-        adminBaseUrl: this.options.adminforth.config.baseUrlSlashed,
+        adminBaseUrl: adminforth.config.baseUrlSlashed,
         emit: input.observability.emit,
         sequenceDebugSink: input.observability.sequenceDebugSink,
       }),

agent/runtime/AgentRuntime.ts

@@ -16,12 +16,12 @@ export class AgentToolProvider {
     );
   }
 
-  async getTools() {
+  async getTools(apiBasedTools: Record<string, ApiBasedTool>) {
     const adminforth = this.getAdminforth();
 
     return createAgentTools(
       adminforth.config.customization.customComponentsDir ?? "custom",
-      this.getApiBasedTools(),
+      apiBasedTools,
       adminforth.activatedPlugins.map((plugin) => plugin.customFolderPath),
     );
   }

agent/runtime/AgentRuntimeFactory.ts

@@ -0,0 +1,36 @@
+import type { AdminUser, IAdminForth } from "adminforth";
+import type { AgentTurnContext, BaseAgentTurnInput } from "./turnTypes.js";
+
+export type UserContextProvider = {
+  getUserTimeZone(adminUser: AdminUser): Promise<string | null | undefined> | string | null | undefined;
+};
+
+export class TurnContextBuilder {
+  constructor(
+    private readonly getAdminforth: () => IAdminForth,
+    private readonly userContextProvider?: UserContextProvider,
+  ) {}
+
+  async build(input: {
+    base: BaseAgentTurnInput;
+    turnId: string;
+  }): Promise<AgentTurnContext> {
+    const adminforth = this.getAdminforth();
+
+    return {
+      adminUser: input.base.adminUser,
+      sessionId: input.base.sessionId,
+      turnId: input.turnId,
+      abortSignal: input.base.abortSignal,
+      currentPage: input.base.currentPage,
+      chatSurface: input.base.chatSurface,
+      userTimeZone:
+        input.base.userTimeZone
+        ?? await this.userContextProvider?.getUserTimeZone(input.base.adminUser)
+        ?? "UTC",
+      adminPublicOrigin:
+        input.base.adminPublicOrigin
+        ?? adminforth.config.baseUrlSlashed,
+    };
+  }
+}

agent/tools/AgentToolProvider.ts

@@ -10,7 +10,7 @@ export type BaseAgentTurnInput = {
   prompt: string;
   sessionId: string;
   modeName?: string | null;
-  userTimeZone: string;
+  userTimeZone?: string;
   currentPage?: CurrentPageContext;
   chatSurface?: string;
   adminPublicOrigin?: string;

agent/turn/TurnContextBuilder.ts

@@ -1,28 +1,21 @@
 import { logger } from "adminforth";
 import { randomUUID } from "crypto";
-import type { IAdminForth } from "adminforth";
-import type { BaseCheckpointSaver } from "@langchain/langgraph";
 import { AgentModelFactory } from "./agent/models/AgentModelFactory.js";
 import { AgentModeResolver } from "./agent/models/AgentModeResolver.js";
 import { createSequenceDebugCollector } from "./agent/middleware/sequenceDebug.js";
-import { AgentRuntimeFactory } from "./agent/runtime/AgentRuntimeFactory.js";
-import { AgentToolProvider } from "./agent/tools/AgentToolProvider.js";
+import { AgentRuntime } from "./agent/runtime/AgentRuntime.js";
+import { TurnContextBuilder } from "./agent/turn/TurnContextBuilder.js";
 import { TurnLifecycleService } from "./agent/turn/TurnLifecycleService.js";
-import { TurnPersistenceService } from "./agent/turn/TurnPersistenceService.js";
 import { TurnPromptBuilder } from "./agent/turn/TurnPromptBuilder.js";
 import { TurnStreamConsumer } from "./agent/turn/TurnStreamConsumer.js";
 import type {
   BaseAgentTurnInput,
-  HandleSpeechTurnInput,
   HandleTurnInput,
   PreparedAgentTurn,
   RunAndPersistAgentResponseInput,
   RunAndPersistAgentResponseResult,
 } from "./agent/turn/turnTypes.js";
 import { getErrorMessage, isAbortError } from "./errors.js";
-import type { AgentSessionStore } from "./sessionStore.js";
-import { SpeechTurnService } from "./agent/speech/SpeechTurnService.js";
-import type { PluginOptions } from "./types.js";
 
 export type {
   BaseAgentTurnInput,
@@ -32,73 +25,32 @@ export type {
   RunAndPersistAgentResponseResult,
 } from "./agent/turn/turnTypes.js";
 
-type AgentTurnServiceOptions = {
-  getAdminforth: () => IAdminForth;
-  getPluginInstanceId: () => string;
-  options: PluginOptions;
-  sessionStore: AgentSessionStore;
-  getCheckpointer: () => BaseCheckpointSaver;
-  getInternalAgentResourceIds: () => string[];
-  getAgentSystemPrompt: () => Promise<string>;
-};
-
 export class AgentTurnService {
-  private readonly modeResolver: AgentModeResolver;
-  private readonly modelFactory: AgentModelFactory;
-  private readonly runtimeFactory: AgentRuntimeFactory;
-  private readonly lifecycle: TurnLifecycleService;
-  private readonly promptBuilder: TurnPromptBuilder;
-  private readonly streamConsumer = new TurnStreamConsumer();
-  private readonly speechTurnService: SpeechTurnService;
-
-  constructor(private readonly serviceOptions: AgentTurnServiceOptions) {
-    const toolProvider = new AgentToolProvider(
-      serviceOptions.getAdminforth,
-      serviceOptions.getInternalAgentResourceIds,
-    );
-    const persistence = new TurnPersistenceService(
-      serviceOptions.getAdminforth,
-      serviceOptions.options,
-    );
-
-    this.modeResolver = new AgentModeResolver(serviceOptions.options);
-    this.modelFactory = new AgentModelFactory(serviceOptions.options.maxTokens ?? 1000);
-    this.runtimeFactory = new AgentRuntimeFactory(
-      serviceOptions.getAdminforth,
-      serviceOptions.getCheckpointer,
-      toolProvider,
-      () => `adminforth-agent-${serviceOptions.getPluginInstanceId()}`,
-    );
-    this.lifecycle = new TurnLifecycleService(serviceOptions.sessionStore, persistence);
-    this.promptBuilder = new TurnPromptBuilder({
-      getAdminforth: serviceOptions.getAdminforth,
-      getAgentSystemPrompt: serviceOptions.getAgentSystemPrompt,
-    });
-    this.speechTurnService = new SpeechTurnService(
-      this.runAndPersistAgentResponse.bind(this),
-    );
-  }
+  constructor(
+    private readonly lifecycle: TurnLifecycleService,
+    private readonly contextBuilder: TurnContextBuilder,
+    private readonly modeResolver: AgentModeResolver,
+    private readonly modelFactory: AgentModelFactory,
+    private readonly promptBuilder: TurnPromptBuilder,
+    private readonly runtime: AgentRuntime,
+    private readonly streamConsumer: TurnStreamConsumer,
+  ) {}
 
   private async prepareTurn(input: BaseAgentTurnInput): Promise<PreparedAgentTurn> {
     const sequenceDebugCollector = createSequenceDebugCollector();
     const { turnId, previousUserMessages } = await this.lifecycle.start(input);
+    const context = await this.contextBuilder.build({
+      base: input,
+      turnId,
+    });
 
     return {
       prompt: input.prompt,
       sessionId: input.sessionId,
       turnId,
       previousUserMessages,
       modeName: input.modeName,
-      context: {
-        adminUser: input.adminUser,
-        userTimeZone: input.userTimeZone,
-        sessionId: input.sessionId,
-        turnId,
-        abortSignal: input.abortSignal,
-        currentPage: input.currentPage,
-        chatSurface: input.chatSurface,
-        adminPublicOrigin: input.adminPublicOrigin,
-      },
+      context,
       observability: {
         emit: undefined,
         sequenceDebugSink: sequenceDebugCollector,
@@ -119,8 +71,7 @@ export class AgentTurnService {
         abortSignal: input.context.abortSignal,
       }),
     ]);
-    const runtime = this.runtimeFactory.create();
-    const stream = await runtime.stream({
+    const stream = await this.runtime.stream({
       models,
       messages,
       context: input.context,
@@ -214,8 +165,4 @@ export class AgentTurnService {
 
     return agentResponse;
   }
-
-  async handleSpeechTurn(input: HandleSpeechTurnInput) {
-    return this.speechTurnService.handle(input);
-  }
 }

agent/turn/turnTypes.ts

@@ -18,6 +18,16 @@ import type { AgentEndpointsContext } from "./endpoints/context.js";
 import { AgentSessionStore } from "./sessionStore.js";
 import { ChatSurfaceService } from "./chatSurfaceService.js";
 import { AgentTurnService } from "./agentTurnService.js";
+import { AgentModelFactory } from "./agent/models/AgentModelFactory.js";
+import { AgentModeResolver } from "./agent/models/AgentModeResolver.js";
+import { AgentRuntime } from "./agent/runtime/AgentRuntime.js";
+import { SpeechTurnService } from "./agent/speech/SpeechTurnService.js";
+import { AgentToolProvider } from "./agent/tools/AgentToolProvider.js";
+import { TurnContextBuilder } from "./agent/turn/TurnContextBuilder.js";
+import { TurnLifecycleService } from "./agent/turn/TurnLifecycleService.js";
+import { TurnPersistenceService } from "./agent/turn/TurnPersistenceService.js";
+import { TurnPromptBuilder } from "./agent/turn/TurnPromptBuilder.js";
+import { TurnStreamConsumer } from "./agent/turn/TurnStreamConsumer.js";
 
 export type { AgentEvent, AgentEventEmitter } from "./agentEvents.js";
 
@@ -27,6 +37,7 @@ export default class AdminForthAgentPlugin extends AdminForthPlugin {
   private checkpointer: BaseCheckpointSaver | null = null;
   private sessionStore: AgentSessionStore;
   private agentTurnService: AgentTurnService;
+  private speechTurnService: SpeechTurnService;
   private chatSurfaceService: ChatSurfaceService;
   private parseBody<T>(
     schema: z.ZodType<T>,
@@ -62,15 +73,32 @@ export default class AdminForthAgentPlugin extends AdminForthPlugin {
     super(options, import.meta.url);
     this.options = options;
     this.sessionStore = new AgentSessionStore(() => this.adminforth, this.options);
-    this.agentTurnService = new AgentTurnService({
+    const toolProvider = new AgentToolProvider(
+      () => this.adminforth,
+      this.getInternalAgentResourceIds.bind(this),
+    );
+    const runtime = new AgentRuntime({
+      name: `adminforth-agent-${this.pluginInstanceId}`,
       getAdminforth: () => this.adminforth,
-      getPluginInstanceId: () => this.pluginInstanceId,
-      options: this.options,
-      sessionStore: this.sessionStore,
       getCheckpointer: this.getCheckpointer.bind(this),
-      getInternalAgentResourceIds: this.getInternalAgentResourceIds.bind(this),
-      getAgentSystemPrompt: () => this.agentSystemPromptPromise,
+      toolProvider,
     });
+    const persistence = new TurnPersistenceService(() => this.adminforth, this.options);
+    this.agentTurnService = new AgentTurnService(
+      new TurnLifecycleService(this.sessionStore, persistence),
+      new TurnContextBuilder(() => this.adminforth),
+      new AgentModeResolver(this.options),
+      new AgentModelFactory(this.options.maxTokens ?? 1000),
+      new TurnPromptBuilder({
+        getAdminforth: () => this.adminforth,
+        getAgentSystemPrompt: () => this.agentSystemPromptPromise,
+      }),
+      runtime,
+      new TurnStreamConsumer(),
+    );
+    this.speechTurnService = new SpeechTurnService(
+      this.agentTurnService.runAndPersistAgentResponse.bind(this.agentTurnService),
+    );
     this.chatSurfaceService = new ChatSurfaceService(
       () => this.adminforth,
       this.options,
@@ -146,7 +174,7 @@ export default class AdminForthAgentPlugin extends AdminForthPlugin {
       options: this.options,
       parseBody: this.parseBody.bind(this),
       handleTurn: this.agentTurnService.handleTurn.bind(this.agentTurnService),
-      handleSpeechTurn: this.agentTurnService.handleSpeechTurn.bind(this.agentTurnService),
+      handleSpeechTurn: this.speechTurnService.handle.bind(this.speechTurnService),
       runAndPersistAgentResponse: this.agentTurnService.runAndPersistAgentResponse.bind(this.agentTurnService),
       getSessionTurns: this.sessionStore.getSessionTurns.bind(this.sessionStore),
       createNewTurn: this.sessionStore.createNewTurn.bind(this.sessionStore),