105 lines
2.5 KiB
Go
105 lines
2.5 KiB
Go
package agent
|
|
|
|
import (
|
|
"encoding/json"
|
|
"io"
|
|
"strings"
|
|
"testing"
|
|
|
|
"ollie/toolsrv"
|
|
)
|
|
|
|
// mockToolServer creates a Conn backed by a goroutine that responds to
|
|
// list_tools with the given ToolInfo slice. Simulates the ollie-remote RPC.
|
|
func mockToolServer(t *testing.T, infos []toolsrv.ToolInfo) *toolsrv.Conn {
|
|
t.Helper()
|
|
cr, sw := io.Pipe() // client reads, server writes
|
|
sr, cw := io.Pipe() // server reads, client writes
|
|
|
|
go func() {
|
|
dec := json.NewDecoder(sr)
|
|
enc := json.NewEncoder(sw)
|
|
for {
|
|
var req struct {
|
|
JSONRPC string `json:"jsonrpc"`
|
|
ID int64 `json:"id"`
|
|
Method string `json:"method"`
|
|
Params json.RawMessage `json:"params,omitempty"`
|
|
}
|
|
if err := dec.Decode(&req); err != nil {
|
|
return
|
|
}
|
|
switch req.Method {
|
|
case "list_tools":
|
|
result, _ := json.Marshal(infos)
|
|
enc.Encode(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": json.RawMessage(result)})
|
|
default:
|
|
enc.Encode(map[string]any{"jsonrpc": "2.0", "id": req.ID, "result": true})
|
|
}
|
|
}
|
|
}()
|
|
|
|
rwc := struct {
|
|
io.Reader
|
|
io.Writer
|
|
io.Closer
|
|
}{cr, cw, cw}
|
|
return toolsrv.NewConn(rwc, nil)
|
|
}
|
|
|
|
func TestRefreshToolListing(t *testing.T) {
|
|
infos := []toolsrv.ToolInfo{
|
|
{
|
|
Name: "new_tool",
|
|
Description: "New description",
|
|
Prompt: "## new_tool\n\nNew description with details.",
|
|
},
|
|
}
|
|
|
|
conn := mockToolServer(t, infos)
|
|
defer conn.Close()
|
|
|
|
p := &Preamble{}
|
|
p.Set(SectionSystem, "# System")
|
|
p.Set(SectionEnv, "# Env")
|
|
p.Set(SectionAgent, "# Agent")
|
|
p.Set(SectionTools, "# Tools\n\n## old_tool\n\nOld description\n\n")
|
|
|
|
ag := &Agent{runtime: &Runtime{Preamble: p, ToolServer: conn}}
|
|
|
|
ag.refreshToolListing()
|
|
|
|
preamble := ag.runtime.PreambleString()
|
|
|
|
if !strings.Contains(preamble, "## new_tool") {
|
|
t.Error("should contain new tool")
|
|
}
|
|
if !strings.Contains(preamble, "New description with details.") {
|
|
t.Error("should contain tool prompt/docs")
|
|
}
|
|
if strings.Contains(preamble, "old_tool") {
|
|
t.Error("should not contain old tool")
|
|
}
|
|
}
|
|
|
|
func TestRefreshToolListingEmpty(t *testing.T) {
|
|
p := &Preamble{}
|
|
p.Set(SectionSystem, "# System")
|
|
p.Set(SectionEnv, "# Env")
|
|
p.Set(SectionAgent, "# Agent")
|
|
p.Set(SectionTools, "# Tools\n\n## tool\n\nDesc\n\n")
|
|
|
|
ag := &Agent{runtime: &Runtime{Preamble: p}}
|
|
|
|
ag.refreshToolListing()
|
|
|
|
preamble := ag.runtime.PreambleString()
|
|
|
|
if strings.Contains(preamble, "# Tools") {
|
|
t.Error("empty listing should remove tools section")
|
|
}
|
|
if !strings.Contains(preamble, "# System") {
|
|
t.Error("system section should be preserved")
|
|
}
|
|
}
|