diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index 4dd3acda..b92c5397 100644 --- a/agent/assistant/agent.go +++ b/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) } diff --git a/agent/assistant/chat.go b/agent/assistant/chat.go index 9feceab7..63e4d8c6 100644 --- a/agent/assistant/chat.go +++ b/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 diff --git a/agent/assistant/next.go b/agent/assistant/next.go index ff2d61f3..1e8aa348 100644 --- a/agent/assistant/next.go +++ b/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) diff --git a/agent/caller/jsapi.go b/agent/caller/jsapi.go index b9077ba8..dcc3c476 100644 --- a/agent/caller/jsapi.go +++ b/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 { diff --git a/agent/caller/orchestrator.go b/agent/caller/orchestrator.go index 1865ef19..5c9958db 100644 --- a/agent/caller/orchestrator.go +++ b/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 { diff --git a/agent/context/context.go b/agent/context/context.go index 98cc3e20..ceac1610 100644 --- a/agent/context/context.go +++ b/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 +} diff --git a/agent/context/output.go b/agent/context/output.go index 1223ea07..7c85f864 100644 --- a/agent/context/output.go +++ b/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() + } +} diff --git a/agent/context/stack.go b/agent/context/stack.go index de816768..5e95e316 100644 --- a/agent/context/stack.go +++ b/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 diff --git a/agent/context/types.go b/agent/context/types.go index 215c0612..6c2b9010 100644 --- a/agent/context/types.go +++ b/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 { diff --git a/agent/output/adapters/cui/writer.go b/agent/output/adapters/cui/writer.go index 07f2088f..343c374c 100644 --- a/agent/output/adapters/cui/writer.go +++ b/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 } diff --git a/agent/output/safe_writer.go b/agent/output/safe_writer.go new file mode 100644 index 00000000..c415a643 --- /dev/null +++ b/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 +} diff --git a/agent/output/safe_writer_test.go b/agent/output/safe_writer_test.go new file mode 100644 index 00000000..3007a8a1 --- /dev/null +++ b/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") + } +}