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

package openai

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

	"google.golang.org/protobuf/proto"
	edgeservice "iop/apps/edge/internal/service"
	iop "iop/proto/gen/iop"
)

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

// openAIRunTerminalError retains the typed terminal failure only inside the
// OpenAI host. Its public error text is deliberately stable: callers must not
// receive provider text or arbitrary Node metadata through a buffered path.
type openAIRunTerminalError struct {
	failure *iop.ExecutionFailure
}

func newOpenAIRunTerminalError(event *iop.RunEvent) error {
	if event == nil || event.GetFailure() == nil {
		if event != nil {
			if message := event.GetError(); message != "" {
				return errors.New(message)
			}
			if message := event.GetMessage(); message != "" {
				return errors.New(message)
			}
		}
		return errors.New("run failed")
	}
	failure, ok := proto.Clone(event.GetFailure()).(*iop.ExecutionFailure)
	if !ok || failure == nil {
		return errors.New("run failed")
	}
	return &openAIRunTerminalError{failure: failure}
}

func (e *openAIRunTerminalError) Error() string { return "run failed" }

func (e *openAIRunTerminalError) executionFailure() *iop.ExecutionFailure {
	if e == nil || e.failure == nil {
		return nil
	}
	failure, _ := proto.Clone(e.failure).(*iop.ExecutionFailure)
	return failure
}

// 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()),
						ReasoningTokens:   int(u.GetReasoningTokens()),
						CachedInputTokens: int(u.GetCachedInputTokens()),
					}
				}
				return contentBuilder.String(), reasoningBuilder.String(), finishReason, toolCalls, usage, isTextToolFallback(event.GetMetadata()), nil
			case "error", "cancelled":
				return "", "", "", nil, nil, false, newOpenAIRunTerminalError(event)
			}
		}
	}
}

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

func extractToolCallNames(toolCalls []any) []string {
	var names []string
	for _, tc := range toolCalls {
		if tc == nil {
			continue
		}
		b, err := json.Marshal(tc)
		if err != nil {
			continue
		}
		var parsed struct {
			Function struct {
				Name string `json:"name"`
			} `json:"function"`
		}
		if err := json.Unmarshal(b, &parsed); err == nil && parsed.Function.Name != "" {
			names = append(names, parsed.Function.Name)
		}
	}
	return names
}