iop/apps/edge/internal/openai/run_result.go

package openai

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	edgeservice "iop/apps/edge/internal/service"
)

const (
	runtimeMetadataOpenAIToolCalls        = "openai_tool_calls"
	runtimeMetadataOpenAITextToolFallback = "openai_text_tool_fallback"
)

// errRunTimedOut is returned by collectRunResult and the direct streaming
// loop when handle.WaitTimeout() elapses without a terminal run event.
var errRunTimedOut = errors.New("run timed out")

// isCancelWorthyRunError reports whether err means the HTTP caller gave up
// (context cancellation/deadline or a WaitTimeout expiry) before the run
// reached a terminal state, so Edge should propagate CancelRun to Node.
// Terminal run outcomes (complete/error/cancelled events, closed stream,
// node disconnect) are not cancel-worthy: Node already reached a terminal
// state or is unreachable, so sending CancelRun would be redundant or futile.
func isCancelWorthyRunError(err error) bool {
	return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) || errors.Is(err, errRunTimedOut)
}

func collectRunResult(ctx context.Context, stream edgeservice.RunStream, timeout time.Duration) (string, string, string, []any, *openAIUsage, bool, error) {
	if stream.Events == nil {
		return "", "", "", nil, nil, false, fmt.Errorf("run stream unavailable")
	}
	var contentBuilder strings.Builder
	var reasoningBuilder strings.Builder
	finishReason := "stop"
	var usage *openAIUsage
	timer := time.NewTimer(timeout)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return "", "", "", nil, nil, false, ctx.Err()
		case <-timer.C:
			return "", "", "", nil, nil, false, errRunTimedOut
		case nodeEvent, ok := <-stream.NodeEvents:
			if !ok {
				stream.NodeEvents = nil
				continue
			}
			if edgeservice.IsNodeDisconnected(nodeEvent) {
				return "", "", "", nil, nil, false, fmt.Errorf("node disconnected")
			}
		case event, ok := <-stream.Events:
			if !ok {
				return "", "", "", nil, nil, false, fmt.Errorf("run stream closed")
			}
			if event == nil {
				continue
			}
			switch event.GetType() {
			case "delta":
				contentBuilder.WriteString(event.GetDelta())
			case "reasoning_delta":
				reasoningBuilder.WriteString(event.GetDelta())
			case "complete":
				if reason := event.GetMetadata()["finish_reason"]; reason != "" {
					finishReason = reason
				}
				toolCalls := toolCallsFromRunMetadata(event.GetMetadata())
				if len(toolCalls) > 0 {
					finishReason = "tool_calls"
				}
				if u := event.GetUsage(); u != nil {
					usage = &openAIUsage{
						PromptTokens:     int(u.GetInputTokens()),
						CompletionTokens: int(u.GetOutputTokens()),
						TotalTokens:      int(u.GetInputTokens() + u.GetOutputTokens()),
					}
				}
				return contentBuilder.String(), reasoningBuilder.String(), finishReason, toolCalls, usage, isTextToolFallback(event.GetMetadata()), nil
			case "error", "cancelled":
				msg := event.GetError()
				if msg == "" {
					msg = event.GetMessage()
				}
				if msg == "" {
					msg = "run failed"
				}
				return "", "", "", nil, nil, false, fmt.Errorf("%s", msg)
			}
		}
	}
}

func toolCallsFromRunMetadata(metadata map[string]string) []any {
	raw := strings.TrimSpace(metadata[runtimeMetadataOpenAIToolCalls])
	if raw == "" {
		return nil
	}
	var toolCalls []any
	if err := json.Unmarshal([]byte(raw), &toolCalls); err != nil {
		return nil
	}
	if len(toolCalls) == 0 {
		return nil
	}
	return toolCalls
}

func isTextToolFallback(metadata map[string]string) bool {
	return strings.EqualFold(strings.TrimSpace(metadata[runtimeMetadataOpenAITextToolFallback]), "true")
}