yao/agent/sandbox/v2/stream.go

package sandboxv2

import (
	"context"
	"errors"
	"fmt"
	"log"
	"time"

	agentContext "github.com/yaoapp/yao/agent/context"
	"github.com/yaoapp/yao/agent/output/message"
	"github.com/yaoapp/yao/agent/sandbox/v2/types"
	infra "github.com/yaoapp/yao/sandbox/v2"
)

// ExecuteRequest consolidates all parameters for ExecuteSandboxStream.
type ExecuteRequest struct {
	Computer  infra.Computer
	Runner    types.Runner
	Config    *types.SandboxConfig
	StreamReq *types.StreamRequest
	Manager   *infra.Manager
}

// ExecuteSandboxStream is the V2 replacement for executeSandboxStream.
// It calls runner.Stream, handles interrupts, and performs cleanup/lifecycle
// in defer.
func ExecuteSandboxStream(
	ctx *agentContext.Context,
	req *ExecuteRequest,
	handler message.StreamFunc,
) (*agentContext.CompletionResponse, error) {

	if req.Runner == nil || req.Computer == nil {
		return nil, fmt.Errorf("runner and computer are required")
	}

	stdCtx := ctx.Context
	panicked := true // Assume panic; set false on normal exit.

	// Resolve stop timeout from config (default 2s).
	stopTimeout := 2 * time.Second
	if req.Config != nil && req.Config.StopTimeout != "" {
		if d, err := time.ParseDuration(req.Config.StopTimeout); err == nil {
			stopTimeout = d
		}
	}

	// Panic recovery (registered first, executes last in LIFO order).
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[sandbox/v2] panic in stream: %v", r)
			cleanCtx, cancel := context.WithTimeout(context.Background(), stopTimeout)
			defer cancel()
			req.Runner.Cleanup(cleanCtx, req.Computer)
			LifecycleAction(cleanCtx, req.Config, req.Computer, req.Manager)
		}
	}()

	// Lifecycle action (registered second, executes second-to-last).
	defer func() {
		if !panicked {
			LifecycleAction(stdCtx, req.Config, req.Computer, req.Manager)
		}
	}()

	// Runner cleanup (registered last, executes first).
	defer func() {
		if !panicked {
			cleanCtx, cancel := context.WithTimeout(context.Background(), stopTimeout)
			defer cancel()
			req.Runner.Cleanup(cleanCtx, req.Computer)
		}
	}()

	// Build a cancellable runnerCtx that bridges agentContext interrupts.
	runnerCtx, cancelRunner := context.WithCancel(stdCtx)
	defer cancelRunner() // Prevent goroutine leak.

	done := make(chan struct{})
	defer close(done)

	go func() {
		ticker := time.NewTicker(500 * time.Millisecond)
		defer ticker.Stop()
		for {
			select {
			case <-done:
				return
			case <-ticker.C:
				if ctx.Interrupt != nil {
					if sig := ctx.Interrupt.Peek(); sig != nil {
						cancelRunner()
						return
					}
					if ctx.Interrupt.IsInterrupted() {
						cancelRunner()
						return
					}
				}
			case <-stdCtx.Done():
				cancelRunner()
				return
			}
		}
	}()

	var textContent []byte
	wrappedHandler := func(chunkType message.StreamChunkType, data []byte) int {
		if chunkType == message.ChunkText {
			textContent = append(textContent, data...)
		}
		if handler != nil {
			return handler(chunkType, data)
		}
		return 0
	}

	err := req.Runner.Stream(runnerCtx, req.StreamReq, wrappedHandler)

	panicked = false // Normal exit reached.

	if err != nil {
		if errors.Is(err, context.Canceled) {
			return nil, err
		}
		return nil, fmt.Errorf("runner.Stream: %w", err)
	}

	resp := &agentContext.CompletionResponse{
		Role:         "assistant",
		FinishReason: agentContext.FinishReasonStop,
	}
	if len(textContent) > 0 {
		resp.Content = string(textContent)
	}
	return resp, nil
}