mirror/LocalAI
1
0 Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2026-10-28 09:08:54 -04:00
Code Issues Packages Projects Releases Wiki Activity

test(addons): move the add-on tool tests to Ginkgo

Browse Source 
The fine-tuning, quantization and add-on host library tests used
t.Errorf, t.Fatal and t.Run, which the forbidigo rule rejects. They now
use Describe, It and Gomega, with a suite bootstrap in each package. The
cases and the assertions are the same.

The tables in the host client and the tool schema tests keep one entry
per case.

Assisted-by: Claude Code:claude-sonnet-5-5 go
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-10-21 02:20:12 +00:00
1 parent 579ed93daf
commit 63c7c72a48
10 files changed
+883 -964
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

addons/fine-tuning/internal/tools/suite_test.go
+14
View File
@@ -0,0 +1,14 @@
// SPDX-License-Identifier: MIT
package tools
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestTools(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Fine-tuning tools suite")
}
addons/fine-tuning/internal/tools/tools_test.go
+399 -487
View File
@@ -4,18 +4,18 @@ package tools
import (
"context"
"encoding/json"
"errors"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"testing"
"time"
"github.com/mudler/LocalAI/addons/internal/hostclient"
"github.com/mudler/LocalAI/addons/internal/protocol"
"github.com/mudler/LocalAI/pkg/addons/contract"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var ctx = context.Background()
@@ -103,199 +103,16 @@ func envelope(tool, opID, resource, args string) contract.Envelope {
const startJSON = `{"name":"my run","model":"org/base","datasetSource":"org/data"}`
// started returns a runtime that already holds the job job-1 for resource res1.
func started(t *testing.T) (contract.Runtime, *fakeEngine) {
t.Helper()
func started() (contract.Runtime, *fakeEngine) {
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, t.TempDir())
rt := New(eng, GinkgoT().TempDir())
rec, err := rt.Apply(ctx, envelope("fine-tune-start", "create-res1", "res1", startJSON), contract.Revision{Schema: 1})
if err != nil || rec.State != contract.ReceiptCommitted {
t.Fatalf("setup: %v %+v", err, rec)
}
Expect(err).NotTo(HaveOccurred(), "setup")
Expect(rec.State).To(Equal(contract.ReceiptCommitted), "setup: %+v", rec)
return rt, eng
}
func TestStartCallsTheEngineOnceAndRemembersTheJob(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, t.TempDir())
rec, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 0})
if err != nil || rec.State != contract.ReceiptCommitted || rec.Revision.Schema == 0 {
t.Fatalf("%v %+v", err, rec)
}
if got := eng.ops(); !reflect.DeepEqual(got, []string{"finetune.start"}) {
t.Fatalf("ops %v", got)
}
// The same operation again does not start a second job.
again, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 1})
if err != nil || again.State != contract.ReceiptCommitted {
t.Fatalf("%v %+v", err, again)
}
if n := eng.count("finetune.start"); n != 1 {
t.Fatalf("started %d jobs", n)
}
}
func TestStartSendsOnlyTheFieldsOfTheEngineOperation(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, t.TempDir())
args := `{"name":"my run","model":"org/base","datasetSource":"org/data","numEpochs":2,"learningRate":0.0002,"adapterRank":8}`
if _, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", args), contract.Revision{Schema: 1}); err != nil {
t.Fatal(err)
}
c := eng.last("finetune.start")
if c.opID != "op1" {
t.Fatalf("the call must carry the operation of the request, got %q", c.opID)
}
for _, banned := range []string{"tool", "name", "resourceId"} {
if _, ok := c.args[banned]; ok {
t.Fatalf("the engine must not receive %q: %v", banned, c.args)
}
}
if c.args["model"] != "org/base" || c.args["numEpochs"] != float64(2) || c.args["adapterRank"] != float64(8) {
t.Fatalf("args %v", c.args)
}
}
func TestPreviewOfStartHasANonzeroRevisionAndCreatesNothing(t *testing.T) {
eng := &fakeEngine{}
rt := New(eng, t.TempDir())
p, err := rt.Preview(ctx, envelope("fine-tune-start", "op1", "res1", startJSON))
if err != nil || p.Revision.Schema == 0 || len(eng.ops()) != 0 {
t.Fatalf("%v %+v %v", err, p, eng.ops())
}
if !json.Valid(p.Summary) {
t.Fatalf("summary %s", p.Summary)
}
if _, err := rt.Preview(ctx, envelope("fine-tune-start", "op1", "res1", `{"name":"x","model":"","datasetSource":"d"}`)); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("an empty model must be refused, got %v", err)
}
}
func TestStatusReadsTheJobOfTheResource(t *testing.T) {
rt, eng := started(t)
eng.result["finetune.get"] = map[string]any{"id": "job-1", "status": "training", "model": "org/base"}
res, err := rt.Read(ctx, envelope("fine-tune-status", "r1", "res1", ""))
if err != nil {
t.Fatal(err)
}
if got := eng.last("finetune.get").args["jobId"]; got != "job-1" {
t.Fatalf("jobId %v", got)
}
var v map[string]any
if err := json.Unmarshal(res.Payload, &v); err != nil {
t.Fatal(err)
}
if v["name"] != "my run" || v["status"] != "training" {
t.Fatalf("payload %v", v)
}
rev, _ := v["revision"].(map[string]any)
if rev["schema"] == nil || rev["schema"] == float64(0) {
t.Fatalf("inspect needs a revision: %v", v)
}
if _, has := v["id"]; has {
t.Fatalf("the engine job ID must not reach the model: %v", v)
}
if _, has := v["progressPercent"]; has {
t.Fatalf("a snapshot percent is unreliable and must not be reported: %v", v)
}
}
func TestCheckpointsReadTheJobOfTheResource(t *testing.T) {
rt, eng := started(t)
eng.result["finetune.checkpoints"] = map[string]any{"checkpoints": []any{map[string]any{"step": 10}}}
res, err := rt.Read(ctx, envelope("fine-tune-checkpoints", "r1", "res1", ""))
if err != nil || !strings.Contains(string(res.Payload), `"step":10`) {
t.Fatalf("%v %s", err, res.Payload)
}
if eng.last("finetune.checkpoints").args["jobId"] != "job-1" {
t.Fatal("wrong job")
}
}
func TestAToolCannotControlAJobItDidNotStart(t *testing.T) {
eng := &fakeEngine{}
rt := New(eng, t.TempDir())
if _, err := rt.Apply(ctx, envelope("fine-tune-stop", "op2", "nope", `{}`), contract.Revision{Schema: 1}); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("stop: %v", err)
}
if _, err := rt.Apply(ctx, envelope("fine-tune-export", "op3", "nope", `{}`), contract.Revision{Schema: 1}); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("export: %v", err)
}
if _, err := rt.Read(ctx, envelope("fine-tune-status", "r1", "nope", "")); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("status: %v", err)
}
if _, err := rt.Preview(ctx, envelope("fine-tune-stop", "op2", "nope", `{}`)); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("preview: %v", err)
}
if n := len(eng.ops()); n != 0 {
t.Fatalf("the engine was called %d times: %v", n, eng.ops())
}
}
func TestStopAndExportUseTheStoredJobAndAdvanceTheRevision(t *testing.T) {
rt, eng := started(t)
eng.result["finetune.stop"] = map[string]any{}
eng.result["finetune.export"] = map[string]any{"model": "my-model"}
p, err := rt.Preview(ctx, envelope("fine-tune-stop", "op2", "res1", `{"saveCheckpoint":true}`))
if err != nil || p.Revision.Schema == 0 {
t.Fatalf("%v %+v", err, p)
}
rec, err := rt.Apply(ctx, envelope("fine-tune-stop", "op2", "res1", `{"saveCheckpoint":true}`), p.Revision)
if err != nil || rec.State != contract.ReceiptCommitted {
t.Fatalf("%v %+v", err, rec)
}
c := eng.last("finetune.stop")
if c.args["jobId"] != "job-1" || c.args["saveCheckpoint"] != true {
t.Fatalf("stop args %v", c.args)
}
if rec.Revision.Data <= p.Revision.Data {
t.Fatalf("the revision must advance: %+v then %+v", p.Revision, rec.Revision)
}
// A repeat of the same operation is answered from the receipt.
if _, err := rt.Apply(ctx, envelope("fine-tune-stop", "op2", "res1", `{"saveCheckpoint":true}`), p.Revision); err != nil || eng.count("finetune.stop") != 1 {
t.Fatalf("stop ran %d times (%v)", eng.count("finetune.stop"), err)
}
// Export names the format and the model, and never a path.
ex, err := rt.Apply(ctx, envelope("fine-tune-export", "op3", "res1", `{"exportFormat":"gguf","name":"my-model"}`), rec.Revision)
if err != nil || ex.State != contract.ReceiptCommitted || !strings.Contains(string(ex.Result.Payload), "my-model") {
t.Fatalf("%v %+v", err, ex)
}
if a := eng.last("finetune.export").args; a["jobId"] != "job-1" || a["exportFormat"] != "gguf" || a["name"] != "my-model" {
t.Fatalf("export args %v", a)
}
}
func TestAStaleRevisionIsRejectedAndTheEngineIsNotCalled(t *testing.T) {
rt, eng := started(t)
rec, err := rt.Apply(ctx, envelope("fine-tune-stop", "op2", "res1", `{}`), contract.Revision{Schema: 1, Data: 99})
if err != nil || rec.State != contract.ReceiptRejected {
t.Fatalf("%v %+v", err, rec)
}
if eng.count("finetune.stop") != 0 {
t.Fatal("the engine was called")
}
}
func TestHostRefusalsMapToProtocolErrors(t *testing.T) {
rt, eng := started(t)
for code, want := range map[string]error{
contract.HostFeature: protocol.ErrForbidden, contract.HostBadRequest: protocol.ErrBadRequest,
contract.HostNotFound: protocol.ErrNotFound, contract.HostConflict: protocol.ErrConflict,
} {
eng.result["finetune.export"] = &contract.HostError{Code: code, Message: "m"}
_, err := rt.Apply(ctx, envelope("fine-tune-export", "op-"+code, "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
if !errors.Is(err, want) {
t.Fatalf("%s: %v, want %v", code, err, want)
}
}
eng.result["finetune.export"] = &contract.HostError{Code: contract.HostUnavailable, Message: "m"}
_, err := rt.Apply(ctx, envelope("fine-tune-export", "op-u", "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
if err == nil || errors.Is(err, protocol.ErrBadRequest) || errors.Is(err, protocol.ErrForbidden) || errors.Is(err, protocol.ErrNotFound) || errors.Is(err, protocol.ErrConflict) {
t.Fatalf("an unavailable engine must stay uncertain, got %v", err)
}
}
func exists(t *testing.T, rt contract.Runtime, res string) (bool, error) {
t.Helper()
func exists(rt contract.Runtime, res string) (bool, error) {
r, err := rt.Read(ctx, envelope("resource-exists", "x-"+res, res, ""))
if err != nil {
return false, err
@@ -307,309 +124,404 @@ func exists(t *testing.T, rt contract.Runtime, res string) (bool, error) {
return v.Exists, nil
}
func TestResourceExistsAnswersFromTheJournalAndTheEngine(t *testing.T) {
rt, eng := started(t)
eng.result["finetune.get"] = map[string]any{"id": "job-1", "status": "training"}
if ok, err := exists(t, rt, "res1"); err != nil || !ok {
t.Fatalf("mapped: %v %v", ok, err)
}
if ok, err := exists(t, rt, "other"); err != nil || ok {
t.Fatalf("unmapped: %v %v", ok, err)
}
// The job was deleted in the UI: the engine says not found.
eng.result["finetune.get"] = &contract.HostError{Code: contract.HostNotFound, Message: "gone"}
if ok, err := exists(t, rt, "res1"); err != nil || ok {
t.Fatalf("deleted job: %v %v", ok, err)
}
// The engine cannot answer: unknown, never false.
eng.result["finetune.get"] = &contract.HostError{Code: contract.HostUnavailable, Message: "down"}
if _, err := exists(t, rt, "res1"); err == nil {
t.Fatal("an unavailable engine must not read as absent")
}
eng.result["finetune.get"] = hostclient.ErrNoHostOps
if _, err := exists(t, rt, "res1"); err == nil {
t.Fatal("no host endpoint must not read as absent")
}
}
func TestJournalSurvivesARestart(t *testing.T) {
dir := t.TempDir()
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}, "finetune.get": map[string]any{"status": "training"}}}
if _, err := New(eng, dir).Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1}); err != nil {
t.Fatal(err)
}
again := New(eng, dir)
if _, err := again.Read(ctx, envelope("fine-tune-status", "r1", "res1", "")); err != nil {
t.Fatal(err)
}
if eng.last("finetune.get").args["jobId"] != "job-1" {
t.Fatal("the restarted runtime lost the job")
}
// The receipt of the first run answers a repeat after the restart.
if _, err := again.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1}); err != nil || eng.count("finetune.start") != 1 {
t.Fatalf("started %d times (%v)", eng.count("finetune.start"), err)
}
if r, err := again.Receipt(ctx, "op1"); err != nil || r.State != contract.ReceiptCommitted {
t.Fatalf("%v %+v", err, r)
}
if _, err := again.Receipt(ctx, "never"); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("an unknown operation must be not found, got %v", err)
}
// The journal is a plain file written through a temporary one.
entries, _ := os.ReadDir(dir)
for _, e := range entries {
if strings.Contains(e.Name(), ".tmp") {
t.Fatalf("a temporary file was left: %s", e.Name())
}
}
if _, err := os.Stat(filepath.Join(dir, "journal.json")); err != nil {
t.Fatal(err)
}
}
func TestWithoutAHostEndpointEveryToolRefuses(t *testing.T) {
var nilClient *hostclient.Client
dir := t.TempDir()
rt := New(nilClient, dir)
if _, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1}); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("start: %v", err)
}
// Seed a job as an earlier run with an endpoint would have.
seed := New(&fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1"}}}, dir)
if _, err := seed.Apply(ctx, envelope("fine-tune-start", "op0", "res1", startJSON), contract.Revision{Schema: 1}); err != nil {
t.Fatal(err)
}
rt = New(nilClient, dir)
for _, tool := range []string{"fine-tune-status", "fine-tune-checkpoints", "resource-exists"} {
if _, err := rt.Read(ctx, envelope(tool, "r-"+tool, "res1", "")); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("%s: %v", tool, err)
}
}
for _, tool := range []string{"fine-tune-stop", "fine-tune-export"} {
if _, err := rt.Apply(ctx, envelope(tool, "a-"+tool, "res1", `{}`), contract.Revision{Schema: 1, Data: 1}); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("%s: %v", tool, err)
}
}
}
func TestAnUnknownToolIsABadRequest(t *testing.T) {
rt := New(&fakeEngine{}, t.TempDir())
if _, err := rt.Read(ctx, envelope("fine-tune-nope", "r1", "res1", "")); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("read: %v", err)
}
if _, err := rt.Apply(ctx, envelope("fine-tune-nope", "a1", "res1", ""), contract.Revision{Schema: 1}); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("apply: %v", err)
}
if _, err := rt.OpenOutput(ctx, envelope("fine-tune-status", "o1", "res1", ""), "h"); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("output: %v", err)
}
}
func loadManifest(t *testing.T) contract.Manifest {
t.Helper()
func loadManifest() contract.Manifest {
b, err := os.ReadFile("../../manifest.json")
if err != nil {
t.Fatal(err)
}
Expect(err).NotTo(HaveOccurred())
m, err := contract.ParseManifest(b)
if err != nil {
t.Fatalf("the manifest must satisfy the contract: %v", err)
}
Expect(err).NotTo(HaveOccurred(), "the manifest must satisfy the contract")
return m
}
func TestManifestDeclaresTheToolsAndTheResourceMapping(t *testing.T) {
m := loadManifest(t)
if m.ID != "fine-tuning" || m.Capabilities.UI.MainPage != "jobs" || !reflect.DeepEqual(m.Capabilities.UI.Pages, []string{"jobs"}) {
t.Fatalf("manifest %+v", m)
}
if m.Capabilities.UI.SDK.Min != 2 || m.Capabilities.UI.SDK.Max != 2 {
t.Fatalf("sdk %+v", m.Capabilities.UI.SDK)
}
if !reflect.DeepEqual(m.Capabilities.Tools, Tools) {
t.Fatalf("manifest tools %v, package tools %v", m.Capabilities.Tools, Tools)
}
r := m.Capabilities.Resources
if r == nil || r.Noun != "fine-tuning job" || r.Create != "fine-tune-start" || r.Inspect != "fine-tune-status" {
t.Fatalf("resources %+v", r)
}
kinds := map[string]string{}
for _, d := range m.Capabilities.ToolDefinitions {
kinds[d.Name] = d.Kind
}
want := map[string]string{"fine-tune-list": "read", "fine-tune-status": "read", "fine-tune-checkpoints": "read", "fine-tune-start": "create", "fine-tune-dataset": "create", "fine-tune-stop": "write", "fine-tune-export": "write"}
if !reflect.DeepEqual(kinds, want) {
t.Fatalf("kinds %v", kinds)
}
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name == "fine-tune-list" && !contract.IsResourceless(d) {
t.Fatal("the list tool must declare no resource")
}
if d.Name != "fine-tune-list" && contract.IsResourceless(d) {
t.Fatalf("%s must not be resource-less", d.Name)
}
if len(d.Description) >= 400 || !strings.Contains(d.Description, "started") {
t.Fatalf("%s: the description must be short and say that only jobs a tool started can be controlled", d.Name)
}
}
}
func TestNoToolArgumentNamesAUserOrAPath(t *testing.T) {
m := loadManifest(t)
for _, d := range m.Capabilities.ToolDefinitions {
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
}
if err := json.Unmarshal(d.InputSchema, &s); err != nil {
t.Fatal(err)
}
for name := range s.Properties {
lower := strings.ToLower(name)
if contract.ForbiddenToolProperty(name) || strings.Contains(lower, "path") || strings.Contains(lower, "dir") {
t.Fatalf("%s: property %q names an identity or a path", d.Name, name)
}
}
}
}
func TestStartSchemaHasTheTask20FieldsAndTheNameTheHostNeeds(t *testing.T) {
m := loadManifest(t)
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name != "fine-tune-start" {
continue
}
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
Required []string `json:"required"`
}
_ = json.Unmarshal(d.InputSchema, &s)
var got []string
for k := range s.Properties {
got = append(got, k)
}
want := []string{"name", "model", "backend", "trainingType", "trainingMethod", "datasetSource", "datasetSplit", "learningRate", "numEpochs", "batchSize", "maxSteps", "adapterRank", "adapterAlpha"}
if len(got) != len(want) {
t.Fatalf("properties %v", got)
}
for _, w := range want {
if _, ok := s.Properties[w]; !ok {
t.Fatalf("missing %s", w)
}
}
return
}
t.Fatal("no start definition")
}
func TestListReadsTheCallersJobsWithoutAResource(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"finetune.list": map[string]any{"jobs": []any{map[string]any{"id": "ui-job", "status": "completed"}}}}}
rt := New(eng, t.TempDir())
res, err := rt.Read(ctx, envelope("fine-tune-list", "r1", "", ""))
if err != nil || !strings.Contains(string(res.Payload), "ui-job") {
t.Fatalf("%v %s", err, res.Payload)
}
if got := eng.ops(); !reflect.DeepEqual(got, []string{"finetune.list"}) {
t.Fatalf("ops %v", got)
}
if c := eng.last("finetune.list"); c.args != nil {
t.Fatalf("the list takes no argument, got %v", c.args)
}
if _, err := rt.Read(ctx, envelope("fine-tune-list", "r2", "", `{"userId":"other"}`)); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("an argument must be refused, got %v", err)
}
var nilClient *hostclient.Client
if _, err := New(nilClient, t.TempDir()).Read(ctx, envelope("fine-tune-list", "r3", "", "")); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("no endpoint: %v", err)
}
}
const datasetJSON = `{"name":"data.jsonl","content":"{\"text\":\"a\"}\n"}`
func TestDatasetIsSentOnceAndItsRefComesBackOnlyThroughTheResource(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"finetune.dataset": map[string]any{"datasetRef": "0000000a-data.jsonl"}}}
rt := New(eng, t.TempDir())
p, err := rt.Preview(ctx, envelope("fine-tune-dataset", "create-d1", "d1", datasetJSON))
if err != nil || p.Revision.Schema == 0 || len(eng.ops()) != 0 {
t.Fatalf("%v %+v %v", err, p, eng.ops())
}
if strings.Contains(string(p.Summary), "text") {
t.Fatalf("the approval card must not carry the content: %s", p.Summary)
}
rec, err := rt.Apply(ctx, envelope("fine-tune-dataset", "create-d1", "d1", datasetJSON), p.Revision)
if err != nil || rec.State != contract.ReceiptCommitted || !strings.Contains(string(rec.Result.Payload), "0000000a-data.jsonl") {
t.Fatalf("%v %+v", err, rec)
}
if _, err := rt.Apply(ctx, envelope("fine-tune-dataset", "create-d1", "d1", datasetJSON), p.Revision); err != nil || eng.count("finetune.dataset") != 1 {
t.Fatalf("sent %d times (%v)", eng.count("finetune.dataset"), err)
}
if a := eng.last("finetune.dataset").args; a["name"] != "data.jsonl" || a["content"] == nil || len(a) != 2 {
t.Fatalf("host operation args %v", a)
}
// The resource is the dataset: its status gives the reference to the owner of the resource.
st, err := rt.Read(ctx, envelope("fine-tune-status", "r1", "d1", ""))
if err != nil || !strings.Contains(string(st.Payload), `"datasetRef":"0000000a-data.jsonl"`) || !strings.Contains(string(st.Payload), `"kind":"dataset"`) {
t.Fatalf("%v %s", err, st.Payload)
}
if n := len(eng.ops()); n != 1 {
t.Fatalf("the status of a dataset must not call the engine: %v", eng.ops())
}
if ok, err := exists(t, rt, "d1"); err != nil || !ok {
t.Fatalf("exists %v %v", ok, err)
}
// A dataset is not a job.
for _, tool := range []string{"fine-tune-checkpoints"} {
if _, err := rt.Read(ctx, envelope(tool, "r2", "d1", "")); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("%s: %v", tool, err)
}
}
if _, err := rt.Apply(ctx, envelope("fine-tune-stop", "s1", "d1", `{}`), contract.Revision{Schema: 1, Data: 1}); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("stop: %v", err)
}
}
var _ = Describe("fine-tuning tools", func() {
It("start calls the engine once and remembers the job", func() {
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, GinkgoT().TempDir())
rec, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 0})
Expect(err).NotTo(HaveOccurred())
Expect(rec.State).To(Equal(contract.ReceiptCommitted))
Expect(rec.Revision.Schema).NotTo(BeZero())
Expect(eng.ops()).To(Equal([]string{"finetune.start"}))
// The same operation again does not start a second job.
again, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 1})
Expect(err).NotTo(HaveOccurred())
Expect(again.State).To(Equal(contract.ReceiptCommitted))
Expect(eng.count("finetune.start")).To(Equal(1))
})
func TestDatasetRefusesWhatTheHostWouldRefuse(t *testing.T) {
rt := New(&fakeEngine{}, t.TempDir())
big := strings.Repeat("a", 900<<10+1)
for name, args := range map[string]string{
"empty": `{"name":"d.jsonl","content":""}`,
"too large": `{"name":"d.jsonl","content":"` + big + `"}`,
"a folder": `{"name":"a/d.jsonl","content":"x"}`,
"parent": `{"name":"../d.jsonl","content":"x"}`,
"wrong type": `{"name":"d.exe","content":"x"}`,
"only the type": `{"name":".jsonl","content":"x"}`,
"extra field": `{"name":"d.jsonl","content":"x","path":"/etc"}`,
} {
if _, err := rt.Preview(ctx, envelope("fine-tune-dataset", "o-"+name, "d9", args)); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("%s: %v", name, err)
It("start sends only the fields of the engine operation", func() {
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, GinkgoT().TempDir())
args := `{"name":"my run","model":"org/base","datasetSource":"org/data","numEpochs":2,"learningRate":0.0002,"adapterRank":8}`
_, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", args), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
c := eng.last("finetune.start")
Expect(c.opID).To(Equal("op1"), "the call must carry the operation of the request")
for _, banned := range []string{"tool", "name", "resourceId"} {
Expect(c.args).NotTo(HaveKey(banned), "the engine must not receive %q", banned)
}
}
// The largest allowed dataset passes.
ok := `{"name":"d.jsonl","content":"` + strings.Repeat("a", 900<<10) + `"}`
if _, err := rt.Preview(ctx, envelope("fine-tune-dataset", "o-ok", "d9", ok)); err != nil {
t.Fatalf("900 KiB: %v", err)
}
}
Expect(c.args["model"]).To(Equal("org/base"))
Expect(c.args["numEpochs"]).To(Equal(float64(2)))
Expect(c.args["adapterRank"]).To(Equal(float64(8)))
})
func TestStartPassesADatasetRefAndKeepsNoJobWhenTheHostRefusesIt(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, t.TempDir())
args := `{"name":"run","model":"org/base","datasetSource":"0000000a-data.jsonl"}`
if _, err := rt.Apply(ctx, envelope("fine-tune-start", "create-r1", "r1", args), contract.Revision{Schema: 1}); err != nil {
t.Fatal(err)
}
if eng.last("finetune.start").args["datasetSource"] != "0000000a-data.jsonl" {
t.Fatalf("args %v", eng.last("finetune.start").args)
}
// The host refuses the reference (the per-user check is proven in the host tests); the add-on only maps the refusal and keeps nothing.
eng.result["finetune.start"] = &contract.HostError{Code: contract.HostBadRequest, Message: "datasetSource is not a dataset reference of this user"}
_, err := rt.Apply(ctx, envelope("fine-tune-start", "create-r2", "r2", `{"name":"run","model":"org/base","datasetSource":"0000000b-theirs.jsonl"}`), contract.Revision{Schema: 1})
if !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("want a bad request, got %v", err)
}
if _, err := rt.Read(ctx, envelope("fine-tune-status", "r3", "r2", "")); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("a refused start must not leave a job: %v", err)
}
if _, err := rt.Receipt(ctx, "create-r2"); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("a refused start must not leave a receipt: %v", err)
}
}
It("previews a start with a nonzero revision and creates nothing", func() {
eng := &fakeEngine{}
rt := New(eng, GinkgoT().TempDir())
p, err := rt.Preview(ctx, envelope("fine-tune-start", "op1", "res1", startJSON))
Expect(err).NotTo(HaveOccurred())
Expect(p.Revision.Schema).NotTo(BeZero())
Expect(eng.ops()).To(BeEmpty())
Expect(json.Valid(p.Summary)).To(BeTrue(), "summary %s", p.Summary)
_, err = rt.Preview(ctx, envelope("fine-tune-start", "op1", "res1", `{"name":"x","model":"","datasetSource":"d"}`))
Expect(err).To(MatchError(protocol.ErrBadRequest), "an empty model must be refused")
})
It("status reads the job of the resource", func() {
rt, eng := started()
eng.result["finetune.get"] = map[string]any{"id": "job-1", "status": "training", "model": "org/base"}
res, err := rt.Read(ctx, envelope("fine-tune-status", "r1", "res1", ""))
Expect(err).NotTo(HaveOccurred())
Expect(eng.last("finetune.get").args["jobId"]).To(Equal("job-1"))
var v map[string]any
Expect(json.Unmarshal(res.Payload, &v)).To(Succeed())
Expect(v["name"]).To(Equal("my run"))
Expect(v["status"]).To(Equal("training"))
rev, _ := v["revision"].(map[string]any)
Expect(rev["schema"]).NotTo(BeNil(), "inspect needs a revision: %v", v)
Expect(rev["schema"]).NotTo(Equal(float64(0)), "inspect needs a revision: %v", v)
Expect(v).NotTo(HaveKey("id"), "the engine job ID must not reach the model")
Expect(v).NotTo(HaveKey("progressPercent"), "a snapshot percent is unreliable and must not be reported")
})
It("checkpoints read the job of the resource", func() {
rt, eng := started()
eng.result["finetune.checkpoints"] = map[string]any{"checkpoints": []any{map[string]any{"step": 10}}}
res, err := rt.Read(ctx, envelope("fine-tune-checkpoints", "r1", "res1", ""))
Expect(err).NotTo(HaveOccurred())
Expect(string(res.Payload)).To(ContainSubstring(`"step":10`))
Expect(eng.last("finetune.checkpoints").args["jobId"]).To(Equal("job-1"), "wrong job")
})
It("does not let a tool control a job it did not start", func() {
eng := &fakeEngine{}
rt := New(eng, GinkgoT().TempDir())
_, err := rt.Apply(ctx, envelope("fine-tune-stop", "op2", "nope", `{}`), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrNotFound), "stop")
_, err = rt.Apply(ctx, envelope("fine-tune-export", "op3", "nope", `{}`), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrNotFound), "export")
_, err = rt.Read(ctx, envelope("fine-tune-status", "r1", "nope", ""))
Expect(err).To(MatchError(protocol.ErrNotFound), "status")
_, err = rt.Preview(ctx, envelope("fine-tune-stop", "op2", "nope", `{}`))
Expect(err).To(MatchError(protocol.ErrNotFound), "preview")
Expect(eng.ops()).To(BeEmpty(), "the engine must not be called")
})
It("stop and export use the stored job and advance the revision", func() {
rt, eng := started()
eng.result["finetune.stop"] = map[string]any{}
eng.result["finetune.export"] = map[string]any{"model": "my-model"}
p, err := rt.Preview(ctx, envelope("fine-tune-stop", "op2", "res1", `{"saveCheckpoint":true}`))
Expect(err).NotTo(HaveOccurred())
Expect(p.Revision.Schema).NotTo(BeZero())
rec, err := rt.Apply(ctx, envelope("fine-tune-stop", "op2", "res1", `{"saveCheckpoint":true}`), p.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(rec.State).To(Equal(contract.ReceiptCommitted))
c := eng.last("finetune.stop")
Expect(c.args["jobId"]).To(Equal("job-1"))
Expect(c.args["saveCheckpoint"]).To(Equal(true))
Expect(rec.Revision.Data).To(BeNumerically(">", p.Revision.Data), "the revision must advance")
// A repeat of the same operation is answered from the receipt.
_, err = rt.Apply(ctx, envelope("fine-tune-stop", "op2", "res1", `{"saveCheckpoint":true}`), p.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(eng.count("finetune.stop")).To(Equal(1))
// Export names the format and the model, and never a path.
ex, err := rt.Apply(ctx, envelope("fine-tune-export", "op3", "res1", `{"exportFormat":"gguf","name":"my-model"}`), rec.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(ex.State).To(Equal(contract.ReceiptCommitted))
Expect(string(ex.Result.Payload)).To(ContainSubstring("my-model"))
a := eng.last("finetune.export").args
Expect(a["jobId"]).To(Equal("job-1"))
Expect(a["exportFormat"]).To(Equal("gguf"))
Expect(a["name"]).To(Equal("my-model"))
})
It("rejects a stale revision and does not call the engine", func() {
rt, eng := started()
rec, err := rt.Apply(ctx, envelope("fine-tune-stop", "op2", "res1", `{}`), contract.Revision{Schema: 1, Data: 99})
Expect(err).NotTo(HaveOccurred())
Expect(rec.State).To(Equal(contract.ReceiptRejected))
Expect(eng.count("finetune.stop")).To(BeZero(), "the engine was called")
})
Describe("host refusals", func() {
for code, want := range map[string]error{
contract.HostFeature: protocol.ErrForbidden, contract.HostBadRequest: protocol.ErrBadRequest,
contract.HostNotFound: protocol.ErrNotFound, contract.HostConflict: protocol.ErrConflict,
} {
It("maps "+code+" to a protocol error", func() {
rt, eng := started()
eng.result["finetune.export"] = &contract.HostError{Code: code, Message: "m"}
_, err := rt.Apply(ctx, envelope("fine-tune-export", "op-"+code, "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
Expect(err).To(MatchError(want))
})
}
It("keeps an unavailable engine uncertain", func() {
rt, eng := started()
eng.result["finetune.export"] = &contract.HostError{Code: contract.HostUnavailable, Message: "m"}
_, err := rt.Apply(ctx, envelope("fine-tune-export", "op-u", "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
Expect(err).To(HaveOccurred())
for _, e := range []error{protocol.ErrBadRequest, protocol.ErrForbidden, protocol.ErrNotFound, protocol.ErrConflict} {
Expect(err).NotTo(MatchError(e), "an unavailable engine must stay uncertain")
}
})
})
It("answers resource-exists from the journal and the engine", func() {
rt, eng := started()
eng.result["finetune.get"] = map[string]any{"id": "job-1", "status": "training"}
ok, err := exists(rt, "res1")
Expect(err).NotTo(HaveOccurred(), "mapped")
Expect(ok).To(BeTrue(), "mapped")
ok, err = exists(rt, "other")
Expect(err).NotTo(HaveOccurred(), "unmapped")
Expect(ok).To(BeFalse(), "unmapped")
// The job was deleted in the UI: the engine says not found.
eng.result["finetune.get"] = &contract.HostError{Code: contract.HostNotFound, Message: "gone"}
ok, err = exists(rt, "res1")
Expect(err).NotTo(HaveOccurred(), "deleted job")
Expect(ok).To(BeFalse(), "deleted job")
// The engine cannot answer: unknown, never false.
eng.result["finetune.get"] = &contract.HostError{Code: contract.HostUnavailable, Message: "down"}
_, err = exists(rt, "res1")
Expect(err).To(HaveOccurred(), "an unavailable engine must not read as absent")
eng.result["finetune.get"] = hostclient.ErrNoHostOps
_, err = exists(rt, "res1")
Expect(err).To(HaveOccurred(), "no host endpoint must not read as absent")
})
It("keeps the journal across a restart", func() {
dir := GinkgoT().TempDir()
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}, "finetune.get": map[string]any{"status": "training"}}}
_, err := New(eng, dir).Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
again := New(eng, dir)
_, err = again.Read(ctx, envelope("fine-tune-status", "r1", "res1", ""))
Expect(err).NotTo(HaveOccurred())
Expect(eng.last("finetune.get").args["jobId"]).To(Equal("job-1"), "the restarted runtime lost the job")
// The receipt of the first run answers a repeat after the restart.
_, err = again.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
Expect(eng.count("finetune.start")).To(Equal(1))
r, err := again.Receipt(ctx, "op1")
Expect(err).NotTo(HaveOccurred())
Expect(r.State).To(Equal(contract.ReceiptCommitted))
_, err = again.Receipt(ctx, "never")
Expect(err).To(MatchError(protocol.ErrNotFound), "an unknown operation must be not found")
// The journal is a plain file written through a temporary one.
entries, _ := os.ReadDir(dir)
for _, e := range entries {
Expect(e.Name()).NotTo(ContainSubstring(".tmp"), "a temporary file was left")
}
_, err = os.Stat(filepath.Join(dir, "journal.json"))
Expect(err).NotTo(HaveOccurred())
})
It("refuses every tool without a host endpoint", func() {
var nilClient *hostclient.Client
dir := GinkgoT().TempDir()
rt := New(nilClient, dir)
_, err := rt.Apply(ctx, envelope("fine-tune-start", "op1", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).To(MatchError(hostclient.ErrNoHostOps), "start")
// Seed a job as an earlier run with an endpoint would have.
seed := New(&fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1"}}}, dir)
_, err = seed.Apply(ctx, envelope("fine-tune-start", "op0", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
rt = New(nilClient, dir)
for _, tool := range []string{"fine-tune-status", "fine-tune-checkpoints", "resource-exists"} {
_, err := rt.Read(ctx, envelope(tool, "r-"+tool, "res1", ""))
Expect(err).To(MatchError(hostclient.ErrNoHostOps), tool)
}
for _, tool := range []string{"fine-tune-stop", "fine-tune-export"} {
_, err := rt.Apply(ctx, envelope(tool, "a-"+tool, "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
Expect(err).To(MatchError(hostclient.ErrNoHostOps), tool)
}
})
It("treats an unknown tool as a bad request", func() {
rt := New(&fakeEngine{}, GinkgoT().TempDir())
_, err := rt.Read(ctx, envelope("fine-tune-nope", "r1", "res1", ""))
Expect(err).To(MatchError(protocol.ErrBadRequest), "read")
_, err = rt.Apply(ctx, envelope("fine-tune-nope", "a1", "res1", ""), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrBadRequest), "apply")
_, err = rt.OpenOutput(ctx, envelope("fine-tune-status", "o1", "res1", ""), "h")
Expect(err).To(MatchError(protocol.ErrNotFound), "output")
})
Describe("manifest", func() {
It("declares the tools and the resource mapping", func() {
m := loadManifest()
Expect(m.ID).To(Equal("fine-tuning"))
Expect(m.Capabilities.UI.MainPage).To(Equal("jobs"))
Expect(m.Capabilities.UI.Pages).To(Equal([]string{"jobs"}))
Expect(m.Capabilities.UI.SDK.Min).To(BeEquivalentTo(2))
Expect(m.Capabilities.UI.SDK.Max).To(BeEquivalentTo(2))
Expect(m.Capabilities.Tools).To(Equal(Tools))
r := m.Capabilities.Resources
Expect(r).NotTo(BeNil())
Expect(r.Noun).To(Equal("fine-tuning job"))
Expect(r.Create).To(Equal("fine-tune-start"))
Expect(r.Inspect).To(Equal("fine-tune-status"))
kinds := map[string]string{}
for _, d := range m.Capabilities.ToolDefinitions {
kinds[d.Name] = d.Kind
}
want := map[string]string{"fine-tune-list": "read", "fine-tune-status": "read", "fine-tune-checkpoints": "read", "fine-tune-start": "create", "fine-tune-dataset": "create", "fine-tune-stop": "write", "fine-tune-export": "write"}
Expect(kinds).To(Equal(want))
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name == "fine-tune-list" {
Expect(contract.IsResourceless(d)).To(BeTrue(), "the list tool must declare no resource")
} else {
Expect(contract.IsResourceless(d)).To(BeFalse(), "%s must not be resource-less", d.Name)
}
Expect(len(d.Description)).To(BeNumerically("<", 400), d.Name)
Expect(d.Description).To(ContainSubstring("started"), "%s: the description must say that only jobs a tool started can be controlled", d.Name)
}
})
It("names no user or path in any tool argument", func() {
m := loadManifest()
for _, d := range m.Capabilities.ToolDefinitions {
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
}
Expect(json.Unmarshal(d.InputSchema, &s)).To(Succeed())
for name := range s.Properties {
lower := strings.ToLower(name)
Expect(contract.ForbiddenToolProperty(name) || strings.Contains(lower, "path") || strings.Contains(lower, "dir")).To(BeFalse(),
"%s: property %q names an identity or a path", d.Name, name)
}
}
})
It("gives the start schema the task 20 fields and the name the host needs", func() {
m := loadManifest()
found := false
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name != "fine-tune-start" {
continue
}
found = true
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
Required []string `json:"required"`
}
_ = json.Unmarshal(d.InputSchema, &s)
want := []string{"name", "model", "backend", "trainingType", "trainingMethod", "datasetSource", "datasetSplit", "learningRate", "numEpochs", "batchSize", "maxSteps", "adapterRank", "adapterAlpha"}
Expect(s.Properties).To(HaveLen(len(want)))
for _, w := range want {
Expect(s.Properties).To(HaveKey(w))
}
break
}
Expect(found).To(BeTrue(), "no start definition")
})
})
It("lists the caller's jobs without a resource", func() {
eng := &fakeEngine{result: map[string]any{"finetune.list": map[string]any{"jobs": []any{map[string]any{"id": "ui-job", "status": "completed"}}}}}
rt := New(eng, GinkgoT().TempDir())
res, err := rt.Read(ctx, envelope("fine-tune-list", "r1", "", ""))
Expect(err).NotTo(HaveOccurred())
Expect(string(res.Payload)).To(ContainSubstring("ui-job"))
Expect(eng.ops()).To(Equal([]string{"finetune.list"}))
Expect(eng.last("finetune.list").args).To(BeNil(), "the list takes no argument")
_, err = rt.Read(ctx, envelope("fine-tune-list", "r2", "", `{"userId":"other"}`))
Expect(err).To(MatchError(protocol.ErrBadRequest), "an argument must be refused")
var nilClient *hostclient.Client
_, err = New(nilClient, GinkgoT().TempDir()).Read(ctx, envelope("fine-tune-list", "r3", "", ""))
Expect(err).To(MatchError(hostclient.ErrNoHostOps), "no endpoint")
})
It("sends a dataset once and returns its ref only through the resource", func() {
eng := &fakeEngine{result: map[string]any{"finetune.dataset": map[string]any{"datasetRef": "0000000a-data.jsonl"}}}
rt := New(eng, GinkgoT().TempDir())
p, err := rt.Preview(ctx, envelope("fine-tune-dataset", "create-d1", "d1", datasetJSON))
Expect(err).NotTo(HaveOccurred())
Expect(p.Revision.Schema).NotTo(BeZero())
Expect(eng.ops()).To(BeEmpty())
Expect(string(p.Summary)).NotTo(ContainSubstring("text"), "the approval card must not carry the content")
rec, err := rt.Apply(ctx, envelope("fine-tune-dataset", "create-d1", "d1", datasetJSON), p.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(rec.State).To(Equal(contract.ReceiptCommitted))
Expect(string(rec.Result.Payload)).To(ContainSubstring("0000000a-data.jsonl"))
_, err = rt.Apply(ctx, envelope("fine-tune-dataset", "create-d1", "d1", datasetJSON), p.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(eng.count("finetune.dataset")).To(Equal(1))
a := eng.last("finetune.dataset").args
Expect(a["name"]).To(Equal("data.jsonl"))
Expect(a["content"]).NotTo(BeNil())
Expect(a).To(HaveLen(2), "host operation args")
// The resource is the dataset: its status gives the reference to the owner of the resource.
st, err := rt.Read(ctx, envelope("fine-tune-status", "r1", "d1", ""))
Expect(err).NotTo(HaveOccurred())
Expect(string(st.Payload)).To(ContainSubstring(`"datasetRef":"0000000a-data.jsonl"`))
Expect(string(st.Payload)).To(ContainSubstring(`"kind":"dataset"`))
Expect(eng.ops()).To(HaveLen(1), "the status of a dataset must not call the engine")
ok, err := exists(rt, "d1")
Expect(err).NotTo(HaveOccurred())
Expect(ok).To(BeTrue())
// A dataset is not a job.
for _, tool := range []string{"fine-tune-checkpoints"} {
_, err := rt.Read(ctx, envelope(tool, "r2", "d1", ""))
Expect(err).To(MatchError(protocol.ErrNotFound), tool)
}
_, err = rt.Apply(ctx, envelope("fine-tune-stop", "s1", "d1", `{}`), contract.Revision{Schema: 1, Data: 1})
Expect(err).To(MatchError(protocol.ErrNotFound), "stop")
})
Describe("dataset validation", func() {
big := strings.Repeat("a", 900<<10+1)
for name, args := range map[string]string{
"empty": `{"name":"d.jsonl","content":""}`,
"too large": `{"name":"d.jsonl","content":"` + big + `"}`,
"a folder": `{"name":"a/d.jsonl","content":"x"}`,
"parent": `{"name":"../d.jsonl","content":"x"}`,
"wrong type": `{"name":"d.exe","content":"x"}`,
"only the type": `{"name":".jsonl","content":"x"}`,
"extra field": `{"name":"d.jsonl","content":"x","path":"/etc"}`,
} {
It("refuses "+name, func() {
rt := New(&fakeEngine{}, GinkgoT().TempDir())
_, err := rt.Preview(ctx, envelope("fine-tune-dataset", "o-"+name, "d9", args))
Expect(err).To(MatchError(protocol.ErrBadRequest))
})
}
It("accepts the largest allowed dataset", func() {
rt := New(&fakeEngine{}, GinkgoT().TempDir())
ok := `{"name":"d.jsonl","content":"` + strings.Repeat("a", 900<<10) + `"}`
_, err := rt.Preview(ctx, envelope("fine-tune-dataset", "o-ok", "d9", ok))
Expect(err).NotTo(HaveOccurred(), "900 KiB")
})
})
It("passes a dataset ref to start and keeps no job when the host refuses it", func() {
eng := &fakeEngine{result: map[string]any{"finetune.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, GinkgoT().TempDir())
args := `{"name":"run","model":"org/base","datasetSource":"0000000a-data.jsonl"}`
_, err := rt.Apply(ctx, envelope("fine-tune-start", "create-r1", "r1", args), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
Expect(eng.last("finetune.start").args["datasetSource"]).To(Equal("0000000a-data.jsonl"))
// The host refuses the reference (the per-user check is proven in the host tests); the add-on only maps the refusal and keeps nothing.
eng.result["finetune.start"] = &contract.HostError{Code: contract.HostBadRequest, Message: "datasetSource is not a dataset reference of this user"}
_, err = rt.Apply(ctx, envelope("fine-tune-start", "create-r2", "r2", `{"name":"run","model":"org/base","datasetSource":"0000000b-theirs.jsonl"}`), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrBadRequest))
_, err = rt.Read(ctx, envelope("fine-tune-status", "r3", "r2", ""))
Expect(err).To(MatchError(protocol.ErrNotFound), "a refused start must not leave a job")
_, err = rt.Receipt(ctx, "create-r2")
Expect(err).To(MatchError(protocol.ErrNotFound), "a refused start must not leave a receipt")
})
})
addons/internal/hostclient/client_test.go
+53 -50
View File
@@ -7,65 +7,68 @@ import (
"errors"
"net/http"
"net/http/httptest"
"testing"
"github.com/mudler/LocalAI/addons/internal/protocol"
"github.com/mudler/LocalAI/pkg/addons/contract"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func startupFor(endpoint string) protocol.Startup {
return protocol.Startup{HostEndpoint: endpoint, HostCredential: "secret", Limits: contract.DefaultLimits()}
}
func TestDoSendsTheCredentialAndTheBody(t *testing.T) {
var gotCred string
var gotReq contract.HostRequest
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotCred = r.Header.Get("X-Addon-Host-Credential")
_ = json.NewDecoder(r.Body).Decode(&gotReq)
_ = json.NewEncoder(w).Encode(contract.HostResponse{Result: json.RawMessage(`{"id":"j1"}`)})
}))
defer srv.Close()
c := New(startupFor(srv.URL))
if c == nil {
t.Fatal("no client for a loopback endpoint")
}
var out struct{ ID string }
if err := c.Do(context.Background(), "op1", "finetune.get", map[string]string{"jobId": "j1"}, &out); err != nil {
t.Fatal(err)
}
if gotCred != "secret" || gotReq.OperationID != "op1" || gotReq.Op != "finetune.get" || string(gotReq.Arguments) != `{"jobId":"j1"}` || out.ID != "j1" {
t.Fatalf("%q %+v %+v", gotCred, gotReq, out)
}
}
var _ = Describe("host client", func() {
It("sends the credential and the body", func() {
var gotCred string
var gotReq contract.HostRequest
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
gotCred = r.Header.Get("X-Addon-Host-Credential")
_ = json.NewDecoder(r.Body).Decode(&gotReq)
_ = json.NewEncoder(w).Encode(contract.HostResponse{Result: json.RawMessage(`{"id":"j1"}`)})
}))
DeferCleanup(srv.Close)
c := New(startupFor(srv.URL))
Expect(c).NotTo(BeNil(), "no client for a loopback endpoint")
var out struct{ ID string }
Expect(c.Do(context.Background(), "op1", "finetune.get", map[string]string{"jobId": "j1"}, &out)).To(Succeed())
Expect(gotCred).To(Equal("secret"))
Expect(gotReq.OperationID).To(Equal("op1"))
Expect(gotReq.Op).To(Equal("finetune.get"))
Expect(string(gotReq.Arguments)).To(Equal(`{"jobId":"j1"}`))
Expect(out.ID).To(Equal("j1"))
})
func TestDoTurnsAnErrorAnswerIntoAHostError(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusForbidden)
_ = json.NewEncoder(w).Encode(contract.HostResponse{Error: &contract.HostError{Code: "scope", Message: "no"}})
}))
defer srv.Close()
err := New(startupFor(srv.URL)).Do(context.Background(), "op1", "finetune.start", nil, nil)
var he *contract.HostError
if !errors.As(err, &he) || he.Code != "scope" {
t.Fatalf("%v", err)
}
}
It("turns an error answer into a host error", func() {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusForbidden)
_ = json.NewEncoder(w).Encode(contract.HostResponse{Error: &contract.HostError{Code: "scope", Message: "no"}})
}))
DeferCleanup(srv.Close)
err := New(startupFor(srv.URL)).Do(context.Background(), "op1", "finetune.start", nil, nil)
var he *contract.HostError
Expect(errors.As(err, &he)).To(BeTrue(), "got %v", err)
Expect(he.Code).To(Equal("scope"))
})
func TestAStartupWithoutAnEndpointGivesANilClient(t *testing.T) {
c := New(protocol.Startup{})
if c != nil {
t.Fatal("expected nil")
}
if err := c.Do(context.Background(), "op1", "finetune.get", nil, nil); !errors.Is(err, ErrNoHostOps) {
t.Fatalf("%v", err)
}
}
It("gives a nil client for a startup without an endpoint", func() {
c := New(protocol.Startup{})
Expect(c).To(BeNil())
Expect(c.Do(context.Background(), "op1", "finetune.get", nil, nil)).To(MatchError(ErrNoHostOps))
})
func TestANonLoopbackEndpointGivesANilClient(t *testing.T) {
for _, e := range []string{"http://10.0.0.1:80", "http://example.com:80", "https://127.0.0.1:80", "http://127.0.0.1", "http://127.0.0.1:80/x", "http://localhost:80", "http://127.0.0.1:0", "http://user@127.0.0.1:80", "http://[::1]:80"} {
if New(startupFor(e)) != nil {
t.Errorf("%s: got a client", e)
}
}
}
DescribeTable("gives a nil client for a non-loopback endpoint",
func(endpoint string) {
Expect(New(startupFor(endpoint))).To(BeNil())
},
Entry(nil, "http://10.0.0.1:80"),
Entry(nil, "http://example.com:80"),
Entry(nil, "https://127.0.0.1:80"),
Entry(nil, "http://127.0.0.1"),
Entry(nil, "http://127.0.0.1:80/x"),
Entry(nil, "http://localhost:80"),
Entry(nil, "http://127.0.0.1:0"),
Entry(nil, "http://user@127.0.0.1:80"),
Entry(nil, "http://[::1]:80"),
)
})
addons/internal/hostclient/suite_test.go
+14
View File
@@ -0,0 +1,14 @@
// SPDX-License-Identifier: MIT
package hostclient
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestHostclient(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Hostclient suite")
}
addons/internal/idle/idle_test.go
+33 -29
View File
@@ -3,39 +3,43 @@ package idle
import (
"context"
"errors"
"testing"
"github.com/mudler/LocalAI/pkg/addons/contract"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var _ contract.Runtime = Runtime{}
func TestEveryOperationRefusesBecauseTheAddonDeclaresNoTools(t *testing.T) {
ctx := context.Background()
rt := Runtime{}
if _, err := rt.Read(ctx, contract.Envelope{}); !errors.Is(err, ErrNoTools) {
t.Fatalf("Read: %v", err)
}
if _, err := rt.Preview(ctx, contract.Envelope{}); !errors.Is(err, ErrNoTools) {
t.Fatalf("Preview: %v", err)
}
if _, err := rt.Apply(ctx, contract.Envelope{}, contract.Revision{}); !errors.Is(err, ErrNoTools) {
t.Fatalf("Apply: %v", err)
}
if _, err := rt.Receipt(ctx, "x"); !errors.Is(err, ErrNoTools) {
t.Fatalf("Receipt: %v", err)
}
if _, err := rt.OpenOutput(ctx, contract.Envelope{}, "x"); !errors.Is(err, ErrNoTools) {
t.Fatalf("OpenOutput: %v", err)
}
}
var _ = Describe("idle runtime", func() {
Context("when the addon declares no tools", func() {
ctx := context.Background()
rt := Runtime{}
func TestRunRefusesWithoutTheProtocolFlag(t *testing.T) {
if got := Run([]string{"voice"}); got != 2 {
t.Fatalf("exit code %d, want 2", got)
}
if got := Run([]string{"voice", "--other"}); got != 2 {
t.Fatalf("exit code %d, want 2", got)
}
}
It("refuses Read", func() {
_, err := rt.Read(ctx, contract.Envelope{})
Expect(err).To(MatchError(ErrNoTools))
})
It("refuses Preview", func() {
_, err := rt.Preview(ctx, contract.Envelope{})
Expect(err).To(MatchError(ErrNoTools))
})
It("refuses Apply", func() {
_, err := rt.Apply(ctx, contract.Envelope{}, contract.Revision{})
Expect(err).To(MatchError(ErrNoTools))
})
It("refuses Receipt", func() {
_, err := rt.Receipt(ctx, "x")
Expect(err).To(MatchError(ErrNoTools))
})
It("refuses OpenOutput", func() {
_, err := rt.OpenOutput(ctx, contract.Envelope{}, "x")
Expect(err).To(MatchError(ErrNoTools))
})
})
It("refuses to run without the protocol flag", func() {
Expect(Run([]string{"voice"})).To(Equal(2))
Expect(Run([]string{"voice", "--other"})).To(Equal(2))
})
})
addons/internal/idle/suite_test.go
+14
View File
@@ -0,0 +1,14 @@
// SPDX-License-Identifier: MIT
package idle
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestIdle(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Idle suite")
}
addons/internal/protocol/run_test.go
+11 -13
View File
@@ -2,20 +2,18 @@
package protocol
import (
"testing"
"github.com/mudler/LocalAI/pkg/addons/contract"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestRunRefusesWithoutTheProtocolFlagAndBuildsNoRuntime(t *testing.T) {
built := false
factory := func(Startup) contract.Runtime { built = true; return nil }
for _, args := range [][]string{{"x"}, {"x", "--other"}, {"x", "--localai-addon-v1", "extra"}} {
if got := Run(args, factory); got != 2 {
t.Fatalf("%v: exit code %d, want 2", args, got)
var _ = Describe("Run", func() {
It("refuses without the protocol flag and builds no runtime", func() {
built := false
factory := func(Startup) contract.Runtime { built = true; return nil }
for _, args := range [][]string{{"x"}, {"x", "--other"}, {"x", "--localai-addon-v1", "extra"}} {
Expect(Run(args, factory)).To(Equal(2), "args %v", args)
}
}
if built {
t.Fatal("the factory ran before the protocol flag was checked")
}
}
Expect(built).To(BeFalse(), "the factory ran before the protocol flag was checked")
})
})
addons/internal/protocol/suite_test.go
+14
View File
@@ -0,0 +1,14 @@
// SPDX-License-Identifier: MIT
package protocol
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestProtocol(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Protocol suite")
}
addons/quantization/internal/tools/suite_test.go
+14
View File
@@ -0,0 +1,14 @@
// SPDX-License-Identifier: MIT
package tools
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestTools(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Quantization tools suite")
}
addons/quantization/internal/tools/tools_test.go
+317 -385
View File
@@ -4,18 +4,17 @@ package tools
import (
"context"
"encoding/json"
"errors"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"testing"
"time"
"github.com/mudler/LocalAI/addons/internal/hostclient"
"github.com/mudler/LocalAI/addons/internal/protocol"
"github.com/mudler/LocalAI/pkg/addons/contract"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var ctx = context.Background()
@@ -103,194 +102,16 @@ func envelope(tool, opID, resource, args string) contract.Envelope {
const startJSON = `{"name":"my quant","model":"org/base","quantizationType":"q4_k_m"}`
// started returns a runtime that already holds the job job-1 for resource res1.
func started(t *testing.T) (contract.Runtime, *fakeEngine) {
t.Helper()
func started() (contract.Runtime, *fakeEngine) {
eng := &fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, t.TempDir())
rt := New(eng, GinkgoT().TempDir())
rec, err := rt.Apply(ctx, envelope("quantize-start", "create-res1", "res1", startJSON), contract.Revision{Schema: 1})
if err != nil || rec.State != contract.ReceiptCommitted {
t.Fatalf("setup: %v %+v", err, rec)
}
Expect(err).NotTo(HaveOccurred(), "setup")
Expect(rec.State).To(Equal(contract.ReceiptCommitted), "setup: %+v", rec)
return rt, eng
}
func TestStartCallsTheEngineOnceAndRemembersTheJob(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, t.TempDir())
rec, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 0})
if err != nil || rec.State != contract.ReceiptCommitted || rec.Revision.Schema == 0 {
t.Fatalf("%v %+v", err, rec)
}
if got := eng.ops(); !reflect.DeepEqual(got, []string{"quantize.start"}) {
t.Fatalf("ops %v", got)
}
again, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 1})
if err != nil || again.State != contract.ReceiptCommitted {
t.Fatalf("%v %+v", err, again)
}
if n := eng.count("quantize.start"); n != 1 {
t.Fatalf("started %d jobs", n)
}
// A different operation on the same resource is a clash, not a second job.
if _, err := rt.Apply(ctx, envelope("quantize-start", "op9", "res1", startJSON), contract.Revision{Schema: 1}); !errors.Is(err, protocol.ErrConflict) {
t.Fatalf("want a conflict, got %v", err)
}
}
func TestStartSendsOnlyTheFieldsOfTheEngineOperation(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, t.TempDir())
args := `{"name":"my quant","model":"org/base","backend":"llama-cpp-quantization","quantizationType":"q5_k_m"}`
if _, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", args), contract.Revision{Schema: 1}); err != nil {
t.Fatal(err)
}
c := eng.last("quantize.start")
if c.opID != "op1" {
t.Fatalf("the call must carry the operation of the request, got %q", c.opID)
}
if len(c.args) != 3 || c.args["model"] != "org/base" || c.args["backend"] != "llama-cpp-quantization" || c.args["quantizationType"] != "q5_k_m" {
t.Fatalf("args %v", c.args)
}
}
func TestPreviewOfStartHasANonzeroRevisionAndCreatesNothing(t *testing.T) {
eng := &fakeEngine{}
rt := New(eng, t.TempDir())
p, err := rt.Preview(ctx, envelope("quantize-start", "op1", "res1", startJSON))
if err != nil || p.Revision.Schema == 0 || len(eng.ops()) != 0 {
t.Fatalf("%v %+v %v", err, p, eng.ops())
}
if !json.Valid(p.Summary) || !strings.Contains(string(p.Summary), "q4_k_m") {
t.Fatalf("summary %s", p.Summary)
}
for name, args := range map[string]string{
"empty model": `{"name":"x","model":""}`,
"bad type": `{"name":"x","model":"m","quantizationType":"Q4 K"}`,
"long name": `{"name":"` + strings.Repeat("n", 81) + `","model":"m"}`,
"a token": `{"name":"x","model":"m","hfToken":"hf_x"}`,
"unknown part": `{"name":"x","model":"m","path":"/etc"}`,
} {
if _, err := rt.Preview(ctx, envelope("quantize-start", "op-"+name, "res1", args)); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("%s: %v", name, err)
}
}
}
func TestStatusReadsTheJobOfTheResource(t *testing.T) {
rt, eng := started(t)
eng.result["quantize.get"] = map[string]any{"id": "job-1", "status": "quantizing", "model": "org/base"}
res, err := rt.Read(ctx, envelope("quantize-status", "r1", "res1", ""))
if err != nil {
t.Fatal(err)
}
if got := eng.last("quantize.get").args["jobId"]; got != "job-1" {
t.Fatalf("jobId %v", got)
}
var v map[string]any
if err := json.Unmarshal(res.Payload, &v); err != nil {
t.Fatal(err)
}
if v["name"] != "my quant" || v["status"] != "quantizing" {
t.Fatalf("payload %v", v)
}
rev, _ := v["revision"].(map[string]any)
if rev["schema"] == nil || rev["schema"] == float64(0) {
t.Fatalf("inspect needs a revision: %v", v)
}
if _, has := v["id"]; has {
t.Fatalf("the engine job ID must not reach the model: %v", v)
}
}
func TestAToolCannotControlAJobItDidNotStart(t *testing.T) {
eng := &fakeEngine{}
rt := New(eng, t.TempDir())
if _, err := rt.Apply(ctx, envelope("quantize-stop", "op2", "nope", `{}`), contract.Revision{Schema: 1}); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("stop: %v", err)
}
if _, err := rt.Apply(ctx, envelope("quantize-import", "op3", "nope", `{}`), contract.Revision{Schema: 1}); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("import: %v", err)
}
if _, err := rt.Read(ctx, envelope("quantize-status", "r1", "nope", "")); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("status: %v", err)
}
if _, err := rt.Preview(ctx, envelope("quantize-stop", "op2", "nope", `{}`)); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("preview: %v", err)
}
if n := len(eng.ops()); n != 0 {
t.Fatalf("the engine was called %d times: %v", n, eng.ops())
}
}
func TestStopAndImportUseTheStoredJobAndAdvanceTheRevision(t *testing.T) {
rt, eng := started(t)
eng.result["quantize.import"] = map[string]any{"model": "my-model"}
p, err := rt.Preview(ctx, envelope("quantize-stop", "op2", "res1", `{}`))
if err != nil || p.Revision.Schema == 0 {
t.Fatalf("%v %+v", err, p)
}
rec, err := rt.Apply(ctx, envelope("quantize-stop", "op2", "res1", `{}`), p.Revision)
if err != nil || rec.State != contract.ReceiptCommitted {
t.Fatalf("%v %+v", err, rec)
}
c := eng.last("quantize.stop")
if len(c.args) != 1 || c.args["jobId"] != "job-1" {
t.Fatalf("stop args %v (a quantization job has no checkpoint to save)", c.args)
}
if rec.Revision.Data <= p.Revision.Data {
t.Fatalf("the revision must advance: %+v then %+v", p.Revision, rec.Revision)
}
if _, err := rt.Apply(ctx, envelope("quantize-stop", "op2", "res1", `{}`), p.Revision); err != nil || eng.count("quantize.stop") != 1 {
t.Fatalf("stop ran %d times (%v)", eng.count("quantize.stop"), err)
}
// Import names the model and never a path.
ex, err := rt.Apply(ctx, envelope("quantize-import", "op3", "res1", `{"name":"my-model"}`), rec.Revision)
if err != nil || ex.State != contract.ReceiptCommitted || !strings.Contains(string(ex.Result.Payload), "my-model") {
t.Fatalf("%v %+v", err, ex)
}
if a := eng.last("quantize.import").args; len(a) != 2 || a["jobId"] != "job-1" || a["name"] != "my-model" {
t.Fatalf("import args %v", a)
}
if _, err := rt.Preview(ctx, envelope("quantize-import", "op4", "res1", `{"name":"../x"}`)); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("a bad model name must be refused, got %v", err)
}
if _, err := rt.Apply(ctx, envelope("quantize-import", "op5", "res1", `{"name":"a b"}`), ex.Revision); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("a bad model name must be refused at apply, got %v", err)
}
}
func TestAStaleRevisionIsRejectedAndTheEngineIsNotCalled(t *testing.T) {
rt, eng := started(t)
rec, err := rt.Apply(ctx, envelope("quantize-stop", "op2", "res1", `{}`), contract.Revision{Schema: 1, Data: 99})
if err != nil || rec.State != contract.ReceiptRejected {
t.Fatalf("%v %+v", err, rec)
}
if eng.count("quantize.stop") != 0 {
t.Fatal("the engine was called")
}
}
func TestHostRefusalsMapToProtocolErrors(t *testing.T) {
rt, eng := started(t)
for code, want := range map[string]error{
contract.HostFeature: protocol.ErrForbidden, contract.HostBadRequest: protocol.ErrBadRequest,
contract.HostNotFound: protocol.ErrNotFound, contract.HostConflict: protocol.ErrConflict,
} {
eng.result["quantize.import"] = &contract.HostError{Code: code, Message: "m"}
_, err := rt.Apply(ctx, envelope("quantize-import", "op-"+code, "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
if !errors.Is(err, want) {
t.Fatalf("%s: %v, want %v", code, err, want)
}
}
eng.result["quantize.import"] = &contract.HostError{Code: contract.HostUnavailable, Message: "m"}
_, err := rt.Apply(ctx, envelope("quantize-import", "op-u", "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
if err == nil || errors.Is(err, protocol.ErrBadRequest) || errors.Is(err, protocol.ErrForbidden) || errors.Is(err, protocol.ErrNotFound) || errors.Is(err, protocol.ErrConflict) {
t.Fatalf("an unavailable engine must stay uncertain, got %v", err)
}
}
func exists(t *testing.T, rt contract.Runtime, res string) (bool, error) {
t.Helper()
func exists(rt contract.Runtime, res string) (bool, error) {
r, err := rt.Read(ctx, envelope("resource-exists", "x-"+res, res, ""))
if err != nil {
return false, err
@@ -302,214 +123,325 @@ func exists(t *testing.T, rt contract.Runtime, res string) (bool, error) {
return v.Exists, nil
}
func TestResourceExistsAnswersFromTheJournalAndTheEngine(t *testing.T) {
rt, eng := started(t)
eng.result["quantize.get"] = map[string]any{"id": "job-1", "status": "quantizing"}
if ok, err := exists(t, rt, "res1"); err != nil || !ok {
t.Fatalf("mapped: %v %v", ok, err)
}
if ok, err := exists(t, rt, "other"); err != nil || ok {
t.Fatalf("unmapped: %v %v", ok, err)
}
eng.result["quantize.get"] = &contract.HostError{Code: contract.HostNotFound, Message: "gone"}
if ok, err := exists(t, rt, "res1"); err != nil || ok {
t.Fatalf("deleted job: %v %v", ok, err)
}
eng.result["quantize.get"] = &contract.HostError{Code: contract.HostUnavailable, Message: "down"}
if _, err := exists(t, rt, "res1"); err == nil {
t.Fatal("an unavailable engine must not read as absent")
}
eng.result["quantize.get"] = hostclient.ErrNoHostOps
if _, err := exists(t, rt, "res1"); err == nil {
t.Fatal("no host endpoint must not read as absent")
}
}
func TestJournalSurvivesARestart(t *testing.T) {
dir := t.TempDir()
eng := &fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1", "status": "queued"}, "quantize.get": map[string]any{"status": "quantizing"}}}
if _, err := New(eng, dir).Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1}); err != nil {
t.Fatal(err)
}
again := New(eng, dir)
if _, err := again.Read(ctx, envelope("quantize-status", "r1", "res1", "")); err != nil {
t.Fatal(err)
}
if eng.last("quantize.get").args["jobId"] != "job-1" {
t.Fatal("the restarted runtime lost the job")
}
if _, err := again.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1}); err != nil || eng.count("quantize.start") != 1 {
t.Fatalf("started %d times (%v)", eng.count("quantize.start"), err)
}
if r, err := again.Receipt(ctx, "op1"); err != nil || r.State != contract.ReceiptCommitted {
t.Fatalf("%v %+v", err, r)
}
if _, err := again.Receipt(ctx, "never"); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("an unknown operation must be not found, got %v", err)
}
entries, _ := os.ReadDir(dir)
for _, e := range entries {
if strings.Contains(e.Name(), ".tmp") {
t.Fatalf("a temporary file was left: %s", e.Name())
}
}
if _, err := os.Stat(filepath.Join(dir, "journal.json")); err != nil {
t.Fatal(err)
}
}
func TestWithoutAHostEndpointEveryToolRefuses(t *testing.T) {
var nilClient *hostclient.Client
dir := t.TempDir()
rt := New(nilClient, dir)
if _, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1}); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("start: %v", err)
}
seed := New(&fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1"}}}, dir)
if _, err := seed.Apply(ctx, envelope("quantize-start", "op0", "res1", startJSON), contract.Revision{Schema: 1}); err != nil {
t.Fatal(err)
}
rt = New(nilClient, dir)
for _, tool := range []string{"quantize-status", "resource-exists"} {
if _, err := rt.Read(ctx, envelope(tool, "r-"+tool, "res1", "")); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("%s: %v", tool, err)
}
}
if _, err := rt.Read(ctx, envelope("quantize-list", "r-list", "", "")); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("list: %v", err)
}
for _, tool := range []string{"quantize-stop", "quantize-import"} {
if _, err := rt.Apply(ctx, envelope(tool, "a-"+tool, "res1", `{}`), contract.Revision{Schema: 1, Data: 1}); !errors.Is(err, hostclient.ErrNoHostOps) {
t.Fatalf("%s: %v", tool, err)
}
}
}
func TestAnUnknownToolIsABadRequest(t *testing.T) {
rt := New(&fakeEngine{}, t.TempDir())
if _, err := rt.Read(ctx, envelope("quantize-nope", "r1", "res1", "")); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("read: %v", err)
}
if _, err := rt.Apply(ctx, envelope("quantize-nope", "a1", "res1", ""), contract.Revision{Schema: 1}); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("apply: %v", err)
}
if _, err := rt.OpenOutput(ctx, envelope("quantize-status", "o1", "res1", ""), "h"); !errors.Is(err, protocol.ErrNotFound) {
t.Fatalf("output: %v", err)
}
}
func TestListReadsTheCallersJobsWithoutAResource(t *testing.T) {
eng := &fakeEngine{result: map[string]any{"quantize.list": map[string]any{"jobs": []any{map[string]any{"id": "ui-job", "status": "completed"}}}}}
rt := New(eng, t.TempDir())
res, err := rt.Read(ctx, envelope("quantize-list", "r1", "", ""))
if err != nil || !strings.Contains(string(res.Payload), "ui-job") {
t.Fatalf("%v %s", err, res.Payload)
}
if got := eng.ops(); !reflect.DeepEqual(got, []string{"quantize.list"}) {
t.Fatalf("ops %v", got)
}
if c := eng.last("quantize.list"); c.args != nil {
t.Fatalf("the list takes no argument, got %v", c.args)
}
if _, err := rt.Read(ctx, envelope("quantize-list", "r2", "", `{"userId":"other"}`)); !errors.Is(err, protocol.ErrBadRequest) {
t.Fatalf("an argument must be refused, got %v", err)
}
}
func loadManifest(t *testing.T) contract.Manifest {
t.Helper()
func loadManifest() contract.Manifest {
b, err := os.ReadFile("../../manifest.json")
if err != nil {
t.Fatal(err)
}
Expect(err).NotTo(HaveOccurred())
m, err := contract.ParseManifest(b)
if err != nil {
t.Fatalf("the manifest must satisfy the contract: %v", err)
}
Expect(err).NotTo(HaveOccurred(), "the manifest must satisfy the contract")
return m
}
func TestManifestDeclaresTheToolsAndTheResourceMapping(t *testing.T) {
m := loadManifest(t)
if m.ID != "quantization" || m.Capabilities.UI.MainPage != "jobs" || !reflect.DeepEqual(m.Capabilities.UI.Pages, []string{"jobs"}) || m.Capabilities.UI.Title == nil || m.Capabilities.UI.Title.Text != "Quantization" {
t.Fatalf("manifest %+v", m)
}
if m.Capabilities.UI.SDK.Min != 2 || m.Capabilities.UI.SDK.Max != 2 {
t.Fatalf("sdk %+v", m.Capabilities.UI.SDK)
}
if !reflect.DeepEqual(m.Capabilities.Tools, Tools) {
t.Fatalf("manifest tools %v, package tools %v", m.Capabilities.Tools, Tools)
}
r := m.Capabilities.Resources
if r == nil || r.Noun != "quantization job" || r.Create != "quantize-start" || r.Inspect != "quantize-status" {
t.Fatalf("resources %+v", r)
}
kinds := map[string]string{}
for _, d := range m.Capabilities.ToolDefinitions {
kinds[d.Name] = d.Kind
}
want := map[string]string{"quantize-list": "read", "quantize-status": "read", "quantize-start": "create", "quantize-stop": "write", "quantize-import": "write"}
if !reflect.DeepEqual(kinds, want) {
t.Fatalf("kinds %v", kinds)
}
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name == "quantize-list" && !contract.IsResourceless(d) {
t.Fatal("the list tool must declare no resource")
}
if d.Name != "quantize-list" && contract.IsResourceless(d) {
t.Fatalf("%s must not be resource-less", d.Name)
}
if len(d.Description) >= 400 || !strings.Contains(d.Description, "started") {
t.Fatalf("%s: the description must be short and say that only jobs a tool started can be controlled", d.Name)
}
}
}
var _ = Describe("quantization tools", func() {
It("start calls the engine once and remembers the job", func() {
eng := &fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, GinkgoT().TempDir())
rec, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 0})
Expect(err).NotTo(HaveOccurred())
Expect(rec.State).To(Equal(contract.ReceiptCommitted))
Expect(rec.Revision.Schema).NotTo(BeZero())
Expect(eng.ops()).To(Equal([]string{"quantize.start"}))
again, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1, Data: 1})
Expect(err).NotTo(HaveOccurred())
Expect(again.State).To(Equal(contract.ReceiptCommitted))
Expect(eng.count("quantize.start")).To(Equal(1))
// A different operation on the same resource is a clash, not a second job.
_, err = rt.Apply(ctx, envelope("quantize-start", "op9", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrConflict))
})
func TestNoToolArgumentNamesAUserAPathOrASecret(t *testing.T) {
m := loadManifest(t)
for _, d := range m.Capabilities.ToolDefinitions {
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
}
if err := json.Unmarshal(d.InputSchema, &s); err != nil {
t.Fatal(err)
}
for name := range s.Properties {
lower := strings.ToLower(name)
if contract.ForbiddenToolProperty(name) || strings.Contains(lower, "path") || strings.Contains(lower, "dir") || strings.Contains(lower, "token") {
t.Fatalf("%s: property %q names an identity, a path or a secret", d.Name, name)
}
}
}
}
It("start sends only the fields of the engine operation", func() {
eng := &fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1", "status": "queued"}}}
rt := New(eng, GinkgoT().TempDir())
args := `{"name":"my quant","model":"org/base","backend":"llama-cpp-quantization","quantizationType":"q5_k_m"}`
_, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", args), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
c := eng.last("quantize.start")
Expect(c.opID).To(Equal("op1"), "the call must carry the operation of the request")
Expect(c.args).To(HaveLen(3))
Expect(c.args["model"]).To(Equal("org/base"))
Expect(c.args["backend"]).To(Equal("llama-cpp-quantization"))
Expect(c.args["quantizationType"]).To(Equal("q5_k_m"))
})
func TestStartSchemaHasTheHostOperationFieldsAndTheNameTheHostNeeds(t *testing.T) {
m := loadManifest(t)
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name != "quantize-start" {
continue
Describe("preview of start", func() {
It("has a nonzero revision and creates nothing", func() {
eng := &fakeEngine{}
rt := New(eng, GinkgoT().TempDir())
p, err := rt.Preview(ctx, envelope("quantize-start", "op1", "res1", startJSON))
Expect(err).NotTo(HaveOccurred())
Expect(p.Revision.Schema).NotTo(BeZero())
Expect(eng.ops()).To(BeEmpty())
Expect(json.Valid(p.Summary)).To(BeTrue(), "summary %s", p.Summary)
Expect(string(p.Summary)).To(ContainSubstring("q4_k_m"))
})
for name, args := range map[string]string{
"empty model": `{"name":"x","model":""}`,
"bad type": `{"name":"x","model":"m","quantizationType":"Q4 K"}`,
"long name": `{"name":"` + strings.Repeat("n", 81) + `","model":"m"}`,
"a token": `{"name":"x","model":"m","hfToken":"hf_x"}`,
"unknown part": `{"name":"x","model":"m","path":"/etc"}`,
} {
It("refuses "+name, func() {
rt := New(&fakeEngine{}, GinkgoT().TempDir())
_, err := rt.Preview(ctx, envelope("quantize-start", "op-"+name, "res1", args))
Expect(err).To(MatchError(protocol.ErrBadRequest))
})
}
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
Required []string `json:"required"`
})
It("status reads the job of the resource", func() {
rt, eng := started()
eng.result["quantize.get"] = map[string]any{"id": "job-1", "status": "quantizing", "model": "org/base"}
res, err := rt.Read(ctx, envelope("quantize-status", "r1", "res1", ""))
Expect(err).NotTo(HaveOccurred())
Expect(eng.last("quantize.get").args["jobId"]).To(Equal("job-1"))
var v map[string]any
Expect(json.Unmarshal(res.Payload, &v)).To(Succeed())
Expect(v["name"]).To(Equal("my quant"))
Expect(v["status"]).To(Equal("quantizing"))
rev, _ := v["revision"].(map[string]any)
Expect(rev["schema"]).NotTo(BeNil(), "inspect needs a revision: %v", v)
Expect(rev["schema"]).NotTo(Equal(float64(0)), "inspect needs a revision: %v", v)
Expect(v).NotTo(HaveKey("id"), "the engine job ID must not reach the model")
})
It("does not let a tool control a job it did not start", func() {
eng := &fakeEngine{}
rt := New(eng, GinkgoT().TempDir())
_, err := rt.Apply(ctx, envelope("quantize-stop", "op2", "nope", `{}`), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrNotFound), "stop")
_, err = rt.Apply(ctx, envelope("quantize-import", "op3", "nope", `{}`), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrNotFound), "import")
_, err = rt.Read(ctx, envelope("quantize-status", "r1", "nope", ""))
Expect(err).To(MatchError(protocol.ErrNotFound), "status")
_, err = rt.Preview(ctx, envelope("quantize-stop", "op2", "nope", `{}`))
Expect(err).To(MatchError(protocol.ErrNotFound), "preview")
Expect(eng.ops()).To(BeEmpty(), "the engine must not be called")
})
It("stop and import use the stored job and advance the revision", func() {
rt, eng := started()
eng.result["quantize.import"] = map[string]any{"model": "my-model"}
p, err := rt.Preview(ctx, envelope("quantize-stop", "op2", "res1", `{}`))
Expect(err).NotTo(HaveOccurred())
Expect(p.Revision.Schema).NotTo(BeZero())
rec, err := rt.Apply(ctx, envelope("quantize-stop", "op2", "res1", `{}`), p.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(rec.State).To(Equal(contract.ReceiptCommitted))
c := eng.last("quantize.stop")
Expect(c.args).To(HaveLen(1), "a quantization job has no checkpoint to save")
Expect(c.args["jobId"]).To(Equal("job-1"))
Expect(rec.Revision.Data).To(BeNumerically(">", p.Revision.Data), "the revision must advance")
_, err = rt.Apply(ctx, envelope("quantize-stop", "op2", "res1", `{}`), p.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(eng.count("quantize.stop")).To(Equal(1))
// Import names the model and never a path.
ex, err := rt.Apply(ctx, envelope("quantize-import", "op3", "res1", `{"name":"my-model"}`), rec.Revision)
Expect(err).NotTo(HaveOccurred())
Expect(ex.State).To(Equal(contract.ReceiptCommitted))
Expect(string(ex.Result.Payload)).To(ContainSubstring("my-model"))
a := eng.last("quantize.import").args
Expect(a).To(HaveLen(2))
Expect(a["jobId"]).To(Equal("job-1"))
Expect(a["name"]).To(Equal("my-model"))
_, err = rt.Preview(ctx, envelope("quantize-import", "op4", "res1", `{"name":"../x"}`))
Expect(err).To(MatchError(protocol.ErrBadRequest), "a bad model name must be refused")
_, err = rt.Apply(ctx, envelope("quantize-import", "op5", "res1", `{"name":"a b"}`), ex.Revision)
Expect(err).To(MatchError(protocol.ErrBadRequest), "a bad model name must be refused at apply")
})
It("rejects a stale revision and does not call the engine", func() {
rt, eng := started()
rec, err := rt.Apply(ctx, envelope("quantize-stop", "op2", "res1", `{}`), contract.Revision{Schema: 1, Data: 99})
Expect(err).NotTo(HaveOccurred())
Expect(rec.State).To(Equal(contract.ReceiptRejected))
Expect(eng.count("quantize.stop")).To(BeZero(), "the engine was called")
})
Describe("host refusals", func() {
for code, want := range map[string]error{
contract.HostFeature: protocol.ErrForbidden, contract.HostBadRequest: protocol.ErrBadRequest,
contract.HostNotFound: protocol.ErrNotFound, contract.HostConflict: protocol.ErrConflict,
} {
It("maps "+code+" to a protocol error", func() {
rt, eng := started()
eng.result["quantize.import"] = &contract.HostError{Code: code, Message: "m"}
_, err := rt.Apply(ctx, envelope("quantize-import", "op-"+code, "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
Expect(err).To(MatchError(want))
})
}
_ = json.Unmarshal(d.InputSchema, &s)
var got []string
for k := range s.Properties {
got = append(got, k)
}
if len(got) != 4 {
t.Fatalf("properties %v", got)
}
for _, w := range []string{"name", "model", "backend", "quantizationType"} {
if _, ok := s.Properties[w]; !ok {
t.Fatalf("missing %s", w)
It("keeps an unavailable engine uncertain", func() {
rt, eng := started()
eng.result["quantize.import"] = &contract.HostError{Code: contract.HostUnavailable, Message: "m"}
_, err := rt.Apply(ctx, envelope("quantize-import", "op-u", "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
Expect(err).To(HaveOccurred())
for _, e := range []error{protocol.ErrBadRequest, protocol.ErrForbidden, protocol.ErrNotFound, protocol.ErrConflict} {
Expect(err).NotTo(MatchError(e), "an unavailable engine must stay uncertain")
}
})
})
It("answers resource-exists from the journal and the engine", func() {
rt, eng := started()
eng.result["quantize.get"] = map[string]any{"id": "job-1", "status": "quantizing"}
ok, err := exists(rt, "res1")
Expect(err).NotTo(HaveOccurred(), "mapped")
Expect(ok).To(BeTrue(), "mapped")
ok, err = exists(rt, "other")
Expect(err).NotTo(HaveOccurred(), "unmapped")
Expect(ok).To(BeFalse(), "unmapped")
eng.result["quantize.get"] = &contract.HostError{Code: contract.HostNotFound, Message: "gone"}
ok, err = exists(rt, "res1")
Expect(err).NotTo(HaveOccurred(), "deleted job")
Expect(ok).To(BeFalse(), "deleted job")
eng.result["quantize.get"] = &contract.HostError{Code: contract.HostUnavailable, Message: "down"}
_, err = exists(rt, "res1")
Expect(err).To(HaveOccurred(), "an unavailable engine must not read as absent")
eng.result["quantize.get"] = hostclient.ErrNoHostOps
_, err = exists(rt, "res1")
Expect(err).To(HaveOccurred(), "no host endpoint must not read as absent")
})
It("keeps the journal across a restart", func() {
dir := GinkgoT().TempDir()
eng := &fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1", "status": "queued"}, "quantize.get": map[string]any{"status": "quantizing"}}}
_, err := New(eng, dir).Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
again := New(eng, dir)
_, err = again.Read(ctx, envelope("quantize-status", "r1", "res1", ""))
Expect(err).NotTo(HaveOccurred())
Expect(eng.last("quantize.get").args["jobId"]).To(Equal("job-1"), "the restarted runtime lost the job")
_, err = again.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
Expect(eng.count("quantize.start")).To(Equal(1))
r, err := again.Receipt(ctx, "op1")
Expect(err).NotTo(HaveOccurred())
Expect(r.State).To(Equal(contract.ReceiptCommitted))
_, err = again.Receipt(ctx, "never")
Expect(err).To(MatchError(protocol.ErrNotFound), "an unknown operation must be not found")
entries, _ := os.ReadDir(dir)
for _, e := range entries {
Expect(e.Name()).NotTo(ContainSubstring(".tmp"), "a temporary file was left")
}
if !reflect.DeepEqual(s.Required, []string{"name", "model"}) {
t.Fatalf("required %v", s.Required)
_, err = os.Stat(filepath.Join(dir, "journal.json"))
Expect(err).NotTo(HaveOccurred())
})
It("refuses every tool without a host endpoint", func() {
var nilClient *hostclient.Client
dir := GinkgoT().TempDir()
rt := New(nilClient, dir)
_, err := rt.Apply(ctx, envelope("quantize-start", "op1", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).To(MatchError(hostclient.ErrNoHostOps), "start")
seed := New(&fakeEngine{result: map[string]any{"quantize.start": map[string]any{"id": "job-1"}}}, dir)
_, err = seed.Apply(ctx, envelope("quantize-start", "op0", "res1", startJSON), contract.Revision{Schema: 1})
Expect(err).NotTo(HaveOccurred())
rt = New(nilClient, dir)
for _, tool := range []string{"quantize-status", "resource-exists"} {
_, err := rt.Read(ctx, envelope(tool, "r-"+tool, "res1", ""))
Expect(err).To(MatchError(hostclient.ErrNoHostOps), tool)
}
return
}
t.Fatal("no start definition")
}
_, err = rt.Read(ctx, envelope("quantize-list", "r-list", "", ""))
Expect(err).To(MatchError(hostclient.ErrNoHostOps), "list")
for _, tool := range []string{"quantize-stop", "quantize-import"} {
_, err := rt.Apply(ctx, envelope(tool, "a-"+tool, "res1", `{}`), contract.Revision{Schema: 1, Data: 1})
Expect(err).To(MatchError(hostclient.ErrNoHostOps), tool)
}
})
It("treats an unknown tool as a bad request", func() {
rt := New(&fakeEngine{}, GinkgoT().TempDir())
_, err := rt.Read(ctx, envelope("quantize-nope", "r1", "res1", ""))
Expect(err).To(MatchError(protocol.ErrBadRequest), "read")
_, err = rt.Apply(ctx, envelope("quantize-nope", "a1", "res1", ""), contract.Revision{Schema: 1})
Expect(err).To(MatchError(protocol.ErrBadRequest), "apply")
_, err = rt.OpenOutput(ctx, envelope("quantize-status", "o1", "res1", ""), "h")
Expect(err).To(MatchError(protocol.ErrNotFound), "output")
})
It("lists the caller's jobs without a resource", func() {
eng := &fakeEngine{result: map[string]any{"quantize.list": map[string]any{"jobs": []any{map[string]any{"id": "ui-job", "status": "completed"}}}}}
rt := New(eng, GinkgoT().TempDir())
res, err := rt.Read(ctx, envelope("quantize-list", "r1", "", ""))
Expect(err).NotTo(HaveOccurred())
Expect(string(res.Payload)).To(ContainSubstring("ui-job"))
Expect(eng.ops()).To(Equal([]string{"quantize.list"}))
Expect(eng.last("quantize.list").args).To(BeNil(), "the list takes no argument")
_, err = rt.Read(ctx, envelope("quantize-list", "r2", "", `{"userId":"other"}`))
Expect(err).To(MatchError(protocol.ErrBadRequest), "an argument must be refused")
})
Describe("manifest", func() {
It("declares the tools and the resource mapping", func() {
m := loadManifest()
Expect(m.ID).To(Equal("quantization"))
Expect(m.Capabilities.UI.MainPage).To(Equal("jobs"))
Expect(m.Capabilities.UI.Pages).To(Equal([]string{"jobs"}))
Expect(m.Capabilities.UI.Title).NotTo(BeNil())
Expect(m.Capabilities.UI.Title.Text).To(Equal("Quantization"))
Expect(m.Capabilities.UI.SDK.Min).To(BeEquivalentTo(2))
Expect(m.Capabilities.UI.SDK.Max).To(BeEquivalentTo(2))
Expect(m.Capabilities.Tools).To(Equal(Tools))
r := m.Capabilities.Resources
Expect(r).NotTo(BeNil())
Expect(r.Noun).To(Equal("quantization job"))
Expect(r.Create).To(Equal("quantize-start"))
Expect(r.Inspect).To(Equal("quantize-status"))
kinds := map[string]string{}
for _, d := range m.Capabilities.ToolDefinitions {
kinds[d.Name] = d.Kind
}
want := map[string]string{"quantize-list": "read", "quantize-status": "read", "quantize-start": "create", "quantize-stop": "write", "quantize-import": "write"}
Expect(kinds).To(Equal(want))
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name == "quantize-list" {
Expect(contract.IsResourceless(d)).To(BeTrue(), "the list tool must declare no resource")
} else {
Expect(contract.IsResourceless(d)).To(BeFalse(), "%s must not be resource-less", d.Name)
}
Expect(len(d.Description)).To(BeNumerically("<", 400), d.Name)
Expect(d.Description).To(ContainSubstring("started"), "%s: the description must say that only jobs a tool started can be controlled", d.Name)
}
})
It("names no user, path or secret in any tool argument", func() {
m := loadManifest()
for _, d := range m.Capabilities.ToolDefinitions {
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
}
Expect(json.Unmarshal(d.InputSchema, &s)).To(Succeed())
for name := range s.Properties {
lower := strings.ToLower(name)
Expect(contract.ForbiddenToolProperty(name) || strings.Contains(lower, "path") || strings.Contains(lower, "dir") || strings.Contains(lower, "token")).To(BeFalse(),
"%s: property %q names an identity, a path or a secret", d.Name, name)
}
}
})
It("gives the start schema the host operation fields and the name the host needs", func() {
m := loadManifest()
found := false
for _, d := range m.Capabilities.ToolDefinitions {
if d.Name != "quantize-start" {
continue
}
found = true
var s struct {
Properties map[string]json.RawMessage `json:"properties"`
Required []string `json:"required"`
}
_ = json.Unmarshal(d.InputSchema, &s)
Expect(s.Properties).To(HaveLen(4))
for _, w := range []string{"name", "model", "backend", "quantizationType"} {
Expect(s.Properties).To(HaveKey(w))
}
Expect(s.Required).To(Equal([]string{"name", "model"}))
break
}
Expect(found).To(BeTrue(), "no start definition")
})
})
})