Agent-Native/runtime/ext/process.go

package ext

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os/exec"
	"sync"

	"agent-native/runtime/llm"
	"agent-native/runtime/tools"
)

// Message is one JSONL frame on the extension pipe.
type Message struct {
	Type        string          `json:"type"`
	ID          string          `json:"id,omitempty"`
	Name        string          `json:"name,omitempty"`
	Description string          `json:"description,omitempty"`
	Parameters  json.RawMessage `json:"parameters,omitempty"`
	Arguments   json.RawMessage `json:"arguments,omitempty"`
	Content     string          `json:"content,omitempty"`
	IsError     bool            `json:"isError,omitempty"`
	Event       string          `json:"event,omitempty"`
	Block       bool            `json:"block,omitempty"`
	Reason      string          `json:"reason,omitempty"`
	Message     string          `json:"message,omitempty"`
	Protocol    string          `json:"protocol,omitempty"`
	BaseURL     string          `json:"baseURL,omitempty"`
	CustomType  string          `json:"customType,omitempty"`
	Data        json.RawMessage `json:"data,omitempty"`
	Command     string          `json:"command,omitempty"`
}

// UIRequest is a question an extension wants a human to answer.
type UIRequest struct {
	ID      string
	Message string
}

// Process is one extension subprocess.
type Process struct {
	cmd      *exec.Cmd
	stdin    io.WriteCloser
	mu       sync.Mutex
	tools    []registeredTool
	commands []string
	providers []llm.Provider
	pending  map[string]chan Message
	UI       chan UIRequest
	OnCustom func(customType string, data json.RawMessage)
	closed   bool
}

type registeredTool struct {
	name        string
	description string
	parameters  json.RawMessage
	proc        *Process
}

func (t registeredTool) Name() string                { return t.name }
func (t registeredTool) Description() string         { return t.description }
func (t registeredTool) Parameters() json.RawMessage { return t.parameters }
func (t registeredTool) Execute(ctx context.Context, args json.RawMessage) (tools.Result, error) {
	return t.proc.callTool(ctx, t.name, args)
}

// Start launches command and reads its registration frames until it goes idle
// or the context ends. command is argv, not a shell line.
func Start(ctx context.Context, command []string) (*Process, error) {
	if len(command) == 0 {
		return nil, fmt.Errorf("ext: empty command")
	}
	cmd := exec.CommandContext(ctx, command[0], command[1:]...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	if err := cmd.Start(); err != nil {
		return nil, err
	}
	p := &Process{
		cmd: cmd, stdin: stdin, pending: map[string]chan Message{},
		UI: make(chan UIRequest, 4),
	}
	go p.read(stdout)
	if err := p.send(Message{Type: "init"}); err != nil {
		_ = p.Close()
		return nil, err
	}
	return p, nil
}

// Tools returns the tools the extension registered.
func (p *Process) Tools() []tools.Tool {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := make([]tools.Tool, len(p.tools))
	for i := range p.tools {
		out[i] = p.tools[i]
	}
	return out
}

// Commands returns slash-command names.
func (p *Process) Commands() []string {
	p.mu.Lock()
	defer p.mu.Unlock()
	return append([]string(nil), p.commands...)
}

// Providers returns endpoints the extension registered.
func (p *Process) Providers() []llm.Provider {
	p.mu.Lock()
	defer p.mu.Unlock()
	return append([]llm.Provider(nil), p.providers...)
}

// Hook tells the extension about an event. A block result stops the host action.
func (p *Process) Hook(ctx context.Context, event, name string, args json.RawMessage) (bool, string) {
	id := name + ":" + event
	ch := make(chan Message, 1)
	p.mu.Lock()
	p.pending[id] = ch
	p.mu.Unlock()
	if err := p.send(Message{Type: "event", ID: id, Event: event, Name: name, Arguments: args}); err != nil {
		return false, ""
	}
	select {
	case <-ctx.Done():
		return false, ctx.Err().Error()
	case msg := <-ch:
		return msg.Block, msg.Reason
	}
}

func (p *Process) callTool(ctx context.Context, name string, args json.RawMessage) (tools.Result, error) {
	id := name
	ch := make(chan Message, 1)
	p.mu.Lock()
	p.pending[id] = ch
	p.mu.Unlock()
	if err := p.send(Message{Type: "tool_call", ID: id, Name: name, Arguments: args}); err != nil {
		return tools.Result{IsError: true, Content: err.Error()}, err
	}
	select {
	case <-ctx.Done():
		return tools.Result{IsError: true, Content: ctx.Err().Error()}, ctx.Err()
	case msg := <-ch:
		return tools.Result{Content: msg.Content, IsError: msg.IsError}, nil
	}
}

func (p *Process) send(msg Message) error {
	raw, err := json.Marshal(msg)
	if err != nil {
		return err
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		return fmt.Errorf("ext: closed")
	}
	_, err = p.stdin.Write(append(raw, '\n'))
	return err
}

func (p *Process) read(r io.Reader) {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
	for sc.Scan() {
		var msg Message
		if err := json.Unmarshal(sc.Bytes(), &msg); err != nil {
			continue
		}
		p.dispatch(msg)
	}
}

func (p *Process) dispatch(msg Message) {
	switch msg.Type {
	case "register_tool":
		p.mu.Lock()
		p.tools = append(p.tools, registeredTool{
			name: msg.Name, description: msg.Description, parameters: msg.Parameters, proc: p,
		})
		p.mu.Unlock()
	case "register_command":
		p.mu.Lock()
		p.commands = append(p.commands, msg.Name)
		p.mu.Unlock()
	case "register_provider":
		p.mu.Lock()
		p.providers = append(p.providers, llm.Provider{
			Name: msg.Name, Protocol: llm.Protocol(msg.Protocol), BaseURL: msg.BaseURL,
		})
		p.mu.Unlock()
	case "tool_result", "hook_result":
		p.mu.Lock()
		ch := p.pending[msg.ID]
		delete(p.pending, msg.ID)
		p.mu.Unlock()
		if ch != nil {
			ch <- msg
		}
	case "ui_request":
		select {
		case p.UI <- UIRequest{ID: msg.ID, Message: msg.Message}:
		default:
		}
	case "custom_entry":
		if p.OnCustom != nil {
			p.OnCustom(msg.CustomType, msg.Data)
		}
	}
}

// Close stops the subprocess.
func (p *Process) Close() error {
	p.mu.Lock()
	p.closed = true
	_ = p.stdin.Close()
	p.mu.Unlock()
	if p.cmd.Process != nil {
		_ = p.cmd.Process.Kill()
	}
	return p.cmd.Wait()
}