Agent-System/internal/agent/remote_result.go

package agent

import (
	"encoding/json"
	"fmt"
	"strings"
)

type remoteRunWire struct {
	OK             bool      `json:"ok"`
	RunID          string    `json:"run_id"`
	FallbackReason string    `json:"fallback_reason"`
	Retryable      bool      `json:"retryable"`
	RunnerProtocol string    `json:"runner_protocol"`
	Capabilities   *struct {
		TypedResult bool `json:"typed_result"`
	} `json:"capabilities"`
	ResultV1 *ResultV1 `json:"result_v1"`
	Patches  []struct {
		RelPath string `json:"rel_path"`
		Content string `json:"content"`
	} `json:"patches"`
}

type decodedRemote struct {
	OK       bool
	Retryable bool
	Reason    string
	Incoming  map[string][]byte
	Typed     *ResultV1
	Legacy    bool
}

func claimsTypedCapability(w remoteRunWire) bool {
	if w.Capabilities != nil && w.Capabilities.TypedResult {
		return true
	}
	return strings.EqualFold(strings.TrimSpace(w.RunnerProtocol), RunnerProtocolTypedV1)
}

func isLegacyHandshake(w remoteRunWire) bool {
	return strings.EqualFold(strings.TrimSpace(w.RunnerProtocol), RunnerProtocolLegacy)
}

func DecodeRemoteRun(body []byte, expectedRunID string, typedRequested bool) (*decodedRemote, error) {
	if len(body) == 0 {
		return nil, fmt.Errorf("%w: empty body", ErrTypedResult)
	}
	var w remoteRunWire
	if err := json.Unmarshal(body, &w); err != nil {
		return nil, fmt.Errorf("failed to decode remote response")
	}
	incoming := make(map[string][]byte, len(w.Patches))
	for _, pt := range w.Patches {
		if len(pt.Content) > MaxResultMemberBytes {
			return nil, fmt.Errorf("%w: member size", ErrTypedResult)
		}
		incoming[pt.RelPath] = []byte(pt.Content)
	}

	out := &decodedRemote{OK: w.OK, Incoming: incoming}

	if w.ResultV1 != nil {
		if err := w.ResultV1.Validate(expectedRunID, incoming); err != nil {
			return nil, err
		}
		out.Typed = w.ResultV1
		out.Retryable = TypedRetryable(*w.ResultV1)
		if w.ResultV1.Outcome == "error" {
			out.OK = false
			out.Reason = w.ResultV1.ErrorCode
		}
		return out, nil
	}

	// typed-present is handled above. Missing typed payload:
	// handshake downgrade — claiming typed capability without result_v1 — fail closed.
	if typedRequested && (claimsTypedCapability(w) || !isLegacyHandshake(w)) {
		return nil, fmt.Errorf("%w: missing result_v1", ErrTypedResult)
	}

	out.Legacy = true
	out.Retryable = w.Retryable
	out.Reason = safeLegacyReason(w.FallbackReason)
	return out, nil
}

func safeLegacyReason(reason string) string {
	r := strings.TrimSpace(reason)
	switch {
	case r == "":
		return "remote pi-agent reported task failure"
	case strings.EqualFold(r, "prompt_empty"):
		return "PROMPT_EMPTY"
	case strings.Contains(strings.ToLower(r), "unauthorized"):
		return "UNAUTHORIZED"
	case strings.Contains(strings.ToLower(r), "timeout"):
		return "TIMEOUT"
	default:
		if strings.ContainsAny(r, "/\\") || strings.Contains(strings.ToLower(r), "token") {
			return "RUN_FAILED"
		}
		if len(r) > 64 {
			return "RUN_FAILED"
		}
		return r
	}
}