Merge pull request #941 from trheyi/main

feat: Enhance context management and agent functionality in SUI

neo/api.go

@@ -232,6 +232,12 @@ func (neo *DSL) handleChat(c *gin.Context) {
 	defer cancel()
 	defer ctx.Release() // Release the context after the request is done
 
+	// Set the assistant ID
+	assistantID := c.Query("assistant_id")
+	if assistantID != "" {
+		ctx = chatctx.WithAssistantID(ctx, assistantID)
+	}
+
 	err := neo.Answer(ctx, content, c)
 
 	// Error handling

neo/context/context.go

@@ -71,6 +71,12 @@ func NewWithCancel(sid, cid, payload string) (Context, context.CancelFunc) {
 	return WithCancel(ctx)
 }
 
+// WithAssistantID set the assistant ID
+func WithAssistantID(ctx Context, assistantID string) Context {
+	ctx.AssistantID = assistantID
+	return ctx
+}
+
 // NewWithTimeout create a new context with timeout
 func NewWithTimeout(sid, cid, payload string, timeout time.Duration) (Context, context.CancelFunc) {
 	ctx := New(sid, cid, payload)

sui/libsui/agent.ts

@@ -103,7 +103,7 @@ class Agent {
    * @param input Text message or input object with text and optional attachments
    * @param args Additional arguments to pass to the agent
    */
-  async Call(input: AgentInput, ...args: any[]) {
+  async Call(input: AgentInput, ...args: any[]): Promise<any> {
     const messages: AgentMessage[] = [];
     let currentContent = "";
     let lastAssistant = {
@@ -302,6 +302,13 @@ class Agent {
 
             messages.push(newMessage);
             messageHandler(newMessage);
+
+            // If the message is done, close the event source
+            if (done) {
+              const doneHandler = this.events["done"] as DoneHandler;
+              doneHandler?.(messages);
+              es.close();
+            }
             return;
           }