@pi-in-go/pigpen-a2a 0.1.0

package/extensions/a2a/survivors_test.go

@@ -0,0 +1,302 @@
+package a2aext
+
+// Tests added after the first mutation run (pigeq mutate --unit): each one kills a mutant that the
+// first suite let survive. The mutant's name is in the test's name or comment.
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"runtime"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/a2aproject/a2a-go/v2/a2a"
+	"github.com/a2aproject/a2a-go/v2/a2asrv"
+)
+
+func TestTokenEnvNotSetIsSaidPlainly(t *testing.T) { // missing-token-env-allowed
+	dir := t.TempDir()
+	writeConfig(t, dir, `{"listen":"127.0.0.1:5555","tokens":[{"name":"a","tokenEnv":"A2A_TOKEN_A"}]}`)
+	_, err := LoadConfig(LoadOptions{ConfigHome: dir, Getenv: envFrom(nil)})
+	if err == nil || !strings.Contains(err.Error(), "not set") {
+		t.Fatalf("%v", err)
+	}
+}
+
+func TestTwoVariablesWithOneValueAreRejected(t *testing.T) { // duplicate-token-values-allowed
+	dir := t.TempDir()
+	writeConfig(t, dir, `{"listen":"127.0.0.1:1","tokens":[{"name":"a","tokenEnv":"V1"},{"name":"b","tokenEnv":"V2"}]}`)
+	same := "0123456789abcdef0123456789abcdef"
+	_, err := LoadConfig(LoadOptions{ConfigHome: dir, Getenv: envFrom(map[string]string{"V1": same, "V2": same})})
+	if err == nil || !strings.Contains(err.Error(), "same value") {
+		t.Fatalf("two callers with one token cannot be told apart: %v", err)
+	}
+}
+
+func TestInsecureNoAuthWithTokensIsAmbiguous(t *testing.T) { // insecure-and-tokens-both
+	dir := t.TempDir()
+	writeConfig(t, dir, `{"listen":"127.0.0.1:1","insecureNoAuth":true,"tokens":[{"name":"a","tokenEnv":"T"}]}`)
+	_, err := LoadConfig(LoadOptions{ConfigHome: dir, Getenv: envFrom(map[string]string{"T": "0123456789abcdef0123456789abcdef"})})
+	if err == nil || !strings.Contains(err.Error(), "insecureNoAuth") {
+		t.Fatalf("%v", err)
+	}
+}
+
+func TestMalformedListenAddressSaysHostPort(t *testing.T) { // listen-not-host-port
+	dir := t.TempDir()
+	writeConfig(t, dir, `{"listen":"not-an-address","insecureNoAuth":true}`)
+	_, err := LoadConfig(LoadOptions{ConfigHome: dir, Getenv: envFrom(nil)})
+	if err == nil || !strings.Contains(err.Error(), "host:port") {
+		t.Fatalf("%v", err)
+	}
+}
+
+func TestOversizedPromptIsRefusedBelowTheBodyLimit(t *testing.T) { // prompt-size-unbounded
+	w := &scriptedWorker{}
+	s := startServer(t, serverConfig(), w)
+	c := a2aClient(t, "http://"+s.Addr(), tokenA)
+	if _, err := c.SendMessage(context.Background(), &a2a.SendMessageRequest{Message: textMessage(strings.Repeat("a", maxPromptBytes+1))}); err == nil {
+		t.Fatal("a prompt over the limit must be refused")
+	}
+	if len(w.Turns()) != 0 {
+		t.Fatal("worker ran")
+	}
+}
+
+func TestBodyLimitAppliesEvenWithASmallPrompt(t *testing.T) { // body-size-unbounded
+	w := &scriptedWorker{}
+	s := startServer(t, serverConfig(), w)
+	body, _ := json.Marshal(map[string]any{"jsonrpc": "2.0", "id": 1, "method": "SendMessage", "params": map[string]any{
+		"message": map[string]any{"messageId": "m", "role": "ROLE_USER", "parts": []any{map[string]any{"text": "small prompt"}},
+			"metadata": map[string]any{"padding": strings.Repeat("x", 2<<20)}}}})
+	req, _ := http.NewRequest(http.MethodPost, "http://"+s.Addr(), bytes.NewReader(body))
+	req.Header.Set("Authorization", "Bearer "+tokenA)
+	req.Header.Set("A2A-Version", "1.0")
+	if resp, err := http.DefaultClient.Do(req); err == nil {
+		resp.Body.Close()
+	}
+	if len(w.Turns()) != 0 {
+		t.Fatal("a 2 MiB request with a small prompt reached the worker; the body cap is the first line of defence")
+	}
+}
+
+func TestTextPlusFilePartIsRefusedNotTrimmed(t *testing.T) { // file-parts-dropped
+	w := &scriptedWorker{}
+	s := startServer(t, serverConfig(), w)
+	c := a2aClient(t, "http://"+s.Addr(), tokenA)
+	m := a2a.NewMessage(a2a.MessageRoleUser, a2a.NewTextPart("summarise this"), a2a.NewRawPart([]byte("attachment")))
+	if _, err := c.SendMessage(context.Background(), &a2a.SendMessageRequest{Message: m}); err == nil {
+		t.Fatal("an attachment the worker cannot see must not be silently dropped")
+	}
+	if len(w.Turns()) != 0 {
+		t.Fatal("worker ran")
+	}
+}
+
+func TestMixedVersionHeaderIsRefused(t *testing.T) { // mixed-version-values-accepted
+	w := &scriptedWorker{}
+	s := startServer(t, serverConfig(), w)
+	params := map[string]any{"message": map[string]any{"messageId": "m1", "role": "ROLE_USER", "parts": []any{map[string]any{"text": "hello there"}}}}
+	_, out := rawRPC(t, "http://"+s.Addr(), tokenA, "1.0, 0.3", "SendMessage", params)
+	if out["error"] == nil {
+		t.Fatalf("a header that also names 0.3 is not 1.0: %v", out)
+	}
+	_, out = rawRPC(t, "http://"+s.Addr(), tokenA, "1.0, 1.0", "SendMessage", params)
+	if out["error"] != nil {
+		t.Fatalf("kagent's client repeats the header; that must work: %v", out)
+	}
+}
+
+func TestExecutorRefusesAForgedUser(t *testing.T) { // principal-not-checked
+	for name, u := range map[string]*a2asrv.User{
+		"nil":             nil,
+		"unauthenticated": {Name: "token:alice", Authenticated: false, Attributes: map[string]any{"name": "alice"}},
+		"key mismatch":    {Name: "tenant:other", Authenticated: true, Attributes: map[string]any{"name": "alice", "tenant": "team-a"}},
+	} {
+		if _, err := principalOf(&a2asrv.ExecutorContext{User: u}); err == nil {
+			t.Errorf("%s: a user the interceptor did not vouch for must not run tasks", name)
+		}
+	}
+	p, err := principalOf(&a2asrv.ExecutorContext{User: a2asrv.NewAuthenticatedUser("tenant:team-a", map[string]any{"name": "alice", "tenant": "team-a"})})
+	if err != nil || p.Name != "alice" || p.Tenant != "team-a" {
+		t.Fatalf("%+v %v", p, err)
+	}
+}
+
+func TestCancelReturnsOnlyAfterTheWorkerExited(t *testing.T) { // cancel-does-not-wait
+	var exited, started = make(chan struct{}), make(chan struct{})
+	w := &scriptedWorker{run: func(ctx context.Context, tn Turn, up func(Update)) (Result, error) {
+		close(started)
+		<-ctx.Done()
+		time.Sleep(300 * time.Millisecond) // the worker takes its time to stop
+		close(exited)
+		return Result{}, ctx.Err()
+	}}
+	s := startServer(t, serverConfig(), w)
+	c := a2aClient(t, "http://"+s.Addr(), tokenA)
+	go func() {
+		for range c.SendStreamingMessage(context.Background(), &a2a.SendMessageRequest{Message: textMessage("long job")}) {
+		}
+	}()
+	waitChan(t, started, "the worker")
+	var id a2a.TaskID
+	waitFor(t, func() bool {
+		list, err := c.ListTasks(context.Background(), &a2a.ListTasksRequest{})
+		if err == nil && len(list.Tasks) == 1 {
+			id = list.Tasks[0].ID
+			return true
+		}
+		return false
+	}, "the task to be listed")
+	if _, err := c.CancelTask(context.Background(), &a2a.CancelTaskRequest{ID: id}); err != nil {
+		t.Fatal(err)
+	}
+	select {
+	case <-exited:
+	default:
+		t.Fatal("CancelTask answered while the worker was still running: a session file would have two writers")
+	}
+}
+
+func TestSecondStartIsAnError(t *testing.T) { // second-start-allowed
+	s := startServer(t, serverConfig(), &scriptedWorker{})
+	if err := s.Start(); err == nil {
+		t.Fatal("starting a running server twice must fail, not leak a listener")
+	}
+}
+
+func TestWorkerAsksForTermBeforeKill(t *testing.T) { // no-terminate
+	if runtime.GOOS == "windows" {
+		t.Skip("POSIX signals")
+	}
+	w, logPath, _ := newTestWorker(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		_, _ = w.Run(ctx, Turn{Principal: alice, ContextID: "c", Prompt: "IGNORE_ABORT"}, func(Update) {})
+		close(done)
+	}()
+	waitFakeLog(t, logPath, func(l fakeLog) bool { return l.Started })
+	cancel()
+	waitChan(t, done, "Run")
+	if !readFakeLog(t, logPath).Terminated {
+		t.Fatal("after the grace period the worker gets SIGTERM before SIGKILL")
+	}
+}
+
+func TestWorkerEscalatesToKillWhenTermIsIgnored(t *testing.T) { // no-kill-after-grace
+	if runtime.GOOS == "windows" {
+		t.Skip("POSIX signals")
+	}
+	w, logPath, _ := newTestWorker(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		_, _ = w.Run(ctx, Turn{Principal: alice, ContextID: "c", Prompt: "IGNORE_TERM"}, func(Update) {})
+		close(done)
+	}()
+	l := waitFakeLog(t, logPath, func(l fakeLog) bool { return l.Started })
+	cancel()
+	select {
+	case <-done:
+	case <-time.After(20 * time.Second):
+		t.Fatal("a worker that ignores abort and SIGTERM must still be killed")
+	}
+	if processAlive(l.PID) {
+		t.Fatal("child survived")
+	}
+}
+
+func TestWorkerFailsATurnThatNeedsInteractiveInput(t *testing.T) { // ui-request-hangs
+	w, _, _ := newTestWorker(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	res, err := w.Run(ctx, Turn{Principal: alice, ContextID: "c", Prompt: "NEEDUI now"}, func(Update) {})
+	if err != nil || !strings.Contains(res.Failure, "interactive") {
+		t.Fatalf("%+v %v", res, err)
+	}
+}
+
+func TestWorkerReportsARejectedPromptWithoutTheReason(t *testing.T) { // rejected-prompt-ignored
+	w, _, _ := newTestWorker(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	res, err := w.Run(ctx, Turn{Principal: alice, ContextID: "c", Prompt: "REJECT me"}, func(Update) {})
+	if err != nil || res.Failure != "PiG rejected the prompt" {
+		t.Fatalf("%+v %v", res, err)
+	}
+}
+
+func TestModelFailureMessageIsExactlyGeneric(t *testing.T) { // model-error-text-copied
+	w, _, _ := newTestWorker(t)
+	res, err := w.Run(context.Background(), Turn{Principal: alice, ContextID: "c", Prompt: "FAIL now"}, func(Update) {})
+	if err != nil || res.Failure != "the model call failed" {
+		t.Fatalf("%+v %v", res, err)
+	}
+}
+
+func TestSendRefusesAnEmptyMessageBeforeAnyRequest(t *testing.T) { // empty-message-sent
+	r := newFakeRemote(t)
+	_, err := remotes(t, r, nil, nil).Send(context.Background(), SendArgs{Agent: "peer", Message: "  "}, nil)
+	if err == nil || !strings.Contains(err.Error(), "empty") {
+		t.Fatalf("%v", err)
+	}
+	r.mu.Lock()
+	n := len(r.headers)
+	r.mu.Unlock()
+	if n != 0 {
+		t.Fatalf("%d requests for an empty message", n)
+	}
+}
+
+func TestAuthFailureIsExplainedWithoutTheToken(t *testing.T) { // auth-error-shows-token-context
+	r := newFakeRemote(t)
+	r.requireAuth = "the-right-one"
+	rs := remotes(t, r, func(ra *RemoteAgent) { ra.BearerTokenEnv = "REMOTE_TOKEN" }, map[string]string{"REMOTE_TOKEN": "wrong-secret-token"})
+	_, err := rs.Send(context.Background(), SendArgs{Agent: "peer", Message: "hello"}, nil)
+	if err == nil || !strings.Contains(err.Error(), "authentication failed") {
+		t.Fatalf("the model needs to know the credentials were refused: %v", err)
+	}
+}
+
+func TestCredentialTransportRefusesAnotherHost(t *testing.T) { // transport-host-not-checked
+	hit := false
+	other := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { hit = true }))
+	defer other.Close()
+	// Same scheme as the other server, so only the host check can refuse it.
+	tr := headerTransport{h: http.Header{"Authorization": {"Bearer secret-value"}}, host: "configured.example:443", scheme: "http", base: http.DefaultTransport}
+	req, _ := http.NewRequest(http.MethodGet, other.URL, nil)
+	if _, err := tr.RoundTrip(req); err == nil || hit {
+		t.Fatalf("credentials went to %s (hit=%v, err=%v)", other.URL, hit, err)
+	}
+}
+
+func TestRemoteRedirectsAreNotFollowed(t *testing.T) { // redirects-followed-with-credentials
+	rs := NewRemotes(map[string]RemoteAgent{"peer": {URL: "http://127.0.0.1:1", BearerTokenEnv: "T"}}, envFrom(map[string]string{"T": "secret-value"}))
+	hc, err := rs.httpClient("peer", rs.cfg["peer"])
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := hc.CheckRedirect(&http.Request{}, nil); got != http.ErrUseLastResponse {
+		t.Fatalf("redirects must not be followed: %v", got)
+	}
+}
+
+func TestSummaryOfAnInputRequiredTaskIsNotTerminal(t *testing.T) { // input-required-called-terminal
+	task := &a2a.Task{ID: "t", ContextID: "c", Status: a2a.TaskStatus{State: a2a.TaskStateInputRequired, Message: a2a.NewMessage(a2a.MessageRoleAgent, a2a.NewTextPart("which file?"))}}
+	s := summarize(task)
+	if s.Terminal || s.State != "input-required" || s.Text != "which file?" {
+		t.Fatalf("%+v", s)
+	}
+	out := formatSummary(s) // unfinished-task-not-flagged
+	if !strings.Contains(out, "not finished") || !strings.Contains(out, "which file?") {
+		t.Fatalf("%s", out)
+	}
+	if strings.Contains(formatSummary(TaskSummary{State: "completed", Terminal: true, Text: "done"}), "not finished") {
+		t.Fatal("a finished task is not flagged")
+	}
+}

package/extensions/a2a/worker.go

@@ -0,0 +1,388 @@
+package a2aext
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"crypto/sha256"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"time"
+)
+
+// Turn is one task's run of PiG.
+type Turn struct {
+	Principal Principal
+	ContextID string
+	TaskID    string
+	Prompt    string
+}
+
+// Update is incremental output from a running turn.
+type Update struct {
+	Text string // assistant text delta
+	Tool string // name of a tool that started
+}
+
+// Result is a finished turn. Text is exactly the concatenation of the streamed text updates.
+type Result struct {
+	Text string
+	// Failure is a short, peer-safe reason the turn failed. Provider error text is never copied into it.
+	Failure string
+}
+
+// Worker runs turns.
+type Worker interface {
+	Run(ctx context.Context, turn Turn, onUpdate func(Update)) (Result, error)
+}
+
+// ProcessWorker runs each turn in a `pig --mode rpc` child process. A context is one PiG
+// session file, so a later task in the same context continues the conversation.
+type ProcessWorker struct {
+	cfg      WorkerConfig
+	stateDir string
+	getenv   func(string) string
+}
+
+const maxFrameBytes = 16 << 20
+
+// baseEnv is what a worker inherits. Provider keys and everything else must be named in passEnv.
+var baseEnv = []string{
+	"PATH", "HOME", "USER", "LANG", "LC_ALL", "TMPDIR", "TERM",
+	"PIG_HOME", "PIG_CODING_AGENT_DIR", "PI_CODING_AGENT_DIR",
+	"XDG_CONFIG_HOME", "XDG_STATE_HOME", "XDG_CACHE_HOME", "XDG_DATA_HOME",
+	"HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "no_proxy",
+	"SSL_CERT_FILE", "SSL_CERT_DIR", "SystemRoot", "USERPROFILE", "APPDATA", "LOCALAPPDATA",
+}
+
+// NewProcessWorker builds a ProcessWorker whose per-principal sessions live under stateDir.
+func NewProcessWorker(cfg WorkerConfig, stateDir string, getenv func(string) string) (*ProcessWorker, error) {
+	if stateDir == "" {
+		return nil, errors.New("a2a: worker needs a state directory")
+	}
+	if getenv == nil {
+		getenv = os.Getenv
+	}
+	if cfg.Command == "" {
+		cfg.Command = getenv("PIG_A2A_PIG")
+	}
+	if cfg.Command == "" {
+		cfg.Command = "pig"
+	}
+	for _, name := range cfg.PassEnv {
+		if !envRE.MatchString(name) {
+			return nil, fmt.Errorf("a2a: worker passEnv %q is not an environment variable name", name)
+		}
+	}
+	if cfg.GraceSeconds <= 0 {
+		cfg.GraceSeconds = defaultGraceSeconds
+	}
+	return &ProcessWorker{cfg: cfg, stateDir: stateDir, getenv: getenv}, nil
+}
+
+func keyDigest(s string) string {
+	sum := sha256.Sum256([]byte(s))
+	return hex.EncodeToString(sum[:])
+}
+
+// SessionID derives the PiG session id for a principal's context. It is hex, so a hostile
+// context id cannot reach a file name, and it depends on the principal's tenant boundary,
+// so two tenants that pick the same context id never share a session.
+func SessionID(p Principal, contextID string) string {
+	return keyDigest("a2a-session\x00" + p.Key() + "\x00" + contextID)[:32]
+}
+
+func (w *ProcessWorker) sessionDir(p Principal) string {
+	return filepath.Join(w.stateDir, "sessions", keyDigest("a2a-dir\x00" + p.Key())[:16])
+}
+
+func (w *ProcessWorker) cwd(p Principal) string {
+	if w.cfg.Cwd != "" {
+		return w.cfg.Cwd
+	}
+	// Without an explicit workspace the worker sees an empty private directory, not the operator's files.
+	return filepath.Join(w.stateDir, "workspaces", keyDigest("a2a-dir\x00" + p.Key())[:16])
+}
+
+func (w *ProcessWorker) args(t Turn) []string {
+	args := []string{"--mode", "rpc", "--offline", "--no-extensions", "--no-skills", "--no-context-files", "--no-prompt-templates", "--no-themes"}
+	if len(w.cfg.Tools) > 0 {
+		args = append(args, "--tools", strings.Join(w.cfg.Tools, ","))
+	} else {
+		// PiG's file tools are not confined to the working directory, so a worker gets no tools unless named.
+		args = append(args, "--no-tools")
+	}
+	if w.cfg.Provider != "" {
+		args = append(args, "--provider", w.cfg.Provider)
+	}
+	if w.cfg.Model != "" {
+		args = append(args, "--model", w.cfg.Model)
+	}
+	args = append(args, "--session-dir", w.sessionDir(t.Principal), "--session-id", SessionID(t.Principal, t.ContextID))
+	return append(args, w.cfg.Args...)
+}
+
+func (w *ProcessWorker) env() []string {
+	var env []string
+	seen := map[string]bool{}
+	for _, name := range append(append([]string{}, baseEnv...), w.cfg.PassEnv...) {
+		if v := w.getenv(name); v != "" && !seen[name] {
+			seen[name] = true
+			env = append(env, name+"="+v)
+		}
+	}
+	return append(env, "PIG_A2A_WORKER=1", "PI_SKIP_VERSION_CHECK=1")
+}
+
+// wireEvent is the part of PiG's RPC JSONL the worker consumes.
+type wireEvent struct {
+	Type                  string `json:"type"`
+	ID                    string `json:"id"`
+	Command               string `json:"command"`
+	Success               bool   `json:"success"`
+	ToolName              string `json:"toolName"`
+	AssistantMessageEvent *struct {
+		Type  string `json:"type"`
+		Delta string `json:"delta"`
+	} `json:"assistantMessageEvent"`
+	Message *struct {
+		Role       string          `json:"role"`
+		StopReason string          `json:"stopReason"`
+		Content    json.RawMessage `json:"content"`
+	} `json:"message"`
+}
+
+type frame struct {
+	event wireEvent
+	err   error
+}
+
+func readFrames(r io.Reader, out chan<- frame) {
+	defer close(out)
+	br := bufio.NewReaderSize(r, 1<<20)
+	for {
+		var line []byte
+		for {
+			chunk, err := br.ReadSlice('\n')
+			line = append(line, chunk...)
+			if len(line) > maxFrameBytes {
+				out <- frame{err: fmt.Errorf("PiG frame exceeds %d bytes", maxFrameBytes)}
+				return
+			}
+			if err == bufio.ErrBufferFull {
+				continue
+			}
+			if err != nil {
+				if len(bytes.TrimSpace(line)) > 0 {
+					out <- frame{err: fmt.Errorf("PiG stream closed inside a record: %w", err)}
+				}
+				return
+			}
+			break
+		}
+		if len(bytes.TrimSpace(line)) == 0 {
+			continue
+		}
+		var ev wireEvent
+		if err := json.Unmarshal(line, &ev); err != nil || ev.Type == "" {
+			out <- frame{err: errors.New("PiG sent a record that is not an event")}
+			return
+		}
+		out <- frame{event: ev}
+	}
+}
+
+type tailBuffer struct{ b []byte }
+
+func (t *tailBuffer) Write(p []byte) (int, error) {
+	t.b = append(t.b, p...)
+	if len(t.b) > 4096 {
+		t.b = t.b[len(t.b)-4096:]
+	}
+	return len(p), nil
+}
+
+// Run implements Worker.
+func (w *ProcessWorker) Run(ctx context.Context, t Turn, onUpdate func(Update)) (Result, error) {
+	if strings.TrimSpace(t.Prompt) == "" {
+		return Result{}, errors.New("a2a: empty prompt")
+	}
+	if err := ctx.Err(); err != nil {
+		return Result{}, err
+	}
+	for _, dir := range []string{w.sessionDir(t.Principal), w.cwd(t.Principal)} {
+		if err := os.MkdirAll(dir, 0o700); err != nil {
+			return Result{}, fmt.Errorf("a2a: prepare worker directory: %w", err)
+		}
+	}
+	cmd := exec.Command(w.cfg.Command, w.args(t)...)
+	cmd.Dir, cmd.Env = w.cwd(t.Principal), w.env()
+	configureProcess(cmd)
+	stderr := &tailBuffer{}
+	cmd.Stderr = stderr
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		return Result{}, err
+	}
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return Result{}, err
+	}
+	if err := cmd.Start(); err != nil {
+		return Result{}, fmt.Errorf("a2a: start worker %q: %w", w.cfg.Command, err)
+	}
+	frames := make(chan frame, 64)
+	go readFrames(stdout, frames)
+	exited := make(chan error, 1)
+	go func() { exited <- cmd.Wait() }()
+	grace := time.Duration(w.cfg.GraceSeconds) * time.Second
+
+	write := func(v map[string]any) error {
+		b, _ := json.Marshal(v)
+		_, err := stdin.Write(append(b, '\n'))
+		return err
+	}
+	var result Result
+	var runErr error
+	defer func() {
+		_ = stdin.Close()
+		select {
+		case <-exited:
+		case <-time.After(grace):
+			_ = terminateTree(cmd.Process)
+			select {
+			case <-exited:
+			case <-time.After(grace):
+				_ = killTree(cmd.Process)
+				<-exited
+			}
+		}
+		// A tool the worker started can outlive it.
+		_ = killTree(cmd.Process)
+		go func() {
+			for range frames { // let the reader finish
+			}
+		}()
+	}()
+
+	if err := write(map[string]any{"id": "prompt", "type": "prompt", "message": t.Prompt}); err != nil {
+		return Result{}, fmt.Errorf("a2a: write prompt: %w", err)
+	}
+
+	var (
+		text       strings.Builder
+		msgText    string // streamed text of the current assistant message
+		msgHasText bool
+		lastStop   string
+		accepted   bool
+		aborting   bool
+		abortTimer <-chan time.Time
+		ctxDone    = ctx.Done()
+	)
+	emit := func(s string) {
+		if s == "" {
+			return
+		}
+		if !msgHasText && text.Len() > 0 {
+			text.WriteString("\n\n")
+			onUpdate(Update{Text: "\n\n"})
+		}
+		msgHasText = true
+		msgText += s
+		text.WriteString(s)
+		onUpdate(Update{Text: s})
+	}
+	for {
+		select {
+		case <-ctxDone:
+			ctxDone = nil
+			aborting = true
+			_ = write(map[string]any{"id": "abort", "type": "abort"})
+			abortTimer = time.After(grace)
+		case <-abortTimer:
+			return Result{}, ctx.Err()
+		case f, ok := <-frames:
+			if !ok {
+				if aborting {
+					return Result{}, ctx.Err()
+				}
+				return Result{}, fmt.Errorf("a2a: worker exited before the turn settled (%s)", strings.TrimSpace(string(stderr.b)))
+			}
+			if f.err != nil {
+				return Result{}, fmt.Errorf("a2a: worker protocol: %w", f.err)
+			}
+			ev := f.event
+			switch ev.Type {
+			case "response":
+				if ev.ID == "prompt" {
+					if !ev.Success {
+						return Result{Failure: "PiG rejected the prompt"}, nil
+					}
+					accepted = true
+				}
+			case "message_start":
+				if ev.Message != nil && ev.Message.Role == "assistant" {
+					msgText, msgHasText, lastStop = "", false, ""
+				}
+			case "message_update":
+				if !aborting && ev.AssistantMessageEvent != nil && ev.AssistantMessageEvent.Type == "text_delta" {
+					emit(ev.AssistantMessageEvent.Delta)
+				}
+			case "message_end":
+				if ev.Message == nil || ev.Message.Role != "assistant" {
+					break
+				}
+				lastStop = ev.Message.StopReason
+				if lastStop == "error" || lastStop == "aborted" || aborting {
+					break // a failed attempt may retry; never copy provider text
+				}
+				var blocks []struct {
+					Type string `json:"type"`
+					Text string `json:"text"`
+				}
+				if json.Unmarshal(ev.Message.Content, &blocks) == nil {
+					var final strings.Builder
+					for _, b := range blocks {
+						if b.Type == "text" {
+							final.WriteString(b.Text)
+						}
+					}
+					if f := final.String(); len(f) > len(msgText) && strings.HasPrefix(f, msgText) {
+						emit(f[len(msgText):])
+					}
+				}
+			case "tool_execution_start":
+				if !aborting && ev.ToolName != "" {
+					onUpdate(Update{Tool: ev.ToolName})
+				}
+			case "error":
+				result.Failure = "PiG reported an error"
+				return result, runErr
+			case "extension_ui_request":
+				result.Failure = "PiG asked for interactive input, which an A2A task cannot provide"
+				return result, runErr
+			case "agent_settled":
+				if aborting {
+					return Result{}, ctx.Err()
+				}
+				if !accepted {
+					return Result{}, errors.New("a2a: worker settled before accepting the prompt")
+				}
+				switch lastStop {
+				case "error":
+					return Result{Text: text.String(), Failure: "the model call failed"}, nil
+				case "aborted":
+					return Result{Text: text.String(), Failure: "the turn was aborted"}, nil
+				}
+				return Result{Text: text.String()}, nil
+			}
+		}
+	}
+}