1package agentrunner23import (4 "bufio"5 "context"6 "encoding/json"7 "errors"8 "os"9 "path/filepath"10 "testing"11 "time"12)1314func TestACPHandshakeAndStream(t *testing.T) {15 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)16 defer cancel()17 result, err := RunProcess(ctx, ProcessOptions{Command: os.Args[0], Args: []string{"-test.run=TestACPHelper"}, Env: append(os.Environ(), "TINY_ACP_HELPER=1"), CWD: t.TempDir(), Protocol: "acp"}, "hello")18 if err != nil || result != "hello from agent" {19 t.Fatalf("result=%q err=%v", result, err)20 }21}2223func TestACPCancellationNotifiesAgentBeforeExit(t *testing.T) {24 marker := filepath.Join(t.TempDir(), "cancelled")25 ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)26 defer cancel()27 _, err := RunProcess(ctx, ProcessOptions{Command: os.Args[0], Args: []string{"-test.run=TestACPHelper"}, Env: append(os.Environ(), "TINY_ACP_HELPER=1", "TINY_ACP_CANCEL_HELPER=1", "TINY_ACP_CANCEL_MARKER="+marker), CWD: t.TempDir(), Protocol: "acp"}, "hello")28 if !errors.Is(err, context.DeadlineExceeded) && !errors.Is(err, context.Canceled) {29 t.Fatalf("err=%v", err)30 }31 if _, err := os.Stat(marker); err != nil {32 t.Fatalf("agent did not receive cancellation: %v", err)33 }34}3536func TestCommandCleansUpDescendantProcess(t *testing.T) {37 marker := filepath.Join(t.TempDir(), "child-survived")38 ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)39 defer cancel()40 _, err := RunProcess(ctx, ProcessOptions{Command: "sh", Args: []string{"-c", "(sleep 0.5; touch \"$1\") & wait", "descendant", marker}, CWD: t.TempDir(), Protocol: "command"}, "hello")41 if err == nil {42 t.Fatal("expected cancellation")43 }44 time.Sleep(200 * time.Millisecond)45 if _, err := os.Stat(marker); err == nil {46 t.Fatal("descendant survived")47 }48}49func TestACPHelper(t *testing.T) {50 if os.Getenv("TINY_ACP_HELPER") != "1" && os.Getenv("TINY_ACP_CANCEL_HELPER") != "1" {51 return52 }53 scan := bufio.NewScanner(os.Stdin)54 encoder := json.NewEncoder(os.Stdout)55 for scan.Scan() {56 var message map[string]any57 if json.Unmarshal(scan.Bytes(), &message) != nil {58 os.Exit(2)59 }60 params, _ := message["params"].(map[string]any)61 switch message["method"] {62 case "initialize":63 if params["protocolVersion"] != float64(1) {64 os.Exit(3)65 }66 encoder.Encode(map[string]any{"jsonrpc": "2.0", "id": message["id"], "result": map[string]any{"protocolVersion": 1}})67 case "session/new":68 if params["cwd"] == nil || params["mcpServers"] == nil {69 os.Exit(4)70 }71 encoder.Encode(map[string]any{"jsonrpc": "2.0", "id": message["id"], "result": map[string]any{"sessionId": "test-session"}})72 case "session/prompt":73 if params["sessionId"] != "test-session" || params["prompt"] == nil {74 os.Exit(5)75 }76 if os.Getenv("TINY_ACP_CANCEL_HELPER") == "1" {77 continue78 }79 encoder.Encode(map[string]any{"jsonrpc": "2.0", "method": "session/update", "params": map[string]any{"sessionId": "test-session", "update": map[string]any{"sessionUpdate": "agent_message_chunk", "content": map[string]any{"type": "text", "text": "hello from agent"}}}})80 encoder.Encode(map[string]any{"jsonrpc": "2.0", "id": message["id"], "result": map[string]any{"stopReason": "end_turn"}})81 case "session/cancel":82 if marker := os.Getenv("TINY_ACP_CANCEL_MARKER"); marker != "" {83 _ = os.WriteFile(marker, []byte("cancelled"), 0600)84 }85 os.Exit(0)86 }87 }88 os.Exit(0)89}