From fd1b676e280bccfe464b6d8b37025877307859fa Mon Sep 17 00:00:00 2001 From: weichen Date: Mon, 21 Sep 2026 14:53:09 +0000 Subject: [PATCH] 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. --- internal/agentx/loop.go | 257 ++++++++++++++++++++++++++ internal/agentx/loop_test.go | 186 +++++++++++++++++++ internal/agentx/session/jsonl.go | 73 ++++++++ internal/agentx/session/jsonl_test.go | 38 ++++ 4 files changed, 554 insertions(+) create mode 100644 internal/agentx/loop.go create mode 100644 internal/agentx/loop_test.go create mode 100644 internal/agentx/session/jsonl.go create mode 100644 internal/agentx/session/jsonl_test.go diff --git a/internal/agentx/loop.go b/internal/agentx/loop.go new file mode 100644 index 00000000..6841a3dc --- /dev/null +++ b/internal/agentx/loop.go @@ -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) } diff --git a/internal/agentx/loop_test.go b/internal/agentx/loop_test.go new file mode 100644 index 00000000..28ebe774 --- /dev/null +++ b/internal/agentx/loop_test.go @@ -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() +} diff --git a/internal/agentx/session/jsonl.go b/internal/agentx/session/jsonl.go new file mode 100644 index 00000000..b2417bca --- /dev/null +++ b/internal/agentx/session/jsonl.go @@ -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() +} diff --git a/internal/agentx/session/jsonl_test.go b/internal/agentx/session/jsonl_test.go new file mode 100644 index 00000000..1cc9f10e --- /dev/null +++ b/internal/agentx/session/jsonl_test.go @@ -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]) + } + } +}