Make two WhatsApp tests independent of the load

[onechat] / subagent_run_test.go

diff --git a/subagent_run_test.go b/subagent_run_test.go
index 5caf0d886ef3a6855ef142f1ed28ca41efec46f9..9b8b8c5b8279b43c8369a1b26d2856c45e9e7ffa 100644 (file)
--- a/subagent_run_test.go
+++ b/subagent_run_test.go
@@ -1,12 +1,14 @@
 package main
 
 import (
 package main
 
 import (

+       "context"

        "io"
        "net/http"
        "os"
        "path/filepath"
        "strconv"
        "strings"
        "io"
        "net/http"
        "os"
        "path/filepath"
        "strconv"
        "strings"

+       "sync/atomic"

        "testing"
        "time"
 
        "testing"
        "time"
 


@@ -414,6 +416,140 @@ func TestCloseEndsTheSubagentsWithNoReport(t *testing.T) {
        require.Len(t, entries, 1)
 }
 
        require.Len(t, entries, 1)
 }
 

+func TestRestartResumesTheCutSubagents(t *testing.T) {
+       server := newModelFake(t)
+       arrived := make(chan struct{}, 1)
+       release := make(chan struct{})
+       t.Cleanup(func() {
+               close(release)
+       })
+
+       var restarted atomic.Bool
+       server.reply = func(w http.ResponseWriter, call ModelFakeCall) {
+               switch {
+               case subagentTestCall(call) && !restarted.Load():
+                       arrived <- struct{}{}
+                       <-release
+               case subagentTestCall(call):
+                       modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("Built.")}, 100, 0, 5))
+               default:
+                       modelWriteEvents(w, modelCompletedEvent([]any{modelMessageItem("Noted.")}, 100, 0, 5))
+               }
+       }
+
+       agent, err := openAgent(t.TempDir())
+       require.NoError(t, err)
+
+       first := &TurnEventRecorder{}
+       loop, _, stop := subagentTestStart(t, agent, server, first)
+       _, err = loop.runs.Spawn(t.Context(), []string{"Build the project."}, nil)
+       require.NoError(t, err)
+
+       <-arrived
+       name := first.Last("subagent")
+       stop()
+       require.Zero(t, first.Count("work"))
+
+       data, err := os.ReadFile(filepath.Join(agent.dir, subagentWorkFileName))
+       require.NoError(t, err)
+       require.JSONEq(t, `{"running": ["`+name+`"]}`, string(data))
+
+       restarted.Store(true)
+       second := &TurnEventRecorder{}
+       loop, memory, stop := subagentTestStart(t, agent, server, second)
+       defer stop()
+
+       second.WaitFor(t, "idle", 1)
+       note := "The daemon restarted, cutting " + name + "; nothing was lost: go on."
+       log := subagentTestChat(t, memory)
+       require.Equal(t, []string{
+               "note: " + note,
+               "work: " + subagentReport(name, []string{"Built."}),
+               "reply: Noted.",
+       }, log[len(log)-3:])
+
+       subagent, err := loop.runs.subagents.Get(t.Context(), name)
+       require.NoError(t, err)
+
+       text, err := subagent.Show()
+       require.NoError(t, err)
+       require.Equal(t, "0|note: subagent\n\n1|user: Build the project.\n\n2|user: "+note+"\n\n3|reply: Built.", text)
+
+       data, err = os.ReadFile(filepath.Join(agent.dir, subagentWorkFileName))
+       require.NoError(t, err)
+       require.JSONEq(t, `{"running": []}`, string(data))
+}
+
+func TestRestartReportsASubagentThatCannotGoOn(t *testing.T) {
+       server := newModelFake(t)
+       agent, err := openAgent(t.TempDir())
+       require.NoError(t, err)
+       require.NoError(t, os.WriteFile(filepath.Join(agent.dir, subagentWorkFileName), []byte(`{"running": ["BlueFox"]}`), 0o600))
+
+       events := &TurnEventRecorder{}
+       _, memory, stop := subagentTestStart(t, agent, server, events)
+       defer stop()
+
+       events.WaitFor(t, "idle", 1)
+       require.Equal(t, []string{
+               "note: The daemon restarted, cutting BlueFox; nothing was lost: go on.",
+               "work: " + subagentReport("BlueFox", []string{"failed: no agent BlueFox"}),
+               "reply: a line",
+       }, subagentTestChat(t, memory))
+       require.Equal(t, subagentFailed, events.LastEvent("work").Status)
+
+       data, err := os.ReadFile(filepath.Join(agent.dir, subagentWorkFileName))
+       require.NoError(t, err)
+       require.JSONEq(t, `{"running": []}`, string(data))
+}
+
+// subagentTestStart runs a loop of agent with its memory and subagents in the
+// agent directory, as the daemon does. The returned function stops the loop
+// and closes the files.
+func subagentTestStart(t *testing.T, agent *Agent, server *ModelFake, events *TurnEventRecorder) (*TurnLoop, *Memory, func()) {
+       t.Helper()
+
+       agent.config.Workdir = t.TempDir()
+       clock := newTestClock()
+       memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer())
+       require.NoError(t, err)
+
+       subagents := newSubagents(filepath.Join(agent.dir, subagentDirName), clock.Now, subagentNameRandom)
+       client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
+       whatsapp, _, _ := whatsappStartTest(t, t.TempDir(), clock)
+       loop, err := newTurnLoop(agent, memory, subagents, whatsapp, client, clock.Now, events.Add)
+       require.NoError(t, err)
+
+       ctx, cancel := context.WithCancel(t.Context())
+       done := make(chan error)
+       go func() {
+               done <- loop.Run(ctx)
+       }()
+
+       return loop, memory, func() {
+               cancel()
+               require.ErrorIs(t, <-done, context.Canceled)
+               require.NoError(t, subagents.Close())
+               require.NoError(t, memory.Close())
+       }
+}
+
+// subagentTestChat returns each message of the chat in memory as
+// "kind: text".
+func subagentTestChat(t *testing.T, memory *Memory) []string {
+       t.Helper()
+
+       var lines []string
+       for id := range memory.Len() {
+               message, err := memory.Message(id)
+               require.NoError(t, err)
+
+               lines = append(lines, message.Kind+": "+message.Text)
+       }
+
+       return lines
+}
+

 // subagentTestCall tells if call is a call of a subagent.
 func subagentTestCall(call ModelFakeCall) bool {
        instructions, _ := call.body["instructions"].(string)
 // subagentTestCall tells if call is a call of a subagent.
 func subagentTestCall(call ModelFakeCall) bool {
        instructions, _ := call.body["instructions"].(string)