@axiom-lattice/gateway 2.1.44 → 2.1.46

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@axiom-lattice/gateway",
-  "version": "2.1.44",
+  "version": "2.1.46",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
   "types": "dist/index.d.ts",
@@ -35,16 +35,19 @@
     "pino": "^9.7.0",
     "pino-pretty": "^13.0.0",
     "pino-roll": "^3.1.0",
+    "pg": "^8.11.0",
     "redis": "^5.0.1",
     "uuid": "^9.0.1",
-    "@axiom-lattice/core": "2.1.38",
-    "@axiom-lattice/protocols": "2.1.20",
-    "@axiom-lattice/queue-redis": "1.0.19"
+    "@axiom-lattice/core": "2.1.40",
+    "@axiom-lattice/pg-stores": "1.0.30",
+    "@axiom-lattice/protocols": "2.1.22",
+    "@axiom-lattice/queue-redis": "1.0.21"
   },
   "devDependencies": {
     "@types/jest": "^29.5.14",
     "@types/lodash": "^4.17.16",
     "@types/node": "^20.19.13",
+    "@types/pg": "^8.11.0",
     "@types/uuid": "^9.0.8",
     "@types/ws": "^8.18.1",
     "@typescript-eslint/eslint-plugin": "^7.2.0",

package/src/controllers/assistant.ts

@@ -1,6 +1,5 @@
 import { FastifyRequest, FastifyReply } from "fastify";
-import { draw_graph } from "../services/agent_service";
-import { getStoreLattice } from "@axiom-lattice/core";
+import { agentInstanceManager, getStoreLattice } from "@axiom-lattice/core";
 import type {
   Assistant,
   CreateAssistantRequest,
@@ -26,7 +25,7 @@ function getTenantId(request: FastifyRequest): string {
   if (userTenantId) {
     return userTenantId;
   }
-  
+
   // Fallback to request headers for backward compatibility
   return (request.headers["x-tenant-id"] as string) || "default";
 }
@@ -321,8 +320,10 @@ export const getAgentGraph = async (
     const { assistantId } = request.params;
     const tenant_id = getTenantId(request);
 
+    const agent = agentInstanceManager.getAgent({ assistant_id: assistantId, tenant_id, thread_id: '' })
+
     // Call drawing service to get image data
-    const imageData = await draw_graph(assistantId, tenant_id);
+    const imageData = await agent.get_draw_graph()
 
     // Set response header and return image data
     reply.header("Content-Type", "application/json").send({

package/src/controllers/memory.ts

@@ -1,5 +1,5 @@
 import { FastifyRequest, FastifyReply } from "fastify";
-import { agent_messages, agent_state } from "../services/agent_service";
+import { agentInstanceManager } from "@axiom-lattice/core";
 
 // 设置内存项
 export const setMemoryItem = async (
@@ -108,11 +108,9 @@ export const getAllMemoryItems = async (
       return;
     }
 
-    const result = await agent_messages({
-      assistant_id: assistantId,
-      thread_id: thread_id,
-      tenant_id: tenant_id,
-    });
+    const agent = agentInstanceManager.getAgent({ assistant_id: assistantId, tenant_id, thread_id: thread_id })
+
+    const result = await agent.getCurrentMessages();
 
     if (!result) {
       reply.status(500).send(result);
@@ -156,19 +154,25 @@ export const getAgentState = async (
     }
 
     const tenant_id = request.headers["x-tenant-id"] as string;
-    
-    const result = await agent_state({
-      assistant_id: assistantId,
-      thread_id: thread_id,
-      tenant_id: tenant_id,
-    });
+    const agent = agentInstanceManager.getAgent({ assistant_id: assistantId, tenant_id, thread_id: thread_id })
+
+    // 1. 获取 LangGraph 状态（现有逻辑）
+    const result = await agent.getCurrentState();
+    const pendingMessages = await agent.getPendingMessages();
+
 
     if (!result) {
       reply.status(500).send(result);
       return;
     }
 
-    reply.send(result);
+    // 3. 合并返回
+    const mergedResult = {
+      ...result,
+      pendingMessages
+    };
+
+    reply.send(mergedResult);
   } catch (error: any) {
     reply.status(500).send({
       success: false,

package/src/controllers/models.ts

@@ -43,6 +43,7 @@ function extractModelKey(tenantId: string, tenantModelKey: string): string {
 interface ModelConfig {
   key: string;
   model: string;
+  displayName?: string;
   provider: "azure" | "openai" | "deepseek" | "siliconcloud" | "volcengine";
   streaming?: boolean;
   apiKey?: string;
@@ -75,10 +76,16 @@ export async function getModels(request: FastifyRequest, reply: FastifyReply) {
       // Extract config from the lattice client
       // Note: This is a simplified approach - you may need to adjust based on actual implementation
       const config = (lattice.client as any).config || {};
+      // Use provided displayName or auto-generate from key
+      const displayName = config.displayName || lattice.key
+        .split('-')
+        .map(word => word.charAt(0).toUpperCase() + word.slice(1))
+        .join(' ');
       return {
         key: lattice.key,
         model: config.model || "",
         provider: config.provider || "openai",
+        displayName: displayName,
         streaming: config.streaming || false,
         apiKey: config.apiKey || "",
         baseURL: config.baseURL || "",
@@ -144,6 +151,7 @@ export async function updateModels(
         const llmConfig: LLMConfig = {
           provider: modelConfig.provider,
           model: modelConfig.model,
+          displayName: modelConfig.displayName,
           streaming: modelConfig.streaming ?? false,
           apiKey: modelConfig.apiKey,
           baseURL: modelConfig.baseURL,

package/src/controllers/run.ts

@@ -1,7 +1,12 @@
 import { FastifyRequest, FastifyReply } from "fastify";
-import * as agentService from "../services/agent_service";
 import { CreateRunRequest } from "../types";
 import { v4 } from "uuid";
+import {
+  Agent,
+  ThreadStatus,
+  agentLatticeManager,
+  agentInstanceManager,
+} from "@axiom-lattice/core";
 
 interface ResumeStreamRequest {
   thread_id: string;
@@ -10,7 +15,7 @@ interface ResumeStreamRequest {
   poll_interval?: number;
 }
 
-// 创建运行
+// Create run
 export const createRun = async (
   request: FastifyRequest,
   reply: FastifyReply
@@ -31,7 +36,7 @@ export const createRun = async (
     const project_id = request.headers["x-project-id"] as string;
     const x_request_id = (request.headers["x-request-id"] as string) || v4();
 
-    // 验证请求数据
+    // Validate request
     if (!assistant_id) {
       reply.status(400).send({
         success: false,
@@ -40,35 +45,21 @@ export const createRun = async (
       return;
     }
 
-    // 如果请求streaming，则agent_stream
-    if (streaming) {
-      // 先检查agent是否存在（在hijack之前）
-      const agentExists = await agentService.checkAgentExists(tenant_id, assistant_id);
-      if (!agentExists) {
-        reply.status(404).send({
-          success: false,
-          error: `Agent ${assistant_id} not found for tenant ${tenant_id}`,
-        });
-        return;
-      }
-      
-      // 开始运行
-      const stream = await agentService.agent_stream({
-        assistant_id: assistant_id,
-        input: input,
-        thread_id: thread_id,
-        command,
-        tenant_id: tenant_id,
-        workspace_id: workspace_id,
-        project_id: project_id,
-        run_id: x_request_id,
-        custom_run_config,
-      });
+    // Get or create Agent instance (ensures single instance per thread)
+    const agent = agentInstanceManager.getAgent({
+      assistant_id,
+      thread_id,
+      tenant_id,
+      workspace_id,
+      project_id,
+      custom_run_config,
+    });
 
-      // 通知 Fastify 我们将手动处理响应
-      reply.hijack();
 
-      // 设置 SSE 响应头
+    // Handle streaming requests
+    if (streaming) {
+      // Setup SSE
+      reply.hijack();
       reply.raw.writeHead(200, {
         "Content-Type": "text/event-stream",
         "Cache-Control": "no-cache",
@@ -77,41 +68,59 @@ export const createRun = async (
       });
 
       try {
-        let chunkCount = 0;
+
+        // Execute agent with streaming
+        const result = await agent.addMessage({
+          input,
+          command,
+          custom_run_config,
+        });
+
+        // const agentStatus = await agent.getRunStatus()
+        // console.log(agentStatus)
+        // if (agentStatus === "busy") {
+        //   reply.status(200).send({
+        //     success: true,
+
+        //   });
+        //   return
+        // }
+
+        const stream = agent.chunkStream((result).messageId)
+
+        // Forward all chunks to SSE
         for await (const chunk of stream) {
-          chunkCount++;
+          console.log(input.message, chunk.data.content);
           const success = reply.raw.write(`data: ${JSON.stringify(chunk)}\n\n`);
           if (!success) {
             await new Promise(resolve => reply.raw.once('drain', resolve));
           }
         }
       } catch (error: any) {
-        // Send error as SSE event before closing (following MessageChunk format)
+        // Send error as SSE event
         const errorEvent = {
           type: 'error',
           data: {
             id: v4(),
-            content: error.message || 'Stream processing error'
-          }
+            content: error.message || 'Stream processing error',
+          },
         };
         reply.raw.write(`data: ${JSON.stringify(errorEvent)}\n\n`);
       } finally {
         reply.raw.end();
       }
     } else {
-      // 后台运行的情况
-      const result = await agentService.agent_invoke({
-        assistant_id: assistant_id,
-        input: input,
-        command: command,
-        thread_id: thread_id,
-        tenant_id: tenant_id,
-        workspace_id: workspace_id,
-        project_id: project_id,
-        run_id: x_request_id,
+      // Non-streaming: use agentExecutor
+      const { message: msg, ...restInputNonStream } = input;
+      const result = await agent.invoke({
+        input: { message: msg, ...restInputNonStream },
+        command,
         custom_run_config,
       });
-      reply.status(200).send(result);
+      reply.status(200).send({
+        success: true,
+        ...result,
+      });
     }
   } catch (error: any) {
     reply.status(500).send({
@@ -151,14 +160,17 @@ export const resumeStream = async (
     });
 
     try {
-      // Get the stream from agent service
-      const stream = await agentService.resume_stream({
+      // Create Agent instance for accessing chunk buffer
+      // Note: assistant_id, tenant_id are not required for chunkStream
+      const agent = new Agent({
+        assistant_id: '',
         thread_id,
-        message_id,
-        known_content: known_content,
-        poll_interval: poll_interval || 100,
+        tenant_id: '',
       });
 
+      // Get the stream from agent's chunk buffer
+      const stream = agent.chunkStream(message_id, known_content);
+
       // Stream the chunks to the client
       for await (const chunk of stream) {
         reply.raw.write(`data: ${JSON.stringify(chunk)}\n\n`);

package/src/controllers/thread_status.ts

@@ -0,0 +1,228 @@
+// import { FastifyRequest, FastifyReply } from "fastify";
+// import { QueueMode } from "@axiom-lattice/core";
+
+
+// /**
+//  * GET /api/threads/:thread_id/status
+//  * Query single thread status
+//  */
+// export async function getThreadStatusHandler(
+//   request: FastifyRequest,
+//   reply: FastifyReply
+// ) {
+//   try {
+//     const { thread_id } = request.params as { thread_id: string };
+//     const tenant_id = request.headers["x-tenant-id"] as string;
+
+//     if (!tenant_id) {
+//       return reply.code(400).send({ error: "Missing x-tenant-id header" });
+//     }
+
+//     const state = lifecycleManager.getThreadStatus(tenant_id, thread_id);
+
+//     if (!state) {
+//       return reply.code(404).send({
+//         error: "Thread not found",
+//         threadId: thread_id,
+//       });
+//     }
+
+//     return reply.send(state);
+//   } catch (error) {
+//     console.error("Error getting thread status:", error);
+//     return reply.code(500).send({
+//       error: "Internal server error",
+//       message: error instanceof Error ? error.message : String(error),
+//     });
+//   }
+// }
+
+// /**
+//  * GET /api/assistants/:assistant_id/threads/status
+//  * Query all threads status for an agent
+//  */
+// export async function getAgentThreadsHandler(
+//   request: FastifyRequest,
+//   reply: FastifyReply
+// ) {
+//   try {
+//     const { assistant_id } = request.params as { assistant_id: string };
+//     const tenant_id = request.headers["x-tenant-id"] as string;
+
+//     if (!tenant_id) {
+//       return reply.code(400).send({ error: "Missing x-tenant-id header" });
+//     }
+
+//     const threads = lifecycleManager.getAgentThreads(tenant_id, assistant_id);
+
+//     return reply.send({
+//       assistantId: assistant_id,
+//       threads,
+//       count: threads.length,
+//     });
+//   } catch (error) {
+//     console.error("Error getting agent threads:", error);
+//     return reply.code(500).send({
+//       error: "Internal server error",
+//       message: error instanceof Error ? error.message : String(error),
+//     });
+//   }
+// }
+
+// /**
+//  * GET /api/threads/active
+//  * Query all active threads (BUSY + INTERRUPTED)
+//  */
+// export async function getActiveThreadsHandler(
+//   request: FastifyRequest,
+//   reply: FastifyReply
+// ) {
+//   try {
+//     const tenant_id = request.headers["x-tenant-id"] as string;
+
+//     if (!tenant_id) {
+//       return reply.code(400).send({ error: "Missing x-tenant-id header" });
+//     }
+
+//     const threads = lifecycleManager.getActiveThreads(tenant_id);
+
+//     return reply.send({
+//       threads,
+//       count: threads.length,
+//     });
+//   } catch (error) {
+//     console.error("Error getting active threads:", error);
+//     return reply.code(500).send({
+//       error: "Internal server error",
+//       message: error instanceof Error ? error.message : String(error),
+//     });
+//   }
+// }
+
+// /**
+//  * POST /api/threads/:thread_id/messages
+//  * Add message to queue or execute immediately (steer mode)
+//  */
+// export async function addMessageHandler(
+//   request: FastifyRequest,
+//   reply: FastifyReply
+// ) {
+//   try {
+//     const { thread_id } = request.params as { thread_id: string };
+//     const tenant_id = request.headers["x-tenant-id"] as string;
+//     const { assistant_id, content, mode } = request.body as {
+//       assistant_id: string;
+//       content: any;
+//       mode?: QueueMode;
+//     };
+
+//     if (!tenant_id) {
+//       return reply.code(400).send({ error: "Missing x-tenant-id header" });
+//     }
+
+//     if (!assistant_id || !content) {
+//       return reply.code(400).send({ error: "Missing assistant_id or content" });
+//     }
+
+//     const result = await lifecycleManager.addMessage(
+//       tenant_id,
+//       assistant_id,
+//       thread_id,
+//       content,
+//       mode
+//     );
+
+//     return reply.send({
+//       success: true,
+//       threadId: thread_id,
+//       ...result,
+//     });
+//   } catch (error) {
+//     console.error("Error adding message:", error);
+//     return reply.code(500).send({
+//       error: "Internal server error",
+//       message: error instanceof Error ? error.message : String(error),
+//     });
+//   }
+// }
+
+// /**
+//  * DELETE /api/threads/:thread_id/messages/:message_id
+//  * Remove pending message
+//  */
+// export async function removeMessageHandler(
+//   request: FastifyRequest,
+//   reply: FastifyReply
+// ) {
+//   try {
+//     const { thread_id, message_id } = request.params as {
+//       thread_id: string;
+//       message_id: string;
+//     };
+//     const tenant_id = request.headers["x-tenant-id"] as string;
+
+//     if (!tenant_id) {
+//       return reply.code(400).send({ error: "Missing x-tenant-id header" });
+//     }
+
+//     const success = await lifecycleManager.removePendingMessage(
+//       tenant_id,
+//       thread_id,
+//       message_id
+//     );
+
+//     if (!success) {
+//       return reply.code(404).send({
+//         error: "Message not found",
+//         messageId: message_id,
+//       });
+//     }
+
+//     return reply.send({
+//       success: true,
+//       messageId: message_id,
+//     });
+//   } catch (error) {
+//     console.error("Error removing message:", error);
+//     return reply.code(500).send({
+//       error: "Internal server error",
+//       message: error instanceof Error ? error.message : String(error),
+//     });
+//   }
+// }
+
+// /**
+//  * PUT /api/threads/:thread_id/queue-config
+//  * Update queue configuration
+//  */
+// export async function updateQueueConfigHandler(
+//   request: FastifyRequest,
+//   reply: FastifyReply
+// ) {
+//   try {
+//     const { thread_id } = request.params as { thread_id: string };
+//     const tenant_id = request.headers["x-tenant-id"] as string;
+//     const config = request.body as {
+//       mode?: QueueMode;
+//       maxSize?: number;
+//     };
+
+//     if (!tenant_id) {
+//       return reply.code(400).send({ error: "Missing x-tenant-id header" });
+//     }
+
+//     await lifecycleManager.setQueueConfig(tenant_id, thread_id, config);
+
+//     return reply.send({
+//       success: true,
+//       threadId: thread_id,
+//       config,
+//     });
+//   } catch (error) {
+//     console.error("Error updating queue config:", error);
+//     return reply.code(500).send({
+//       error: "Internal server error",
+//       message: error instanceof Error ? error.message : String(error),
+//     });
+//   }
+// }

package/src/index.ts

@@ -116,7 +116,7 @@ app.addHook("onResponse", (request, reply, done) => {
     "x-tenant-id": getHeaderValue(request.headers["x-tenant-id"]),
     "x-request-id": getHeaderValue(request.headers["x-request-id"]),
   };
-  loggerLattice.info(`${request.method} ${request.url} - ${reply.statusCode}`);
+  //loggerLattice.info(`${request.method} ${request.url} - ${reply.statusCode}`);
   done();
 });
 
@@ -244,6 +244,13 @@ const start = async (config?: LatticeGatewayConfig) => {
     await app.listen({ port: target_port, host: "0.0.0.0" });
     logger.info(`Lattice Gateway is running on port: ${target_port}`);
 
+    // Initialize AgentLifecycleManager
+    try {
+      logger.info("AgentLifecycleManager initialized");
+    } catch (error) {
+      logger.warn("Failed to initialize AgentLifecycleManager", { error });
+    }
+
     // Configure queue service
     const queueServiceConfig = config?.queueServiceConfig;
     if (queueServiceConfig) {

package/src/routes/index.ts

@@ -38,6 +38,14 @@ import { registerMcpServerConfigRoutes } from "../controllers/mcp-configs";
 import { registerUserRoutes } from "../controllers/users";
 import { registerTenantRoutes } from "../controllers/tenants";
 import { registerAuthRoutes } from "../controllers/auth";
+// import {
+//   getThreadStatusHandler,
+//   getAgentThreadsHandler,
+//   getActiveThreadsHandler,
+//   addMessageHandler,
+//   removeMessageHandler,
+//   updateQueueConfigHandler,
+// } from "../controllers/thread_status";
 
 export const registerLatticeRoutes = (app: FastifyInstance): void => {
   // 运行路由
@@ -328,4 +336,14 @@ export const registerLatticeRoutes = (app: FastifyInstance): void => {
     autoApproveUsers: process.env.AUTO_APPROVE_USERS !== "false",
     allowTenantRegistration: process.env.ALLOW_TENANT_REGISTRATION !== "false",
   });
+
+  // // Thread 状态查询路由
+  // app.get("/api/threads/:thread_id/status", getThreadStatusHandler);
+  // app.get("/api/assistants/:assistant_id/threads/status", getAgentThreadsHandler);
+  // app.get("/api/threads/active", getActiveThreadsHandler);
+
+  // // Thread 消息队列路由
+  // app.post("/api/threads/:thread_id/messages", addMessageHandler);
+  // app.delete("/api/threads/:thread_id/messages/:message_id", removeMessageHandler);
+  // app.put("/api/threads/:thread_id/queue-config", updateQueueConfigHandler);
 };