Agent-System/internal/agent/piagent.go

package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"chat-agent-system/internal/config"
)

type PiAgentBackend struct {
	cfg *config.Config
}

func (p *PiAgentBackend) Name() string { return "pi-agent" }

func (p *PiAgentBackend) Execute(ctx context.Context, req Request) (*Result, error) {
	if len(req.Targets) == 0 {
		req.Targets = []string{"output.md"}
	}

	workDir := req.WorkDir
	if !filepath.IsAbs(workDir) {
		wd, _ := os.Getwd()
		workDir = filepath.Join(wd, workDir)
	}

	payload, _ := json.Marshal(p.runnerPayload(req, workDir, false, 0))

	bridgeBin := p.cfg.PiAgentBridgeBin
	if !filepath.IsAbs(bridgeBin) {
		wd, _ := os.Getwd()
		bridgeBin = filepath.Join(wd, bridgeBin)
	}

	timeout := req.Timeout
	if timeout <= 0 {
		timeout = p.cfg.AgentTimeout
	}
	runCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	if runnerURL := strings.TrimRight(strings.TrimSpace(p.cfg.PiAgentRunnerURL), "/"); runnerURL != "" {
		return p.executeRemote(runCtx, runnerURL, req, workDir, timeout)
	}

	cmd := exec.CommandContext(runCtx, bridgeBin)
	cmd.Dir = workDir
	cmd.Stdin = bytes.NewReader(payload)
	cmd.Env = append(os.Environ(),
		"PI_AGENT_BIN="+p.cfg.PiAgentBin,
		"PI_AGENT_MODEL="+p.cfg.PiAgentModel,
		"PI_AGENT_PROVIDER="+p.cfg.PiAgentProvider,
		"PI_AGENT_SYSTEM_PROMPT_FILE=AGENTS.md",
	)

	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = io.Discard

	err := cmd.Run()
	if err != nil {
		res := &Result{
			Success:     false,
			ErrorReason: "pi execution error",
		}
		res.Retryable = AssessStepFailure(StepFailureSignals{Result: res, WorkDir: workDir}).Class == FailureTransient
		return res, nil
	}

	type bridgeResponse struct {
		OK             bool   `json:"ok"`
		FallbackReason string `json:"fallback_reason"`
		Retryable      bool   `json:"retryable"`
	}
	var bResp bridgeResponse
	lines := strings.Split(strings.TrimSpace(stdout.String()), "\n")
	if len(lines) > 0 {
		_ = json.Unmarshal([]byte(lines[len(lines)-1]), &bResp)
	}
	logDigest := "stdout_digest=" + DigestBytes(stdout.Bytes())

	if !bResp.OK {
		errMsg := safeLegacyReason(bResp.FallbackReason)
		if errMsg == "" {
			errMsg = "pi bridge reported task failed"
		}
		res := &Result{
			Success:     false,
			ErrorReason: errMsg,
			Logs:        logDigest,
			Retryable:   bResp.Retryable,
		}
		if !res.Retryable {
			res.Retryable = AssessStepFailure(StepFailureSignals{Result: res, WorkDir: workDir}).Class == FailureTransient
		}
		return res, nil
	}

	patches, err := MaterializeTargets(workDir, req.Targets, nil)
	if err != nil {
		res := &Result{
			Success:     false,
			ErrorReason: err.Error(),
			Patches:     patches,
			Logs:        logDigest,
		}
		res.Retryable = AssessStepFailure(StepFailureSignals{
			Result:        res,
			WorkDir:       workDir,
			TargetMissing: true,
		}).Class == FailureTransient
		return res, nil
	}

	return &Result{
		Success: true,
		Patches: patches,
		Logs:    logDigest,
	}, nil
}

func (p *PiAgentBackend) runnerPayload(req Request, workDir string, remote bool, timeout time.Duration) map[string]any {
	profile := strings.TrimSpace(req.ToolProfile)
	if profile == "" {
		profile = ToolProfileWorkspace
	}
	body := map[string]any{
		"workspace":    workDir,
		"model":        p.cfg.PiAgentModel,
		"provider":     p.cfg.PiAgentProvider,
		"prompt":       req.Prompt,
		"instructions": req.Instructions,
		"targets":      req.Targets,
		"tool_profile": profile,
	}
	if remote {
		runID := filepath.Base(workDir)
		if runID == "" || runID == "." || runID == string(filepath.Separator) {
			runID = "unknown"
		}
		body["timeout_ms"] = timeout.Milliseconds()
		body["run_id"] = runID
		body["typed_result"] = true
	}
	return body
}

func (p *PiAgentBackend) executeRemote(ctx context.Context, runnerURL string, req Request, workDir string, timeout time.Duration) (*Result, error) {
	expectedRunID := filepath.Base(workDir)
	if expectedRunID == "" || expectedRunID == "." || expectedRunID == string(filepath.Separator) {
		expectedRunID = "unknown"
	}
	reqBody, _ := json.Marshal(p.runnerPayload(req, workDir, true, timeout))

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, runnerURL+"/v1/run", bytes.NewReader(reqBody))
	if err != nil {
		return &Result{
			Success:     false,
			ErrorReason: fmt.Sprintf("failed to build remote request: %v", err),
		}, nil
	}
	httpReq.Header.Set("Content-Type", "application/json")
	ApplyRunnerAuth(httpReq, p.cfg.PiAgentRunnerAuthToken)

	client := &http.Client{Timeout: timeout}
	resp, err := client.Do(httpReq)
	if err != nil {
		return &Result{
			Success:     false,
			ErrorReason: fmt.Sprintf("pi remote runner call failed: %v", err),
			Retryable:   true,
		}, nil
	}
	defer resp.Body.Close()
	limited, err := io.ReadAll(io.LimitReader(resp.Body, MaxResultBodyBytes+1))
	if err != nil {
		return &Result{Success: false, ErrorReason: "failed to read remote response", Retryable: true}, nil
	}
	if len(limited) > MaxResultBodyBytes {
		return &Result{Success: false, ErrorReason: "typed result rejected"}, nil
	}

	decoded, decErr := DecodeRemoteRun(limited, expectedRunID, true)
	if decErr != nil {
		return &Result{Success: false, ErrorReason: "typed result rejected"}, nil
	}

	if !decoded.OK {
		reason := decoded.Reason
		if reason == "" {
			reason = "remote pi-agent reported task failure"
		}
		res := &Result{
			Success:     false,
			ErrorReason: reason,
			Retryable:   decoded.Retryable,
		}
		if !res.Retryable {
			res.Retryable = AssessStepFailure(StepFailureSignals{Result: res, WorkDir: workDir}).Class == FailureTransient
		}
		return res, nil
	}

	patches, matErr := MaterializeTargets(workDir, req.Targets, decoded.Incoming)
	if matErr != nil {
		res := &Result{
			Success:     false,
			ErrorReason: matErr.Error(),
			Patches:     patches,
		}
		res.Retryable = decoded.Retryable || AssessStepFailure(StepFailureSignals{
			Result:        res,
			WorkDir:       workDir,
			TargetMissing: true,
		}).Class == FailureTransient
		return res, nil
	}

	return &Result{
		Success: true,
		Patches: patches,
	}, nil
}