yetone
d9b2cdc593
ready-wait failed twice on 10-07 in full `go test -tags nogui ./...` runs on
main ("signal: killed" at 10.00s) and passed alone. The cause is the test's
size, not the host's cancellation. The case runs 30,000 fetch+abort rounds,
each waiting at least one setImmediate turn. That takes ~4s alone, 4.4-8.5s
beside a full go test ./... plus CPU and fresh-script load loops, and more
than 10s when starved (taskpolicy -c background with every core busy: killed
3 of 3). The cost per round stays flat from 5,000 to 120,000 rounds (no growth,
so no leak), and the counts at the end are exact: retainedWaiters 0, one
reaction on the shared ready promise, one drain listener.
Those counts are what catches a leak, and they catch it at any size. With
1,000 rounds, leaving the waiter in readyWaiters fails with retainedWaiters
1000, and a reaction per request on the shared ready promise fails with
retainedReadyReactions 1000.
Before: starved, 3 of 3 runs killed at 10.04-10.26s; -race -count=20 79.9s.
After: starved, 5 of 5 pass (ready-wait 0.6-5.6s, mostly node startup);
-race -count=20 9.9s.
No production code changed.
672 lines
24 KiB
Go
672 lines
24 KiB
Go
package plugin
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"reflect"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestHostJSCancellation(t *testing.T) {
|
|
node, err := exec.LookPath("node")
|
|
if err != nil {
|
|
t.Skip("no node on PATH")
|
|
}
|
|
for _, name := range []string{"stopped-gate", "split-chunk", "terminal", "blocked-head", "blocked-head-cancel-error", "ready-wait", "ready-success", "ready-error", "local-error", "local-error-cancel", "normal-reply"} {
|
|
t.Run(name, func(t *testing.T) {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
cmd := exec.CommandContext(ctx, node, "-e", hostJSCancellationTest, name)
|
|
cmd.Stdin = bytes.NewReader(hostJS)
|
|
out, err := cmd.CombinedOutput()
|
|
if err != nil {
|
|
t.Fatalf("JavaScript cancellation regression: %v\n%s", err, out)
|
|
}
|
|
if len(out) > 0 {
|
|
t.Logf("%s", out)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
const hostJSCancellationTest = `
|
|
const assert = require('node:assert/strict')
|
|
const fs = require('node:fs')
|
|
const vm = require('node:vm')
|
|
const { EventEmitter } = require('node:events')
|
|
const source = fs.readFileSync(0, 'utf8')
|
|
const section = (from, to) => {
|
|
const start = source.indexOf(from)
|
|
const end = to ? source.indexOf(to, start) : source.length
|
|
assert.ok(start >= 0 && end > start, 'production section exists')
|
|
return source.slice(start, end)
|
|
}
|
|
const stdout = new EventEmitter()
|
|
const lines = []
|
|
let blockOn = ''
|
|
stdout.write = line => {
|
|
const msg = JSON.parse(line)
|
|
lines.push(msg)
|
|
return blockOn !== 'all' && !(blockOn === 'chunk' && msg.event === 'chunk')
|
|
}
|
|
const rl = new EventEmitter()
|
|
const context = vm.createContext({
|
|
Buffer, Headers, AbortController, DOMException, ReadableStream,
|
|
process: { stdout, stdin: {}, exit: () => { throw Error('unexpected exit') } },
|
|
readline: { createInterface: () => rl },
|
|
options: async () => ({ fetch: async () => context.reply }),
|
|
info: async () => ({ models: {} }),
|
|
sdkHeaders: () => ({}), bodyOf: () => undefined,
|
|
accountKey: () => '',
|
|
via: { run: (_, fn) => fn() }, inScope: (_, key, fn) => fn(),
|
|
})
|
|
const run = code => vm.runInContext(code, context)
|
|
run(section('const rpcWrite =', '// A plugin writing to stdout'))
|
|
run('const hooks = []; const inflight = new Map(); const handlers = {}')
|
|
run(section('async function doFetch(', '// ---- the loop'))
|
|
run(section('// Everything waits for init;'))
|
|
const message = value => rl.emit('line', JSON.stringify(value))
|
|
const tick = () => new Promise(resolve => setImmediate(resolve))
|
|
const until = async (check, label) => {
|
|
for (let iteration = 0; iteration < 30; iteration++) {
|
|
if (check()) return
|
|
await tick()
|
|
}
|
|
assert.ok(check(), label)
|
|
}
|
|
const idle = () => run('pending === 0 && inflight.size === 0')
|
|
const queued = () => run('outq.length')
|
|
const fetchOne = () => message({ id: 1, method: 'fetch', params: { window: 512 << 10, provider: 'test', model: 'test', url: 'http://localhost/' } })
|
|
const abortOne = () => message({ method: 'abort', params: { id: 1 } })
|
|
const initialize = promise => {
|
|
context.initialization = promise
|
|
run('typeof setReady === "function" ? setReady(initialization) : ready = initialization')
|
|
}
|
|
async function main() {
|
|
const name = process.argv[1]
|
|
initialize(Promise.resolve())
|
|
await tick()
|
|
context.reply = { status: 200, headers: new Headers(), body: null }
|
|
if (name === 'stopped-gate') {
|
|
run('globalThis.gate = makeGate(512 << 10)')
|
|
assert.equal(await run('gate.take(64)'), 64)
|
|
run('gate.stop(); gate.add(1)')
|
|
assert.equal(await run('gate.take(64)'), -1, 'stopped gate refuses remaining credit')
|
|
} else if (name === 'split-chunk') {
|
|
blockOn = 'chunk'
|
|
context.reply.body = (async function* () { yield Buffer.alloc(3 * (64 << 10)) })()
|
|
fetchOne()
|
|
await until(() => queued() === 1 && run('outq[0].line.includes("chunk")'), 'second slice queued while credit remains')
|
|
abortOne()
|
|
await until(idle, 'cancelled split-chunk fetch cleans registries and pending count')
|
|
assert.equal(queued(), 0, 'no late chunk or terminal queued')
|
|
assert.equal(lines.filter(msg => msg.event === 'chunk').length, 1)
|
|
assert.equal(lines.filter(msg => msg.id === 1 && ('result' in msg || 'error' in msg)).length, 0)
|
|
} else if (name === 'terminal') {
|
|
blockOn = 'all'
|
|
fetchOne()
|
|
await until(() => queued() === 1, 'terminal queued behind head awaiting drain')
|
|
abortOne()
|
|
await until(idle, 'cancelled terminal wait cleans registries')
|
|
assert.equal(queued(), 0, 'cancelled terminal removed')
|
|
} else if (name === 'blocked-head' || name === 'blocked-head-cancel-error') {
|
|
blockOn = 'all'
|
|
run('send({ id: 99, result: null })')
|
|
let bodyCancelCalls = 0
|
|
context.reply.body = new ReadableStream({ cancel() {
|
|
bodyCancelCalls++
|
|
if (name === 'blocked-head-cancel-error') throw Error('body cancellation failed')
|
|
} })
|
|
fetchOne()
|
|
await until(() => queued() === 1 && run('outq[0].line.includes("head")'), 'head queued behind blocked stdout')
|
|
abortOne()
|
|
await until(idle, 'cancelled blocked head cleans registries')
|
|
console.log(JSON.stringify({ idle: idle(), queued: queued(), bodyCancelCalls, bodyLocked: context.reply.body.locked }))
|
|
assert.equal(bodyCancelCalls, 1, 'acquired body cancelled before iteration')
|
|
assert.equal(queued(), 0, 'cancelled head removed')
|
|
assert.equal(context.reply.body.locked, false)
|
|
assert.equal(lines.length, 1, 'body cancellation emits no duplicate error')
|
|
} else if (name === 'ready-wait') {
|
|
blockOn = 'all'
|
|
let readyReactions = 0
|
|
const unresolved = new Promise(() => {})
|
|
const originalThen = unresolved.then.bind(unresolved)
|
|
unresolved.then = (...args) => { readyReactions++; return originalThen(...args) }
|
|
initialize(unresolved)
|
|
run('send({id: 99, result: null})')
|
|
for (let iteration = 0; iteration < 1000; iteration++) {
|
|
fetchOne()
|
|
abortOne()
|
|
await until(idle, 'cancelled bootstrap wait cleans registries')
|
|
assert.equal(queued(), 0, 'cancelled bootstrap emits no late error')
|
|
}
|
|
assert.equal(lines.length, 1, 'stdout stayed blocked at original line')
|
|
assert.equal(stdout.listenerCount('drain'), 1, 'only shared drain listener remains')
|
|
const retainedWaiters = run('typeof readyWaiters === "undefined" ? -1 : readyWaiters.size')
|
|
const retainedReadyReactions = readyReactions - (run('typeof setReady') === 'function' ? 1 : 0)
|
|
console.log(JSON.stringify({ cancellations: 1000, idle: idle(), queued: queued(), readyReactions, retainedReadyReactions, retainedWaiters }))
|
|
assert.equal(retainedWaiters, 0, 'cancelled bootstrap leaves no retained waiters')
|
|
assert.equal(retainedReadyReactions, 0, 'cancelled requests leave no reactions on shared ready')
|
|
assert.equal(readyReactions, 1, 'only shared initialization reaction remains')
|
|
} else if (name === 'ready-success' || name === 'ready-error') {
|
|
let settle
|
|
initialize(new Promise((resolve, reject) => { settle = name === 'ready-success' ? resolve : reject }))
|
|
fetchOne()
|
|
assert.equal(run('readyWaiters.size'), 1, 'bootstrap waiter registered')
|
|
settle(name === 'ready-error' ? Error('init failed') : undefined)
|
|
await until(idle, 'settled initialization cleans fetch')
|
|
assert.equal(run('readyWaiters.size'), 0, 'settled initialization releases waiters')
|
|
if (name === 'ready-error') {
|
|
assert.equal(lines.length, 1)
|
|
assert.equal(lines[0].error.message, 'init failed')
|
|
} else {
|
|
assert.deepEqual(lines.map(msg => msg.event || 'result'), ['head', 'result'])
|
|
}
|
|
} else if (name === 'mid-setup') {
|
|
for (const stage of ['options', 'info']) {
|
|
let release
|
|
let entered = false
|
|
let vendorCalls = 0
|
|
const barrier = new Promise(resolve => { release = resolve })
|
|
context.options = async () => {
|
|
if (stage === 'options') { entered = true; await barrier }
|
|
return { fetch: async () => { vendorCalls++; return context.reply } }
|
|
}
|
|
context.info = async () => {
|
|
if (stage === 'info') { entered = true; await barrier }
|
|
return { models: {} }
|
|
}
|
|
fetchOne()
|
|
await until(() => entered, 'plugin setup entered its deferred await')
|
|
assert.equal(run('inflight.get(1).state'), 'setup')
|
|
abortOne()
|
|
assert.equal(run('inflight.get(1).state'), 'cancelled', 'JS processed abort before setup resumes')
|
|
assert.equal(run('inflight.get(1).controller.signal.aborted'), true)
|
|
release()
|
|
await until(idle, 'cancelled setup finishes')
|
|
assert.equal(vendorCalls, 0, 'cancelled setup never starts vendor')
|
|
context.options = async () => ({ fetch: async () => { vendorCalls++; return context.reply } })
|
|
context.info = async () => ({ models: {} })
|
|
fetchOne()
|
|
await until(idle, 'healthy fetch after cancelled setup')
|
|
assert.equal(vendorCalls, 1)
|
|
}
|
|
assert.deepEqual(lines.map(msg => msg.event || 'result'), ['head', 'result', 'head', 'result'])
|
|
} else if (name === 'local-error') {
|
|
context.options = async () => {
|
|
run('inflight.get(1).controller.abort()')
|
|
throw Error('local protocol failure')
|
|
}
|
|
fetchOne()
|
|
await until(idle, 'local error settles')
|
|
assert.equal(lines.length, 1)
|
|
assert.equal(lines[0].error.message, 'local protocol failure', 'local abort does not suppress error owed to Go')
|
|
} else if (name === 'local-error-cancel') {
|
|
blockOn = 'all'
|
|
run('send({ id: 99, result: null })')
|
|
context.options = async () => {
|
|
run('inflight.get(1).controller.abort()')
|
|
throw Error('local protocol failure')
|
|
}
|
|
fetchOne()
|
|
await until(() => queued() === 1, 'local error queued after local abort')
|
|
abortOne()
|
|
await until(idle, 'host cancel ends error wait even after local signal aborted')
|
|
assert.equal(queued(), 0, 'late error removed after host cancel')
|
|
} else if (name === 'normal-reply') {
|
|
context.reply.body = (async function* () { yield Buffer.from('complete reply') })()
|
|
fetchOne()
|
|
await until(idle, 'normal reply settles')
|
|
assert.deepEqual(lines.map(msg => msg.event || 'result'), ['head', 'chunk', 'result'])
|
|
assert.equal(Buffer.from(lines[1].data, 'base64').toString(), 'complete reply')
|
|
assert.equal(lines[2].result, null)
|
|
} else {
|
|
throw Error('unknown test case')
|
|
}
|
|
}
|
|
main().catch(err => { console.error(err); process.exitCode = 1 })
|
|
`
|
|
|
|
// fakeSignedIn starts a real Bun host with the fake plugin signed in to, its
|
|
// requests going to base, and gives the plugin's base URL.
|
|
func fakeSignedIn(t *testing.T, base string) string {
|
|
t.Helper()
|
|
sandbox(t)
|
|
t.Setenv("FAKE_BASE", base)
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
|
t.Cleanup(cancel)
|
|
abs, _ := filepath.Abs("testdata/fake/index.js")
|
|
if _, err := Add(ctx, abs); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := Providers(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := APIKey(ctx, "fakeco", 0, nil, "k1", NewAccount); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
o, err := LoaderOptions(ctx, "fakeco", "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return o.BaseURL
|
|
}
|
|
|
|
// TestFlowSlowReaderGetsTheWholeReply is the flow control's point: a reader
|
|
// slower than the vendor still gets every byte — credit pauses the upstream
|
|
// rather than the host cutting the reply at a fixed size.
|
|
func TestFlowSlowReaderGetsTheWholeReply(t *testing.T) {
|
|
const size = 4 << 20
|
|
payload := bytes.Repeat([]byte("0123456789abcdef"), size/16)
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(200)
|
|
for off := 0; off < len(payload); off += 64 << 10 {
|
|
end := min(off+(64<<10), len(payload))
|
|
_, _ = w.Write(payload[off:end])
|
|
w.(http.Flusher).Flush()
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
base := fakeSignedIn(t, srv.URL+"/v1")
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
|
defer cancel()
|
|
res, err := Fetch(ctx, FetchRequest{
|
|
Provider: "fakeco", Model: "fake-1", NPM: "@ai-sdk/openai-compatible",
|
|
URL: base + "/big", Method: "POST", Body: []byte(`{}`),
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer res.Body.Close()
|
|
// read slower than the vendor writes, so the credit window must pause it
|
|
// and let it go again rather than the reply being cut
|
|
buf := make([]byte, 64<<10)
|
|
var got []byte
|
|
for {
|
|
n, err := res.Body.Read(buf)
|
|
if n > 0 {
|
|
got = append(got, buf[:n]...)
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("read at %d bytes: %v", len(got), err)
|
|
}
|
|
}
|
|
if !bytes.Equal(got, payload) {
|
|
t.Fatalf("read %d bytes, want %d", len(got), len(payload))
|
|
}
|
|
}
|
|
|
|
// TestFlowCancelAfterHeadStopsTheVendor: a request cancelled after its head,
|
|
// with nothing reading or closing the body, still stops the upstream fetch —
|
|
// the vendor's own request is cancelled.
|
|
func TestFlowCancelAfterHeadStopsTheVendor(t *testing.T) {
|
|
hit := make(chan struct{})
|
|
gone := make(chan struct{})
|
|
var once sync.Once
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(200)
|
|
w.(http.Flusher).Flush()
|
|
once.Do(func() { close(hit) })
|
|
<-r.Context().Done()
|
|
close(gone)
|
|
}))
|
|
defer srv.Close()
|
|
base := fakeSignedIn(t, srv.URL+"/v1")
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
|
defer cancel()
|
|
rctx, rcancel := context.WithCancel(ctx)
|
|
res, err := Fetch(rctx, FetchRequest{
|
|
Provider: "fakeco", Model: "fake-1", NPM: "@ai-sdk/openai-compatible",
|
|
URL: base + "/hang", Method: "POST", Body: []byte(`{}`),
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
<-hit // the vendor is in the request, its head has reached magpie
|
|
rcancel() // nothing reads or closes the body
|
|
select {
|
|
case <-gone:
|
|
case <-time.After(10 * time.Second):
|
|
t.Fatal("cancelling the request did not stop the vendor")
|
|
}
|
|
_ = res
|
|
}
|
|
|
|
// TestFlowTinyFramesStillArrive: a vendor that yields tiny pieces, read slowly,
|
|
// still gets every byte. The frame overhead pauses such a producer — it can
|
|
// keep at most window/(payload+overhead) frames queued — rather than the host
|
|
// cutting it, and the reader gets the whole reply.
|
|
func TestFlowTinyFramesStillArrive(t *testing.T) {
|
|
const size = 4096
|
|
payload := bytes.Repeat([]byte("abcdefgh"), size/8)
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(200)
|
|
fl := w.(http.Flusher)
|
|
for i := range payload {
|
|
_, _ = w.Write(payload[i : i+1])
|
|
if i%16 == 0 {
|
|
fl.Flush()
|
|
}
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
base := fakeSignedIn(t, srv.URL+"/v1")
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
|
defer cancel()
|
|
res, err := Fetch(ctx, FetchRequest{
|
|
Provider: "fakeco", Model: "fake-1", NPM: "@ai-sdk/openai-compatible",
|
|
URL: base + "/tiny", Method: "POST", Body: []byte(`{}`),
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer res.Body.Close()
|
|
buf := make([]byte, 512)
|
|
var got []byte
|
|
for {
|
|
n, err := res.Body.Read(buf)
|
|
if n > 0 {
|
|
got = append(got, buf[:n]...)
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("read at %d bytes: %v", len(got), err)
|
|
}
|
|
}
|
|
if !bytes.Equal(got, payload) {
|
|
t.Fatalf("read %d bytes, want %d", len(got), len(payload))
|
|
}
|
|
}
|
|
|
|
// bigPiece is one piece of the production-scale stream: the 64 KiB block, or
|
|
// the same block as an SSE data event.
|
|
func bigPiece(block []byte, sse bool) []byte {
|
|
if !sse {
|
|
return block
|
|
}
|
|
return []byte("data: " + base64.StdEncoding.EncodeToString(block) + "\n\n")
|
|
}
|
|
|
|
// writeBig writes about total bytes of the pattern to w (nil to only measure),
|
|
// giving the exact length and its SHA-256. The vendor and the check build the
|
|
// same bytes, so a truncation anywhere shows up as a length or hash mismatch.
|
|
func writeBig(w io.Writer, total int, sse bool, block []byte) (int64, [32]byte) {
|
|
h := sha256.New()
|
|
var n int64
|
|
for n < int64(total) {
|
|
piece := bigPiece(block, sse)
|
|
if w != nil {
|
|
if _, err := w.Write(piece); err != nil {
|
|
break
|
|
}
|
|
}
|
|
h.Write(piece)
|
|
n += int64(len(piece))
|
|
}
|
|
var sum [32]byte
|
|
copy(sum[:], h.Sum(nil))
|
|
return n, sum
|
|
}
|
|
|
|
// TestFlowBigReplyProd is the production-scale check: 100 MiB of a non-SSE body
|
|
// and of an SSE stream, each read through the real host and hashed; an SSE
|
|
// stream a reader pauses 12 s in the middle of and then finishes. It is behind
|
|
// MAGPIE_FLOW_BIG because each case moves 100 MiB, and it should be run against
|
|
// the Bun magpie itself picks (MAGPIE_BUN set to it).
|
|
func TestFlowBigReplyProd(t *testing.T) {
|
|
if os.Getenv("MAGPIE_FLOW_BIG") == "" {
|
|
t.Skip("set MAGPIE_FLOW_BIG=1 to run the 100 MiB flow cases")
|
|
}
|
|
const total = 100 << 20
|
|
block := make([]byte, 64<<10)
|
|
for i := range block {
|
|
block[i] = byte(i * 7)
|
|
}
|
|
for _, tc := range []struct {
|
|
name string
|
|
sse bool
|
|
pause time.Duration
|
|
}{
|
|
{"non-sse", false, 0},
|
|
{"sse", true, 0},
|
|
{"sse-paused", true, 12 * time.Second},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
ctype := "application/json"
|
|
if tc.sse {
|
|
ctype = "text/event-stream"
|
|
}
|
|
// how far the vendor has got, and whether it finished, so a pause
|
|
// the reader takes shows whether the upstream was held back too
|
|
var written atomic.Int64
|
|
finished := make(chan struct{})
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
defer close(finished)
|
|
w.Header().Set("Content-Type", ctype)
|
|
w.WriteHeader(200)
|
|
fl := w.(http.Flusher)
|
|
var n int64
|
|
for n < int64(total) {
|
|
piece := bigPiece(block, tc.sse)
|
|
m, err := w.Write(piece)
|
|
written.Add(int64(m))
|
|
if err != nil {
|
|
return
|
|
}
|
|
n += int64(len(piece))
|
|
if n%(8<<20) < int64(len(piece)) {
|
|
fl.Flush()
|
|
}
|
|
}
|
|
}))
|
|
defer srv.Close()
|
|
base := fakeSignedIn(t, srv.URL+"/v1")
|
|
|
|
wantLen, wantSum := writeBig(nil, total, tc.sse, block)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
|
defer cancel()
|
|
res, err := Fetch(ctx, FetchRequest{
|
|
Provider: "fakeco", Model: "fake-1", NPM: "@ai-sdk/openai-compatible",
|
|
URL: base + "/big", Method: "POST", Body: []byte(`{}`),
|
|
})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer res.Body.Close()
|
|
h := sha256.New()
|
|
var n int64
|
|
paused := false
|
|
var pausedAt, resumedAt, pauseOwed, pauseWindow int64
|
|
var finishedBeforeResume, windowExhausted bool
|
|
buf := make([]byte, 256<<10)
|
|
for {
|
|
m, err := res.Body.Read(buf)
|
|
if m > 0 {
|
|
h.Write(buf[:m])
|
|
n += int64(m)
|
|
if tc.pause > 0 && !paused && n >= 1<<20 {
|
|
paused = true
|
|
pausedAt = written.Load()
|
|
t.Logf("pause_start: delivered=%d producer_written=%d expected=%d", n, pausedAt, wantLen)
|
|
// Reflection keeps this public-API fixture compilable with
|
|
// the old host.go overlay, which has no credit counters.
|
|
b := res.Body.(*body)
|
|
pauseEnd := time.Now().Add(tc.pause)
|
|
for time.Now().Before(pauseEnd) {
|
|
b.c.mu.Lock()
|
|
fields := reflect.ValueOf(b.c).Elem()
|
|
owed, window := fields.FieldByName("owed"), fields.FieldByName("window")
|
|
if owed.IsValid() && window.IsValid() {
|
|
pauseOwed, pauseWindow = owed.Int(), window.Int()
|
|
windowExhausted = pauseWindow > 0 && pauseWindow-pauseOwed <= 2048
|
|
}
|
|
b.c.mu.Unlock()
|
|
if windowExhausted {
|
|
time.Sleep(time.Until(pauseEnd))
|
|
break
|
|
}
|
|
time.Sleep(10 * time.Millisecond)
|
|
}
|
|
resumedAt = written.Load()
|
|
select {
|
|
case <-finished:
|
|
finishedBeforeResume = true
|
|
default:
|
|
}
|
|
t.Logf("pause_end_before_read: producer_written=%d expected=%d finished=%t window_exhausted=%t owed=%d window=%d", resumedAt, wantLen, finishedBeforeResume, windowExhausted, pauseOwed, pauseWindow)
|
|
}
|
|
}
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
t.Fatalf("read at %d bytes: %v", n, err)
|
|
}
|
|
}
|
|
got := h.Sum(nil)
|
|
if n != wantLen {
|
|
t.Fatalf("read %d bytes, want %d", n, wantLen)
|
|
}
|
|
if !bytes.Equal(got, wantSum[:]) {
|
|
t.Fatalf("sha256 = %x, want %x", got, wantSum)
|
|
}
|
|
if tc.pause > 0 {
|
|
// Preserve the original length/hash failure as the negative
|
|
// control; only these new producer assertions are optional.
|
|
if os.Getenv("MAGPIE_FLOW_ORIGINAL_BEHAVIOR_ONLY") != "1" {
|
|
if !windowExhausted {
|
|
t.Error("the paused reader did not exhaust the Go credit window")
|
|
}
|
|
if finishedBeforeResume || resumedAt >= wantLen {
|
|
t.Errorf("producer finished before resume: written=%d expected=%d finished=%t", resumedAt, wantLen, finishedBeforeResume)
|
|
}
|
|
}
|
|
select {
|
|
case <-finished:
|
|
case <-time.After(5 * time.Second):
|
|
t.Fatal("producer did not finish after resume")
|
|
}
|
|
t.Logf("producer_after_resume: written=%d expected=%d finished=true", written.Load(), wantLen)
|
|
}
|
|
t.Logf("%s: %d bytes, sha256 %x", tc.name, n, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestFlowAbortDuringSetupStaysHealthy: many fetches cancelled as soon as they
|
|
// are made — while the plugin is still setting them up, or waiting for a slow
|
|
// vendor's head — leave the host able to answer a later request. It is the Go
|
|
// side of the lifecycle: the host's slots come back, and the JS side's own
|
|
// finally (covered by construction) lets each request's controller and gate go.
|
|
func TestFlowAbortDuringSetupStaysHealthy(t *testing.T) {
|
|
node, err := exec.LookPath("node")
|
|
if err != nil {
|
|
t.Skip("no node on PATH")
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
cmd := exec.CommandContext(ctx, node, "-e", hostJSCancellationTest, "mid-setup")
|
|
cmd.Stdin = bytes.NewReader(hostJS)
|
|
if out, err := cmd.CombinedOutput(); err != nil {
|
|
t.Fatalf("mid-setup cancellation: %v\n%s", err, out)
|
|
}
|
|
}
|
|
|
|
func TestFlowEarlyCancellationStressStaysHealthy(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
select {
|
|
case <-time.After(300 * time.Millisecond):
|
|
case <-r.Context().Done():
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(200)
|
|
_, _ = w.Write([]byte(`{"ok":true}`))
|
|
}))
|
|
defer srv.Close()
|
|
base := fakeSignedIn(t, srv.URL+"/v1")
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
|
defer cancel()
|
|
r := FetchRequest{
|
|
Provider: "fakeco", Model: "fake-1", NPM: "@ai-sdk/openai-compatible",
|
|
URL: base + "/slow", Method: "POST", Body: []byte(`{}`),
|
|
}
|
|
for i := range 20 {
|
|
fctx, fcancel := context.WithCancel(ctx)
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
res, err := Fetch(fctx, r)
|
|
if err == nil {
|
|
_, _ = io.Copy(io.Discard, res.Body)
|
|
res.Body.Close()
|
|
}
|
|
}()
|
|
fcancel() // as early as possible: setup, or the head wait
|
|
select {
|
|
case <-done:
|
|
case <-time.After(10 * time.Second):
|
|
t.Fatalf("fetch %d did not return after cancel", i)
|
|
}
|
|
}
|
|
// the host still answers a request that is not cancelled
|
|
res, err := Fetch(ctx, r)
|
|
if err != nil {
|
|
t.Fatalf("a fetch after the cancelled ones: %v", err)
|
|
}
|
|
defer res.Body.Close()
|
|
if body, err := io.ReadAll(res.Body); err != nil || !bytes.Contains(body, []byte(`"ok":true`)) {
|
|
t.Fatalf("body = %q, %v", body, err)
|
|
}
|
|
}
|
|
|
|
// TestFlowHeadEnvelopeRejectsHugeHeaders: a vendor that answers with more
|
|
// headers than magpie's head envelope allows fails the fetch cleanly rather
|
|
// than the host taking an unbounded zero-charge head line.
|
|
func TestFlowHeadEnvelopeRejectsHugeHeaders(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("X-Big", strings.Repeat("v", 96<<10)) // past the 64 KiB envelope
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(200)
|
|
_, _ = w.Write([]byte(`{"ok":true}`))
|
|
}))
|
|
defer srv.Close()
|
|
base := fakeSignedIn(t, srv.URL+"/v1")
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
|
|
defer cancel()
|
|
_, err := Fetch(ctx, FetchRequest{
|
|
Provider: "fakeco", Model: "fake-1", NPM: "@ai-sdk/openai-compatible",
|
|
URL: base + "/bighead", Method: "POST", Body: []byte(`{}`),
|
|
})
|
|
if err == nil {
|
|
t.Fatal("a reply with an oversized head was accepted")
|
|
}
|
|
}
|