309 lines
8.6 KiB
Go
309 lines
8.6 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"chat-agent-system/internal/config"
|
|
"chat-agent-system/internal/workspace"
|
|
)
|
|
|
|
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{workspace.DeliverableKey}
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
// runnerDataRoots caches a successful /health data_dir per runner URL.
|
|
// A failed probe is not cached: the next call may try again, and this call skips the check.
|
|
var runnerDataRoots sync.Map
|
|
|
|
func probeRunnerDataRoot(ctx context.Context, runnerURL string) (string, bool) {
|
|
if v, ok := runnerDataRoots.Load(runnerURL); ok {
|
|
root, _ := v.(string)
|
|
return root, root != ""
|
|
}
|
|
probeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
|
|
defer cancel()
|
|
httpReq, err := http.NewRequestWithContext(probeCtx, http.MethodGet, runnerURL+"/health", nil)
|
|
if err != nil {
|
|
return "", false
|
|
}
|
|
client := &http.Client{Timeout: 2 * time.Second}
|
|
resp, err := client.Do(httpReq)
|
|
if err != nil {
|
|
return "", false
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
return "", false
|
|
}
|
|
var body struct {
|
|
DataDir string `json:"data_dir"`
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&body); err != nil {
|
|
return "", false
|
|
}
|
|
root := filepath.Clean(strings.TrimSpace(body.DataDir))
|
|
if root == "" || root == "." {
|
|
return "", false
|
|
}
|
|
runnerDataRoots.Store(runnerURL, root)
|
|
log.Printf("pi runner data_dir check passed")
|
|
return root, true
|
|
}
|
|
|
|
func workDirWithinRunnerRoot(workDir, root string) bool {
|
|
rel, err := filepath.Rel(filepath.Clean(root), filepath.Clean(workDir))
|
|
if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (p *PiAgentBackend) executeRemote(ctx context.Context, runnerURL string, req Request, workDir string, timeout time.Duration) (*Result, error) {
|
|
if root, ok := probeRunnerDataRoot(ctx, runnerURL); ok && !workDirWithinRunnerRoot(workDir, root) {
|
|
// Do not include either path: the caller classifies this phrase, and paths leak host layout.
|
|
log.Printf("pi runner data root mismatch; refusing /v1/run")
|
|
return &Result{Success: false, ErrorReason: "pi runner data root mismatch"}, nil
|
|
}
|
|
|
|
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: "runner result body too large"}, nil
|
|
}
|
|
|
|
decoded, decErr := DecodeRemoteRun(limited, expectedRunID, true)
|
|
if decErr != nil {
|
|
reason := "typed result rejected"
|
|
msg := decErr.Error()
|
|
switch {
|
|
case strings.Contains(msg, "schema"):
|
|
reason = "typed result rejected: schema"
|
|
case strings.Contains(msg, "run mismatch"):
|
|
reason = "typed result rejected: run mismatch"
|
|
case strings.Contains(msg, "digest"):
|
|
reason = "typed result rejected: digest"
|
|
case strings.Contains(msg, "missing result_v1"):
|
|
reason = "typed result rejected: missing result_v1"
|
|
case strings.Contains(msg, "member"):
|
|
reason = "typed result rejected: member"
|
|
case strings.Contains(msg, "outcome"), strings.Contains(msg, "error_code"):
|
|
reason = "typed result rejected: outcome"
|
|
case strings.Contains(msg, "empty body"):
|
|
reason = "typed result rejected: empty body"
|
|
case strings.Contains(msg, "failed to decode"):
|
|
reason = "typed result rejected: decode"
|
|
}
|
|
return &Result{Success: false, ErrorReason: reason}, 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
|
|
}
|