Agent-Native/rpc/server.go

package rpc

import (
	"context"
	"encoding/json"
	"fmt"

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

// Request is one RPC command. Type matches the pi headless command set.
type Request struct {
	ID            string `json:"id,omitempty"`
	Type          string `json:"type"`
	Message       string `json:"message,omitempty"`
	Model         string `json:"model,omitempty"`
	ThinkingLevel string `json:"thinkingLevel,omitempty"`
	Path          string `json:"path,omitempty"`
	Name          string `json:"name,omitempty"`
	Leaf          string `json:"leaf,omitempty"`
	Keep          int    `json:"keep,omitempty"`
}

// Response is the synchronous result of one command.
type Response struct {
	ID     string `json:"id,omitempty"`
	Type   string `json:"type"`
	OK     bool   `json:"ok"`
	Text   string `json:"text,omitempty"`
	Error  string `json:"error,omitempty"`
	Result any    `json:"result,omitempty"`
}

// Server serves RPC commands against one agent session and its transcript.
type Server struct {
	Agent   *runtime.AgentSession
	Store   *session.Store
	NoStore bool
}

// Handle runs one command and returns the JSON response line.
func (s *Server) Handle(ctx context.Context, line []byte) ([]byte, error) {
	var req Request
	if err := json.Unmarshal(line, &req); err != nil {
		return encode(Response{Type: "response", OK: false, Error: err.Error()})
	}
	resp := s.dispatch(ctx, req)
	resp.ID = req.ID
	if resp.Type == "" {
		resp.Type = "response"
	}
	return encode(resp)
}

func (s *Server) dispatch(ctx context.Context, req Request) Response {
	switch req.Type {
	case "prompt":
		if s.Agent == nil {
			return fail("no agent")
		}
		res, err := s.Agent.Prompt(ctx, llm.Message{Role: llm.RoleUser, Content: req.Message})
		if err != nil {
			return fail(err.Error())
		}
		if res.Err != nil {
			return fail(res.Err.Error())
		}
		return Response{OK: true, Text: res.Text}
	case "steer":
		s.Agent.Steer(llm.Message{Role: llm.RoleUser, Content: req.Message})
		return Response{OK: true}
	case "follow_up":
		s.Agent.FollowUp(llm.Message{Role: llm.RoleUser, Content: req.Message})
		return Response{OK: true}
	case "abort":
		s.Agent.Abort()
		return Response{OK: true}
	case "clear_queue":
		s.Agent.ClearQueue()
		return Response{OK: true}
	case "set_model":
		if err := s.Agent.SetModel(req.Model, req.ThinkingLevel); err != nil {
			return fail(err.Error())
		}
		return Response{OK: true}
	case "compact":
		keep := req.Keep
		if keep <= 0 {
			keep = 4
		}
		done, err := s.Agent.Compact(keep, func(messages []llm.Message) (string, error) {
			return fmt.Sprintf("%d messages", len(messages)), nil
		})
		if err != nil {
			return fail(err.Error())
		}
		return Response{OK: true, Result: done}
	case "get_state":
		leaf, _ := s.Store.Leaf()
		return Response{OK: true, Result: map[string]any{"leaf": leaf, "name": s.Store.Name(), "path": s.Store.Path()}}
	case "get_tree":
		entries, err := s.Store.Entries()
		if err != nil {
			return fail(err.Error())
		}
		return Response{OK: true, Result: entries}
	case "get_entries":
		entries, err := s.Store.Entries()
		if err != nil {
			return fail(err.Error())
		}
		return Response{OK: true, Result: entries}
	case "get_last_assistant_text":
		msgs, err := s.Store.Load()
		if err != nil {
			return fail(err.Error())
		}
		text := ""
		for _, m := range msgs {
			if m.Role == llm.RoleAssistant {
				text = m.Content
			}
		}
		return Response{OK: true, Text: text}
	case "set_session_name":
		if err := s.Store.SetName(req.Name); err != nil {
			return fail(err.Error())
		}
		return Response{OK: true}
	case "fork":
		dest := req.Path
		if dest == "" {
			return fail("path is required")
		}
		if _, err := s.Store.Fork(dest, req.Leaf); err != nil {
			return fail(err.Error())
		}
		return Response{OK: true, Text: dest}
	case "clone":
		if req.Path == "" {
			return fail("path is required")
		}
		if _, err := s.Store.Clone(req.Path); err != nil {
			return fail(err.Error())
		}
		return Response{OK: true, Text: req.Path}
	case "switch_session":
		if req.Path == "" {
			return fail("path is required")
		}
		st, err := session.Open(req.Path)
		if err != nil {
			return fail(err.Error())
		}
		s.Store = st
		return Response{OK: true}
	case "new_session":
		if req.Path == "" {
			return fail("path is required")
		}
		st, err := session.Open(req.Path)
		if err != nil {
			return fail(err.Error())
		}
		s.Store = st
		return Response{OK: true, Text: req.Path}
	case "export_html":
		if req.Path == "" {
			return fail("path is required")
		}
		if err := s.Store.ExportHTML(req.Path); err != nil {
			return fail(err.Error())
		}
		return Response{OK: true, Text: req.Path}
	case "get_session_stats":
		st := s.Store.Stats()
		return Response{OK: true, Result: st}
	default:
		return fail("unknown command " + req.Type)
	}
}

func fail(msg string) Response { return Response{OK: false, Error: msg} }

func encode(resp Response) ([]byte, error) {
	raw, err := json.Marshal(resp)
	if err != nil {
		return nil, err
	}
	return append(raw, '\n'), nil
}