| 1 | package agentrunner |
| 2 | |
| 3 | import ( |
| 4 | "bufio" |
| 5 | "context" |
| 6 | "encoding/json" |
| 7 | "errors" |
| 8 | "os" |
| 9 | "path/filepath" |
| 10 | "testing" |
| 11 | "time" |
| 12 | ) |
| 13 | |
| 14 | func 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 | } |
| 22 | |
| 23 | func 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 | } |
| 35 | |
| 36 | func 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 | } |
| 49 | func TestACPHelper(t *testing.T) { |
| 50 | if os.Getenv("TINY_ACP_HELPER" ) != "1" && os.Getenv("TINY_ACP_CANCEL_HELPER" ) != "1" { |
| 51 | return |
| 52 | } |
| 53 | scan := bufio.NewScanner(os.Stdin) |
| 54 | encoder := json.NewEncoder(os.Stdout) |
| 55 | for scan.Scan() { |
| 56 | var message map[string]any |
| 57 | 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 | continue |
| 78 | } |
| 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 | } |