feat(addons): move connector file bytes by handle

A connector add-on that declares files.transfer no longer carries file
bytes in the connector envelope. The host puts a random handle in the
request of the op media (write) or send_file (read), and the add-on
streams the file with two new host operations, connector.file.put and
connector.file.get, in chunks of at most 256 KiB.

The process that runs the add-on registers the handle when the call
arrives and ends it when the call returns. A handle is bound to the
add-on, so another add-on gets the same answer as for an unknown
handle. An unfinished inbound file is dropped with its partial bytes,
and a handle with no call for two minutes expires. Inbound bytes go
to the blob store of the workspace, outbound bytes come from it, so a
cluster needs the shared object store and no shared disk. The file is
never held in memory. The host overwrites the size and hash of an
inbound file with its own record.

The change is additive: the field files.transfer needs manifest API 2,
a kind without it keeps the inline path and its 1 MiB cap, and the host
never sends it a handle. A call that carries a handle may run up to 15
minutes. The ceiling of a file also stays under the largest file of a
workspace.

Assisted-by: Claude:claude-sonnet-5-5 go test
This commit is contained in:
Ettore Di Giacinto committed 2026-11-03 15:54:45 +00:00
1 parent f5f4b7ea58
commit 4ba1b7b2bd
21 files changed
+1453 -14

No files matched your search

+80 -6
View File
@@ -13,6 +13,7 @@ import (
"time"
"github.com/mudler/LocalAI/core/services/addons"
"github.com/mudler/LocalAI/core/services/addons/hostops"
"github.com/mudler/LocalAI/core/services/studio/connectors"
"github.com/mudler/LocalAI/pkg/addons/contract"
)
@@ -51,7 +52,7 @@ func (g *Gateway) connectorSource() ConnectorSource {
func (g *Gateway) ConnectorKinds(ctx context.Context) ([]connectors.Kind, error) {
if g.Forward != nil {
var out []connectors.Kind
err := g.Forward.connectorRequest(ctx, "GET", "/addons/connectors/kinds", nil, &out)
err := g.Forward.connectorRequest(ctx, "GET", "/addons/connectors/kinds", nil, &out, connectorTimeout)
return out, err
}
return g.localConnectorKinds(ctx)
@@ -77,13 +78,17 @@ func (g *Gateway) ConnectorCall(ctx context.Context, s connectors.ConnectorCallS
req.Token = ""
}
var out contract.ConnectorResponse
err := g.Forward.connectorRequest(ctx, "POST", "/addons/connectors/call", connectorCall{s, req}, &out)
err := g.Forward.connectorRequest(ctx, "POST", "/addons/connectors/call", connectorCall{s, req}, &out, req.CallTimeout())
return out, err
}
return g.localConnectorCall(ctx, s, req)
}
func (g *Gateway) localConnectorCall(ctx context.Context, s connectors.ConnectorCallScope, req contract.ConnectorRequest) (contract.ConnectorResponse, error) {
ctx, cancel := context.WithTimeout(ctx, connectorTimeout)
timeout := connectorTimeout
if req.Transfer != nil {
timeout = req.CallTimeout()
}
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
source := g.connectorSource()
if source == nil || g.ConnectorAuthority == nil {
@@ -109,7 +114,18 @@ func (g *Gateway) localConnectorCall(ctx context.Context, s connectors.Connector
}
return contract.ConnectorResponse{}, ErrForbidden
}
if req.Transfer != nil && (kinds[i].Files == nil || !kinds[i].Files.Transfer) {
// A handle only goes to an add-on that declared that it takes one.
return contract.ConnectorResponse{}, ErrForbidden
}
end, err := g.beginTransfer(s, req)
if err != nil {
return contract.ConnectorResponse{}, err
}
out, err := source.Connector(ctx, s.AddonID, req)
if end != nil {
out = finishTransfer(end, out)
}
if err != nil {
// A definite refusal of the add-on keeps its type; anything else could
// not tell, so the caller treats it as an outage, with the cause kept.
@@ -120,8 +136,8 @@ func (g *Gateway) localConnectorCall(ctx context.Context, s connectors.Connector
}
return out, nil
}
func (c *Carrier) connectorRequest(ctx context.Context, method, path string, body any, out any) error {
ctx, cancel := context.WithTimeout(ctx, connectorTimeout)
func (c *Carrier) connectorRequest(ctx context.Context, method, path string, body any, out any, timeout time.Duration) error {
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
var raw []byte
if body != nil {
@@ -136,7 +152,13 @@ func (c *Carrier) connectorRequest(ctx context.Context, method, path string, bod
return ErrUnavailable
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.client.Do(req)
// A call that moves a file runs longer than the client's own timeout; the
// context carries the bound.
client := c.client
if timeout > connectorTimeout {
client = c.long
}
resp, err := client.Do(req)
if err != nil {
return ErrUnavailable
}
@@ -172,6 +194,10 @@ func (g *Gateway) serveConnectors(w http.ResponseWriter, r *http.Request) bool {
typedError(w, ErrForbidden)
return true
}
if call.Request.Transfer != nil {
// The server's write timeout is sized for short calls.
_ = http.NewResponseController(w).SetWriteDeadline(time.Now().Add(contract.ConnectorTransferTimeout + time.Minute))
}
out, err := g.localConnectorCall(r.Context(), call.Scope, call.Request)
if err != nil {
typedError(w, err)
@@ -193,3 +219,51 @@ func decodeConnector(r io.Reader, out any) error {
}
return decode(bytes.NewReader(raw), out)
}
// beginTransfer registers the handle of a call that carries one, on the
// process that runs the add-on. The frontend made the handle and checked the
// chat's access; this process binds it to the add-on of the call, so no other
// add-on can use it, and to the content or the bound the scope names.
func (g *Gateway) beginTransfer(s connectors.ConnectorCallScope, req contract.ConnectorRequest) (func() hostops.TransferOutcome, error) {
if req.Transfer == nil {
return nil, nil
}
grant := hostops.TransferGrant{Handle: req.Transfer.Handle, AddonID: s.AddonID, BindingID: s.BindingID}
switch req.Op {
case contract.ConnectorOpMedia:
grant.Direction, grant.MaxBytes = hostops.TransferIn, req.Transfer.MaxBytes
case contract.ConnectorOpSendFile:
if s.Content == nil || req.File != nil {
return nil, ErrForbidden
}
grant.Direction, grant.SHA, grant.Size = hostops.TransferOut, s.Content.SHA, s.Content.Size
default:
return nil, ErrForbidden
}
if g.Transfers == nil {
return nil, ErrUnavailable
}
end, err := g.Transfers.Begin(grant)
if err != nil {
if errors.Is(err, hostops.ErrBadRequest) || errors.Is(err, hostops.ErrConflict) {
return nil, ErrForbidden
}
return nil, ErrUnavailable
}
return end, nil
}
// finishTransfer ends the handle after the add-on answered. For an inbound
// file it writes the size and the hash from the host's own record into the
// answer; an add-on that did not finish the file gets no Received.
func finishTransfer(end func() hostops.TransferOutcome, out contract.ConnectorResponse) contract.ConnectorResponse {
o := end()
if out.Received != nil {
if o.Direction != hostops.TransferIn || o.Result != hostops.TransferDone || o.SHA == "" {
out.Received = nil
} else {
out.Received.Size, out.Received.SHA256 = o.Bytes, o.SHA
}
}
return out
}
@@ -69,7 +69,7 @@ var _ = Describe("Forwarded connector answers", func() {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(status) }))
defer server.Close()
c := &Carrier{config: CarrierConfig{OwnerAddress: server.URL}, client: server.Client()}
return c.connectorRequest(context.Background(), "GET", "/addons/connectors/kinds", nil, &[]connectors.Kind{})
return c.connectorRequest(context.Background(), "GET", "/addons/connectors/kinds", nil, &[]connectors.Kind{}, connectorTimeout)
}
It("keeps a refusal apart from an outage", func() {
@@ -0,0 +1,111 @@
// SPDX-License-Identifier: MIT
package gateway
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"net/http"
"net/http/httptest"
"github.com/mudler/LocalAI/core/services/addons"
"github.com/mudler/LocalAI/core/services/addons/hostops"
"github.com/mudler/LocalAI/core/services/studio/connectors"
"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
"github.com/mudler/LocalAI/pkg/addons/contract"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
// handleSource is an add-on that takes handles: it runs fn during the call.
type handleSource struct {
transfer bool
fn func(contract.ConnectorRequest) contract.ConnectorResponse
}
func (s handleSource) ConnectorKinds(context.Context) ([]addons.ConnectorKind, error) {
k := contract.ConnectorKind{ID: "slack", Delivery: contract.DeliveryPull}
if s.transfer {
k.Files = &contract.ConnectorFiles{Receive: &contract.FileAbility{MaxBytes: 1 << 30}, Transfer: true}
}
return []addons.ConnectorKind{{AddonID: "slack", ConnectorKind: k}}, nil
}
func (s handleSource) Connector(_ context.Context, _ string, r contract.ConnectorRequest) (contract.ConnectorResponse, error) {
return s.fn(r), nil
}
var _ = Describe("Connector calls that carry a file handle", func() {
scope := connectors.ConnectorCallScope{AddonID: "slack", Kind: "slack", BindingID: "b1"}
var (
t *hostops.Transfers
host http.Handler
)
BeforeEach(func() {
blobs, err := blob.NewFS(GinkgoT().TempDir())
Expect(err).NotTo(HaveOccurred())
t = hostops.NewTransfers(blobs)
b := hostops.New(hostops.NewInFlight(), func(string, string) bool { return true }, contract.DefaultLimits())
Expect(hostops.RegisterConnectorTransfer(b, t)).To(Succeed())
host = b.Handler("slack", "k", []string{"connector."})
})
put := func(h string, off int64, data []byte, final bool) int {
raw, _ := json.Marshal(contract.ConnectorFilePut{Handle: h, Offset: off, Data: data, Final: final})
body, _ := json.Marshal(contract.HostRequest{OperationID: "o", Op: contract.ConnectorFilePutOp, Arguments: raw})
req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body))
req.Header.Set(contract.HostCredentialHeader, "k")
rec := httptest.NewRecorder()
host.ServeHTTP(rec, req)
return rec.Code
}
media := func(s handleSource, tr *hostops.Transfers) (contract.ConnectorResponse, error) {
g := &Gateway{Connectors: s, ConnectorAuthority: allowAllConnectors{}, Transfers: tr}
return g.ConnectorCall(context.Background(), scope, contract.ConnectorRequest{Op: contract.ConnectorOpMedia, Kind: "slack", Token: "t",
Transfer: &contract.ConnectorTransfer{Handle: "hh", MaxBytes: 100}})
}
It("registers the handle for the call and gives the host's size and hash, not the add-on's", func() {
src := handleSource{transfer: true, fn: func(r contract.ConnectorRequest) contract.ConnectorResponse {
Expect(t.Live()).To(Equal(1))
Expect(put(r.Transfer.Handle, 0, []byte("hello"), true)).To(Equal(http.StatusOK))
return contract.ConnectorResponse{Received: &contract.ConnectorReceived{Name: "a.txt", Size: 999, SHA256: "forged"}}
}}
out, err := media(src, t)
Expect(err).NotTo(HaveOccurred())
sum := sha256.Sum256([]byte("hello"))
Expect(out.Received.SHA256).To(Equal(hex.EncodeToString(sum[:])))
Expect(out.Received.Size).To(BeEquivalentTo(5))
Expect(t.Live()).To(BeZero(), "the handle ends with the call")
})
It("drops an answer for a file the add-on did not finish", func() {
src := handleSource{transfer: true, fn: func(r contract.ConnectorRequest) contract.ConnectorResponse {
Expect(put(r.Transfer.Handle, 0, []byte("half"), false)).To(Equal(http.StatusOK))
return contract.ConnectorResponse{Received: &contract.ConnectorReceived{Name: "a.txt", Size: 4, SHA256: "forged"}}
}}
out, err := media(src, t)
Expect(err).NotTo(HaveOccurred())
Expect(out.Received).To(BeNil())
Expect(t.Live()).To(BeZero())
})
It("refuses a handle for a kind that did not declare transfer, and without a table", func() {
src := handleSource{fn: func(contract.ConnectorRequest) contract.ConnectorResponse { return contract.ConnectorResponse{} }}
_, err := media(src, t)
Expect(err).To(MatchError(ErrForbidden))
src.transfer = true
_, err = media(src, nil)
Expect(err).To(MatchError(ErrUnavailable))
})
It("refuses a handle on an op that moves no file, and an outbound handle with no content", func() {
src := handleSource{transfer: true, fn: func(contract.ConnectorRequest) contract.ConnectorResponse { return contract.ConnectorResponse{} }}
g := &Gateway{Connectors: src, ConnectorAuthority: allowAllConnectors{}, Transfers: t}
_, err := g.ConnectorCall(context.Background(), scope, contract.ConnectorRequest{Op: contract.ConnectorOpSend, Kind: "slack", Transfer: &contract.ConnectorTransfer{Handle: "x"}})
Expect(err).To(MatchError(ErrForbidden))
_, err = g.ConnectorCall(context.Background(), scope, contract.ConnectorRequest{Op: contract.ConnectorOpSendFile, Kind: "slack", Transfer: &contract.ConnectorTransfer{Handle: "x", Size: 1}})
Expect(err).To(MatchError(ErrForbidden), "the frontend names the content in the scope")
})
})
+5
View File
@@ -11,6 +11,7 @@ import (
"sync"
"github.com/mudler/LocalAI/core/services/addons"
"github.com/mudler/LocalAI/core/services/addons/hostops"
"github.com/mudler/LocalAI/pkg/addons/contract"
)
@@ -25,6 +26,10 @@ var ErrNotSent = fmt.Errorf("the add-on is not running now and nothing was sent:
// Gateway deliberately admits lifecycle commands only until resource admission
// and proposal confirmation are installed. Installation admin is not data access.
type Gateway struct {
// Transfers holds the handles of connector file transfers on the process
// that runs the add-ons. Without it a call that carries a handle is
// unavailable.
Transfers *hostops.Transfers
Connectors ConnectorSource
ConnectorAuthority ConnectorAuthority
// HostSettings is owner-only durable policy; frontends forward, never fall back.
+1 -1
View File
@@ -54,5 +54,5 @@ func Requests(req *contract.HostOperations, registered, grants []string) Request
// must not mistake that startup ordering for unsupported operations. The broker
// still requires a registered handler and revalidates every request's authority.
func SupportedOperations() []string {
return strings.Fields(`finetune.start finetune.get finetune.list finetune.checkpoints finetune.dataset finetune.export finetune.progress finetune.stop quantize.start quantize.get quantize.list quantize.import quantize.progress quantize.stop workspace.materialize workspace.commit workspace.export workspace.import workspace.changes workspace.admin egress.check egress.audit connector.inbound`)
return strings.Fields(`finetune.start finetune.get finetune.list finetune.checkpoints finetune.dataset finetune.export finetune.progress finetune.stop quantize.start quantize.get quantize.list quantize.import quantize.progress quantize.stop workspace.materialize workspace.commit workspace.export workspace.import workspace.changes workspace.admin egress.check egress.audit connector.inbound connector.file.put connector.file.get`)
}
@@ -37,6 +37,7 @@ func registeredBuiltinOperations() []string {
})).To(Succeed())
Expect(RegisterEgress(b, &fakeEgress{})).To(Succeed())
Expect(RegisterConnector(b, admissionInbox{})).To(Succeed())
Expect(RegisterConnectorTransfer(b, NewTransfers(nil))).To(Succeed())
// Each registration error is checked above, so a duplicate cannot be hidden
// by the broker's map. Ops includes dynamic progress/stop and inbound specs.
return b.Ops()
+4
View File
@@ -327,6 +327,10 @@ func (b *Broker) serveInbound(w http.ResponseWriter, r *http.Request, spec Spec,
status, code, msg = http.StatusForbidden, contract.HostScope, "not allowed"
case errors.Is(err, ErrNotFound):
status, code, msg = http.StatusNotFound, contract.HostNotFound, "not found"
case errors.Is(err, ErrTransferTooLarge):
status, code, msg = http.StatusRequestEntityTooLarge, contract.HostTooLarge, "the file is larger than the limit"
case errors.Is(err, ErrConflict):
status, code, msg = http.StatusConflict, contract.HostConflict, "conflict"
case errors.Is(err, ErrBadRequest):
status, code, msg = http.StatusBadRequest, contract.HostBadRequest, err.Error()
}
+450
View File
@@ -0,0 +1,450 @@
// SPDX-License-Identifier: MIT
package hostops
import (
"bytes"
"context"
"crypto/subtle"
"encoding/json"
"errors"
"fmt"
"io"
"sync"
"time"
"github.com/mudler/LocalAI/core/services/studio/workspace"
"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
"github.com/mudler/LocalAI/pkg/addons/contract"
)
// Direction of a transfer, from the host.
const (
// TransferIn moves a file from an add-on to the host: the add-on puts.
TransferIn = "in"
// TransferOut moves a file from the host to an add-on: the add-on gets.
TransferOut = "out"
)
// Results of a transfer, as the audit shows them.
const (
TransferDone = "done"
TransferTooLarge = "too_large"
TransferExpired = "expired"
TransferRefused = "refused"
TransferIncomplete = "incomplete"
)
// Default times of a handle.
const (
// DefaultTransferIdle ends a handle that has seen no call for this long.
DefaultTransferIdle = 2 * time.Minute
// DefaultTransferLife ends a handle this long after it was made, whatever
// it does. A connector call is shorter than this.
DefaultTransferLife = time.Hour
)
// ErrTransferTooLarge means a put went past the bound of its handle.
var ErrTransferTooLarge = errors.New("the file is larger than the limit of the handle")
// TransferGrant is what the host that issued a handle records about it. The
// owner of the add-ons registers it when a connector call that carries a
// ConnectorTransfer arrives.
type TransferGrant struct {
Handle string
Direction string
// AddonID is the only add-on that may use the handle.
AddonID string
// BindingID is for the audit.
BindingID string
// MaxBytes bounds an inbound file.
MaxBytes int64
// SHA and Size name the outbound content, a blob of the workspace store.
SHA string
Size int64
}
// TransferOutcome is what a handle ended with. It never holds a name or any
// content.
type TransferOutcome struct {
Direction string
BindingID string
Bytes int64
Result string
Duration time.Duration
// SHA is the hash of a finished inbound file.
SHA string
}
// FileScan is the hook after the last byte of an inbound file. An error
// rejects the file. It is a no-op until an owner sets one.
type FileScan func(ctx context.Context, sha string, size int64) error
// Transfers is the table of live handles of one process, and the bytes behind
// them. Inbound bytes stream into the blob store of the workspace, which in a
// cluster is the shared object store, so no node needs a shared disk. Outbound
// bytes stream out of the same store. A file is never held in memory: a put
// or a get moves one chunk.
type Transfers struct {
// Blobs is the blob store of the workspace. UseBlobs sets it when the
// files start, which can be after the table is made.
Blobs blob.Store
Scan FileScan
// Audit gets the outcome of every handle when it ends. Nil: none.
Audit func(TransferOutcome)
// Idle and Life default to DefaultTransferIdle and DefaultTransferLife.
Idle, Life time.Duration
mu sync.Mutex
m map[string]*xfer
}
type xfer struct {
t *Transfers
g TransferGrant
start time.Time
mu sync.Mutex // one call at a time per handle
received int64
pw *io.PipeWriter
done chan struct{} // closed when the blob Put of an inbound handle returned
sha string
size int64
putErr error
final bool
rd io.ReadSeekCloser
cancel context.CancelFunc
idle *time.Timer
life *time.Timer
ended bool
result string
}
// NewTransfers builds a table over the blob store of the workspace.
func NewTransfers(blobs blob.Store) *Transfers {
return &Transfers{Blobs: blobs, m: map[string]*xfer{}}
}
// UseBlobs sets the blob store. Handles that begin before it is set are
// refused as unavailable.
func (t *Transfers) UseBlobs(b blob.Store) {
t.mu.Lock()
t.Blobs = b
t.mu.Unlock()
}
func (t *Transfers) store() blob.Store {
t.mu.Lock()
defer t.mu.Unlock()
return t.Blobs
}
func (t *Transfers) idle() time.Duration {
if t.Idle > 0 {
return t.Idle
}
return DefaultTransferIdle
}
func (t *Transfers) life() time.Duration {
if t.Life > 0 {
return t.Life
}
return DefaultTransferLife
}
// Begin registers a handle. The returned function ends it: it drops a file
// that is not finished and reports the outcome. For a finished inbound file the
// outcome holds the hash and the size. Calling it again returns the same
// outcome. A handle that is already registered is refused, so a handle is
// never used for two calls.
func (t *Transfers) Begin(g TransferGrant) (end func() TransferOutcome, err error) {
if t == nil || t.store() == nil {
return nil, errors.New("file transfer is not available here")
}
if g.Handle == "" || len(g.Handle) > contract.MaxTransferHandle || g.AddonID == "" {
return nil, fmt.Errorf("%w: bad handle", ErrBadRequest)
}
switch g.Direction {
case TransferIn:
if g.MaxBytes <= 0 {
return nil, fmt.Errorf("%w: no bound", ErrBadRequest)
}
case TransferOut:
if g.SHA == "" || g.Size <= 0 {
return nil, fmt.Errorf("%w: no content", ErrBadRequest)
}
default:
return nil, fmt.Errorf("%w: bad direction", ErrBadRequest)
}
x := &xfer{t: t, g: g, start: time.Now()}
t.mu.Lock()
if t.m == nil {
t.m = map[string]*xfer{}
}
if _, dup := t.m[g.Handle]; dup {
t.mu.Unlock()
return nil, fmt.Errorf("%w: handle in use", ErrConflict)
}
t.m[g.Handle] = x
t.mu.Unlock()
x.idle = time.AfterFunc(t.idle(), func() { x.finish(TransferExpired) })
x.life = time.AfterFunc(t.life(), func() { x.finish(TransferExpired) })
return func() TransferOutcome { return x.finish("") }, nil
}
// Live reports how many handles exist; a test and the sweep use it.
func (t *Transfers) Live() int {
t.mu.Lock()
defer t.mu.Unlock()
return len(t.m)
}
func (t *Transfers) lookup(addonID, handle string) (*xfer, error) {
t.mu.Lock()
x := t.m[handle]
t.mu.Unlock()
// The same answer for a handle that does not exist and one of another
// add-on: a handle cannot be probed.
if x == nil || subtle.ConstantTimeCompare([]byte(x.g.AddonID), []byte(addonID)) != 1 {
return nil, ErrNotFound
}
return x, nil
}
// finish ends the handle once. why is the result when the handle did not end
// by its own end function.
func (x *xfer) finish(why string) TransferOutcome {
x.idle.Stop()
x.life.Stop()
// A put or a get in progress holds mu; an expiry must not wait for a
// stuck one, so it cancels first.
if x.cancel != nil {
x.cancel()
}
if x.pw != nil {
_ = x.pw.CloseWithError(context.Canceled)
}
x.mu.Lock()
defer x.mu.Unlock()
if x.ended {
return x.outcome()
}
x.ended = true
x.t.mu.Lock()
delete(x.t.m, x.g.Handle)
x.t.mu.Unlock()
if x.rd != nil {
_ = x.rd.Close()
}
if x.done != nil {
<-x.done // the blob Put drops its temp file when its reader fails
}
switch {
case x.result != "":
case why != "":
x.result = why
case x.g.Direction == TransferIn && !x.final:
x.result = TransferIncomplete
default:
x.result = TransferDone
}
o := x.outcome()
if x.t.Audit != nil {
x.t.Audit(o)
}
return o
}
func (x *xfer) outcome() TransferOutcome {
o := TransferOutcome{Direction: x.g.Direction, BindingID: x.g.BindingID, Bytes: x.received, Result: x.result, Duration: time.Since(x.start)}
if x.g.Direction == TransferOut {
o.Bytes = x.size
}
if x.final && x.result != TransferTooLarge && x.result != TransferRefused {
o.SHA = x.sha
}
return o
}
func (x *xfer) touch() { x.idle.Reset(x.t.idle()) }
// put writes the next chunk of an inbound file.
func (x *xfer) put(ctx context.Context, p contract.ConnectorFilePut) (int64, error) {
if x.g.Direction != TransferIn {
return 0, ErrNotFound
}
x.mu.Lock()
defer x.mu.Unlock()
if x.ended || x.final {
return 0, ErrConflict
}
x.touch()
if len(p.Data) > contract.TransferChunkBytes {
return 0, fmt.Errorf("%w: chunk too large", ErrBadRequest)
}
if p.Offset != x.received {
// The add-on resumes where the host is.
return x.received, nil
}
if x.received+int64(len(p.Data)) > x.g.MaxBytes {
x.result = TransferTooLarge
x.abortLocked()
return x.received, ErrTransferTooLarge
}
if x.pw == nil {
pr, pw := io.Pipe()
bctx, cancel := context.WithCancel(context.Background())
x.pw, x.cancel, x.done = pw, cancel, make(chan struct{})
go func() {
defer close(x.done)
sha, size, err := x.t.store().Put(bctx, pr, x.g.MaxBytes)
x.sha, x.size, x.putErr = sha, size, err
_ = pr.CloseWithError(errors.New("closed"))
}()
}
if len(p.Data) > 0 {
if _, err := io.Copy(x.pw, bytes.NewReader(p.Data)); err != nil {
return x.received, x.failedLocked(err)
}
x.received += int64(len(p.Data))
}
if !p.Final {
return x.received, nil
}
if x.received == 0 {
x.result = TransferRefused
x.abortLocked()
return 0, fmt.Errorf("%w: an empty file", ErrBadRequest)
}
_ = x.pw.Close()
<-x.done
switch {
case errors.Is(x.putErr, workspace.ErrTooLarge):
x.result = TransferTooLarge
return x.received, ErrTransferTooLarge
case x.putErr != nil:
return x.received, x.putErr
case x.size != x.received:
return x.received, errors.New("the stored size differs from the bytes received")
}
if x.t.Scan != nil {
if err := x.t.Scan(ctx, x.sha, x.size); err != nil {
x.result = TransferRefused
return x.received, fmt.Errorf("%w: the file was refused", ErrBadRequest)
}
}
x.final = true
return x.received, nil
}
// failedLocked turns a failed write into the error of the blob store, which
// ended the pipe, so that a bound reached inside the store is too large.
func (x *xfer) failedLocked(err error) error {
<-x.done
if errors.Is(x.putErr, workspace.ErrTooLarge) {
x.result = TransferTooLarge
return ErrTransferTooLarge
}
if x.putErr != nil {
return x.putErr
}
return err
}
func (x *xfer) abortLocked() {
if x.pw != nil {
_ = x.pw.CloseWithError(context.Canceled)
if x.cancel != nil {
x.cancel()
}
<-x.done
}
}
// get reads the next chunk of an outbound file.
func (x *xfer) get(ctx context.Context, g contract.ConnectorFileGet) (contract.ConnectorFileGetResult, error) {
if x.g.Direction != TransferOut {
return contract.ConnectorFileGetResult{}, ErrNotFound
}
x.mu.Lock()
defer x.mu.Unlock()
if x.ended {
return contract.ConnectorFileGetResult{}, ErrConflict
}
x.touch()
if g.Offset < 0 || g.Offset > x.g.Size || g.Length < 1 || g.Length > contract.TransferChunkBytes {
return contract.ConnectorFileGetResult{}, fmt.Errorf("%w: offset or length", ErrBadRequest)
}
if x.rd == nil {
r, err := x.t.store().Open(ctx, x.g.SHA)
if err != nil {
return contract.ConnectorFileGetResult{}, err
}
n, err := r.Seek(0, io.SeekEnd)
if err != nil || n != x.g.Size {
_ = r.Close()
return contract.ConnectorFileGetResult{}, errors.New("the file changed since the handle was made")
}
x.rd, x.size = r, n
}
if _, err := x.rd.Seek(g.Offset, io.SeekStart); err != nil {
return contract.ConnectorFileGetResult{}, err
}
buf := make([]byte, min(g.Length, x.g.Size-g.Offset))
n, err := io.ReadFull(x.rd, buf)
if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) && !errors.Is(err, io.EOF) {
return contract.ConnectorFileGetResult{}, err
}
buf = buf[:n]
return contract.ConnectorFileGetResult{Data: buf, EOF: g.Offset+int64(n) >= x.g.Size}, nil
}
// RegisterConnectorTransfer registers the two operations with which a
// connector add-on moves file bytes by handle. Both are inbound class: the
// credential of the live incarnation authorizes them, and the handle binds
// them to one connector call. The arguments name no user.
func RegisterConnectorTransfer(b *Broker, t *Transfers) error {
if t == nil {
return errors.New("connector file transfer needs a table")
}
decode := func(raw json.RawMessage, v any) error {
d := json.NewDecoder(bytes.NewReader(raw))
d.DisallowUnknownFields()
if err := d.Decode(v); err != nil {
return ErrBadRequest
}
return nil
}
if err := b.Register(Spec{Op: contract.ConnectorFilePutOp, Scope: ScopeInbound, Timeout: 2 * time.Minute, Handler: func(ctx context.Context, c Call) (json.RawMessage, error) {
var p contract.ConnectorFilePut
if err := decode(c.Arguments, &p); err != nil {
return nil, err
}
x, err := t.lookup(c.AddonID, p.Handle)
if err != nil {
return nil, err
}
got, err := x.put(ctx, p)
if err != nil {
return nil, err
}
return json.Marshal(contract.ConnectorFilePutResult{Received: got})
}}); err != nil {
return err
}
return b.Register(Spec{Op: contract.ConnectorFileGetOp, Scope: ScopeInbound, Timeout: 2 * time.Minute, Handler: func(ctx context.Context, c Call) (json.RawMessage, error) {
var g contract.ConnectorFileGet
if err := decode(c.Arguments, &g); err != nil {
return nil, err
}
x, err := t.lookup(c.AddonID, g.Handle)
if err != nil {
return nil, err
}
r, err := x.get(ctx, g)
if err != nil {
return nil, err
}
return json.Marshal(r)
}})
}
@@ -0,0 +1,284 @@
// SPDX-License-Identifier: MIT
package hostops
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"io"
"math/rand"
"net/http"
"net/http/httptest"
"path/filepath"
"runtime"
"time"
"github.com/mudler/LocalAI/core/services/storage"
"github.com/mudler/LocalAI/core/services/studio/workspace/blob"
"github.com/mudler/LocalAI/pkg/addons/contract"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
type xferFixture struct {
t *Transfers
h http.Handler
outs []TransferOutcome
blobs blob.Store
dir string
shared *storage.FilesystemStore
}
func newXferFixture() *xferFixture {
dir := GinkgoT().TempDir()
shared, err := storage.NewFilesystemStore(filepath.Join(dir, "objects"))
Expect(err).NotTo(HaveOccurred())
// The owner node has its own cache disk; the object store is the only
// thing the nodes share.
blobs := blob.NewObjectStore(shared, filepath.Join(dir, "owner-cache"))
f := &xferFixture{blobs: blobs, dir: dir, shared: shared}
f.t = NewTransfers(blobs)
f.t.Audit = func(o TransferOutcome) { f.outs = append(f.outs, o) }
b := New(NewInFlight(), always, contract.DefaultLimits())
Expect(RegisterConnectorTransfer(b, f.t)).To(Succeed())
f.h = b.Handler("slack", "secret", []string{"connector."})
return f
}
// call posts one host operation as the add-on addon.
func (f *xferFixture) call(addon, op string, args any) (int, contract.HostResponse) {
raw, _ := json.Marshal(args)
body, _ := json.Marshal(contract.HostRequest{OperationID: "x", Op: op, Arguments: raw})
req := httptest.NewRequest(http.MethodPost, "/host", bytes.NewReader(body))
req.Header.Set(contract.HostCredentialHeader, "secret")
rec := httptest.NewRecorder()
h := f.h
if addon != "slack" {
b := New(NewInFlight(), always, contract.DefaultLimits())
Expect(RegisterConnectorTransfer(b, f.t)).To(Succeed())
h = b.Handler(addon, "secret", []string{"connector."})
}
h.ServeHTTP(rec, req)
var out contract.HostResponse
_ = json.Unmarshal(rec.Body.Bytes(), &out)
return rec.Code, out
}
func (f *xferFixture) put(handle string, off int64, data []byte, final bool) (int, int64) {
code, r := f.call("slack", contract.ConnectorFilePutOp, contract.ConnectorFilePut{Handle: handle, Offset: off, Data: data, Final: final})
var res contract.ConnectorFilePutResult
if r.Result != nil {
_ = json.Unmarshal(r.Result, &res)
}
return code, res.Received
}
func (f *xferFixture) get(handle string, off, n int64) (int, contract.ConnectorFileGetResult) {
code, r := f.call("slack", contract.ConnectorFileGetOp, contract.ConnectorFileGet{Handle: handle, Offset: off, Length: n})
var res contract.ConnectorFileGetResult
if r.Result != nil {
_ = json.Unmarshal(r.Result, &res)
}
return code, res
}
// chunks streams n bytes of a seeded sequence, chunk bytes at a time, and
// never holds more than one chunk.
func chunks(n, chunk int64, each func(off int64, b []byte, last bool)) string {
rng := rand.New(rand.NewSource(7))
h := sha256.New()
buf := make([]byte, chunk)
for off := int64(0); off < n; off += chunk {
b := buf[:min(chunk, n-off)]
_, _ = rng.Read(b)
h.Write(b)
each(off, b, off+int64(len(b)) >= n)
}
return hex.EncodeToString(h.Sum(nil))
}
func heap() uint64 {
var m runtime.MemStats
runtime.GC()
runtime.ReadMemStats(&m)
return m.HeapAlloc
}
var _ = Describe("connector file transfer by handle", func() {
const mb = 1 << 20
It("moves 20 MB in and 20 MB out with bounded memory", func() {
f := newXferFixture()
end, err := f.t.Begin(TransferGrant{Handle: "h1", Direction: TransferIn, AddonID: "slack", BindingID: "b1", MaxBytes: 50 * mb})
Expect(err).NotTo(HaveOccurred())
base, peak := heap(), uint64(0)
var count int
want := chunks(20*mb, contract.TransferChunkBytes, func(off int64, b []byte, last bool) {
code, got := f.put("h1", off, append([]byte(nil), b...), last)
Expect(code).To(Equal(http.StatusOK))
Expect(got).To(Equal(off + int64(len(b))))
if count++; count%16 == 0 {
peak = max(peak, heap())
}
})
out := end()
Expect(out.Result).To(Equal(TransferDone))
Expect(out.Bytes).To(Equal(int64(20 * mb)))
Expect(out.SHA).To(Equal(want))
Expect(peak).To(BeNumerically("<", base+16*mb), "the file is never held in memory")
Expect(f.t.Live()).To(BeZero())
// Another node: same object store, its own cache disk.
node2 := blob.NewObjectStore(f.shared, filepath.Join(f.dir, "frontend-cache"))
r, err := node2.Open(GinkgoT().Context(), out.SHA)
Expect(err).NotTo(HaveOccurred())
DeferCleanup(func() { _ = r.Close() })
h := sha256.New()
n, err := io.Copy(h, r)
Expect(err).NotTo(HaveOccurred())
Expect(n).To(Equal(int64(20 * mb)))
Expect(hex.EncodeToString(h.Sum(nil))).To(Equal(want))
// Out: the add-on pulls the same file by a new handle.
end2, err := f.t.Begin(TransferGrant{Handle: "h2", Direction: TransferOut, AddonID: "slack", SHA: out.SHA, Size: n})
Expect(err).NotTo(HaveOccurred())
hh := sha256.New()
var off int64
base, peak = heap(), 0
for i := 0; ; i++ {
code, res := f.get("h2", off, contract.TransferChunkBytes)
Expect(code).To(Equal(http.StatusOK))
hh.Write(res.Data)
off += int64(len(res.Data))
if i%16 == 0 {
peak = max(peak, heap())
}
if res.EOF {
break
}
}
Expect(off).To(Equal(n))
Expect(hex.EncodeToString(hh.Sum(nil))).To(Equal(want))
Expect(peak).To(BeNumerically("<", base+16*mb))
Expect(end2().Result).To(Equal(TransferDone))
Expect(f.outs).To(HaveLen(2))
})
It("resumes at the offset the host has after a gap", func() {
f := newXferFixture()
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: mb})
code, got := f.put("h", 0, []byte("abcd"), false)
Expect(code).To(Equal(http.StatusOK))
Expect(got).To(BeEquivalentTo(4))
_, got = f.put("h", 10, []byte("zz"), false)
Expect(got).To(BeEquivalentTo(4), "a gap answers where the host is and writes nothing")
_, got = f.put("h", 4, []byte("efgh"), true)
Expect(got).To(BeEquivalentTo(8))
o := end()
sum := sha256.Sum256([]byte("abcdefgh"))
Expect(o.SHA).To(Equal(hex.EncodeToString(sum[:])))
})
It("refuses an oversize file and drops the partial one", func() {
f := newXferFixture()
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 10})
code, _ := f.put("h", 0, []byte("12345678"), false)
Expect(code).To(Equal(http.StatusOK))
code, _ = f.put("h", 8, []byte("9012"), true)
Expect(code).To(Equal(http.StatusRequestEntityTooLarge))
o := end()
Expect(o.Result).To(Equal(TransferTooLarge))
Expect(o.SHA).To(BeEmpty())
code, _ = f.put("h", 0, []byte("x"), false)
Expect(code).To(Equal(http.StatusNotFound), "the handle ended")
})
It("refuses a file whose last chunk says it is short of the bytes sent, an empty file, and a write after final", func() {
f := newXferFixture()
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
code, _ := f.put("h", 0, nil, true)
Expect(code).To(Equal(http.StatusBadRequest), "empty")
Expect(end().Result).To(Equal(TransferRefused))
end, _ = f.t.Begin(TransferGrant{Handle: "g", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
code, _ = f.put("g", 0, []byte("ab"), true)
Expect(code).To(Equal(http.StatusOK))
code, _ = f.put("g", 2, []byte("cd"), false)
Expect(code).To(Equal(http.StatusConflict), "write after final")
Expect(end().Result).To(Equal(TransferDone))
})
It("refuses another add-on, an unknown handle and the wrong direction with one answer", func() {
f := newXferFixture()
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
defer end()
other, _ := f.call("telegram", contract.ConnectorFilePutOp, contract.ConnectorFilePut{Handle: "h", Data: []byte("x")})
unknown, _ := f.call("slack", contract.ConnectorFilePutOp, contract.ConnectorFilePut{Handle: "nope", Data: []byte("x")})
Expect(other).To(Equal(http.StatusNotFound))
Expect(unknown).To(Equal(other))
code, _ := f.get("h", 0, 10)
Expect(code).To(Equal(http.StatusNotFound), "a write handle cannot read")
})
It("does not let a handle be used for two calls", func() {
f := newXferFixture()
end, err := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
Expect(err).NotTo(HaveOccurred())
_, err = f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
Expect(err).To(HaveOccurred())
end()
end() // idempotent
Expect(f.outs).To(HaveLen(1))
_, err = f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
Expect(err).NotTo(HaveOccurred(), "a handle of an ended call is free; a new call makes its own random one")
})
It("expires a handle that sees no call and removes the partial file", func() {
f := newXferFixture()
f.t.Idle = 80 * time.Millisecond
_, err := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
Expect(err).NotTo(HaveOccurred())
code, _ := f.put("h", 0, []byte("partial"), false)
Expect(code).To(Equal(http.StatusOK))
Eventually(f.t.Live, "2s", "20ms").Should(BeZero())
code, _ = f.put("h", 7, []byte("x"), false)
Expect(code).To(Equal(http.StatusNotFound))
Expect(f.outs).To(HaveLen(1))
Expect(f.outs[0].Result).To(Equal(TransferExpired))
tmp, _ := filepath.Glob(filepath.Join(f.dir, "owner-cache", "tmp", "*"))
Expect(tmp).To(BeEmpty(), "no partial file is left")
})
It("refuses a file the scan hook rejects", func() {
f := newXferFixture()
f.t.Scan = func(_ context.Context, _ string, _ int64) error { return io.ErrClosedPipe }
end, _ := f.t.Begin(TransferGrant{Handle: "h", Direction: TransferIn, AddonID: "slack", MaxBytes: 100})
code, _ := f.put("h", 0, []byte("virus"), true)
Expect(code).To(Equal(http.StatusBadRequest))
o := end()
Expect(o.Result).To(Equal(TransferRefused))
Expect(o.SHA).To(BeEmpty())
})
It("refuses an outbound handle whose content changed size and an offset past the end", func() {
f := newXferFixture()
sha, size, err := f.blobs.Put(GinkgoT().Context(), bytes.NewReader([]byte("hello world")), 0)
Expect(err).NotTo(HaveOccurred())
bad, _ := f.t.Begin(TransferGrant{Handle: "bad", Direction: TransferOut, AddonID: "slack", SHA: sha, Size: size + 1})
code, _ := f.get("bad", 0, 5)
Expect(code).To(Equal(http.StatusServiceUnavailable))
bad()
end, _ := f.t.Begin(TransferGrant{Handle: "ok", Direction: TransferOut, AddonID: "slack", SHA: sha, Size: size})
defer end()
code, res := f.get("ok", 6, contract.TransferChunkBytes+1)
Expect(code).To(Equal(http.StatusBadRequest), "a length above the chunk is refused")
code, res = f.get("ok", 6, 5)
Expect(code).To(Equal(http.StatusOK))
Expect(string(res.Data)).To(Equal("world"))
Expect(res.EOF).To(BeTrue())
code, _ = f.get("ok", size+1, 1)
Expect(code).To(Equal(http.StatusBadRequest))
})
})
+1 -1
View File
@@ -266,7 +266,7 @@ func (c *Client) Connector(ctx context.Context, v contract.ConnectorRequest) (co
if c.registration.Capabilities.Connector == nil {
return r, ErrNotSent
}
e := c.call(ctx, "/connector", nil, v, &r, contract.ConnectorTimeout)
e := c.call(ctx, "/connector", nil, v, &r, v.CallTimeout())
return r, e
}
func (c *Client) Preview(ctx context.Context, v contract.Envelope) (contract.Preview, error) {
+10
View File
@@ -12,6 +12,16 @@ type ConnectorCallScope struct {
AgentID string `json:"agentId,omitempty"`
AddonID string `json:"addonId"`
Kind string `json:"kind"`
// Content names the file an outbound transfer reads: a blob of the
// workspace store. The frontend that checked the chat's access sets it;
// it never reaches the add-on, which gets only the handle.
Content *TransferContent `json:"content,omitempty"`
}
// TransferContent is the stored content behind an outbound handle.
type TransferContent struct {
SHA string `json:"sha"`
Size int64 `json:"size"`
}
const (