1  package tools
  2  
  3  import (
  4  	"bufio"
  5  	"context"
  6  	"encoding/json"
  7  	"fmt"
  8  	"io"
  9  	"net"
 10  	"os"
 11  	"os/exec"
 12  	"path/filepath"
 13  	"runtime"
 14  	"sync"
 15  	"sync/atomic"
 16  	"syscall"
 17  	"time"
 18  )
 19  
 20  // AXRequest is a JSON-RPC request sent to ax_server.
 21  type AXRequest struct {
 22  	ID     int64  `json:"id"`
 23  	Method string `json:"method"`
 24  	Params any    `json:"params"`
 25  }
 26  
 27  // AXResponse is a JSON-RPC response from ax_server.
 28  type AXResponse struct {
 29  	ID     int64            `json:"id"`
 30  	Result json.RawMessage  `json:"result,omitempty"`
 31  	Error  *AXError         `json:"error,omitempty"`
 32  }
 33  
 34  // AXError is an error returned by ax_server.
 35  type AXError struct {
 36  	Code    int    `json:"code"`
 37  	Message string `json:"message"`
 38  }
 39  
 40  // SharedAXClient returns the process-wide singleton AXClient.
 41  // Both the tools (accessibility, computer, wait) and daemon permission
 42  // endpoints must use the same instance, because the socket server
 43  // accepts only one client at a time.
 44  func SharedAXClient() *AXClient {
 45  	sharedOnce.Do(func() {
 46  		sharedInstance = &AXClient{}
 47  	})
 48  	return sharedInstance
 49  }
 50  
 51  var (
 52  	sharedOnce     sync.Once
 53  	sharedInstance *AXClient
 54  )
 55  
 56  // AXClient manages a persistent ax_server process and multiplexes
 57  // requests by ID. Multiple goroutines can call Call() concurrently.
 58  //
 59  // Two transport modes:
 60  // - Bundled: ax_server is inside a .app bundle, launched via LaunchServices
 61  //   (`open -a`), communicates over a Unix domain socket. Required for TCC
 62  //   permission attribution on macOS.
 63  // - Fallback: ax_server is a bare binary, launched via exec.Command,
 64  //   communicates over stdin/stdout pipes. Used for dev, npm, and CLI.
 65  type AXClient struct {
 66  	mu      sync.Mutex // guards process lifecycle (start/restart)
 67  	writeMu sync.Mutex // guards writes to ax_server
 68  
 69  	// Transport-agnostic I/O
 70  	writer io.WriteCloser
 71  	nextID atomic.Int64
 72  
 73  	// Process management
 74  	cmd        *exec.Cmd // non-nil in fallback mode
 75  	conn       net.Conn  // non-nil in bundled mode
 76  	bundlePID  int       // ax_server PID in bundled mode (for cleanup)
 77  	started    bool
 78  
 79  	pendingMu sync.Mutex
 80  	pending   map[int64]chan AXResponse
 81  }
 82  
 83  // Ensure starts the ax_server process if not already running.
 84  func (c *AXClient) Ensure(ctx context.Context) error {
 85  	c.mu.Lock()
 86  	defer c.mu.Unlock()
 87  
 88  	if c.started {
 89  		return nil
 90  	}
 91  
 92  	binPath, bundlePath, err := AXServerPaths()
 93  	if err != nil {
 94  		return err
 95  	}
 96  
 97  	c.pending = make(map[int64]chan AXResponse)
 98  
 99  	if bundlePath != "" {
100  		return c.startBundled(ctx, bundlePath)
101  	}
102  	return c.startFallback(binPath)
103  }
104  
105  // startBundled launches ax_server via LaunchServices and connects over Unix socket.
106  func (c *AXClient) startBundled(ctx context.Context, bundlePath string) error {
107  	socketPath := AXSocketPath()
108  
109  	// Try connecting to an existing socket first — ax_server may already be running
110  	// (e.g. from a previous open -a that's still alive).
111  	conn, err := net.Dial("unix", socketPath)
112  	if err != nil {
113  		// Not running or stale socket — clean up and launch fresh
114  		os.Remove(socketPath)
115  
116  		// Launch via open(1) with -n (new instance) — gives ax_server its own TCC
117  		// identity and avoids reusing a stale instance with different --args.
118  		cmd := exec.CommandContext(ctx, "open", "-n", "-a", bundlePath, "--args", "--socket", socketPath)
119  		cmd.Stderr = os.Stderr
120  		if err := cmd.Run(); err != nil {
121  			return fmt.Errorf("ax_server launch: %w", err)
122  		}
123  
124  		// Wait for socket to appear
125  		deadline := time.Now().Add(10 * time.Second)
126  		for time.Now().Before(deadline) {
127  			conn, err = net.Dial("unix", socketPath)
128  			if err == nil {
129  				break
130  			}
131  			time.Sleep(100 * time.Millisecond)
132  		}
133  	}
134  	if conn == nil {
135  		return fmt.Errorf("ax_server: socket not available after 10s at %s", socketPath)
136  	}
137  
138  	c.conn = conn
139  	c.writer = conn
140  	c.started = true
141  
142  	// Find the ax_server PID for cleanup on Close().
143  	// pgrep -f matches the socket path in the command line args.
144  	if out, err := exec.Command("pgrep", "-f", socketPath).Output(); err == nil {
145  		var pid int
146  		if _, err := fmt.Sscanf(string(out), "%d", &pid); err == nil {
147  			c.bundlePID = pid
148  		}
149  	}
150  
151  	// Reader goroutine dispatches responses by ID.
152  	go c.readLoop(conn)
153  
154  	return nil
155  }
156  
157  // startFallback launches ax_server via exec.Command with stdin/stdout pipes.
158  func (c *AXClient) startFallback(binPath string) error {
159  	// Use exec.Command (not CommandContext) — the process lifecycle is managed
160  	// independently of any single request's context.
161  	c.cmd = exec.Command(binPath)
162  	var pipeErr error
163  	c.writer, pipeErr = c.cmd.StdinPipe()
164  	if pipeErr != nil {
165  		return fmt.Errorf("ax_server stdin pipe: %w", pipeErr)
166  	}
167  	stdout, pipeErr := c.cmd.StdoutPipe()
168  	if pipeErr != nil {
169  		return fmt.Errorf("ax_server stdout pipe: %w", pipeErr)
170  	}
171  	c.cmd.Stderr = os.Stderr
172  
173  	if err := c.cmd.Start(); err != nil {
174  		return fmt.Errorf("ax_server start: %w", err)
175  	}
176  	c.started = true
177  
178  	// Reader goroutine dispatches responses by ID.
179  	go func() {
180  		c.readLoop(stdout)
181  		// Wait for process exit and mark as dead so next Ensure() restarts it.
182  		c.cmd.Wait()
183  		c.mu.Lock()
184  		c.started = false
185  		c.mu.Unlock()
186  	}()
187  
188  	return nil
189  }
190  
191  // readLoop reads NDJSON responses and dispatches them to pending callers.
192  func (c *AXClient) readLoop(reader io.Reader) {
193  	scanner := bufio.NewScanner(reader)
194  	scanner.Buffer(make([]byte, 0, 1024*1024), 1024*1024)
195  	for scanner.Scan() {
196  		var resp AXResponse
197  		if err := json.Unmarshal(scanner.Bytes(), &resp); err != nil {
198  			continue
199  		}
200  		c.pendingMu.Lock()
201  		ch, ok := c.pending[resp.ID]
202  		if ok {
203  			delete(c.pending, resp.ID)
204  		}
205  		c.pendingMu.Unlock()
206  		if ok {
207  			ch <- resp
208  		}
209  	}
210  	// EOF: ax_server died or disconnected — unblock all pending callers
211  	c.pendingMu.Lock()
212  	for id, ch := range c.pending {
213  		ch <- AXResponse{ID: id, Error: &AXError{Code: -1, Message: "ax_server: unexpected EOF"}}
214  		delete(c.pending, id)
215  	}
216  	c.pendingMu.Unlock()
217  
218  	// Mark as not started so next Ensure() reconnects
219  	c.mu.Lock()
220  	c.started = false
221  	c.mu.Unlock()
222  }
223  
224  // Call sends a request and waits for the response.
225  func (c *AXClient) Call(ctx context.Context, method string, params any) (json.RawMessage, error) {
226  	if runtime.GOOS != "darwin" {
227  		return nil, fmt.Errorf("ax_server is macOS-only")
228  	}
229  
230  	if err := c.Ensure(ctx); err != nil {
231  		return nil, err
232  	}
233  
234  	id := c.nextID.Add(1)
235  	req := AXRequest{ID: id, Method: method, Params: params}
236  
237  	// Register pending channel BEFORE writing
238  	ch := make(chan AXResponse, 1)
239  	c.pendingMu.Lock()
240  	c.pending[id] = ch
241  	c.pendingMu.Unlock()
242  
243  	data, _ := json.Marshal(req)
244  	data = append(data, '\n')
245  
246  	c.writeMu.Lock()
247  	n, writeErr := c.writer.Write(data)
248  	if writeErr == nil && n < len(data) {
249  		writeErr = io.ErrShortWrite
250  	}
251  	c.writeMu.Unlock()
252  
253  	if writeErr != nil {
254  		c.pendingMu.Lock()
255  		delete(c.pending, id)
256  		c.pendingMu.Unlock()
257  		return nil, fmt.Errorf("ax_server write: %w", writeErr)
258  	}
259  
260  	select {
261  	case resp := <-ch:
262  		if resp.Error != nil {
263  			return nil, fmt.Errorf("ax_server: %s", resp.Error.Message)
264  		}
265  		return resp.Result, nil
266  	case <-ctx.Done():
267  		c.pendingMu.Lock()
268  		delete(c.pending, id)
269  		c.pendingMu.Unlock()
270  		return nil, ctx.Err()
271  	}
272  }
273  
274  // Close terminates the ax_server process and cleans up resources.
275  func (c *AXClient) Close() {
276  	c.mu.Lock()
277  	defer c.mu.Unlock()
278  	if c.conn != nil {
279  		// Bundled mode: c.writer == c.conn, so only close once.
280  		c.conn.Close()
281  		c.conn = nil
282  		c.writer = nil
283  		os.Remove(AXSocketPath())
284  	} else if c.writer != nil {
285  		// Fallback mode: close stdin pipe.
286  		c.writer.Close()
287  	}
288  	// Bundled mode: kill the LaunchServices-launched ax_server process.
289  	// Closing the socket causes ax_server to exit on its own (it exits
290  	// after the sole client disconnects), but SIGTERM is a safety net.
291  	if c.bundlePID > 0 {
292  		if proc, err := os.FindProcess(c.bundlePID); err == nil {
293  			proc.Signal(syscall.SIGTERM)
294  		}
295  		c.bundlePID = 0
296  	}
297  	// Fallback mode: kill the subprocess
298  	if c.cmd != nil && c.cmd.Process != nil {
299  		c.cmd.Process.Kill()
300  		c.cmd.Wait()
301  	}
302  	c.started = false
303  }
304  
305  // AXSocketPath returns the Unix socket path for bundled mode.
306  func AXSocketPath() string {
307  	tmpDir := os.TempDir()
308  	return filepath.Join(tmpDir, fmt.Sprintf("run.shannon.shanclaw.ax-server.%d.sock", os.Getpid()))
309  }
310  
311  // AXServerPaths returns the binary path and (optionally) the .app bundle path.
312  // If bundlePath is non-empty, use LaunchServices + socket mode.
313  // If bundlePath is empty, use exec.Command + stdin/stdout with binPath.
314  func AXServerPaths() (binPath, bundlePath string, err error) {
315  	exe, exeErr := os.Executable()
316  	if exeErr == nil {
317  		dir := filepath.Dir(exe)
318  
319  		// Bundled: nested app inside engine helper's Helpers/
320  		bp := filepath.Join(dir, "..", "Helpers", "Kocoro AX.app")
321  		bin := filepath.Join(bp, "Contents", "MacOS", "ax_server")
322  		if _, err := os.Stat(bin); err == nil {
323  			return bin, bp, nil
324  		}
325  
326  		// Flat: same directory as shan binary
327  		p := filepath.Join(dir, "ax_server")
328  		if _, err := os.Stat(p); err == nil {
329  			return p, "", nil
330  		}
331  
332  		// npm: bin/ax_server
333  		p = filepath.Join(dir, "bin", "ax_server")
334  		if _, err := os.Stat(p); err == nil {
335  			return p, "", nil
336  		}
337  	}
338  
339  	// Development: relative to working directory
340  	p := filepath.Join("internal", "tools", "axserver", ".build", "debug", "ax_server")
341  	if _, err := os.Stat(p); err == nil {
342  		return p, "", nil
343  	}
344  
345  	return "", "", fmt.Errorf("ax_server binary not found")
346  }