Implement Agent-to-Agent Call Enhancements and Context Management

- Introduced mechanisms to handle agent-to-agent (A2A) calls, including automatic history skipping for forked calls and proper source tracking.
- Enhanced context management with the addition of ForkParentInfo to facilitate child stack creation without race conditions.
- Updated JSAPI methods to ensure correct handling of sub-agent calls, maintaining output isolation and preventing history pollution.
- Improved documentation to clarify the behavior of A2A calls and context management in concurrent scenarios.

agent/assistant/agent.go

@@ -58,6 +58,17 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 	_, _, done := context.EnterStack(ctx, ast.ID, opts)
 	defer done()
 
+	// Auto-skip history for forked Agent-to-Agent calls (ctx.agent.Call/All/Any/Race)
+	// This ensures forked A2A messages don't pollute chat history.
+	// Delegate calls (RefererAgent) still save history as they are part of the main conversation flow.
+	// Note: Output is NOT skipped - sub-agents output normally with ThreadID for UI separation.
+	if ctx.IsForkedA2ACall() {
+		if opts == nil {
+			opts = &context.Options{}
+		}
+		opts.ForceA2A()
+	}
+
 	// ================================================
 	// Initialize Chat Buffer (for root stack only)
 	// Buffer is flushed in defer block at the end
@@ -133,6 +144,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
 
 	// Buffer user input messages (use cleaned input without overlap)
 	// Skip if History is disabled in options (for internal calls like needsearch)
+	// Note: For A2A calls, ForceA2A() sets skip.history = true, so this will be skipped
 	if opts == nil || opts.Skip == nil || !opts.Skip.History {
 		ast.BufferUserInput(ctx, historyResult.InputMessages)
 	}

agent/assistant/chat.go

@@ -254,6 +254,10 @@ func (ast *Assistant) FlushBuffer(ctx *agentcontext.Context, finalStatus string,
 			}
 		}
 	}
+
+	// 4. Close SafeWriter to flush remaining writes (root stack only)
+	// This ensures all pending SSE messages are sent before the response completes
+	ctx.CloseSafeWriter()
 }
 
 // convertBufferedMessages converts BufferedMessage slice to store Message slice

agent/assistant/next.go

@@ -51,6 +51,9 @@ func (ast *Assistant) handleDelegation(
 		return nil, fmt.Errorf("failed to load delegated assistant '%s': %w", delegate.AgentID, err)
 	}
 
+	// Mark this as an agent-to-agent call for proper source tracking
+	ctx.Referer = agentContext.RefererAgent
+
 	// Call the delegated assistant with the same context
 	// The delegated assistant's Stream method will:
 	// 1. Call EnterStack() to push itself onto the Stack (creating parent-child relationship)

agent/caller/jsapi.go

@@ -27,12 +27,13 @@ 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.
+// Note: For sub-agent calls, skip.history = true is automatically set
+// to prevent A2A messages from being saved to chat history.
+// Sub-agents output normally with ThreadID for SSE stream isolation.
 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)
+	// Force skip options for sub-agent calls
+	api.forceSkipForSubAgent(req)
 	result := api.orchestrator.callAgent(req)
 	return result
 }
@@ -75,12 +76,14 @@ 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.
+// Note: For sub-agent calls, skip.history = true is automatically set.
+// Sub-agents output normally with ThreadID. Use the handler callback
+// to receive streaming 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)
+	// Force skip options for sub-agent calls
+	api.forceSkipForSubAgent(req)
 	result := api.orchestrator.callAgent(req)
 	return result
 }
@@ -106,24 +109,31 @@ 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.
+// forceSkipForSubAgent ensures proper A2A call behavior:
+// - skip.history = true: A2A messages are not saved to chat history
+// - skip.output = false: Sub-agents output normally with ThreadID for SSE stream isolation
+//
+// IMPORTANT: skip.output is explicitly set to false to override any user settings.
+// This ensures ThreadID mechanism works correctly for concurrent sub-agent calls.
 // Users can use the onChunk callback to receive streaming messages if needed.
-func (api *JSAPI) forceSkipOutput(req *Request) {
+func (api *JSAPI) forceSkipForSubAgent(req *Request) {
 	if req.Options == nil {
 		req.Options = &CallOptions{}
 	}
 	if req.Options.Skip == nil {
 		req.Options.Skip = &agentContext.Skip{}
 	}
-	req.Options.Skip.Output = true
+	req.Options.Skip.History = true
+	// Force output to be enabled - this overrides any user settings
+	// Sub-agents MUST output with ThreadID for proper SSE stream isolation
+	req.Options.Skip.Output = false
 }
 
 // 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.
+// For all calls, this automatically sets:
+// - skip.history = true: prevents A2A messages from being saved to chat history
+// - skip.output = false: ensures sub-agents output with ThreadID (overrides user settings)
 func (api *JSAPI) parseRequestsWithHandlers(requests []interface{}, globalHandler agentContext.BatchOnMessageFunc) []*Request {
 	reqs := make([]*Request, 0, len(requests))
 
@@ -154,7 +164,7 @@ 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)
+		api.forceSkipForSubAgent(req)
 
 		// Check for per-request handler first (takes precedence)
 		if handler, ok := reqMap["_handler"].(agentContext.OnMessageFunc); ok && handler != nil {

agent/caller/orchestrator.go

@@ -241,6 +241,10 @@ func (o *Orchestrator) callAgentWithContext(ctx *agentContext.Context, req *Requ
 		return result
 	}
 
+	// Mark this as an agent-to-agent fork call for proper source tracking
+	// RefererAgentFork distinguishes ctx.agent.Call from delegate calls
+	ctx.Referer = agentContext.RefererAgentFork
+
 	// Build context options for the call
 	var ctxOpts *agentContext.Options
 	if req.Options != nil {

agent/context/context.go

@@ -112,6 +112,10 @@ func (ctx *Context) Release() {
 	// Clear current stack reference
 	ctx.Stack = nil
 
+	// Close SafeWriter if exists (must be before setting Writer to nil)
+	// This ensures the background goroutine is properly stopped
+	ctx.CloseSafeWriter()
+
 	// Clear writer reference
 	ctx.Writer = nil
 
@@ -181,8 +185,11 @@ func (ctx *Context) Fork() *Context {
 		// Child stacks will be added here by EnterStack
 		Stacks: ctx.Stacks,
 
+		// Stack is nil for forked contexts - will be set by EnterStack
+		// ForkParent stores parent stack info so EnterStack can create child stack
+		Stack: nil,
+
 		// Create independent resources to avoid race conditions
-		Stack:           nil, // Will be set by EnterStack
 		IDGenerator:     message.NewIDGenerator(),
 		Logger:          NewRequestLogger(ctx.AssistantID, ctx.ChatID, childID),
 		messageMetadata: newMessageMetadataStore(),
@@ -204,6 +211,17 @@ func (ctx *Context) Fork() *Context {
 		trace:     nil, // Trace will be inherited via TraceID in Stack
 	}
 
+	// Set ForkParent info if parent has a Stack
+	// This allows EnterStack to create a child stack instead of root stack
+	if ctx.Stack != nil {
+		child.ForkParent = &ForkParentInfo{
+			StackID: ctx.Stack.ID,
+			TraceID: ctx.Stack.TraceID,
+			Depth:   ctx.Stack.Depth,
+			Path:    append([]string{}, ctx.Stack.Path...), // Copy path slice
+		}
+	}
+
 	return child
 }
 
@@ -507,3 +525,19 @@ func (ctx *Context) shouldSkipHistory() bool {
 	}
 	return ctx.Stack.Options.Skip.History
 }
+
+// IsA2ACall returns true if this is any Agent-to-Agent call (delegate or fork)
+// A2A calls are identified by Referer being "agent" or "agent_fork":
+// - ctx.agent.Call/All/Any/Race uses RefererAgentFork (forked context, skips history)
+// - delegate uses RefererAgent (same context flow, saves history)
+func (ctx *Context) IsA2ACall() bool {
+	return ctx.Referer == RefererAgent || ctx.Referer == RefererAgentFork
+}
+
+// IsForkedA2ACall returns true if this is a forked A2A call (ctx.agent.Call/All/Any/Race)
+// Forked calls use RefererAgentFork, while delegate calls use RefererAgent.
+// This is used to skip history saving for forked sub-agent calls,
+// while allowing delegate calls to save history normally.
+func (ctx *Context) IsForkedA2ACall() bool {
+	return ctx.Referer == RefererAgentFork
+}

agent/context/output.go

@@ -516,6 +516,8 @@ func (ctx *Context) sendRaw(msg *message.Message) error {
 
 // getWriter gets the effective Writer for the current context
 // Priority: Skip.Output > Stack.Options.Writer > ctx.Writer
+// Note: The Writer returned is always a SafeWriter (wrapped at context creation)
+// to ensure thread-safe concurrent writes for SSE streaming.
 func (ctx *Context) getWriter() Writer {
 	// Check if output is explicitly skipped (for internal A2A calls)
 	if ctx.Stack != nil && ctx.Stack.Options != nil && ctx.Stack.Options.Skip != nil && ctx.Stack.Options.Skip.Output {
@@ -537,10 +539,33 @@ func (ctx *Context) getOutput() (*output.Output, error) {
 		return ctx.Stack.output, nil
 	}
 
+	// Ensure Writer is wrapped in SafeWriter for concurrent-safe SSE writes
+	// This is essential for ctx.agent.All where multiple sub-agents
+	// write to the same SSE stream concurrently.
+	// We wrap once at the context level so all forked contexts share the same SafeWriter.
+	writer := ctx.getWriter()
+	if writer != nil {
+		// Check if it's already a SafeWriter
+		if _, ok := writer.(*output.SafeWriter); !ok {
+			// Wrap in SafeWriter with context for automatic cleanup on client disconnect
+			// This prevents goroutine leaks in enterprise applications
+			var safeWriter *output.SafeWriter
+			if ctx.Context != nil {
+				// Use request context to detect client disconnection
+				safeWriter = output.NewSafeWriterWithContext(ctx.Context, writer)
+			} else {
+				// Fallback to basic SafeWriter if no context available
+				safeWriter = output.NewSafeWriter(writer)
+			}
+			ctx.Writer = safeWriter
+			writer = safeWriter
+		}
+	}
+
 	trace, _ := ctx.Trace()
 	var options message.Options = message.Options{
 		BaseURL: "/",
-		Writer:  ctx.getWriter(), // Use getWriter() to resolve Writer priority
+		Writer:  writer,
 		Trace:   trace,
 		Locale:  ctx.Locale,
 		Accept:  string(ctx.Accept),
@@ -564,3 +589,15 @@ func (ctx *Context) getOutput() (*output.Output, error) {
 
 	return out, nil
 }
+
+// CloseSafeWriter closes the SafeWriter if one was created
+// This should be called at the end of the root request to flush any pending writes
+// and stop the background goroutine.
+func (ctx *Context) CloseSafeWriter() {
+	if ctx.Writer == nil {
+		return
+	}
+	if sw, ok := ctx.Writer.(*output.SafeWriter); ok {
+		sw.Close()
+	}
+}

agent/context/stack.go

@@ -55,6 +55,32 @@ func (s *Stack) NewChildStack(assistantID, referer string, opts *Options) *Stack
 	}
 }
 
+// NewChildStackFromForkParent creates a child stack from ForkParentInfo
+// This is used by forked contexts (ctx.agent.Call) to create a child stack
+// without sharing the actual Stack reference (which would cause race conditions)
+func NewChildStackFromForkParent(parent *ForkParentInfo, assistantID, referer string, opts *Options) *Stack {
+	stackID := uuid.New().String()
+	now := time.Now().UnixMilli()
+
+	// Build path by appending parent's path with new ID
+	path := make([]string, len(parent.Path)+1)
+	copy(path, parent.Path)
+	path[len(parent.Path)] = stackID
+
+	return &Stack{
+		ID:          stackID,
+		TraceID:     parent.TraceID, // Inherit trace ID from parent
+		AssistantID: assistantID,
+		Referer:     referer,
+		Depth:       parent.Depth + 1,
+		ParentID:    parent.StackID, // Use parent's stack ID
+		Path:        path,
+		Options:     opts,
+		CreatedAt:   now,
+		Status:      StackStatusRunning,
+	}
+}
+
 // Complete marks the stack as completed and calculates duration
 func (s *Stack) Complete() {
 	now := time.Now().UnixMilli()
@@ -185,13 +211,22 @@ func EnterStack(ctx *Context, assistantID string, opts *Options) (*Stack, string
 	}
 
 	if ctx.Stack == nil {
-		// Create root stack for this assistant call (entry point)
-		// Generate a new trace ID for root
-		traceID = trace.GenTraceID()
-		stack = NewStack(traceID, assistantID, referer, opts)
-		ctx.Stack = stack
+		// Check if this is a forked context with parent stack info
+		if ctx.ForkParent != nil {
+			// Create child stack using ForkParent info
+			// This is for forked contexts (ctx.agent.Call) to have proper ThreadID
+			traceID = ctx.ForkParent.TraceID
+			stack = NewChildStackFromForkParent(ctx.ForkParent, assistantID, referer, opts)
+			ctx.Stack = stack
+		} else {
+			// Create root stack for this assistant call (entry point)
+			// Generate a new trace ID for root
+			traceID = trace.GenTraceID()
+			stack = NewStack(traceID, assistantID, referer, opts)
+			ctx.Stack = stack
+		}
 	} else {
-		// Create child stack for nested agent call
+		// Create child stack for nested agent call (delegate)
 		// Inherit trace ID from parent
 		parentStack = ctx.Stack
 		traceID = parentStack.TraceID

agent/context/types.go

@@ -62,9 +62,12 @@ const (
 	// RefererJSSDK request from JavaScript SDK
 	RefererJSSDK = "jssdk"
 
-	// RefererAgent request from agent-to-agent recursive call (assistant calling another assistant)
+	// RefererAgent request from agent-to-agent delegate call (same context, saves history)
 	RefererAgent = "agent"
 
+	// RefererAgentFork request from agent-to-agent fork call (ctx.agent.Call/All/Any/Race, skips history)
+	RefererAgentFork = "agent_fork"
+
 	// RefererTool request from tool/function execution
 	RefererTool = "tool"
 
@@ -83,16 +86,17 @@ const (
 
 // ValidReferers is the map of valid referer types
 var ValidReferers = map[string]bool{
-	RefererAPI:      true,
-	RefererProcess:  true,
-	RefererMCP:      true,
-	RefererJSSDK:    true,
-	RefererAgent:    true,
-	RefererTool:     true,
-	RefererHook:     true,
-	RefererSchedule: true,
-	RefererScript:   true,
-	RefererInternal: true,
+	RefererAPI:       true,
+	RefererProcess:   true,
+	RefererMCP:       true,
+	RefererJSSDK:     true,
+	RefererAgent:     true,
+	RefererAgentFork: true,
+	RefererTool:      true,
+	RefererHook:      true,
+	RefererSchedule:  true,
+	RefererScript:    true,
+	RefererInternal:  true,
 }
 
 const (
@@ -235,6 +239,11 @@ type Context struct {
 	IDGenerator *message.IDGenerator `json:"-"`  // ID generator for this context (chunk, message, block, thread IDs)
 	Logger      *RequestLogger       `json:"-"`  // Request-scoped async logger
 
+	// ForkParent stores parent stack info for forked contexts (set by Fork())
+	// This allows EnterStack to create a child stack instead of root stack
+	// without sharing the actual Stack reference (which would cause race conditions)
+	ForkParent *ForkParentInfo `json:"-"`
+
 	// Chat buffer for batch saving messages and resume steps
 	Buffer *ChatBuffer `json:"-"` // Chat buffer for batch saving at end of Stream()
 
@@ -315,11 +324,34 @@ type Options struct {
 	OnMessage OnMessageFunc `json:"-"`
 }
 
+// ForceA2A sets the options for Agent-to-Agent (A2A) calls.
+// For A2A calls:
+// - Output is NOT skipped - sub-agents output normally with ThreadID
+// - History IS skipped - A2A messages should not be saved to chat history
+// If Skip is nil, it creates a new Skip instance.
+func (opts *Options) ForceA2A() {
+	if opts.Skip == nil {
+		opts.Skip = &Skip{}
+	}
+	opts.Skip.History = true
+	// Note: skip.output is NOT set - sub-agents output normally with ThreadID
+}
+
 // OnMessageFunc is a callback function for receiving output messages
 // Called for each message sent via ctx.Send() - same as SSE messages to client
 // Returns: 0 = continue, non-zero = stop sending
 type OnMessageFunc func(msg *message.Message) int
 
+// ForkParentInfo stores parent stack information for forked contexts
+// This is used by EnterStack to create a child stack with proper inheritance
+// without sharing the actual Stack reference (which would cause race conditions in parallel calls)
+type ForkParentInfo struct {
+	StackID string   // Parent stack ID (used as ParentID for child stack)
+	TraceID string   // Parent trace ID (inherited by child stack)
+	Depth   int      // Parent depth (child depth = parent depth + 1)
+	Path    []string // Parent path (child path = parent path + child ID)
+}
+
 // Stack represents the call stack node for tracing agent-to-agent calls
 // Uses a flat structure to avoid circular references and memory overhead
 type Stack struct {

agent/output/adapters/cui/writer.go

@@ -18,6 +18,8 @@ type Writer struct {
 }
 
 // NewWriter creates a new CUI writer
+// The Writer should already be wrapped in SafeWriter by the context layer
+// to ensure thread-safe concurrent writes for SSE streaming.
 func NewWriter(options message.Options) (*Writer, error) {
 	return &Writer{
 		Writer:  options.Writer,
@@ -83,6 +85,7 @@ func (w *Writer) Flush() error {
 // Close closes the writer and cleans up resources
 func (w *Writer) Close() error {
 	// Nothing to clean up for CUI writer
+	// SafeWriter cleanup is handled by the context layer
 	return nil
 }
 

agent/output/safe_writer.go

@@ -0,0 +1,193 @@
+package output
+
+import (
+	"context"
+	"net/http"
+	"sync"
+)
+
+// SafeWriter wraps http.ResponseWriter with a channel-based queue
+// to serialize concurrent SSE writes and prevent "short write" errors.
+//
+// When multiple goroutines (e.g., concurrent sub-agents via ctx.agent.All)
+// write to the same SSE stream, direct writes can cause data corruption
+// or "short write" errors. SafeWriter solves this by:
+//
+// 1. Accepting write requests via a buffered channel
+// 2. Processing writes sequentially in a dedicated goroutine
+// 3. Providing non-blocking writes with overflow protection
+// 4. Automatic cleanup when context is cancelled (client disconnect)
+type SafeWriter struct {
+	ch        chan writeRequest
+	writer    http.ResponseWriter
+	done      chan struct{}
+	ctx       context.Context    // For detecting client disconnection
+	cancel    context.CancelFunc // To signal run() to stop
+	closeOnce sync.Once
+	closed    bool
+	mu        sync.RWMutex
+}
+
+// writeRequest represents a single write request
+type writeRequest struct {
+	data []byte
+}
+
+// QueueCapacity is the default buffer size for the write queue
+// Large enough to handle high concurrency without blocking
+const QueueCapacity = 10000
+
+// NewSafeWriter creates a new SafeWriter that wraps an http.ResponseWriter
+// and starts a background goroutine to process writes sequentially.
+// The context should be the HTTP request context to detect client disconnection.
+func NewSafeWriter(w http.ResponseWriter) *SafeWriter {
+	// Create internal context for graceful shutdown
+	ctx, cancel := context.WithCancel(context.Background())
+	sw := &SafeWriter{
+		ch:     make(chan writeRequest, QueueCapacity),
+		writer: w,
+		done:   make(chan struct{}),
+		ctx:    ctx,
+		cancel: cancel,
+	}
+	go sw.run()
+	return sw
+}
+
+// NewSafeWriterWithContext creates a SafeWriter that respects the given context.
+// When the context is cancelled (e.g., client disconnects), the run() goroutine exits.
+// This prevents goroutine leaks in enterprise applications with many concurrent requests.
+func NewSafeWriterWithContext(ctx context.Context, w http.ResponseWriter) *SafeWriter {
+	// Derive a cancellable context from the parent
+	childCtx, cancel := context.WithCancel(ctx)
+	sw := &SafeWriter{
+		ch:     make(chan writeRequest, QueueCapacity),
+		writer: w,
+		done:   make(chan struct{}),
+		ctx:    childCtx,
+		cancel: cancel,
+	}
+	go sw.run()
+	return sw
+}
+
+// run processes write requests from the channel sequentially
+// Exits when channel is closed OR context is cancelled (client disconnect)
+func (sw *SafeWriter) run() {
+	defer close(sw.done)
+
+	for {
+		select {
+		case req, ok := <-sw.ch:
+			if !ok {
+				// Channel closed, exit gracefully
+				return
+			}
+			if sw.writer != nil {
+				sw.writer.Write(req.data)
+				// Flush after each write to ensure SSE data is sent immediately
+				if flusher, ok := sw.writer.(http.Flusher); ok {
+					flusher.Flush()
+				}
+			}
+		case <-sw.ctx.Done():
+			// Context cancelled (client disconnected or explicit close)
+			// Continue reading from channel until it's closed to avoid blocking senders
+			// and to process any remaining messages that were already queued
+			sw.drainUntilClosed()
+			return
+		}
+	}
+}
+
+// drainUntilClosed reads from channel until it's closed
+// This prevents senders from blocking after context cancellation
+func (sw *SafeWriter) drainUntilClosed() {
+	for range sw.ch {
+		// Discard messages - context is cancelled so we don't write them
+	}
+}
+
+// Write implements io.Writer interface
+// Queues the data for sequential writing by the background goroutine
+func (sw *SafeWriter) Write(data []byte) (int, error) {
+	sw.mu.RLock()
+	if sw.closed {
+		sw.mu.RUnlock()
+		return 0, nil // Silently ignore writes after close
+	}
+	sw.mu.RUnlock()
+
+	// Make a copy of data since the caller may reuse the buffer
+	dataCopy := make([]byte, len(data))
+	copy(dataCopy, data)
+
+	// Non-blocking send with overflow protection
+	select {
+	case sw.ch <- writeRequest{data: dataCopy}:
+		return len(data), nil
+	default:
+		// Channel full - this shouldn't happen with 10000 capacity
+		// but if it does, drop the message rather than block
+		// Note: In production, this indicates either:
+		// 1. Extremely high concurrency (>10000 pending writes)
+		// 2. The underlying writer is blocked/slow
+		// Consider increasing QueueCapacity if this occurs frequently
+		return len(data), nil
+	}
+}
+
+// Header returns the header map from the underlying ResponseWriter
+func (sw *SafeWriter) Header() http.Header {
+	if sw.writer == nil {
+		return http.Header{}
+	}
+	return sw.writer.Header()
+}
+
+// WriteHeader sends an HTTP response header with the provided status code
+func (sw *SafeWriter) WriteHeader(statusCode int) {
+	if sw.writer != nil {
+		sw.writer.WriteHeader(statusCode)
+	}
+}
+
+// Flush implements http.Flusher interface
+// Note: Actual flushing happens in the run() goroutine after each write
+func (sw *SafeWriter) Flush() {
+	// Flushing is handled automatically in run() after each write
+	// This method exists to satisfy the http.Flusher interface
+}
+
+// Close closes the write channel and waits for all pending writes to complete
+// This is safe to call multiple times (idempotent via sync.Once)
+func (sw *SafeWriter) Close() error {
+	sw.closeOnce.Do(func() {
+		// First close channel to signal run() to stop and process remaining messages
+		close(sw.ch)
+
+		// Wait for run() to finish processing all queued messages
+		<-sw.done
+
+		// Then mark as closed and cancel context
+		sw.mu.Lock()
+		sw.closed = true
+		sw.mu.Unlock()
+
+		sw.cancel()
+	})
+	return nil
+}
+
+// IsClosed returns whether the SafeWriter has been closed
+func (sw *SafeWriter) IsClosed() bool {
+	sw.mu.RLock()
+	defer sw.mu.RUnlock()
+	return sw.closed
+}
+
+// Underlying returns the underlying http.ResponseWriter
+// Use with caution - direct writes bypass the queue
+func (sw *SafeWriter) Underlying() http.ResponseWriter {
+	return sw.writer
+}

agent/output/safe_writer_test.go

@@ -0,0 +1,401 @@
+package output
+
+import (
+	"bytes"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+	"time"
+)
+
+// mockResponseWriter is a thread-safe mock for testing
+type mockResponseWriter struct {
+	mu       sync.Mutex
+	buf      bytes.Buffer
+	header   http.Header
+	flushed  int
+	writeErr error
+}
+
+func newMockResponseWriter() *mockResponseWriter {
+	return &mockResponseWriter{
+		header: make(http.Header),
+	}
+}
+
+func (m *mockResponseWriter) Header() http.Header {
+	return m.header
+}
+
+func (m *mockResponseWriter) Write(data []byte) (int, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	if m.writeErr != nil {
+		return 0, m.writeErr
+	}
+	return m.buf.Write(data)
+}
+
+func (m *mockResponseWriter) WriteHeader(statusCode int) {}
+
+func (m *mockResponseWriter) Flush() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.flushed++
+}
+
+func (m *mockResponseWriter) String() string {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.buf.String()
+}
+
+func (m *mockResponseWriter) FlushCount() int {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.flushed
+}
+
+func TestSafeWriter_BasicWrite(t *testing.T) {
+	mock := newMockResponseWriter()
+	sw := NewSafeWriter(mock)
+	defer sw.Close()
+
+	// Write some data
+	n, err := sw.Write([]byte("hello"))
+	if err != nil {
+		t.Errorf("Write error: %v", err)
+	}
+	if n != 5 {
+		t.Errorf("Expected 5 bytes written, got %d", n)
+	}
+
+	// Wait for async write to complete
+	time.Sleep(10 * time.Millisecond)
+
+	// Verify data was written
+	if got := mock.String(); got != "hello" {
+		t.Errorf("Expected 'hello', got '%s'", got)
+	}
+}
+
+func TestSafeWriter_ConcurrentWrites(t *testing.T) {
+	mock := newMockResponseWriter()
+	sw := NewSafeWriter(mock)
+
+	// Number of concurrent goroutines
+	numGoroutines := 100
+	// Number of writes per goroutine
+	writesPerGoroutine := 100
+
+	var wg sync.WaitGroup
+	wg.Add(numGoroutines)
+
+	// Launch concurrent writes
+	for i := 0; i < numGoroutines; i++ {
+		go func(id int) {
+			defer wg.Done()
+			for j := 0; j < writesPerGoroutine; j++ {
+				sw.Write([]byte("X"))
+			}
+		}(i)
+	}
+
+	// Wait for all goroutines to complete
+	wg.Wait()
+
+	// Close and wait for all writes to be processed
+	sw.Close()
+
+	// Verify all data was written (no data loss)
+	expectedLen := numGoroutines * writesPerGoroutine
+	if got := len(mock.String()); got != expectedLen {
+		t.Errorf("Expected %d bytes, got %d", expectedLen, got)
+	}
+
+	// Verify flush was called (at least once per write)
+	if mock.FlushCount() < expectedLen {
+		t.Errorf("Expected at least %d flushes, got %d", expectedLen, mock.FlushCount())
+	}
+}
+
+func TestSafeWriter_NoDataCorruption(t *testing.T) {
+	mock := newMockResponseWriter()
+	sw := NewSafeWriter(mock)
+
+	// Use exactly 26 goroutines (one per letter A-Z) to avoid duplicates
+	numGoroutines := 26
+	// Message to write (with unique content per goroutine)
+	msgLen := 100
+
+	var wg sync.WaitGroup
+	wg.Add(numGoroutines)
+
+	// Launch concurrent writes with different content
+	for i := 0; i < numGoroutines; i++ {
+		go func(id int) {
+			defer wg.Done()
+			// Create a message with repeating character (unique per goroutine)
+			char := byte('A' + id)
+			msg := bytes.Repeat([]byte{char}, msgLen)
+			sw.Write(msg)
+		}(i)
+	}
+
+	// Wait for all goroutines to complete
+	wg.Wait()
+
+	// Close and wait for all writes to be processed
+	sw.Close()
+
+	// Verify total length
+	result := mock.String()
+	expectedLen := numGoroutines * msgLen
+	if len(result) != expectedLen {
+		t.Errorf("Expected %d bytes, got %d", expectedLen, len(result))
+	}
+
+	// Verify no interleaving (each message should be contiguous)
+	// Check that we have exactly numGoroutines distinct blocks
+	blocks := make(map[byte]int)
+	for i := 0; i < len(result); i += msgLen {
+		if i+msgLen > len(result) {
+			t.Errorf("Unexpected data at end of result")
+			break
+		}
+		block := result[i : i+msgLen]
+		// Verify block is homogeneous (all same character)
+		firstChar := block[0]
+		for j, c := range []byte(block) {
+			if c != firstChar {
+				t.Errorf("Data corruption detected at position %d: expected %c, got %c", i+j, firstChar, c)
+				break
+			}
+		}
+		blocks[firstChar]++
+	}
+
+	// Each character should appear exactly once (one block per goroutine)
+	for char, count := range blocks {
+		if count != 1 {
+			t.Errorf("Character %c appeared %d times, expected 1", char, count)
+		}
+	}
+
+	// Verify we got all 26 letters
+	if len(blocks) != numGoroutines {
+		t.Errorf("Expected %d distinct blocks, got %d", numGoroutines, len(blocks))
+	}
+}
+
+func TestSafeWriter_CloseWaitsForPendingWrites(t *testing.T) {
+	mock := newMockResponseWriter()
+	sw := NewSafeWriter(mock)
+
+	// Write a large number of messages
+	numWrites := 1000
+	for i := 0; i < numWrites; i++ {
+		sw.Write([]byte("X"))
+	}
+
+	// Close should wait for all writes to complete
+	sw.Close()
+
+	// Verify all data was written
+	if got := len(mock.String()); got != numWrites {
+		t.Errorf("Expected %d bytes after close, got %d", numWrites, got)
+	}
+}
+
+func TestSafeWriter_WriteAfterClose(t *testing.T) {
+	mock := newMockResponseWriter()
+	sw := NewSafeWriter(mock)
+
+	sw.Write([]byte("before"))
+	sw.Close()
+
+	// Write after close should be silently ignored
+	n, err := sw.Write([]byte("after"))
+	if err != nil {
+		t.Errorf("Write after close should not error: %v", err)
+	}
+	if n != 0 {
+		t.Errorf("Write after close should return 0, got %d", n)
+	}
+
+	// Verify only "before" was written
+	if got := mock.String(); got != "before" {
+		t.Errorf("Expected 'before', got '%s'", got)
+	}
+}
+
+func TestSafeWriter_ImplementsHTTPInterfaces(t *testing.T) {
+	mock := newMockResponseWriter()
+	sw := NewSafeWriter(mock)
+	defer sw.Close()
+
+	// Verify it implements http.ResponseWriter
+	var _ http.ResponseWriter = sw
+
+	// Verify it implements http.Flusher
+	var _ http.Flusher = sw
+
+	// Test Header()
+	sw.Header().Set("Content-Type", "text/plain")
+	if got := mock.Header().Get("Content-Type"); got != "text/plain" {
+		t.Errorf("Expected Content-Type 'text/plain', got '%s'", got)
+	}
+}
+
+// BenchmarkSafeWriter_ConcurrentWrites benchmarks concurrent write performance
+func BenchmarkSafeWriter_ConcurrentWrites(b *testing.B) {
+	mock := newMockResponseWriter()
+	sw := NewSafeWriter(mock)
+	defer sw.Close()
+
+	data := []byte("benchmark data for SSE streaming")
+
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			sw.Write(data)
+		}
+	})
+}
+
+// TestSafeWriter_RealHTTPServer tests SafeWriter with a real HTTP server
+func TestSafeWriter_RealHTTPServer(t *testing.T) {
+	// This test verifies SafeWriter works correctly with httptest.ResponseRecorder
+	// which is commonly used in testing HTTP handlers
+
+	recorder := httptest.NewRecorder()
+	sw := NewSafeWriter(recorder)
+
+	// Simulate concurrent SSE writes from multiple sub-agents
+	var wg sync.WaitGroup
+	numAgents := 10
+	messagesPerAgent := 10
+
+	wg.Add(numAgents)
+	for i := 0; i < numAgents; i++ {
+		go func(agentID int) {
+			defer wg.Done()
+			for j := 0; j < messagesPerAgent; j++ {
+				// Simulate SSE message format
+				msg := []byte("data: {\"agent\":" + string(rune('0'+agentID)) + "}\n\n")
+				sw.Write(msg)
+			}
+		}(i)
+	}
+
+	wg.Wait()
+	sw.Close()
+
+	// Verify response contains all messages (no data loss)
+	body := recorder.Body.String()
+	expectedMsgs := numAgents * messagesPerAgent
+
+	// Count number of "data: " prefixes
+	count := 0
+	for i := 0; i < len(body); i++ {
+		if i+6 <= len(body) && body[i:i+6] == "data: " {
+			count++
+		}
+	}
+
+	if count != expectedMsgs {
+		t.Errorf("Expected %d messages, found %d", expectedMsgs, count)
+	}
+}
+
+// TestSafeWriter_ContextCancellation tests that SafeWriter handles context cancellation
+// This is critical for enterprise applications to prevent goroutine leaks
+func TestSafeWriter_ContextCancellation(t *testing.T) {
+	mock := newMockResponseWriter()
+
+	// Create a cancellable context
+	ctx, cancel := context.WithCancel(context.Background())
+
+	sw := NewSafeWriterWithContext(ctx, mock)
+
+	// Write some data and wait for it to be processed
+	sw.Write([]byte("before"))
+	time.Sleep(20 * time.Millisecond)
+
+	// Verify "before" was written
+	if got := mock.String(); got != "before" {
+		t.Errorf("Expected 'before' before cancel, got '%s'", got)
+	}
+
+	// Cancel context (simulates client disconnect)
+	cancel()
+
+	// Write after context cancellation - these may or may not be written
+	// depending on timing (select may pick ctx.Done() first)
+	sw.Write([]byte("after_cancel"))
+
+	// Close properly cleans up
+	sw.Close()
+
+	// After close, run() has exited
+	select {
+	case <-sw.done:
+		// Good - run() has exited
+	default:
+		t.Error("run() should have exited after Close()")
+	}
+
+	// The key guarantee: run() goroutine exits cleanly, no leak
+	// Data written before cancel is preserved
+	got := mock.String()
+	if len(got) < 6 { // At least "before" should be there
+		t.Errorf("Expected at least 'before', got '%s'", got)
+	}
+}
+
+// TestSafeWriter_GoroutineLeak tests that SafeWriter doesn't leak goroutines
+func TestSafeWriter_GoroutineLeak(t *testing.T) {
+	// Create many SafeWriters and ensure they all clean up properly
+	numWriters := 100
+
+	var wg sync.WaitGroup
+	wg.Add(numWriters)
+
+	for i := 0; i < numWriters; i++ {
+		go func() {
+			defer wg.Done()
+
+			mock := newMockResponseWriter()
+			ctx, cancel := context.WithCancel(context.Background())
+			sw := NewSafeWriterWithContext(ctx, mock)
+
+			// Write some data
+			sw.Write([]byte("test"))
+
+			// Randomly either close normally or cancel context
+			if time.Now().UnixNano()%2 == 0 {
+				cancel()
+				time.Sleep(5 * time.Millisecond)
+				sw.Close()
+			} else {
+				sw.Close()
+				cancel() // Cancel after close is safe
+			}
+		}()
+	}
+
+	// All goroutines should complete
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// All completed successfully
+	case <-time.After(5 * time.Second):
+		t.Error("Timeout waiting for goroutines to complete - possible goroutine leak")
+	}
+}