From 4f132859a4ba4567ec2460cf186023ba13dfaaf9 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 26 Jan 2026 10:48:00 +0800 Subject: [PATCH] Enhance JSAPI and Context Management for Sub-Agent Calls - Implemented automatic disabling of SSE output for sub-agent calls to prevent message corruption and client disconnection. - Introduced a `forceSkipOutput` method to ensure `skip.output` is set to true for all sub-agent requests. - Updated the `Fork` method in the context to create independent Memory instances, preventing state sharing during concurrent executions. - Enhanced documentation to clarify the behavior of SSE output and context management in batch operations. --- agent/caller/jsapi.go | 26 +++++++++++++++++++ agent/context/JSAPI.md | 8 ++++++ agent/context/context.go | 27 +++++++++++++++++--- agent/memory/memory.go | 54 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 111 insertions(+), 4 deletions(-) diff --git a/agent/caller/jsapi.go b/agent/caller/jsapi.go index 2edc023c..b9077ba8 100644 --- a/agent/caller/jsapi.go +++ b/agent/caller/jsapi.go @@ -27,8 +27,12 @@ func NewJSAPI(ctx *agentContext.Context) *JSAPI { // Call executes a single agent call // Usage: ctx.agent.Call("assistant-id", messages, options?) // Returns: { agent_id, response, content, error } +// Note: SSE output is automatically disabled to prevent sub-agent from writing to +// the same SSE stream as the parent agent, which would cause client issues. func (api *JSAPI) Call(agentID string, messages []interface{}, opts map[string]interface{}) interface{} { req := api.buildRequest(agentID, messages, opts) + // Force skip.output = true for sub-agent calls + api.forceSkipOutput(req) result := api.orchestrator.callAgent(req) return result } @@ -71,9 +75,12 @@ func (api *JSAPI) Race(requests []interface{}) []interface{} { // ============================================================================ // CallWithHandler executes a single agent call with an OnMessage handler +// Note: SSE output is automatically disabled; use the handler callback to receive messages. func (api *JSAPI) CallWithHandler(agentID string, messages []interface{}, opts map[string]interface{}, handler agentContext.OnMessageFunc) interface{} { req := api.buildRequest(agentID, messages, opts) req.Handler = handler + // Force skip.output = true for sub-agent calls + api.forceSkipOutput(req) result := api.orchestrator.callAgent(req) return result } @@ -99,8 +106,24 @@ func (api *JSAPI) RaceWithHandler(requests []interface{}, globalHandler agentCon return api.convertResults(results) } +// forceSkipOutput ensures SSE output is disabled for sub-agent calls +// This prevents sub-agents from writing to the same SSE stream as the parent, +// which would cause client disconnection and message corruption. +// Users can use the onChunk callback to receive streaming messages if needed. +func (api *JSAPI) forceSkipOutput(req *Request) { + if req.Options == nil { + req.Options = &CallOptions{} + } + if req.Options.Skip == nil { + req.Options.Skip = &agentContext.Skip{} + } + req.Options.Skip.Output = true +} + // parseRequestsWithHandlers parses requests and attaches handlers // It checks for per-request _handler fields and wraps globalHandler with agentID/index +// For all calls, this automatically sets skip.output = true to prevent sub-agents +// from writing to the same SSE stream as the parent, which would cause client issues. func (api *JSAPI) parseRequestsWithHandlers(requests []interface{}, globalHandler agentContext.BatchOnMessageFunc) []*Request { reqs := make([]*Request, 0, len(requests)) @@ -130,6 +153,9 @@ func (api *JSAPI) parseRequestsWithHandlers(requests []interface{}, globalHandle req := api.buildRequest(agentID, messages, opts) + // Force skip.output = true for all sub-agent calls + api.forceSkipOutput(req) + // Check for per-request handler first (takes precedence) if handler, ok := reqMap["_handler"].(agentContext.OnMessageFunc); ok && handler != nil { req.Handler = handler diff --git a/agent/context/JSAPI.md b/agent/context/JSAPI.md index eaef7d1e..4103b244 100644 --- a/agent/context/JSAPI.md +++ b/agent/context/JSAPI.md @@ -2044,6 +2044,13 @@ Common message types: The parallel methods allow calling multiple agents concurrently, similar to JavaScript Promise patterns. +> **Important: SSE Output is Automatically Disabled** +> +> For all batch calls (`All`, `Any`, `Race`), SSE output is **automatically disabled** (`skip.output = true`). +> This prevents multiple agents from writing to the same SSE stream simultaneously, which would cause +> client disconnection and message corruption. Use the `onChunk` callback to receive streaming messages +> if needed. + #### `ctx.agent.All(requests, options?)` Executes all agent calls and waits for all to complete (like `Promise.all`). @@ -2064,6 +2071,7 @@ interface AgentRequest { // Note: Per-request onChunk is NOT supported in batch calls. // Use the global onChunk callback in the second argument instead. +// Note: skip.output is automatically set to true for all batch calls. ``` **Example:** diff --git a/agent/context/context.go b/agent/context/context.go index b89b4ff6..98cc3e20 100644 --- a/agent/context/context.go +++ b/agent/context/context.go @@ -133,17 +133,34 @@ func (ctx *Context) GetAuthorizedMap() map[string]interface{} { } // Fork creates a child context for concurrent agent/LLM calls -// The forked context shares read-only resources (Memory, Authorized, Cache, Writer) -// but has its own independent Stack and Logger to avoid race conditions +// The forked context shares read-only resources (Authorized, Cache, Writer) +// but has its own independent Stack, Logger, and Memory.Context namespace +// to avoid race conditions and state sharing issues. // // This is essential for batch operations (All/Any/Race) where multiple goroutines -// need to execute concurrently without interfering with each other's Stack state. +// need to execute concurrently without interfering with each other's state. +// +// Key behavior: +// - Memory.User, Memory.Team, Memory.Chat are shared (cross-request state) +// - Memory.Context is INDEPENDENT (request-scoped state, isolated per fork) // // The forked context does NOT need to be released separately - the parent context // manages shared resources. However, the child's Stack will be collected in parent's Stacks map. func (ctx *Context) Fork() *Context { childID := generateContextID() + // Fork memory with independent Context namespace + // This prevents parallel sub-agents from sharing ctx.memory.context state + var forkedMemory *memory.Memory + if ctx.Memory != nil { + var err error + forkedMemory, err = ctx.Memory.Fork(childID) + if err != nil { + // Fallback to shared memory if fork fails (log warning) + forkedMemory = ctx.Memory + } + } + child := &Context{ // Inherit parent's standard context Context: ctx.Context, @@ -151,8 +168,10 @@ func (ctx *Context) Fork() *Context { // New unique ID for this forked context ID: childID, + // Memory with independent Context namespace (see above) + Memory: forkedMemory, + // Share read-only/thread-safe resources with parent - Memory: ctx.Memory, // Memory is designed to be shared Cache: ctx.Cache, // Cache store is thread-safe Writer: ctx.Writer, // Output writer is thread-safe (output module handles concurrency) Authorized: ctx.Authorized, // Read-only auth info diff --git a/agent/memory/memory.go b/agent/memory/memory.go index 7a40d54a..b4db35b0 100644 --- a/agent/memory/memory.go +++ b/agent/memory/memory.go @@ -183,3 +183,57 @@ func (m *Memory) Clear() { m.Context.Clear() } } + +// Fork creates a new Memory instance with an independent Context namespace +// but sharing the User, Team, and Chat namespaces with the parent. +// This is used for parallel agent calls (ctx.agent.All/Any/Race) to prevent +// context state from being shared between concurrent sub-agent executions. +// +// The new Context namespace uses the provided newContextID. +// If newContextID is empty, returns a shallow copy with shared Context. +func (m *Memory) Fork(newContextID string) (*Memory, error) { + if m == nil { + return nil, nil + } + + // If no new context ID provided, share everything (shallow copy) + if newContextID == "" { + return &Memory{ + UserID: m.UserID, + TeamID: m.TeamID, + ChatID: m.ChatID, + ContextID: m.ContextID, + User: m.User, + Team: m.Team, + Chat: m.Chat, + Context: m.Context, + Config: m.Config, + }, nil + } + + // Create new Memory with independent Context namespace + forked := &Memory{ + UserID: m.UserID, + TeamID: m.TeamID, + ChatID: m.ChatID, + ContextID: newContextID, + User: m.User, // Shared + Team: m.Team, // Shared + Chat: m.Chat, // Shared + Context: nil, // Will be created below + Config: m.Config, + } + + // Create new Context namespace with independent ID + storeID := "" + if m.Config != nil { + storeID = m.Config.Context + } + ns, err := newNamespace(SpaceContext, newContextID, storeID, DefaultContextTTL) + if err != nil { + return nil, fmt.Errorf("failed to create forked context namespace: %w", err) + } + forked.Context = ns + + return forked, nil +}