From 73766b6b2a40aa8bf104ab74732e9040c7267fab Mon Sep 17 00:00:00 2001
From: "J. Zane Cook" <jzanecook@z90.studio>
Date: Tue, 2 Apr 2024 12:28:17 -0500
Subject: [PATCH] fix(js-sdk): Got testing suite working with current sdk

- Tasks are paginated and return the entire task rather than just the ID
- Add additional_input to task input
---
 packages/sdk/js/examples/minimal.ts |  2 +-
 packages/sdk/js/src/agent.ts        | 45 +++++++++++++++++++++++++----
 packages/sdk/js/src/index.ts        |  4 +--
 packages/sdk/js/src/models.ts       |  3 ++
 4 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/packages/sdk/js/examples/minimal.ts b/packages/sdk/js/examples/minimal.ts
index d3a7f7f2..34e17e9e 100644
--- a/packages/sdk/js/examples/minimal.ts
+++ b/packages/sdk/js/examples/minimal.ts
@@ -18,4 +18,4 @@ async function taskHandler(taskInput: TaskInput | null): Promise<StepHandler> {
   return stepHandler
 }
 
-Agent.handleTask(taskHandler).start()
+Agent.handleTask(taskHandler, {}).start()
diff --git a/packages/sdk/js/src/agent.ts b/packages/sdk/js/src/agent.ts
index b575a344..9487ce4c 100644
--- a/packages/sdk/js/src/agent.ts
+++ b/packages/sdk/js/src/agent.ts
@@ -65,6 +65,7 @@ export const createAgentTask = async (
   const task: Task = {
     task_id: uuid(),
     input: body?.input ?? null,
+    additional_input: body?.additional_input ?? {},
     artifacts: [],
   }
   const stepHandler = await taskHandler(task.task_id, body?.input ?? null)
@@ -89,15 +90,47 @@ const registerCreateAgentTask: RouteRegisterFn = (router: Router) => {
  * Lists all tasks that have been created for the agent.
  * @returns Promise<string[]>
  */
-export const listAgentTaskIDs = async (): Promise<string[]> => {
-  return tasks.map(([task, _]) => task.task_id)
+export const listAgentTasks = async (): Promise<Task[]> => {
+  return tasks.map(([task, _]) => task)
 }
-const registerListAgentTaskIDs: RouteRegisterFn = (router: Router) => {
+const registerListAgentTasks: RouteRegisterFn = (router: Router) => {
   router.get('/agent/tasks', (req, res) => {
     void (async () => {
       try {
-        const ids = await listAgentTaskIDs()
-        res.status(200).json(ids)
+        const tasks = await listAgentTasks()
+
+        let currentPage = 1
+        let pageSize = 10
+        if (
+          req.query.current_page !== undefined &&
+          !isNaN(Number(req.query.current_page))
+        ) {
+          currentPage = Number(req.query.current_page)
+        }
+
+        if (
+          req.query.page_size !== undefined &&
+          !isNaN(Number(req.query.page_size))
+        ) {
+          pageSize = Number(req.query.page_size)
+        }
+
+        const totalItems = tasks.length
+        const totalPages = Math.ceil(totalItems / pageSize)
+
+        const start = (currentPage - 1) * pageSize
+        const end = start + pageSize
+        const pagedTasks = tasks.slice(start, end)
+
+        res.status(200).json({
+          tasks: pagedTasks,
+          pagination: {
+            total_items: totalItems,
+            total_pages: totalPages,
+            current_page: currentPage,
+            page_size: pageSize,
+          }
+        })
       } catch (err: Error | any) {
         console.error(err)
         res.status(500).json({ error: err.message })
@@ -471,7 +504,7 @@ export class Agent {
       port: port ?? this.config.port ?? defaultAgentConfig.port,
       routes: [
         registerCreateAgentTask,
-        registerListAgentTaskIDs,
+        registerListAgentTasks,
         registerGetAgentTask,
         registerListAgentTaskSteps,
         registerExecuteAgentTaskStep,
diff --git a/packages/sdk/js/src/index.ts b/packages/sdk/js/src/index.ts
index d52b14fd..7a146daa 100644
--- a/packages/sdk/js/src/index.ts
+++ b/packages/sdk/js/src/index.ts
@@ -13,7 +13,7 @@ import {
   type TaskHandler,
   StepResultWithDefaults,
   createAgentTask,
-  listAgentTaskIDs,
+  listAgentTasks,
   getAgentTask,
   listAgentTaskSteps,
   executeAgentTaskStep,
@@ -34,7 +34,7 @@ export {
   type TaskHandler,
   StepResultWithDefaults as StepResult,
   createAgentTask,
-  listAgentTaskIDs,
+  listAgentTasks,
   getAgentTask,
   listAgentTaskSteps,
   executeAgentTaskStep,
diff --git a/packages/sdk/js/src/models.ts b/packages/sdk/js/src/models.ts
index 740f06b7..f54919f4 100644
--- a/packages/sdk/js/src/models.ts
+++ b/packages/sdk/js/src/models.ts
@@ -2,6 +2,7 @@
  * Input parameters for the task. Any value is allowed.
  */
 export type TaskInput = any
+export type AdditionalInput = any
 
 /**
  * Artifact that the task has produced. Any value is allowed.
@@ -87,6 +88,7 @@ export interface StepResult {
 
 export interface Task {
   input?: TaskInput
+  additional_input?: AdditionalInput
   /**
    * The ID of the task.
    */
@@ -99,4 +101,5 @@ export interface Task {
 
 export interface TaskRequestBody {
   input?: TaskInput
+  additional_input?: AdditionalInput
 }