From d67ef3f6d98de2378efa0316997c9e6490bbf3cb Mon Sep 17 00:00:00 2001
From: bsorrentino <bartolomeo.sorrentino@gmail.com>
Date: Fri, 19 Jul 2024 00:11:14 +0200
Subject: [PATCH] feat(server-jetty): add completion of async context

---
 .../langgraph4j/LangGraphStreamingServer.java | 53 ++++++++++---------
 1 file changed, 29 insertions(+), 24 deletions(-)

diff --git a/server-jetty/src/main/java/org/bsc/langgraph4j/LangGraphStreamingServer.java b/server-jetty/src/main/java/org/bsc/langgraph4j/LangGraphStreamingServer.java
index 16c5e41..e38300d 100644
--- a/server-jetty/src/main/java/org/bsc/langgraph4j/LangGraphStreamingServer.java
+++ b/server-jetty/src/main/java/org/bsc/langgraph4j/LangGraphStreamingServer.java
@@ -32,8 +32,7 @@
  * of LangGraph.
  * Implementations of this interface can be used to create a web server
  * that exposes an API for interacting with compiled language graphs.
- *
- */
+     */
 public interface LangGraphStreamingServer {
 
     Logger log = LoggerFactory.getLogger(LangGraphStreamingServer.class);
@@ -46,7 +45,7 @@ static Builder builder() {
 
     class Builder {
         private int port = 8080;
-        private Map<String,ArgumentMetadata> inputArgs = new HashMap<>();
+        private Map<String, ArgumentMetadata> inputArgs = new HashMap<>();
         private String title = null;
         private ObjectMapper objectMapper;
 
@@ -59,16 +58,19 @@ public Builder objectMapper(ObjectMapper objectMapper) {
             this.objectMapper = objectMapper;
             return this;
         }
+
         public Builder title(String title) {
             this.title = title;
             return this;
         }
+
         public Builder addInputStringArg(String name, boolean required) {
-            inputArgs.put(name, new ArgumentMetadata("string", required) );
+            inputArgs.put(name, new ArgumentMetadata("string", required));
             return this;
         }
+
         public Builder addInputStringArg(String name) {
-            inputArgs.put(name, new ArgumentMetadata("string", true) );
+            inputArgs.put(name, new ArgumentMetadata("string", true));
             return this;
         }
 
@@ -91,14 +93,14 @@ public <State extends AgentState> LangGraphStreamingServer build(CompiledGraph<S
 
             ServletContextHandler context = new ServletContextHandler(ServletContextHandler.SESSIONS);
 
-            if( objectMapper == null ) {
+            if (objectMapper == null) {
                 objectMapper = new ObjectMapper();
             }
             // context.setContextPath("/");
             // Add the streaming servlet
             context.addServlet(new ServletHolder(new GraphExecutionServlet<State>(compiledGraph, objectMapper)), "/stream");
 
-            InitData initData = new InitData( title, inputArgs );
+            InitData initData = new InitData(title, inputArgs);
             context.addServlet(new ServletHolder(new GraphInitServlet<State>(compiledGraph, initData)), "/init");
 
             Handler.Sequence handlerList = new Handler.Sequence(resourceHandler, context);
@@ -146,7 +148,7 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
         });
 
         // Start asynchronous processing
-        request.startAsync();
+        var asyncContext = request.startAsync();
 
         try {
             compiledGraph.stream(dataMap)
@@ -154,24 +156,25 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
                         try {
 
                             writer.print("{");
-                            writer.printf( "\"node\": \"%s\"", s.node() );
+                            writer.printf("\"node\": \"%s\"", s.node());
                             try {
                                 var stateAsString = objectMapper.writeValueAsString(s.state().data());
-                                writer.printf( ",\"state\": %s" , stateAsString );
-                            }
-                            catch( IOException e ) {
+                                writer.printf(",\"state\": %s", stateAsString);
+                            } catch (IOException e) {
                                 LangGraphStreamingServer.log.info("error serializing state", e);
-                                writer.printf( ",\"state\": {}" );
+                                writer.printf(",\"state\": {}");
                             }
                             writer.print("}");
                             writer.flush();
                             TimeUnit.SECONDS.sleep(1);
-                        } catch (InterruptedException  e) {
+                        } catch (InterruptedException e) {
                             throw new RuntimeException(e);
                         }
 
                     })
-                    .thenAccept(v -> writer.close() );
+                    .thenAccept(v -> writer.close())
+                    .thenAccept(v -> asyncContext.complete())
+            ;
 
         } catch (Exception e) {
             throw new RuntimeException(e);
@@ -179,13 +182,15 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response)
     }
 }
 
-record ArgumentMetadata (
+record ArgumentMetadata(
         String type,
-        boolean required ) {}
+        boolean required) {
+}
 
 record InitData(
         String title,
-        Map<String, ArgumentMetadata> args ) {}
+        Map<String, ArgumentMetadata> args) {
+}
 
 /**
  * return the graph representation in mermaid format
@@ -196,14 +201,14 @@ class GraphInitServlet<State extends AgentState> extends HttpServlet {
     final ObjectMapper objectMapper = new ObjectMapper();
     final InitData initData;
 
-    record Result (
-        String graph,
-        String title,
-        Map<String, ArgumentMetadata> args
+    record Result(
+            String graph,
+            String title,
+            Map<String, ArgumentMetadata> args
     ) {
 
-        public Result(GraphRepresentation graph, InitData initData ) {
-            this( graph.getContent(), initData.title(), initData.args() ); // graph.getContent();
+        public Result(GraphRepresentation graph, InitData initData) {
+            this(graph.getContent(), initData.title(), initData.args()); // graph.getContent();
         }
     }