yao/agent/sandbox/v2/stream.go

package sandboxv2

import (
	"context"
	"errors"
	"fmt"
	"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 runs runner.Stream and bridges agentContext interrupts.
//
// Cleanup (runner.Cleanup + LifecycleAction) is NOT performed here; the caller
// (agent.go sandboxCleanup closure) is responsible for all lifecycle management
// so that cleanup happens exactly once regardless of code path.
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

	// 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.ChunkExecute || 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)
	}

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