Merge branch 'fix/g06' into fix/integrated
# Conflicts: # cmd/mavend/memoryeval.go
This commit is contained in:
@@ -0,0 +1,48 @@
|
||||
package email
|
||||
|
||||
// windows-1251 (and its ASCII-compatible low half) is decoded here rather than
|
||||
// pulled in from x/text.
|
||||
//
|
||||
// The alternative was returning an error for the charset, which ParseMessage
|
||||
// turns into a subject-only message. That is a live gap and not a small one:
|
||||
// cp1251 is still what older Russian senders emit, and subject-only means those
|
||||
// mails can never produce a task candidate. The whole of x/text/encoding is a
|
||||
// large dependency for the most privacy-sensitive path in the tree, and
|
||||
// windows-1251 is a 128-entry table.
|
||||
//
|
||||
// Only cp1251 is added. Guessing at an unknown charset stays forbidden: mojibake
|
||||
// is worse than nothing, because the model extracts a task from it happily.
|
||||
|
||||
// cp1251High — the 0x80..0xFF half of windows-1251. The low half is ASCII.
|
||||
var cp1251High = [128]rune{
|
||||
0x0402, 0x0403, 0x201A, 0x0453, 0x201E, 0x2026, 0x2020, 0x2021,
|
||||
0x20AC, 0x2030, 0x0409, 0x2039, 0x040A, 0x040C, 0x040B, 0x040F,
|
||||
0x0452, 0x2018, 0x2019, 0x201C, 0x201D, 0x2022, 0x2013, 0x2014,
|
||||
0xFFFD, 0x2122, 0x0459, 0x203A, 0x045A, 0x045C, 0x045B, 0x045F,
|
||||
0x00A0, 0x040E, 0x045E, 0x0408, 0x00A4, 0x0490, 0x00A6, 0x00A7,
|
||||
0x0401, 0x00A9, 0x0404, 0x00AB, 0x00AC, 0x00AD, 0x00AE, 0x0407,
|
||||
0x00B0, 0x00B1, 0x0406, 0x0456, 0x0491, 0x00B5, 0x00B6, 0x00B7,
|
||||
0x0451, 0x2116, 0x0454, 0x00BB, 0x0458, 0x0405, 0x0455, 0x0457,
|
||||
0x0410, 0x0411, 0x0412, 0x0413, 0x0414, 0x0415, 0x0416, 0x0417,
|
||||
0x0418, 0x0419, 0x041A, 0x041B, 0x041C, 0x041D, 0x041E, 0x041F,
|
||||
0x0420, 0x0421, 0x0422, 0x0423, 0x0424, 0x0425, 0x0426, 0x0427,
|
||||
0x0428, 0x0429, 0x042A, 0x042B, 0x042C, 0x042D, 0x042E, 0x042F,
|
||||
0x0430, 0x0431, 0x0432, 0x0433, 0x0434, 0x0435, 0x0436, 0x0437,
|
||||
0x0438, 0x0439, 0x043A, 0x043B, 0x043C, 0x043D, 0x043E, 0x043F,
|
||||
0x0440, 0x0441, 0x0442, 0x0443, 0x0444, 0x0445, 0x0446, 0x0447,
|
||||
0x0448, 0x0449, 0x044A, 0x044B, 0x044C, 0x044D, 0x044E, 0x044F,
|
||||
}
|
||||
|
||||
// decodeCP1251 maps each byte through the table. Every byte has a defined
|
||||
// meaning in this charset, so decoding cannot fail.
|
||||
func decodeCP1251(b []byte) string {
|
||||
out := make([]rune, 0, len(b))
|
||||
for _, c := range b {
|
||||
if c < 0x80 {
|
||||
out = append(out, rune(c))
|
||||
continue
|
||||
}
|
||||
out = append(out, cp1251High[c-0x80])
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
package email
|
||||
|
||||
import "time"
|
||||
|
||||
// WithDial sets the connection seam for a test. It lives in a _test.go file so
|
||||
// the seam has no linker symbol in the shipped binary: no code outside this
|
||||
// package can hand FetchSince a dialer, and therefore no code outside this
|
||||
// package can point the mail reader at a cleartext transport and give it his
|
||||
// password. The compiler is what enforces that, which is the whole reason the
|
||||
// field is unexported.
|
||||
func (f FetchSince) WithDial(d func(addr string, timeout time.Duration) (*Conn, error)) FetchSince {
|
||||
f.dial = d
|
||||
return f
|
||||
}
|
||||
@@ -35,6 +35,14 @@ import (
|
||||
// search input" — mail is the same class), and Evidence keeps only the subject
|
||||
// line, so the review page shows him where a candidate came from without the
|
||||
// store growing a copy of his mailbox.
|
||||
//
|
||||
// One constraint for whoever adds task context to a prompt later: a candidate's
|
||||
// text is a model paraphrase of the content of his mail, and it lives in
|
||||
// tasks.text. "Maven never sends his mail anywhere" holds today because nothing
|
||||
// assembles a context block out of live tasks. The moment something does, mail
|
||||
// content reaches whatever that block is sent to, and an outbound search would
|
||||
// be sending his mailbox out a paraphrase at a time. Tasks sourced "email:" have
|
||||
// to be excluded there, not here.
|
||||
|
||||
// MaxCandidates — at most this many candidates per message, enforced by the
|
||||
// grammar. A mail with four tasks in it is a mail he has to read himself; a
|
||||
@@ -81,6 +89,12 @@ func NewExtractor(c Completer, max int, contextBlock func() string) *Extractor {
|
||||
return &Extractor{llm: c, max: max, contextBlock: contextBlock}
|
||||
}
|
||||
|
||||
// Max — the normalised candidate bound. Exported so the daemon logs what it will
|
||||
// actually allow rather than what the config file said: 0 in the config means
|
||||
// MaxCandidates here, and logging the raw value said "max 0" and then wrote
|
||||
// three.
|
||||
func (e *Extractor) Max() int { return e.max }
|
||||
|
||||
// extractGrammar — GBNF pinning the answer to a bounded array of fixed-shape
|
||||
// candidates. Same reasoning as memeval's evalGrammar and the router's
|
||||
// routeGrammar: the shape and the length bound are what keep a small model from
|
||||
@@ -190,9 +204,13 @@ func renderForModel(msg Message) string {
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// parseCandidates decodes the grammar-constrained reply, tolerating the
|
||||
// wrappers a Thinking model sometimes leaves around it (a fenced block, or
|
||||
// leading reasoning before the array).
|
||||
// parseCandidates decodes the reply and trims a fenced block or stray prose
|
||||
// around the array.
|
||||
//
|
||||
// Through Extract that tolerance is unreachable: extractGrammar pins the first
|
||||
// token to "[", so the model cannot emit reasoning before it. It is kept for
|
||||
// callers that pass a raw reply from an ungrammared path, and the note is here
|
||||
// so the next reader does not conclude that thinking output is expected.
|
||||
func parseCandidates(raw string) ([]Candidate, error) {
|
||||
s := strings.TrimSpace(raw)
|
||||
if i := strings.Index(s, "["); i > 0 {
|
||||
|
||||
+56
-15
@@ -1,6 +1,7 @@
|
||||
package email
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
@@ -26,26 +27,53 @@ type FetchSince struct {
|
||||
Since time.Time
|
||||
Max int
|
||||
Skip func(uid uint32) bool
|
||||
// OnSearch, when set, is handed the whole SEARCH result before anything is
|
||||
// fetched, ascending, seen UIDs included. It is how the poller learns which
|
||||
// UIDs are still inside the lookback window: anything below the lowest one
|
||||
// can never be searched for again, and therefore can never be read again.
|
||||
// Without that the poller cannot tell a UID it has not got to yet from one
|
||||
// that has aged out of the window.
|
||||
OnSearch func(uids []uint32)
|
||||
|
||||
// dial — the connection seam, unexported on purpose: see dialer(). Tests
|
||||
// inside this package set it through export_test.go; nothing outside can.
|
||||
dial func(addr string, timeout time.Duration) (*Conn, error)
|
||||
}
|
||||
|
||||
// dial is the connection seam. nil means Dial (implicit TLS); the tests set it
|
||||
// through export_test.go. It is unexported and there is no exported wrapper
|
||||
// that takes a dialer, so no code outside this package can point the reader at
|
||||
// a non-TLS transport and hand it the password. That is a property the compiler
|
||||
// enforces, not a claim about the callers that happen to exist today.
|
||||
func (f FetchSince) dialer() func(addr string, timeout time.Duration) (*Conn, error) {
|
||||
if f.dial != nil {
|
||||
return f.dial
|
||||
}
|
||||
return Dial
|
||||
}
|
||||
|
||||
// Run performs one read. password is passed here, not stored in the struct, so
|
||||
// the configuration of a mailbox and the secret for it are never the same value
|
||||
// sitting in the same place.
|
||||
//
|
||||
// One message that cannot be read does not abandon the poll: the rest of the
|
||||
// mailbox is still worth reading, and returning early meant one oversized or
|
||||
// unreadable message permanently hid every older message behind it, poll after
|
||||
// poll. The returned error joins whatever failed, and the messages that did
|
||||
// come back come back with it.
|
||||
func (f FetchSince) Run(password string) ([]Message, error) {
|
||||
return f.RunWith(password, nil)
|
||||
}
|
||||
|
||||
// RunWith is Run with an explicit connection function, which is how the reader
|
||||
// daemon and the tests substitute an in-process server. nil ⇒ Dial, i.e.
|
||||
// implicit TLS with certificate verification; there is no configuration path
|
||||
// that reaches this, so no deployment can end up talking cleartext IMAP.
|
||||
func (f FetchSince) RunWith(password string, dial func(addr string, timeout time.Duration) (*Conn, error)) ([]Message, error) {
|
||||
if f.Addr == "" || f.User == "" || f.Mailbox == "" {
|
||||
return nil, fmt.Errorf("email: mailbox not configured (addr/user/mailbox)")
|
||||
}
|
||||
if dial == nil {
|
||||
dial = Dial
|
||||
// Timeout is validated like the other three fields. Zero disables every
|
||||
// deadline in the path — net.Dialer{Timeout: 0} and a Conn that never calls
|
||||
// SetDeadline — so a dead server parks the poller forever on a socket read,
|
||||
// with his credential live in a TLS state. That is the exact thing the
|
||||
// connect-read-drop shape exists to avoid.
|
||||
if f.Timeout <= 0 {
|
||||
return nil, fmt.Errorf("email: timeout must be positive")
|
||||
}
|
||||
dial := f.dialer()
|
||||
c, err := dial(f.Addr, f.Timeout)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -63,6 +91,10 @@ func (f FetchSince) RunWith(password string, dial func(addr string, timeout time
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if f.OnSearch != nil {
|
||||
f.OnSearch(uids)
|
||||
}
|
||||
|
||||
// Newest UIDs first — IMAP hands them back ascending, and when Max clips the
|
||||
// list the recent mail is what matters.
|
||||
wanted := make([]uint32, 0, len(uids))
|
||||
@@ -77,13 +109,22 @@ func (f FetchSince) RunWith(password string, dial func(addr string, timeout time
|
||||
}
|
||||
|
||||
out := make([]Message, 0, len(wanted))
|
||||
var failed []error
|
||||
for _, uid := range wanted {
|
||||
raw, err := c.Fetch(uid)
|
||||
if errors.Is(err, ErrMessageTooLarge) {
|
||||
// Too big to read is a permanent verdict, not a failure to retry:
|
||||
// the message will be the same size next poll. Carried as bulk so
|
||||
// the poller marks it seen and stops fetching it, exactly like a
|
||||
// newsletter. Nothing is sent to the model.
|
||||
out = append(out, Message{UID: uid, Junk: true, JunkReason: "oversize"})
|
||||
continue
|
||||
}
|
||||
if err != nil {
|
||||
// One unreadable message does not abandon the poll; the rest of the
|
||||
// mailbox is still worth reading. The error names the UID, not the
|
||||
// message.
|
||||
return out, fmt.Errorf("email: fetch uid %d: %w", uid, err)
|
||||
// The error names the UID, never the message. Collected rather than
|
||||
// returned, so the messages behind this one are still read.
|
||||
failed = append(failed, fmt.Errorf("email: fetch uid %d: %w", uid, err))
|
||||
continue
|
||||
}
|
||||
if len(raw) == 0 {
|
||||
continue // vanished between SEARCH and FETCH
|
||||
@@ -94,5 +135,5 @@ func (f FetchSince) RunWith(password string, dial func(addr string, timeout time
|
||||
}
|
||||
out = append(out, msg)
|
||||
}
|
||||
return out, nil
|
||||
return out, errors.Join(failed...)
|
||||
}
|
||||
|
||||
@@ -22,11 +22,7 @@ func TestFetchSinceRun(t *testing.T) {
|
||||
Max: 2,
|
||||
Skip: func(uid uint32) bool { return uid == 3 },
|
||||
}
|
||||
msgs, err := fs.RunWith("secret", func(addr string, timeout time.Duration) (*Conn, error) {
|
||||
cli, srv := net.Pipe()
|
||||
go f.serve(t, srv)
|
||||
return NewConn(cli, timeout)
|
||||
})
|
||||
msgs, err := fs.WithDial(dialer(t, f)).Run("secret")
|
||||
if err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
@@ -47,3 +43,75 @@ func TestFetchSinceRequiresConfig(t *testing.T) {
|
||||
t.Fatal("an unconfigured mailbox must not be read")
|
||||
}
|
||||
}
|
||||
|
||||
// Timeout zero disables the dial timeout AND every socket deadline, so a dead
|
||||
// server parks the poller forever with his credential live in a TLS state.
|
||||
func TestFetchSinceRejectsZeroTimeout(t *testing.T) {
|
||||
fs := FetchSince{Addr: "mail.example:993", User: "kami", Mailbox: "INBOX"}
|
||||
if _, err := fs.Run("secret"); err == nil {
|
||||
t.Fatal("a zero timeout must be rejected like an empty address")
|
||||
}
|
||||
}
|
||||
|
||||
// dialer wires a client Conn to an in-process fake over net.Pipe.
|
||||
func dialer(t *testing.T, f *fakeIMAP) func(string, time.Duration) (*Conn, error) {
|
||||
t.Helper()
|
||||
return func(addr string, timeout time.Duration) (*Conn, error) {
|
||||
cli, srv := net.Pipe()
|
||||
go f.serve(t, srv)
|
||||
return NewConn(cli, timeout)
|
||||
}
|
||||
}
|
||||
|
||||
// One message that cannot be read must not hide the older ones behind it. The
|
||||
// old code returned on the first failure, so an oversized or unreadable UID
|
||||
// blocked every message below it on every poll, forever.
|
||||
func TestFetchSinceContinuesPastABadMessage(t *testing.T) {
|
||||
mk := func(subject string) string {
|
||||
return "Subject: " + subject + "\r\nContent-Type: text/plain; charset=utf-8\r\n\r\nbody\r\n"
|
||||
}
|
||||
f := &fakeIMAP{
|
||||
uids: []uint32{1, 2, 3},
|
||||
msgs: map[uint32]string{1: mk("one"), 3: mk("three")},
|
||||
quoted: map[uint32]bool{2: true},
|
||||
}
|
||||
fs := FetchSince{
|
||||
Addr: "mail.example:993", User: "kami", Mailbox: "INBOX",
|
||||
Timeout: 5 * time.Second,
|
||||
Since: time.Date(2026, 7, 30, 0, 0, 0, 0, time.UTC),
|
||||
}
|
||||
msgs, err := fs.WithDial(dialer(t, f)).Run("secret")
|
||||
if err == nil {
|
||||
t.Fatal("the unreadable UID must still be reported")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "uid 2") {
|
||||
t.Errorf("error should name the UID: %v", err)
|
||||
}
|
||||
if len(msgs) != 2 {
|
||||
t.Fatalf("got %d messages, want the two readable ones: %+v", len(msgs), msgs)
|
||||
}
|
||||
if msgs[0].Subject != "three" || msgs[1].Subject != "one" {
|
||||
t.Errorf("subjects = %q,%q, want three,one", msgs[0].Subject, msgs[1].Subject)
|
||||
}
|
||||
}
|
||||
|
||||
// An oversized message is retired as bulk rather than retried: it will be the
|
||||
// same size next poll, and the poller marks bulk seen without a model call.
|
||||
func TestFetchSinceRetiresOversizedMessage(t *testing.T) {
|
||||
f := &fakeIMAP{uids: []uint32{7}, oversize: map[uint32]int{7: MaxMessageBytes + 1}}
|
||||
fs := FetchSince{
|
||||
Addr: "mail.example:993", User: "kami", Mailbox: "INBOX",
|
||||
Timeout: 5 * time.Second,
|
||||
Since: time.Date(2026, 7, 30, 0, 0, 0, 0, time.UTC),
|
||||
}
|
||||
msgs, err := fs.WithDial(dialer(t, f)).Run("secret")
|
||||
if err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
if len(msgs) != 1 || !msgs[0].Junk || msgs[0].JunkReason != "oversize" {
|
||||
t.Fatalf("want one oversize-bulk message, got %+v", msgs)
|
||||
}
|
||||
if msgs[0].Body != "" || msgs[0].Subject != "" {
|
||||
t.Error("nothing from an oversized message may be kept")
|
||||
}
|
||||
}
|
||||
|
||||
+117
-8
@@ -3,6 +3,7 @@ package email
|
||||
import (
|
||||
"bufio"
|
||||
"crypto/tls"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
@@ -32,6 +33,31 @@ import (
|
||||
// option to get it wrong once.
|
||||
const DefaultIMAPPort = "993"
|
||||
|
||||
// MaxMessageBytes — the largest message this client will read into memory.
|
||||
//
|
||||
// The literal size comes off the wire, so an unbounded read is an allocation
|
||||
// the server picks: "{2147483647}" is a 2GB make() before a single byte
|
||||
// arrives, and one ordinary mail with a 60MB attachment is a 60MB peak RSS on a
|
||||
// box already holding a 1.7B model resident. All of it would then be thrown
|
||||
// away, because plaintextBody skips attachments and the body is truncated to
|
||||
// MaxBodyBytes anyway.
|
||||
//
|
||||
// 2 MiB is well above what prose plus quoted history plus base64 HTML needs and
|
||||
// well below what hurts. A larger message is drained and reported as
|
||||
// ErrMessageTooLarge rather than read.
|
||||
const MaxMessageBytes = 2 << 20
|
||||
|
||||
// readChunk — how much of a literal is read between deadline refreshes. The
|
||||
// per-connection timeout must stay an IDLE timeout: with one deadline around
|
||||
// the whole read it becomes a whole-message budget, and a healthy but slow
|
||||
// uplink then fails the same message on every poll forever.
|
||||
const readChunk = 64 << 10
|
||||
|
||||
// ErrMessageTooLarge — the server announced a literal above MaxMessageBytes.
|
||||
// The connection stays usable (the bytes are drained), and the caller decides
|
||||
// what to do with the UID. FetchSince retires it rather than retrying it.
|
||||
var ErrMessageTooLarge = errors.New("email: message larger than the read cap")
|
||||
|
||||
// Conn — one authenticated IMAP connection. Not safe for concurrent use; the
|
||||
// poller drives one connection at a time.
|
||||
type Conn struct {
|
||||
@@ -78,6 +104,18 @@ func (c *Conn) Close() error { return c.rwc.Close() }
|
||||
// never stored on the Conn: nothing in this package keeps a credential alive
|
||||
// past the command that uses it, so no struct dump or panic trace can carry it.
|
||||
func (c *Conn) Login(user, pass string) error {
|
||||
// A credential with a line break in it is rejected here, not silently
|
||||
// repaired. quote() strips CR and LF so a stray newline can never become a
|
||||
// second command, but stripping alone means a password file that picked up
|
||||
// a newline authenticates as a DIFFERENT string and comes back as the
|
||||
// server's generic NO, which is a long debugging session. This error names
|
||||
// the problem and cannot leak the value.
|
||||
if strings.ContainsAny(user, "\r\n") {
|
||||
return fmt.Errorf("email: login: username contains a line break")
|
||||
}
|
||||
if strings.ContainsAny(pass, "\r\n") {
|
||||
return fmt.Errorf("email: login: password contains a line break")
|
||||
}
|
||||
// The command line itself is never logged (see exec) — a LOGIN line IS the
|
||||
// credential.
|
||||
if _, err := c.exec(fmt.Sprintf("LOGIN %s %s", quote(user), quote(pass))); err != nil {
|
||||
@@ -130,34 +168,60 @@ var literalSize = regexp.MustCompile(`\{(\d+)\}$`)
|
||||
// Fetch returns the raw RFC 5322 bytes of one message, by UID.
|
||||
//
|
||||
// Returns (nil, nil) when the UID no longer exists — a message he deleted
|
||||
// between SEARCH and FETCH is normal, not an error.
|
||||
// between SEARCH and FETCH is normal, not an error. A FETCH response that came
|
||||
// back with no literal in it is NOT that case and is an error, so a message
|
||||
// that exists and was readable is never dropped without a log line.
|
||||
//
|
||||
// A literal above MaxMessageBytes is drained without being kept and reported as
|
||||
// ErrMessageTooLarge.
|
||||
func (c *Conn) Fetch(uid uint32) ([]byte, error) {
|
||||
tag := c.nextTag()
|
||||
if err := c.send(fmt.Sprintf("%s UID FETCH %d (BODY.PEEK[])", tag, uid)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var raw []byte
|
||||
var sawFetch, tooLarge bool
|
||||
for {
|
||||
line, err := c.readLine()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("email: fetch %d: %w", uid, err)
|
||||
}
|
||||
if done, err := c.tagged(tag, line); done {
|
||||
if err != nil {
|
||||
switch {
|
||||
case err != nil:
|
||||
return nil, fmt.Errorf("email: fetch %d: %w", uid, err)
|
||||
case tooLarge:
|
||||
return nil, fmt.Errorf("email: fetch %d: %w", uid, ErrMessageTooLarge)
|
||||
case sawFetch && raw == nil:
|
||||
// The server answered for this UID but not with a literal (a
|
||||
// quoted string, say). Silently skipping it would look exactly
|
||||
// like a vanished message.
|
||||
return nil, fmt.Errorf("email: fetch %d: no message literal in the FETCH response", uid)
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
if strings.HasPrefix(line, "* ") && strings.Contains(line, " FETCH ") {
|
||||
sawFetch = true
|
||||
}
|
||||
m := literalSize.FindStringSubmatch(strings.TrimSpace(line))
|
||||
if m == nil {
|
||||
continue
|
||||
}
|
||||
n, err := strconv.Atoi(m[1])
|
||||
if err != nil {
|
||||
if err != nil || n < 0 {
|
||||
continue
|
||||
}
|
||||
buf := make([]byte, n)
|
||||
if _, err := io.ReadFull(c.r, buf); err != nil {
|
||||
if n > MaxMessageBytes {
|
||||
// Drained rather than read: the stream has to stay aligned for the
|
||||
// tagged completion, but nothing is allocated and nothing is parsed.
|
||||
tooLarge = true
|
||||
if err := c.discard(int64(n)); err != nil {
|
||||
return nil, fmt.Errorf("email: fetch %d: drain literal: %w", uid, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
buf, err := c.readN(n)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("email: fetch %d: literal: %w", uid, err)
|
||||
}
|
||||
if raw == nil {
|
||||
@@ -203,8 +267,8 @@ func (c *Conn) exec(cmd string) ([]string, error) {
|
||||
// send asks for one outside Fetch, but skip it if it appears so the
|
||||
// stream stays aligned.
|
||||
if m := literalSize.FindStringSubmatch(strings.TrimSpace(line)); m != nil {
|
||||
if n, err := strconv.Atoi(m[1]); err == nil {
|
||||
if _, err := io.CopyN(io.Discard, c.r, int64(n)); err != nil {
|
||||
if n, err := strconv.Atoi(m[1]); err == nil && n > 0 {
|
||||
if err := c.discard(int64(n)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
@@ -238,6 +302,43 @@ func (c *Conn) send(line string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// readN reads exactly n bytes, refreshing the deadline every readChunk so the
|
||||
// timeout stays an idle timeout rather than a budget for the whole literal.
|
||||
func (c *Conn) readN(n int) ([]byte, error) {
|
||||
buf := make([]byte, n)
|
||||
for off := 0; off < n; {
|
||||
end := off + readChunk
|
||||
if end > n {
|
||||
end = n
|
||||
}
|
||||
c.setDeadline()
|
||||
got, err := io.ReadFull(c.r, buf[off:end])
|
||||
off += got
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
// discard throws away n bytes of literal, same chunked deadline refresh as
|
||||
// readN and no allocation proportional to n.
|
||||
func (c *Conn) discard(n int64) error {
|
||||
for n > 0 {
|
||||
chunk := int64(readChunk)
|
||||
if chunk > n {
|
||||
chunk = n
|
||||
}
|
||||
c.setDeadline()
|
||||
got, err := io.CopyN(io.Discard, c.r, chunk)
|
||||
n -= got
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Conn) readLine() (string, error) {
|
||||
c.setDeadline()
|
||||
line, err := c.r.ReadString('\n')
|
||||
@@ -259,6 +360,10 @@ func (c *Conn) setDeadline() {
|
||||
}
|
||||
|
||||
// untagged splits "* SEARCH 1 2 3" into its payload when the key matches.
|
||||
//
|
||||
// The key must be the whole word: a prefix test would also match a future
|
||||
// extension's "* SEARCHRES", and reading its payload as UIDs is the kind of
|
||||
// thing that ages badly next to an IMAP capability nobody asked for.
|
||||
func untagged(line, key string) (string, bool) {
|
||||
if !strings.HasPrefix(line, "* ") {
|
||||
return "", false
|
||||
@@ -267,7 +372,11 @@ func untagged(line, key string) (string, bool) {
|
||||
if !strings.HasPrefix(rest, key) {
|
||||
return "", false
|
||||
}
|
||||
return strings.TrimSpace(rest[len(key):]), true
|
||||
rest = rest[len(key):]
|
||||
if rest != "" && rest[0] != ' ' && rest[0] != '\t' {
|
||||
return "", false
|
||||
}
|
||||
return strings.TrimSpace(rest), true
|
||||
}
|
||||
|
||||
// quote renders an IMAP quoted string. Passwords routinely contain characters
|
||||
|
||||
@@ -2,6 +2,7 @@ package email
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"strconv"
|
||||
@@ -18,6 +19,11 @@ type fakeIMAP struct {
|
||||
uids []uint32
|
||||
cmds []string
|
||||
failOn string // substring of a command to answer NO
|
||||
// oversize — UIDs answered with a literal of this many bytes, which the
|
||||
// server then actually sends. Used to exercise the read cap.
|
||||
oversize map[uint32]int
|
||||
// quoted — UIDs answered with a quoted string instead of a literal.
|
||||
quoted map[uint32]bool
|
||||
}
|
||||
|
||||
func (f *fakeIMAP) serve(t *testing.T, c net.Conn) {
|
||||
@@ -54,7 +60,20 @@ func (f *fakeIMAP) serve(t *testing.T, c net.Conn) {
|
||||
fmt.Fprintf(c, "%s OK search done\r\n", tag)
|
||||
case strings.HasPrefix(upper, "UID FETCH"):
|
||||
uid64, _ := strconv.ParseUint(strings.Fields(cmd)[2], 10, 32)
|
||||
raw, ok := f.msgs[uint32(uid64)]
|
||||
uid := uint32(uid64)
|
||||
if n, big := f.oversize[uid]; big {
|
||||
fmt.Fprintf(c, "* 1 FETCH (UID %d BODY[] {%d}\r\n", uid64, n)
|
||||
fmt.Fprint(c, strings.Repeat("x", n))
|
||||
fmt.Fprint(c, ")\r\n")
|
||||
fmt.Fprintf(c, "%s OK fetch done\r\n", tag)
|
||||
continue
|
||||
}
|
||||
if f.quoted[uid] {
|
||||
fmt.Fprintf(c, "* 1 FETCH (UID %d BODY[] \"short\")\r\n", uid64)
|
||||
fmt.Fprintf(c, "%s OK fetch done\r\n", tag)
|
||||
continue
|
||||
}
|
||||
raw, ok := f.msgs[uid]
|
||||
if ok {
|
||||
fmt.Fprintf(c, "* 1 FETCH (UID %d BODY[] {%d}\r\n", uid64, len(raw))
|
||||
fmt.Fprint(c, raw)
|
||||
@@ -160,3 +179,59 @@ func TestQuoteStripsNewlines(t *testing.T) {
|
||||
t.Errorf("quote kept a line break: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A credential with a line break in it is a broken password file, not a
|
||||
// password. Stripping it silently authenticates as a different string and the
|
||||
// server answers its generic NO.
|
||||
func TestLoginRejectsCredentialWithNewline(t *testing.T) {
|
||||
f := &fakeIMAP{}
|
||||
c := dialFake(t, f)
|
||||
err := c.Login("kami", "s3cr3t\nA1 LOGOUT")
|
||||
if err == nil {
|
||||
t.Fatal("a password with a line break must be rejected")
|
||||
}
|
||||
if strings.Contains(err.Error(), "s3cr3t") {
|
||||
t.Errorf("error leaks the password: %v", err)
|
||||
}
|
||||
if len(f.cmds) != 0 {
|
||||
t.Errorf("nothing should have been sent, got %v", f.cmds)
|
||||
}
|
||||
}
|
||||
|
||||
// The literal size comes off the wire. Without a cap the server picks the
|
||||
// allocation, and one 60MB attachment is 60MB of peak RSS on a box holding a
|
||||
// 1.7B model, all of it thrown away by plaintextBody afterwards.
|
||||
func TestFetchRefusesOversizedLiteral(t *testing.T) {
|
||||
f := &fakeIMAP{uids: []uint32{1}, oversize: map[uint32]int{1: MaxMessageBytes + 1}}
|
||||
c := dialFake(t, f)
|
||||
raw, err := c.Fetch(1)
|
||||
if !errors.Is(err, ErrMessageTooLarge) {
|
||||
t.Fatalf("fetch err = %v, want ErrMessageTooLarge", err)
|
||||
}
|
||||
if raw != nil {
|
||||
t.Errorf("an oversized message must not be kept, got %d bytes", len(raw))
|
||||
}
|
||||
// The stream stayed aligned: the connection is still usable.
|
||||
if err := c.Select("INBOX"); err != nil {
|
||||
t.Errorf("connection unusable after draining: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// A FETCH that answered without a literal is not a vanished message, and must
|
||||
// not be skipped as silently as one.
|
||||
func TestFetchNonLiteralResponseIsAnError(t *testing.T) {
|
||||
f := &fakeIMAP{uids: []uint32{1}, quoted: map[uint32]bool{1: true}}
|
||||
c := dialFake(t, f)
|
||||
if _, err := c.Fetch(1); err == nil {
|
||||
t.Fatal("a FETCH response with no literal must be reported, not dropped")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUntaggedMatchesWholeKeyOnly(t *testing.T) {
|
||||
if _, ok := untagged("* SEARCHRES 1 2 3", "SEARCH"); ok {
|
||||
t.Error("SEARCH must not match SEARCHRES")
|
||||
}
|
||||
if rest, ok := untagged("* SEARCH 1 2 3", "SEARCH"); !ok || rest != "1 2 3" {
|
||||
t.Errorf("untagged = (%q, %v), want (\"1 2 3\", true)", rest, ok)
|
||||
}
|
||||
}
|
||||
|
||||
+12
-21
@@ -21,8 +21,18 @@ import (
|
||||
// - Auto-Submitted other than "no" (RFC 3834) — generated by a machine.
|
||||
// - X-Spam-Flag: YES, X-Spam-Status: Yes — the spam filter upstream already
|
||||
// decided; we do not second-guess it in the other direction.
|
||||
// - X-GM-LABELS / X-Gmail-Labels containing a Gmail category — Gmail's own
|
||||
// Promotions/Social/Forums/Spam classification, when the server sends it.
|
||||
//
|
||||
// There is deliberately NO Gmail-category rule. One was written and removed:
|
||||
// it matched X-GM-LABELS and X-Gmail-Labels against the parsed header block,
|
||||
// and neither is a header. X-GM-LABELS is a Gmail FETCH data item, requested as
|
||||
// "UID FETCH n (X-GM-LABELS)" and never present in the message source;
|
||||
// X-Gmail-Labels only exists in a Takeout mbox export. This client asks for
|
||||
// BODY.PEEK[] and nothing else, so the rule could not fire against a real
|
||||
// mailbox while its doc comment promised a Promotions filter. Gmail's promotion
|
||||
// mail carries List-Unsubscribe in practice and is caught by the rule above.
|
||||
// Bringing the category rule back means adding the FETCH item and carrying the
|
||||
// labels into classifyJunk out of band, not matching a header that never
|
||||
// arrives.
|
||||
//
|
||||
// Deliberately NOT here: sender allow/deny lists and subject keyword matching.
|
||||
// Both are configuration that ages badly and both would be a place for his
|
||||
@@ -35,16 +45,6 @@ import (
|
||||
// junkHeaders — headers whose mere presence marks bulk mail.
|
||||
var junkPresence = []string{"List-Unsubscribe", "List-Id", "List-Post"}
|
||||
|
||||
// gmailCategories — Gmail's category labels, lowercased as they appear in
|
||||
// X-GM-LABELS. "important" and "inbox" are labels too, and are NOT categories.
|
||||
// Matching is by these exact tokens (substring is fine — they are namespaced
|
||||
// and cannot appear in a hand-made label by accident), so a user label named
|
||||
// "Social Club" is not mistaken for Gmail's Social category.
|
||||
var gmailCategories = []string{
|
||||
"category_promotions", "category_social", "category_forums", "category_updates",
|
||||
`\spam`, `\junk`,
|
||||
}
|
||||
|
||||
// classifyJunk returns whether the message is bulk/automated and why. The
|
||||
// reason is a short header name, safe to log — it names the marker, never the
|
||||
// sender or the subject.
|
||||
@@ -67,14 +67,5 @@ func classifyJunk(h mail.Header) (bool, string) {
|
||||
if v := strings.ToLower(strings.TrimSpace(h.Get("X-Spam-Status"))); strings.HasPrefix(v, "yes") {
|
||||
return true, "x-spam-status"
|
||||
}
|
||||
labels := strings.ToLower(h.Get("X-GM-LABELS") + " " + h.Get("X-Gmail-Labels"))
|
||||
for _, c := range gmailCategories {
|
||||
if c == "" {
|
||||
continue
|
||||
}
|
||||
if strings.Contains(labels, c) {
|
||||
return true, "gmail-category"
|
||||
}
|
||||
}
|
||||
return false, ""
|
||||
}
|
||||
|
||||
@@ -31,8 +31,10 @@ func TestClassifyJunk(t *testing.T) {
|
||||
{"spam flag", "From: a@b.c\nX-Spam-Flag: YES", true, "x-spam-flag"},
|
||||
{"spam status", "From: a@b.c\nX-Spam-Status: Yes, score=9.1", true, "x-spam-status"},
|
||||
{"spam status no", "From: a@b.c\nX-Spam-Status: No, score=0.1", false, ""},
|
||||
{"gmail promo", "From: a@b.c\nX-Gmail-Labels: Inbox,CATEGORY_PROMOTIONS", true, "gmail-category"},
|
||||
{"user label", "From: a@b.c\nX-Gmail-Labels: Social Club,Important", false, ""},
|
||||
// X-GM-LABELS is a Gmail FETCH data item, not a header, so it never
|
||||
// reaches classifyJunk through this client. The rule that matched it was
|
||||
// removed rather than left claiming a Promotions filter that never ran.
|
||||
{"gmail label header is not a rule", "From: a@b.c\nX-Gmail-Labels: Inbox,CATEGORY_PROMOTIONS", false, ""},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
|
||||
+66
-27
@@ -94,7 +94,14 @@ func plaintextBody(contentType, encoding string, body io.Reader) (string, error)
|
||||
if boundary == "" {
|
||||
return "", fmt.Errorf("email: multipart without boundary")
|
||||
}
|
||||
return multipartText(multipart.NewReader(body, boundary))
|
||||
plain, html, err := multipartText(multipart.NewReader(body, boundary))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.TrimSpace(plain) != "" {
|
||||
return plain, nil
|
||||
}
|
||||
return html, nil
|
||||
case mediaType == "text/html":
|
||||
raw, err := decodeBody(body, encoding, params["charset"])
|
||||
if err != nil {
|
||||
@@ -109,10 +116,15 @@ func plaintextBody(contentType, encoding string, body io.Reader) (string, error)
|
||||
}
|
||||
}
|
||||
|
||||
// multipartText reads one multipart level, recursing into nested multiparts.
|
||||
// Returns the plain part if any part yielded one, else the stripped HTML.
|
||||
func multipartText(mr *multipart.Reader) (string, error) {
|
||||
var plain, html string
|
||||
// multipartText reads one multipart level, recursing into nested multiparts,
|
||||
// and returns the two buckets separately: real text/plain, and text stripped
|
||||
// out of HTML.
|
||||
//
|
||||
// The buckets stay separate all the way up because a nested multipart can
|
||||
// contribute either kind. Folding a nested level's answer into one string put
|
||||
// HTML-derived text in the plain bucket, and a real text/plain sibling later in
|
||||
// the message was then thrown away by the "plain is already set" guard.
|
||||
func multipartText(mr *multipart.Reader) (plain, html string, err error) {
|
||||
for {
|
||||
part, err := mr.NextPart()
|
||||
if err == io.EOF {
|
||||
@@ -127,35 +139,47 @@ func multipartText(mr *multipart.Reader) (string, error) {
|
||||
continue // attachment
|
||||
}
|
||||
ct := part.Header.Get("Content-Type")
|
||||
mediaType, _, _ := mime.ParseMediaType(ct)
|
||||
text, err := plaintextBody(ct, part.Header.Get("Content-Transfer-Encoding"), part)
|
||||
part.Close()
|
||||
if err != nil || strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
if mediaType == "text/html" && !strings.HasPrefix(mediaType, "multipart/") {
|
||||
if html == "" {
|
||||
html = text
|
||||
mediaType, params, _ := mime.ParseMediaType(ct)
|
||||
switch {
|
||||
case strings.HasPrefix(mediaType, "multipart/"):
|
||||
var np, nh string
|
||||
if b := params["boundary"]; b != "" {
|
||||
np, nh, _ = multipartText(multipart.NewReader(part, b))
|
||||
}
|
||||
part.Close()
|
||||
if plain == "" {
|
||||
plain = np
|
||||
}
|
||||
if html == "" {
|
||||
html = nh
|
||||
}
|
||||
default:
|
||||
text, terr := plaintextBody(ct, part.Header.Get("Content-Transfer-Encoding"), part)
|
||||
part.Close()
|
||||
if terr != nil || strings.TrimSpace(text) == "" {
|
||||
continue
|
||||
}
|
||||
if mediaType == "text/html" {
|
||||
if html == "" {
|
||||
html = text
|
||||
}
|
||||
continue
|
||||
}
|
||||
if plain == "" {
|
||||
plain = text
|
||||
}
|
||||
continue
|
||||
}
|
||||
if plain == "" {
|
||||
plain = text
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(plain) != "" {
|
||||
return plain, nil
|
||||
}
|
||||
return html, nil
|
||||
return plain, html, nil
|
||||
}
|
||||
|
||||
// decodeBody applies the transfer encoding, then the charset.
|
||||
//
|
||||
// Charset support is UTF-8 (and ASCII, its subset) only, on purpose: x/text's
|
||||
// encoding tables are not vendored here, and guessing at windows-1251 bytes
|
||||
// would feed the model mojibake it would happily extract a task from. An
|
||||
// unsupported charset returns an error, which ParseMessage turns into an empty
|
||||
// body — subject-only, which is honest.
|
||||
// Charset support is UTF-8 (and ASCII, its subset) plus windows-1251, which is
|
||||
// decoded from a table in charset.go — see the reasoning there. Everything else
|
||||
// returns an error, which ParseMessage turns into an empty body: subject-only,
|
||||
// which is honest. Guessing at unknown bytes would feed the model mojibake it
|
||||
// would happily extract a task from.
|
||||
func decodeBody(r io.Reader, encoding, charset string) (string, error) {
|
||||
switch strings.ToLower(strings.TrimSpace(encoding)) {
|
||||
case "quoted-printable":
|
||||
@@ -170,6 +194,8 @@ func decodeBody(r io.Reader, encoding, charset string) (string, error) {
|
||||
switch cs := strings.ToLower(strings.TrimSpace(charset)); cs {
|
||||
case "", "utf-8", "utf8", "us-ascii", "ascii":
|
||||
return string(b), nil
|
||||
case "windows-1251", "cp1251", "windows1251", "x-cp1251":
|
||||
return decodeCP1251(b), nil
|
||||
default:
|
||||
return "", fmt.Errorf("email: unsupported charset %q", cs)
|
||||
}
|
||||
@@ -181,6 +207,19 @@ func decodeBody(r io.Reader, encoding, charset string) (string, error) {
|
||||
// shown to him as evidence.
|
||||
func decodeHeader(v string) string {
|
||||
dec := new(mime.WordDecoder)
|
||||
// Same charset support as the body: an old Russian sender encodes the
|
||||
// subject in windows-1251 too, and a subject is often the whole task.
|
||||
dec.CharsetReader = func(charset string, r io.Reader) (io.Reader, error) {
|
||||
switch strings.ToLower(strings.TrimSpace(charset)) {
|
||||
case "windows-1251", "cp1251", "windows1251", "x-cp1251":
|
||||
b, err := io.ReadAll(io.LimitReader(r, 1<<16))
|
||||
if err != nil && len(b) == 0 {
|
||||
return nil, err
|
||||
}
|
||||
return strings.NewReader(decodeCP1251(b)), nil
|
||||
}
|
||||
return nil, fmt.Errorf("email: unsupported charset %q", charset)
|
||||
}
|
||||
out, err := dec.DecodeHeader(v)
|
||||
if err != nil {
|
||||
return collapse(v)
|
||||
|
||||
@@ -70,13 +70,28 @@ func TestParsePrefersPlainAndSkipsAttachments(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// An unsupported charset must degrade to headers-only rather than to mojibake
|
||||
// the model would then extract a task from.
|
||||
func TestParseUnsupportedCharsetKeepsHeaders(t *testing.T) {
|
||||
// windows-1251 is what older Russian senders still emit. Subject-only for those
|
||||
// mails meant they could never produce a task candidate.
|
||||
func TestParseCP1251(t *testing.T) {
|
||||
msg, err := ParseMessage(3, fixture(t, "cp1251.eml"))
|
||||
if err != nil {
|
||||
t.Fatalf("parse: %v", err)
|
||||
}
|
||||
if want := "Счёт за интернет"; msg.Subject != want {
|
||||
t.Errorf("subject = %q, want %q", msg.Subject, want)
|
||||
}
|
||||
if want := "Оплати счёт до пятницы."; !strings.Contains(msg.Body, want) {
|
||||
t.Errorf("body = %q, want it to contain %q", msg.Body, want)
|
||||
}
|
||||
}
|
||||
|
||||
// A charset with no table here must degrade to headers-only rather than to
|
||||
// mojibake the model would then extract a task from.
|
||||
func TestParseUnsupportedCharsetKeepsHeaders(t *testing.T) {
|
||||
msg, err := ParseMessage(3, fixture(t, "koi8r.eml"))
|
||||
if err != nil {
|
||||
t.Fatalf("parse: %v", err)
|
||||
}
|
||||
if msg.Subject != "Legacy" {
|
||||
t.Errorf("subject = %q, want Legacy", msg.Subject)
|
||||
}
|
||||
@@ -85,6 +100,31 @@ func TestParseUnsupportedCharsetKeepsHeaders(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// A nested multipart/alternative that only had HTML must not fill the plain
|
||||
// bucket: a real text/plain sibling later in the message is the better text and
|
||||
// used to be discarded.
|
||||
func TestParseNestedHTMLDoesNotShadowLaterPlain(t *testing.T) {
|
||||
raw := "Subject: nested\r\n" +
|
||||
"Content-Type: multipart/mixed; boundary=OUT\r\n\r\n" +
|
||||
"--OUT\r\n" +
|
||||
"Content-Type: multipart/alternative; boundary=IN\r\n\r\n" +
|
||||
"--IN\r\n" +
|
||||
"Content-Type: text/html; charset=utf-8\r\n\r\n" +
|
||||
"<p>from the html part</p>\r\n" +
|
||||
"--IN--\r\n" +
|
||||
"--OUT\r\n" +
|
||||
"Content-Type: text/plain; charset=utf-8\r\n\r\n" +
|
||||
"the real plain text\r\n" +
|
||||
"--OUT--\r\n"
|
||||
msg, err := ParseMessage(5, []byte(raw))
|
||||
if err != nil {
|
||||
t.Fatalf("parse: %v", err)
|
||||
}
|
||||
if got := strings.TrimSpace(msg.Body); got != "the real plain text" {
|
||||
t.Errorf("body = %q, want the text/plain part to win", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseTruncatesLongBody(t *testing.T) {
|
||||
var b strings.Builder
|
||||
b.WriteString("Subject: long\r\nContent-Type: text/plain; charset=utf-8\r\n\r\n")
|
||||
|
||||
Vendored
+2
-2
@@ -1,7 +1,7 @@
|
||||
From: legacy@example.org
|
||||
To: kami@example.org
|
||||
Subject: Legacy
|
||||
Subject: =?windows-1251?B?0fe48iDn4CDo7fLl8O3l8g==?=
|
||||
Date: Fri, 01 Aug 2026 05:00:00 +0400
|
||||
Content-Type: text/plain; charset="windows-1251"
|
||||
|
||||
Ï
|
||||
Îïëàòè ñ÷¸ò äî ïÿòíèöû.
|
||||
|
||||
Vendored
+7
@@ -0,0 +1,7 @@
|
||||
From: legacy@example.org
|
||||
To: kami@example.org
|
||||
Subject: Legacy
|
||||
Date: Fri, 01 Aug 2026 05:00:00 +0400
|
||||
Content-Type: text/plain; charset="koi8-r"
|
||||
|
||||
ïÐÌÁÔÉ ÓÞ£Ô.
|
||||
@@ -185,6 +185,12 @@ type IngestMailReq struct {
|
||||
// mailbox dedupes to Created=0). Skipped is set when nothing was asked of the
|
||||
// model at all — junk, or an empty message.
|
||||
//
|
||||
// Created == 0 && !Skipped therefore means the model WAS consulted and found no
|
||||
// task, which is the common answer. A reader deciding whether to mark a UID
|
||||
// seen should treat that the same as a success: asking again would spend the
|
||||
// resident model on the same negative answer. Skipped means the same for a
|
||||
// different reason. Only an error means "not read yet".
|
||||
//
|
||||
// Nothing here echoes the mail back. The reader logs counts.
|
||||
type IngestMailResp struct {
|
||||
TaskIDs []int64 `json:"task_ids,omitempty"`
|
||||
|
||||
@@ -23,6 +23,32 @@ type Client struct {
|
||||
mu sync.RWMutex
|
||||
base string
|
||||
http *http.Client
|
||||
|
||||
// gate / background — priority on the single llama-server slot. Set once
|
||||
// at wiring time (SetGate), read on every request. nil gate ⇒ no gating,
|
||||
// which is what every test and every non-daemon caller gets.
|
||||
gate *Gate
|
||||
background bool
|
||||
}
|
||||
|
||||
// SetGate gives this client a priority on the shared llama-server slot. Call it
|
||||
// immediately after New, before the client is handed to anything: the fields are
|
||||
// read under the same lock as base, but the intent is one-time wiring, not a
|
||||
// knob to turn at runtime.
|
||||
//
|
||||
// background = false means "he is waiting for this" and never blocks.
|
||||
// background = true means the request yields to voice turns and runs one at a
|
||||
// time. See Gate.
|
||||
func (c *Client) SetGate(g *Gate, background bool) {
|
||||
c.mu.Lock()
|
||||
c.gate, c.background = g, background
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
func (c *Client) gateFor() (*Gate, bool) {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
return c.gate, c.background
|
||||
}
|
||||
|
||||
func New(baseURL string, timeout time.Duration) *Client {
|
||||
@@ -81,6 +107,17 @@ type resp struct {
|
||||
}
|
||||
|
||||
func (c *Client) Complete(ctx context.Context, r Req) (string, error) {
|
||||
if g, background := c.gateFor(); g != nil {
|
||||
if background {
|
||||
release, err := g.AcquireBackground(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer release()
|
||||
} else {
|
||||
defer g.Foreground()()
|
||||
}
|
||||
}
|
||||
b, _ := json.Marshal(body{
|
||||
Messages: []msg{{Role: "system", Content: r.System}, {Role: "user", Content: r.User}},
|
||||
MaxTokens: r.MaxTokens,
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Gate — priority access to the one llama-server slot.
|
||||
//
|
||||
// llama-server is started without -np, so it serves one request at a time and
|
||||
// everything else queues. That is fine while every caller is a voice turn, and
|
||||
// it stops being fine the moment a background job joins: mail extraction reads
|
||||
// up to 4000 characters on a Thinking 1.7B with a two minute budget, and a turn
|
||||
// that arrives during one waits for however much of that budget is left. The
|
||||
// router degrades to the classifier cascade on error, so he would get the 36.8%
|
||||
// floor while his mail is being read, and the phraser has no floor at all and
|
||||
// simply waits.
|
||||
//
|
||||
// So background work asks the gate first:
|
||||
//
|
||||
// - at most ONE background request is in flight, whatever the queue depth
|
||||
// upstream. A first poll of a mailbox with 40 unseen messages cannot
|
||||
// serialise 40 extractions ahead of anything.
|
||||
// - a background request waits while any foreground request is in flight, and
|
||||
// for Quiet after the last one finished. The quiet window is what stops an
|
||||
// extraction starting in the gap between the router call and the phraser
|
||||
// call of the same turn.
|
||||
//
|
||||
// Foreground requests never wait. This is not a fair queue and must not become
|
||||
// one: the point is that the thing he is waiting for wins every time.
|
||||
//
|
||||
// It bounds only what goes through an *llm.Client built with SetGate. The
|
||||
// phraser's own HTTP path is not gated, and a turn that reaches the phraser
|
||||
// without touching the router is not marked. Every real turn routes first, so
|
||||
// the marking is good enough to keep extraction out of the way; it is a
|
||||
// courtesy gate, not a scheduler.
|
||||
type Gate struct {
|
||||
mu sync.Mutex
|
||||
// fg — foreground requests in flight.
|
||||
fg int
|
||||
// last — when a foreground request last started or finished.
|
||||
last time.Time
|
||||
// bg — one token, so only one background request runs at a time.
|
||||
bg chan struct{}
|
||||
|
||||
quiet time.Duration
|
||||
poll time.Duration
|
||||
now func() time.Time
|
||||
}
|
||||
|
||||
// NewGate returns a gate that holds background work back for quiet after the
|
||||
// last foreground request. quiet <= 0 means "wait only while one is in flight".
|
||||
func NewGate(quiet time.Duration) *Gate {
|
||||
return &Gate{
|
||||
bg: make(chan struct{}, 1),
|
||||
quiet: quiet,
|
||||
poll: 50 * time.Millisecond,
|
||||
now: time.Now,
|
||||
}
|
||||
}
|
||||
|
||||
// Foreground marks a request as the thing he is waiting for. It never blocks.
|
||||
// The returned function must be called when the request finishes.
|
||||
func (g *Gate) Foreground() func() {
|
||||
if g == nil {
|
||||
return func() {}
|
||||
}
|
||||
g.mu.Lock()
|
||||
g.fg++
|
||||
g.last = g.now()
|
||||
g.mu.Unlock()
|
||||
return func() {
|
||||
g.mu.Lock()
|
||||
g.fg--
|
||||
g.last = g.now()
|
||||
g.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// AcquireBackground blocks until the slot is free enough for background work,
|
||||
// or ctx is done. The returned release function must be called when the request
|
||||
// finishes; it is nil on error.
|
||||
func (g *Gate) AcquireBackground(ctx context.Context) (func(), error) {
|
||||
if g == nil {
|
||||
return func() {}, nil
|
||||
}
|
||||
select {
|
||||
case g.bg <- struct{}{}:
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
release := func() { <-g.bg }
|
||||
for {
|
||||
if g.clear() {
|
||||
return release, nil
|
||||
}
|
||||
t := time.NewTimer(g.poll)
|
||||
select {
|
||||
case <-t.C:
|
||||
case <-ctx.Done():
|
||||
t.Stop()
|
||||
release()
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// clear reports whether no foreground request is in flight and the quiet window
|
||||
// since the last one has passed.
|
||||
func (g *Gate) clear() bool {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
if g.fg > 0 {
|
||||
return false
|
||||
}
|
||||
if g.quiet <= 0 || g.last.IsZero() {
|
||||
return true
|
||||
}
|
||||
return g.now().Sub(g.last) >= g.quiet
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Background work must not start while he is waiting on a turn. llama-server
|
||||
// serves one request at a time, so an extraction that starts first holds the
|
||||
// slot for its whole budget.
|
||||
func TestGateBackgroundWaitsForForeground(t *testing.T) {
|
||||
g := NewGate(0)
|
||||
g.poll = time.Millisecond
|
||||
done := g.Foreground()
|
||||
|
||||
started := make(chan struct{})
|
||||
go func() {
|
||||
release, err := g.AcquireBackground(context.Background())
|
||||
if err != nil {
|
||||
t.Errorf("acquire: %v", err)
|
||||
return
|
||||
}
|
||||
close(started)
|
||||
release()
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-started:
|
||||
t.Fatal("background work started while a foreground request was in flight")
|
||||
case <-time.After(20 * time.Millisecond):
|
||||
}
|
||||
done()
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("background work never started after the foreground request finished")
|
||||
}
|
||||
}
|
||||
|
||||
// Only one background request at a time, whatever the queue depth upstream. A
|
||||
// first poll of a mailbox with 40 unseen messages must not put 40 extractions
|
||||
// on the slot.
|
||||
func TestGateOneBackgroundAtATime(t *testing.T) {
|
||||
g := NewGate(0)
|
||||
g.poll = time.Millisecond
|
||||
first, err := g.AcquireBackground(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("first: %v", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if _, err := g.AcquireBackground(ctx); err == nil {
|
||||
t.Fatal("a second background request ran alongside the first")
|
||||
}
|
||||
first()
|
||||
second, err := g.AcquireBackground(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("second after release: %v", err)
|
||||
}
|
||||
second()
|
||||
}
|
||||
|
||||
// The quiet window covers the gap between the router call and the phraser call
|
||||
// of one turn, so an extraction cannot slip in mid-turn.
|
||||
func TestGateQuietWindow(t *testing.T) {
|
||||
now := time.Now()
|
||||
g := NewGate(time.Minute)
|
||||
g.poll = time.Millisecond
|
||||
g.now = func() time.Time { return now }
|
||||
g.Foreground()()
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if _, err := g.AcquireBackground(ctx); err == nil {
|
||||
t.Fatal("background work started inside the quiet window")
|
||||
}
|
||||
now = now.Add(2 * time.Minute)
|
||||
release, err := g.AcquireBackground(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("acquire after the quiet window: %v", err)
|
||||
}
|
||||
release()
|
||||
}
|
||||
|
||||
// Foreground never waits, whatever else is in flight.
|
||||
func TestGateForegroundNeverBlocks(t *testing.T) {
|
||||
g := NewGate(time.Minute)
|
||||
release, err := g.AcquireBackground(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("acquire: %v", err)
|
||||
}
|
||||
defer release()
|
||||
done := make(chan struct{})
|
||||
go func() { g.Foreground()(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("a foreground request waited behind background work")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user