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/i18n"
	"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
	LoadingMsgID string
}

// 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
			}
		}
	}()

	if req.LoadingMsgID != "" {
		waitMsg := &message.Message{
			MessageID:   req.LoadingMsgID,
			Delta:       true,
			DeltaAction: message.DeltaReplace,
			Type:        message.TypeLoading,
			Props: map[string]any{
				"message": i18n.T(ctx.Locale, "sandbox.waiting_response"),
			},
		}
		ctx.Send(waitMsg)
	}

	var textContent []byte
	loadingClosed := false
	wrappedHandler := func(chunkType message.StreamChunkType, data []byte) int {
		if !loadingClosed && req.LoadingMsgID != "" {
			if chunkType == message.ChunkText || chunkType == message.ChunkToolCall || chunkType == message.ChunkMessageStart {
				closeLoading(ctx, req.LoadingMsgID)
				loadingClosed = true
			}
		}
		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)

	if !loadingClosed && req.LoadingMsgID != "" {
		closeLoading(ctx, req.LoadingMsgID)
	}

	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
}

func closeLoading(ctx *agentContext.Context, loadingMsgID string) {
	if loadingMsgID == "" || ctx == nil {
		return
	}
	msg := &message.Message{
		MessageID:   loadingMsgID,
		Delta:       true,
		DeltaAction: message.DeltaReplace,
		Type:        message.TypeLoading,
		Props: map[string]any{
			"done":    true,
			"message": "",
		},
	}
	ctx.Send(msg)
}