feat(agentx): add the turn loop and JSONL transcripts

Persist each message as it happens and keep going after a tool failure, including steering and follow-up messages between turns.
2026-09-21 14:53:09 +00:00
parent 738bbe1579
commit fd1b676e28
4 changed files with 554 additions and 0 deletions
@@ -0,0 +1,257 @@
+package agentx
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"sync"
+
+	"chat-agent-system/internal/agentx/llm"
+	"chat-agent-system/internal/agentx/session"
+	"chat-agent-system/internal/agentx/tools"
+)
+
+const (
+	ExecSequential = "sequential"
+	ExecParallel   = "parallel"
+)
+
+// Event is a lifecycle notice for logs and UI.
+type Event struct {
+	Type    string
+	Message llm.Message
+	Tool    string
+	Text    string
+}
+
+// Config drives one agent session.
+type Config struct {
+	Model         string
+	SystemPrompt  string
+	Tools         *tools.Registry
+	Stream        llm.StreamFunc
+	Session       *session.Store
+	MaxTurns      int
+	ToolExecution string
+	GetSteering   func() []llm.Message
+	GetFollowUp   func() []llm.Message
+	OnEvent       func(Event)
+}
+
+// Result is the final assistant text and the transcript produced by one run.
+type Result struct {
+	Text       string
+	Messages   []llm.Message
+	StopReason llm.StopReason
+	Usage      llm.Usage
+	Err        error
+	Retryable  bool
+}
+
+// Run executes the turn loop until the model stops, follow-ups drain, or the turn cap is hit.
+func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) {
+	if cfg.Stream == nil {
+		return Result{StopReason: llm.StopError, Err: errString("stream func is required")}, nil
+	}
+	if cfg.MaxTurns <= 0 {
+		cfg.MaxTurns = 12
+	}
+	if cfg.ToolExecution == "" {
+		cfg.ToolExecution = ExecSequential
+	}
+	var messages []llm.Message
+	if cfg.Session != nil {
+		loaded, err := cfg.Session.Load()
+		if err != nil {
+			return Result{}, err
+		}
+		messages = loaded
+	}
+	if len(messages) == 0 || messages[0].Role != llm.RoleSystem {
+		sys := llm.Message{Role: llm.RoleSystem, Content: systemText(cfg)}
+		messages = append([]llm.Message{sys}, messages...)
+		if err := persist(cfg, sys); err != nil {
+			return Result{}, err
+		}
+	}
+	if err := persist(cfg, user); err != nil {
+		return Result{}, err
+	}
+	messages = append(messages, user)
+	emit(cfg, Event{Type: "agent_start"})
+
+	var last llm.Message
+	var usage llm.Usage
+	for turn := 0; turn < cfg.MaxTurns; turn++ {
+		if ctx.Err() != nil {
+			return Result{Text: last.Content, Messages: messages, StopReason: llm.StopAborted, Err: ctx.Err()}, nil
+		}
+		emit(cfg, Event{Type: "turn_start"})
+		end, errEv := consume(ctx, cfg, messages)
+		if errEv != nil {
+			return Result{Text: end.Message.Content, Messages: messages, StopReason: end.Message.StopReason, Err: errEv, Retryable: end.Retryable}, nil
+		}
+		last = end.Message
+		usage.PromptTokens += end.Usage.PromptTokens
+		usage.CompletionTokens += end.Usage.CompletionTokens
+		usage.TotalTokens += end.Usage.TotalTokens
+		if err := persist(cfg, last); err != nil {
+			return Result{}, err
+		}
+		messages = append(messages, last)
+		emit(cfg, Event{Type: "message_end", Message: last})
+
+		if last.StopReason == llm.StopToolCalls && len(last.ToolCalls) > 0 {
+			results := executeTools(ctx, cfg, last)
+			for _, res := range results {
+				if err := persist(cfg, res); err != nil {
+					return Result{}, err
+				}
+				messages = append(messages, res)
+			}
+			if steered := drain(cfg.GetSteering); len(steered) > 0 {
+				for _, m := range steered {
+					if err := persist(cfg, m); err != nil {
+						return Result{}, err
+					}
+					messages = append(messages, m)
+				}
+			}
+			emit(cfg, Event{Type: "turn_end", Message: last})
+			continue
+		}
+		emit(cfg, Event{Type: "turn_end", Message: last})
+		if follow := drain(cfg.GetFollowUp); len(follow) > 0 {
+			for _, m := range follow {
+				if err := persist(cfg, m); err != nil {
+					return Result{}, err
+				}
+				messages = append(messages, m)
+			}
+			continue
+		}
+		emit(cfg, Event{Type: "agent_end", Message: last})
+		return Result{Text: last.Content, Messages: messages, StopReason: last.StopReason, Usage: usage}, nil
+	}
+	return Result{
+		Text:       last.Content,
+		Messages:   messages,
+		StopReason: llm.StopError,
+		Usage:      usage,
+		Err:        errString("max turns exceeded"),
+	}, nil
+}
+
+func systemText(cfg Config) string {
+	var b strings.Builder
+	b.WriteString(cfg.SystemPrompt)
+	if cfg.Tools == nil {
+		return b.String()
+	}
+	decls := cfg.Tools.Decls()
+	if len(decls) == 0 {
+		return b.String()
+	}
+	if b.Len() > 0 {
+		b.WriteString("\n\n")
+	}
+	b.WriteString("Tools available in this workspace:")
+	for _, d := range decls {
+		b.WriteString("\n- ")
+		b.WriteString(d.Name)
+		b.WriteString(": ")
+		b.WriteString(d.Description)
+	}
+	return b.String()
+}
+
+func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event, error) {
+	var decls []llm.ToolDecl
+	if cfg.Tools != nil {
+		for _, d := range cfg.Tools.Decls() {
+			decls = append(decls, llm.ToolDecl{Name: d.Name, Description: d.Description, Parameters: d.Parameters})
+		}
+	}
+	var end llm.Event
+	sawEnd := false
+	for ev := range cfg.Stream(ctx, cfg.Model, messages, decls, llm.StreamOptions{}) {
+		switch ev.Kind {
+		case llm.KindTextDelta:
+			emit(cfg, Event{Type: "text_delta", Text: ev.Text})
+		case llm.KindError:
+			return ev, ev.Err
+		case llm.KindMessageEnd:
+			end = ev
+			sawEnd = true
+		}
+	}
+	if !sawEnd {
+		return llm.Event{Message: llm.Message{StopReason: llm.StopError}}, errString("stream ended without a message")
+	}
+	return end, nil
+}
+
+func executeTools(ctx context.Context, cfg Config, assistant llm.Message) []llm.Message {
+	if cfg.ToolExecution == ExecParallel && len(assistant.ToolCalls) > 1 {
+		return executeParallel(ctx, cfg, assistant)
+	}
+	out := make([]llm.Message, 0, len(assistant.ToolCalls))
+	for _, call := range assistant.ToolCalls {
+		out = append(out, runOne(ctx, cfg, call))
+	}
+	return out
+}
+
+func executeParallel(ctx context.Context, cfg Config, assistant llm.Message) []llm.Message {
+	out := make([]llm.Message, len(assistant.ToolCalls))
+	var wg sync.WaitGroup
+	for i, call := range assistant.ToolCalls {
+		wg.Add(1)
+		go func(i int, call llm.ToolCall) {
+			defer wg.Done()
+			out[i] = runOne(ctx, cfg, call)
+		}(i, call)
+	}
+	wg.Wait()
+	return out
+}
+
+func runOne(ctx context.Context, cfg Config, call llm.ToolCall) llm.Message {
+	emit(cfg, Event{Type: "tool_execution_start", Tool: call.Name})
+	res := tools.Result{Content: "no tools configured", IsError: true}
+	if cfg.Tools != nil {
+		res = cfg.Tools.Execute(ctx, call.Name, json.RawMessage(call.Arguments))
+	}
+	emit(cfg, Event{Type: "tool_execution_end", Tool: call.Name, Text: res.Content})
+	return llm.Message{
+		Role:       llm.RoleTool,
+		Name:       call.Name,
+		ToolCallID: call.ID,
+		Content:    res.Content,
+		IsError:    res.IsError,
+	}
+}
+
+func persist(cfg Config, msg llm.Message) error {
+	if cfg.Session == nil {
+		return nil
+	}
+	return cfg.Session.Append(msg)
+}
+
+func drain(fn func() []llm.Message) []llm.Message {
+	if fn == nil {
+		return nil
+	}
+	return fn()
+}
+
+func emit(cfg Config, ev Event) {
+	if cfg.OnEvent != nil {
+		cfg.OnEvent(ev)
+	}
+}
+
+type errString string
+
+func (e errString) Error() string { return string(e) }
@@ -0,0 +1,186 @@
+package agentx
+
+import (
+	"context"
+	"path/filepath"
+	"strings"
+	"sync"
+	"testing"
+
+	"chat-agent-system/internal/agentx/llm"
+	"chat-agent-system/internal/agentx/session"
+	"chat-agent-system/internal/agentx/tools"
+)
+
+func script(turns ...llm.Message) llm.StreamFunc {
+	var mu sync.Mutex
+	i := 0
+	return func(ctx context.Context, model string, messages []llm.Message, tools []llm.ToolDecl, opts llm.StreamOptions) <-chan llm.Event {
+		ch := make(chan llm.Event, 4)
+		go func() {
+			defer close(ch)
+			mu.Lock()
+			defer mu.Unlock()
+			if i >= len(turns) {
+				ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{Role: llm.RoleAssistant, Content: "done", StopReason: llm.StopEnd}}
+				return
+			}
+			msg := turns[i]
+			i++
+			if msg.Content != "" {
+				ch <- llm.Event{Kind: llm.KindTextDelta, Text: msg.Content}
+			}
+			ch <- llm.Event{Kind: llm.KindMessageEnd, Message: msg}
+		}()
+		return ch
+	}
+}
+
+func TestLoopToolThenContinueAfterToolError(t *testing.T) {
+	dir := t.TempDir()
+	reg := tools.WorkspaceTools(dir)
+	res, err := Run(context.Background(), Config{
+		Model: "fake",
+		Tools: reg,
+		Stream: script(
+			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
+				ID: "c1", Name: "read_file", Arguments: `{"path":"missing.md"}`,
+			}}},
+			llm.Message{Role: llm.RoleAssistant, Content: "recovered", StopReason: llm.StopEnd},
+		),
+	}, llm.Message{Role: llm.RoleUser, Content: "read it"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if res.Text != "recovered" || res.StopReason != llm.StopEnd {
+		t.Fatalf("%#v", res)
+	}
+	sawToolErr := false
+	for _, m := range res.Messages {
+		if m.Role == llm.RoleTool && m.IsError {
+			sawToolErr = true
+		}
+	}
+	if !sawToolErr {
+		t.Fatal("expected tool error to continue the loop")
+	}
+}
+
+func TestLoopSteeringAndFollowUp(t *testing.T) {
+	steered := false
+	followed := false
+	res, err := Run(context.Background(), Config{
+		Model: "fake",
+		Stream: script(
+			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{{
+				ID: "c1", Name: "list_dir", Arguments: `{}`,
+			}}},
+			llm.Message{Role: llm.RoleAssistant, Content: "after-steer", StopReason: llm.StopEnd},
+			llm.Message{Role: llm.RoleAssistant, Content: "after-follow", StopReason: llm.StopEnd},
+		),
+		Tools: tools.WorkspaceTools(t.TempDir()),
+		GetSteering: func() []llm.Message {
+			if steered {
+				return nil
+			}
+			steered = true
+			return []llm.Message{{Role: llm.RoleUser, Content: "add a footer"}}
+		},
+		GetFollowUp: func() []llm.Message {
+			if followed {
+				return nil
+			}
+			followed = true
+			return []llm.Message{{Role: llm.RoleUser, Content: "one more"}}
+		},
+	}, llm.Message{Role: llm.RoleUser, Content: "go"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if res.Text != "after-follow" {
+		t.Fatalf("text %q", res.Text)
+	}
+	joined := transcriptText(res.Messages)
+	if !strings.Contains(joined, "add a footer") || !strings.Contains(joined, "one more") {
+		t.Fatalf("transcript %s", joined)
+	}
+}
+
+func TestLoopTranscriptReplay(t *testing.T) {
+	path := filepath.Join(t.TempDir(), ".agentx", "transcript.jsonl")
+	st, err := session.Open(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = Run(context.Background(), Config{
+		Model:        "fake",
+		SystemPrompt: "be brief",
+		Session:      st,
+		Stream:       script(llm.Message{Role: llm.RoleAssistant, Content: "pong", StopReason: llm.StopEnd}),
+	}, llm.Message{Role: llm.RoleUser, Content: "ping"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	loaded, err := st.Load()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(loaded) != 3 || loaded[0].Role != llm.RoleSystem || loaded[1].Content != "ping" || loaded[2].Content != "pong" {
+		t.Fatalf("%#v", loaded)
+	}
+	again, err := session.Open(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	replay, err := again.Load()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(replay) != len(loaded) || replay[2].Content != loaded[2].Content {
+		t.Fatal("replay diverged")
+	}
+}
+
+func TestLoopParallelReads(t *testing.T) {
+	dir := t.TempDir()
+	reg := tools.WorkspaceTools(dir)
+	if reg.Execute(context.Background(), "write_file", []byte(`{"path":"a.txt","content":"A"}`)).IsError {
+		t.Fatal("seed a")
+	}
+	if reg.Execute(context.Background(), "write_file", []byte(`{"path":"b.txt","content":"B"}`)).IsError {
+		t.Fatal("seed b")
+	}
+	res, err := Run(context.Background(), Config{
+		Model:         "fake",
+		Tools:         reg,
+		ToolExecution: ExecParallel,
+		Stream: script(
+			llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopToolCalls, ToolCalls: []llm.ToolCall{
+				{ID: "1", Name: "read_file", Arguments: `{"path":"a.txt"}`},
+				{ID: "2", Name: "read_file", Arguments: `{"path":"b.txt"}`},
+			}},
+			llm.Message{Role: llm.RoleAssistant, Content: "AB", StopReason: llm.StopEnd},
+		),
+	}, llm.Message{Role: llm.RoleUser, Content: "read both"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := map[string]string{}
+	for _, m := range res.Messages {
+		if m.Role == llm.RoleTool {
+			got[m.ToolCallID] = m.Content
+		}
+	}
+	if got["1"] != "A" || got["2"] != "B" {
+		t.Fatalf("parallel results %#v", got)
+	}
+}
+
+func transcriptText(messages []llm.Message) string {
+	var b strings.Builder
+	for _, m := range messages {
+		b.WriteString(m.Content)
+		b.WriteByte('\n')
+	}
+	return b.String()
+}
@@ -0,0 +1,73 @@
+package session
+
+import (
+	"bufio"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"sync"
+
+	"chat-agent-system/internal/agentx/llm"
+)
+
+// Store is an append-only JSONL transcript. Each line is one llm.Message.
+type Store struct {
+	path string
+	mu   sync.Mutex
+}
+
+func Open(path string) (*Store, error) {
+	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
+		return nil, err
+	}
+	f, err := os.OpenFile(path, os.O_CREATE, 0o644)
+	if err != nil {
+		return nil, err
+	}
+	_ = f.Close()
+	return &Store{path: path}, nil
+}
+
+func (s *Store) Path() string { return s.path }
+
+func (s *Store) Append(msg llm.Message) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	raw, err := json.Marshal(msg)
+	if err != nil {
+		return err
+	}
+	f, err := os.OpenFile(s.path, os.O_APPEND|os.O_WRONLY, 0o644)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	raw = append(raw, '\n')
+	_, err = f.Write(raw)
+	return err
+}
+
+func (s *Store) Load() ([]llm.Message, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	f, err := os.Open(s.path)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	var out []llm.Message
+	sc := bufio.NewScanner(f)
+	sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
+	for sc.Scan() {
+		line := sc.Bytes()
+		if len(line) == 0 {
+			continue
+		}
+		var msg llm.Message
+		if err := json.Unmarshal(line, &msg); err != nil {
+			return nil, err
+		}
+		out = append(out, msg)
+	}
+	return out, sc.Err()
+}
@@ -0,0 +1,38 @@
+package session
+
+import (
+	"path/filepath"
+	"testing"
+
+	"chat-agent-system/internal/agentx/llm"
+)
+
+func TestAppendAndReplay(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "transcript.jsonl")
+	st, err := Open(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	msgs := []llm.Message{
+		{Role: llm.RoleSystem, Content: "sys"},
+		{Role: llm.RoleUser, Content: "hi"},
+		{Role: llm.RoleAssistant, Content: "ok", StopReason: llm.StopEnd},
+	}
+	for _, m := range msgs {
+		if err := st.Append(m); err != nil {
+			t.Fatal(err)
+		}
+	}
+	got, err := st.Load()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(got) != len(msgs) {
+		t.Fatalf("len %d", len(got))
+	}
+	for i := range msgs {
+		if got[i].Role != msgs[i].Role || got[i].Content != msgs[i].Content || got[i].StopReason != msgs[i].StopReason {
+			t.Fatalf("msg %d %#v", i, got[i])
+		}
+	}
+}