diff --git a/packages/server/src/types.ts b/packages/server/src/types.ts
index 5b2676c486bc8f1aaa296af2ce968f2314c8db95..233198eb39822cf1466946c5a748e63d51c68fa7 100644
--- a/packages/server/src/types.ts
+++ b/packages/server/src/types.ts
@@ -1,27 +1,16 @@
 import {
   AgentWorkflow,
   Workflow,
-  type ChatMessage,
-  type ChatResponseChunk,
+  type AgentInputData,
+  type AgentWorkflowContext,
 } from "llamaindex";
 import type next from "next";
 
 /**
- * The input for an AgentWorkflow
- * userInput is the last message content from the user
- * chatHistory is the previous chat history (not including the last message)
- */
-export type AgentInput = {
-  userInput: string;
-  chatHistory: ChatMessage[];
-};
-
-/**
- * ServerWorkflow can be either a normal Workflow or an AgentWorkflow
+ * ServerWorkflow can be either a custom Workflow or an AgentWorkflow
  */
 export type ServerWorkflow =
-  | Workflow<null, AgentInput, ChatResponseChunk>
-  // | Workflow<AgentWorkflowContext, AgentInputData, string>
+  | Workflow<AgentWorkflowContext, AgentInputData, string>
   | AgentWorkflow;
 
 /**
diff --git a/packages/server/src/utils/workflow.ts b/packages/server/src/utils/workflow.ts
index ed893fbf38944dbc61ea0354dbf910e94f5296f8..5abcebe78de38e3f8abe75fd1b7450ba3e296ec7 100644
--- a/packages/server/src/utils/workflow.ts
+++ b/packages/server/src/utils/workflow.ts
@@ -5,26 +5,28 @@ import {
   EngineResponse,
   StopEvent,
   Workflow,
+  type AgentInputData,
+  type AgentWorkflowContext,
   type ChatResponseChunk,
 } from "llamaindex";
 import { ReadableStream } from "stream/web";
-import type { AgentInput, ServerWorkflow } from "../types";
+import type { ServerWorkflow } from "../types";
 
 export async function runWorkflow(
   workflow: ServerWorkflow,
-  agentInput: AgentInput,
+  agentInput: AgentInputData,
 ) {
   if (workflow instanceof AgentWorkflow) {
     return runAgentWorkflow(workflow, agentInput);
   }
-  return runNormalWorkflow(workflow, agentInput);
+  return runCustomWorkflow(workflow, agentInput);
 }
 
 async function runAgentWorkflow(
   workflow: AgentWorkflow,
-  agentInput: AgentInput,
+  agentInput: AgentInputData,
 ) {
-  const { userInput, chatHistory } = agentInput;
+  const { userInput = "", chatHistory = [] } = agentInput;
   const context = workflow.run(userInput, { chatHistory });
 
   const dataStream = new StreamData();
@@ -50,9 +52,9 @@ async function runAgentWorkflow(
   return LlamaIndexAdapter.toDataStreamResponse(stream, { data: dataStream });
 }
 
-async function runNormalWorkflow(
-  workflow: Workflow<null, AgentInput, ChatResponseChunk>,
-  agentInput: AgentInput,
+async function runCustomWorkflow(
+  workflow: Workflow<AgentWorkflowContext, AgentInputData, string>,
+  agentInput: AgentInputData,
 ) {
   const context = workflow.run(agentInput);
   const dataStream = new StreamData();
diff --git a/packages/workflow/src/index.ts b/packages/workflow/src/index.ts
index 1d29b42507cc626d2f7b93fd38776a56c8c5d988..320e24f33340fa9e26e988dcd558b928c26f8bce 100644
--- a/packages/workflow/src/index.ts
+++ b/packages/workflow/src/index.ts
@@ -1,3 +1,4 @@
+export { type AgentWorkflowContext } from "./agent/base.js";
 export {
   WorkflowContext,
   type HandlerContext,