Cyrene_v2/internal/runtime/loop/flow_test.go

package loop

import (
	"context"
	"sync"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/trace"
)

// TestRun_PublishesFlowGraph_Trunk wires a FlowPublisher with a capturing publish
// function and a live tracer, runs one tool round and asserts ws.graph_update was
// published with the input→llm→tool→output trunk (design/19 §1 P1).
func TestRun_PublishesFlowGraph_Trunk(t *testing.T) {
	var mu sync.Mutex
	var seen []*trace.FlowGraph
	flow := NewFlowPublisher(func(subject string, payload any) {
		if subject != graphUpdateSubject {
			return
		}
		if g, ok := payload.(*trace.FlowGraph); ok {
			mu.Lock()
			seen = append(seen, g)
			mu.Unlock()
		}
	}, func(string) bool { return true })

	call := `{"tool_calls":[{"name":"echo","args":{"q":"hi"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "Echoed!"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)
	l := New(Options{Agent: a, MaxRounds: 8, Flow: flow, Tracer: trace.New(nil)})

	res, err := l.Run(context.Background(), "web", []model.Message{model.TextMessage(model.RoleUser, "echo this")})
	if err != nil {
		t.Fatalf("Run: %v", err)
	}
	if res.Text != "Echoed!" {
		t.Fatalf("Text = %q, want Echoed!", res.Text)
	}

	mu.Lock()
	defer mu.Unlock()
	if len(seen) == 0 {
		t.Fatal("no ws.graph_update published")
	}
	last := seen[len(seen)-1]
	if last.SessionID != "web" {
		t.Fatalf("SessionID = %q, want web", last.SessionID)
	}

	kinds := make(map[string]string)
	for _, n := range last.Nodes {
		kinds[n.ID] = n.Kind
	}
	nodeKind := map[string]bool{"input": false, "llm": false, "tool": false, "output": false}
	inputID, outputID := "", ""
	for _, n := range last.Nodes {
		if n.Kind == "input" {
			inputID = n.ID
		}
		if n.Kind == "output" {
			outputID = n.ID
		}
		if _, ok := nodeKind[n.Kind]; ok {
			nodeKind[n.Kind] = true
		}
	}
	for _, want := range []string{"input", "llm", "tool", "output"} {
		if !nodeKind[want] {
			t.Fatalf("trunk incomplete: no %q node; nodes=%+v", want, last.Nodes)
		}
	}

	// Walk the chain from the input node; it must reach output through llm/tool.
	chain := []string{}
	cur := inputID
	for cur != "" {
		chain = append(chain, kinds[cur])
		if cur == outputID {
			break
		}
		cur = nextNode(last, cur)
	}
	hasLLM, hasTool := false, false
	for _, k := range chain {
		if k == "llm" {
			hasLLM = true
		}
		if k == "tool" {
			hasTool = true
		}
	}
	if chain[len(chain)-1] != "output" || !hasLLM || !hasTool {
		t.Fatalf("trunk chain broken: %v", chain)
	}
	// The final frame should show the tool completed (done), not still running.
	for _, n := range last.Nodes {
		if n.Kind == "tool" && n.Status != trace.FlowStatusDone {
			t.Fatalf("tool status = %q, want done", n.Status)
		}
	}
}

// nextNode returns the node that cur points to, or "" when cur has no outgoing
// edge (the terminal output node). A linear chain has at most one successor.
func nextNode(g *trace.FlowGraph, cur string) string {
	for _, e := range g.Edges {
		if e.From == cur {
			return e.To
		}
	}
	return ""
}