Agent-System/internal/agent/piagent.go

package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"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(map[string]any{
		"workspace":    workDir,
		"model":        p.cfg.PiAgentModel,
		"provider":     p.cfg.PiAgentProvider,
		"prompt":       req.Prompt,
		"instructions": req.Instructions,
		"targets":      req.Targets,
	})

	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, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		res := &Result{
			Success:     false,
			ErrorReason: fmt.Sprintf("pi execution error: %v, stderr: %s", err, stderr.String()),
		}
		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"`
		Stderr         string `json:"stderr"`
		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)
	}

	if !bResp.OK {
		errMsg := bResp.FallbackReason
		if bResp.Stderr != "" {
			errMsg += " (" + bResp.Stderr + ")"
		}
		if errMsg == "" {
			errMsg = "pi bridge reported task failed"
		}
		res := &Result{
			Success:     false,
			ErrorReason: errMsg,
			Logs:        stdout.String(),
			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:        stdout.String(),
		}
		res.Retryable = AssessStepFailure(StepFailureSignals{
			Result:        res,
			WorkDir:       workDir,
			TargetMissing: true,
		}).Class == FailureTransient
		return res, nil
	}

	return &Result{
		Success: true,
		Patches: patches,
		Logs:    stdout.String(),
	}, nil
}

func (p *PiAgentBackend) executeRemote(ctx context.Context, runnerURL string, req Request, workDir string, timeout time.Duration) (*Result, error) {
	reqBody, _ := json.Marshal(map[string]any{
		"workspace":    workDir,
		"model":        p.cfg.PiAgentModel,
		"provider":     p.cfg.PiAgentProvider,
		"prompt":       req.Prompt,
		"instructions": req.Instructions,
		"targets":      req.Targets,
		"timeout_ms":   timeout.Milliseconds(),
	})

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

	var remoteResp struct {
		OK             bool   `json:"ok"`
		FallbackReason string `json:"fallback_reason"`
		Stderr         string `json:"stderr"`
		Retryable      bool   `json:"retryable"`
		Patches        []struct {
			RelPath string `json:"rel_path"`
			Content string `json:"content"`
		} `json:"patches"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&remoteResp); err != nil {
		return &Result{
			Success:     false,
			ErrorReason: fmt.Sprintf("failed to decode remote response: %v", err),
		}, nil
	}

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

	incoming := make(map[string][]byte, len(remoteResp.Patches))
	for _, pt := range remoteResp.Patches {
		incoming[pt.RelPath] = []byte(pt.Content)
	}
	patches, matErr := MaterializeTargets(workDir, req.Targets, incoming)
	if matErr != nil {
		res := &Result{
			Success:     false,
			ErrorReason: matErr.Error(),
			Patches:     patches,
		}
		res.Retryable = remoteResp.Retryable || AssessStepFailure(StepFailureSignals{
			Result:        res,
			WorkDir:       workDir,
			TargetMissing: true,
		}).Class == FailureTransient
		return res, nil
	}

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