feat(agent): run jobs from the master and resend results until acked (#7)

2026-10-12 18:45:44 +02:00
parent e313365686
commit 601c6da0d4
2 changed files with 517 additions and 47 deletions
@@ -1,65 +1,97 @@
 // Package conn keeps an agent connected to the master: it dials over pinned
-// TLS, announces itself, sends its inventory, and streams changes until the
-// connection drops, then reconnects with backoff.
+// TLS, announces itself, sends its inventory, runs the jobs the master sends
+// and streams changes and job results until the connection drops, then
+// reconnects with backoff.
 package conn

 import (
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"net/http"
+	"sync"
 	"time"

 	"github.com/coder/websocket"
 	"github.com/coder/websocket/wsjson"

 	"github.com/shcizo/package-updater/internal/agent/inventory"
+	"github.com/shcizo/package-updater/internal/agent/updater"
 	"github.com/shcizo/package-updater/internal/pin"
 	"github.com/shcizo/package-updater/internal/protocol"
 )

-// readLimit caps a single incoming message. The library default (32 KiB)
-// is too small for the master's future job messages; matches the master.
+// readLimit caps a single incoming message. Matches the master.
 const readLimit = 1 << 20

 // dialTimeout bounds the TLS + WebSocket handshake so a master that accepts
 // TCP but never answers does not hang the agent forever.
 const dialTimeout = 30 * time.Second

+// writeTimeout bounds one websocket write so a stalled peer ends the session
+// instead of blocking the writer forever.
+const writeTimeout = 10 * time.Second
+
+// defaultProgressInterval is how often job output is batched to the master.
+const defaultProgressInterval = 500 * time.Millisecond
+
 // Source supplies what the agent reports. *inventory.Collector satisfies it.
 type Source interface {
 	Hello(ctx context.Context) (protocol.Hello, error)
 	Snapshot(ctx context.Context) ([]protocol.Service, error)
 }

+// Jobs is the agent's job queue as the connection sees it. *updater.Runner
+// satisfies it.
+type Jobs interface {
+	// Submit queues a job; a job ID already held is ignored.
+	Submit(id string, start protocol.JobStart)
+	// Ack forgets a finished job the master has stored.
+	Ack(id string)
+	// Known lists every job ID held: queued, running, or finished unacked.
+	Known() []string
+	// TakeProgress returns output produced since the previous call.
+	TakeProgress() []updater.Progress
+	// Pending returns results the master has not acknowledged yet.
+	Pending() []updater.Finished
+	// Changed receives when there is new progress or a new result.
+	Changed() <-chan struct{}
+}
+
 // Config holds connection settings. Token is sent as a bearer header and must
 // never be logged.
 type Config struct {
-	URL         string // wss://host:port/agent
-	Token       string // AGENT_TOKEN
-	Fingerprint string // MASTER_KEY
-	Version     string
-	MinBackoff  time.Duration // default 1s
-	MaxBackoff  time.Duration // default 30s
+	URL              string // wss://host:port/agent
+	Token            string // AGENT_TOKEN
+	Fingerprint      string // MASTER_KEY
+	Version          string
+	MinBackoff       time.Duration // default 1s
+	MaxBackoff       time.Duration // default 30s
+	ProgressInterval time.Duration // default 500ms
 }

 // Agent is the agent's side of the master connection.
 type Agent struct {
 	cfg     Config
 	src     Source
+	jobs    Jobs
 	log     *slog.Logger
 	changes chan struct{}
 }

 // New returns an Agent. Call Run to connect.
-func New(cfg Config, src Source, log *slog.Logger) *Agent {
+func New(cfg Config, src Source, jobs Jobs, log *slog.Logger) *Agent {
 	if cfg.MinBackoff == 0 {
 		cfg.MinBackoff = time.Second
 	}
 	if cfg.MaxBackoff == 0 {
 		cfg.MaxBackoff = 30 * time.Second
 	}
-	return &Agent{cfg: cfg, src: src, log: log, changes: make(chan struct{}, 1)}
+	if cfg.ProgressInterval == 0 {
+		cfg.ProgressInterval = defaultProgressInterval
+	}
+	return &Agent{cfg: cfg, src: src, jobs: jobs, log: log, changes: make(chan struct{}, 1)}
 }

 // Notify tells the agent its inventory may have changed. It never blocks;
@@ -94,6 +126,9 @@ func (a *Agent) Run(ctx context.Context) error {

 // session runs one connection until it fails. connected reports whether the
 // hello was delivered, so Run resets its backoff only after real progress.
+//
+// Two goroutines serve the connection: a reader, which handles what the
+// master sends, and a writer, the only goroutine that writes to the socket.
 func (a *Agent) session(parent context.Context) (connected bool, err error) {
 	httpClient := &http.Client{Transport: &http.Transport{TLSClientConfig: pin.ClientConfig(a.cfg.Fingerprint)}}
 	// The dial context only bounds the handshake; coder/websocket does not
@@ -107,59 +142,140 @@ func (a *Agent) session(parent context.Context) (connected bool, err error) {
 	if err != nil {
 		return false, fmt.Errorf("dial master: %w", err)
 	}
-	defer c.CloseNow()
 	c.SetReadLimit(readLimit)

-	ctx, cancel := context.WithCancel(parent)
-	defer cancel()
-	// The master sends nothing yet in this phase, but reading is how a
-	// dropped connection is noticed, so a reader runs for the whole session.
-	go func() {
-		defer cancel()
-		for {
-			if _, _, err := c.Read(ctx); err != nil {
-				return
-			}
+	ctx, cancel := context.WithCancelCause(parent)
+	out := make(chan protocol.Envelope, 64)
+	var wg sync.WaitGroup
+	defer func() {
+		cancel(nil)
+		if parent.Err() != nil {
+			_ = c.Close(websocket.StatusNormalClosure, "")
 		}
+		c.CloseNow()
+		wg.Wait()
 	}()

+	wg.Add(2)
+	go func() { defer wg.Done(); a.writeLoop(ctx, cancel, c, out) }()
+	go func() { defer wg.Done(); a.readLoop(ctx, cancel, c) }()
+
+	// send queues an envelope for the writer. It gives up when the session ends.
+	send := func(typ string, payload any) error {
+		env, err := protocol.Encode(typ, payload)
+		if err != nil {
+			return err
+		}
+		return enqueue(ctx, out, env)
+	}
+	sendJob := func(typ, id string, payload any) error {
+		env, err := protocol.EncodeJob(typ, id, payload)
+		if err != nil {
+			return err
+		}
+		return enqueue(ctx, out, env)
+	}
+	fail := func(connected bool, err error) (bool, error) {
+		if cause := context.Cause(ctx); cause != nil && !errors.Is(cause, context.Canceled) {
+			return connected, cause
+		}
+		return connected, err
+	}
+
 	hello, err := a.src.Hello(ctx)
 	if err != nil {
-		return false, err
+		return fail(false, err)
 	}
 	hello.AgentVersion = a.cfg.Version
-	if err := send(ctx, c, protocol.TypeHello, hello); err != nil {
-		return false, err
+	hello.ProtocolVersion = protocol.Version
+	hello.KnownJobs = a.jobs.Known()
+	if err := send(protocol.TypeHello, hello); err != nil {
+		return fail(false, err)
 	}

 	last, err := a.src.Snapshot(ctx)
 	if err != nil {
-		return true, err
+		return fail(true, err)
 	}
-	if err := send(ctx, c, protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil {
-		return true, err
+	if err := send(protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil {
+		return fail(true, err)
 	}
-	a.log.Info("master_connected", "services", len(last))
+	a.log.Info("master_connected", "services", len(last), "known_jobs", len(hello.KnownJobs))

+	sent := make(map[string]bool) // job results already sent on this connection
+	flush := func() error {
+		for _, p := range a.jobs.TakeProgress() {
+			if err := sendJob(protocol.TypeJobProgress, p.ID, protocol.JobProgress{Output: p.Output}); err != nil {
+				return err
+			}
+		}
+		for _, f := range a.jobs.Pending() {
+			if sent[f.ID] {
+				continue
+			}
+			if err := sendJob(protocol.TypeJobResult, f.ID, f.Result); err != nil {
+				return err
+			}
+			sent[f.ID] = true
+		}
+		return nil
+	}
+
+	// flush is throttled to one call per ProgressInterval; a signal that
+	// arrives early arms the timer for the remainder instead of being lost.
+	lastFlush := time.Time{}
+	var timer *time.Timer
+	var timerC <-chan time.Time
+	defer func() {
+		if timer != nil {
+			timer.Stop()
+		}
+	}()
+	maybeFlush := func() error {
+		wait := a.cfg.ProgressInterval - time.Since(lastFlush)
+		if wait > 0 {
+			if timerC == nil {
+				timer = time.NewTimer(wait)
+				timerC = timer.C
+			}
+			return nil
+		}
+		lastFlush = time.Now()
+		return flush()
+	}
+	if err := maybeFlush(); err != nil {
+		return fail(true, err)
+	}
+
+	jobsChanged := a.jobs.Changed()
 	for {
 		select {
 		case <-ctx.Done():
-			_ = c.Close(websocket.StatusNormalClosure, "")
-			return true, ctx.Err()
+			return fail(true, ctx.Err())
+		case <-jobsChanged:
+			if err := maybeFlush(); err != nil {
+				return fail(true, err)
+			}
+		case <-timerC:
+			timerC = nil
+			lastFlush = time.Now()
+			if err := flush(); err != nil {
+				return fail(true, err)
+			}
 		case <-a.changes:
 			next, err := a.src.Snapshot(ctx)
 			if err != nil {
-				return true, err
+				return fail(true, err)
 			}
 			upserts, removes := inventory.Diff(last, next)
 			for _, s := range upserts {
-				if err := send(ctx, c, protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil {
-					return true, err
+				if err := send(protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil {
+					return fail(true, err)
 				}
 			}
 			for _, k := range removes {
-				if err := send(ctx, c, protocol.TypeRemove, protocol.Remove{Key: k}); err != nil {
-					return true, err
+				if err := send(protocol.TypeRemove, protocol.Remove{Key: k}); err != nil {
+					return fail(true, err)
 				}
 			}
 			last = next
@@ -167,13 +283,58 @@ func (a *Agent) session(parent context.Context) (connected bool, err error) {
 	}
 }

-func send(ctx context.Context, c *websocket.Conn, typ string, payload any) error {
-	env, err := protocol.Encode(typ, payload)
-	if err != nil {
-		return err
+// enqueue hands env to the writer, or gives up when the session ends.
+func enqueue(ctx context.Context, out chan<- protocol.Envelope, env protocol.Envelope) error {
+	select {
+	case out <- env:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// writeLoop is the only goroutine that writes to c. A failed or timed-out
+// write ends the session.
+func (a *Agent) writeLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn, out <-chan protocol.Envelope) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case env := <-out:
+			wctx, wcancel := context.WithTimeout(ctx, writeTimeout)
+			err := wsjson.Write(wctx, c, env)
+			wcancel()
+			if err != nil {
+				cancel(fmt.Errorf("send %s: %w", env.Type, err))
+				return
+			}
+		}
+	}
+}
+
+// readLoop handles what the master sends. A message it cannot parse ends the
+// session; a message type it does not know is logged and ignored, so a newer
+// master can add message types without breaking this agent.
+func (a *Agent) readLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn) {
+	for {
+		var env protocol.Envelope
+		if err := wsjson.Read(ctx, c, &env); err != nil {
+			cancel(fmt.Errorf("read: %w", err))
+			return
+		}
+		switch env.Type {
+		case protocol.TypeJobStart:
+			start, err := protocol.Decode[protocol.JobStart](env)
+			if err != nil || env.ID == "" {
+				cancel(fmt.Errorf("malformed job.start (id %q): %v", env.ID, err))
+				return
+			}
+			a.jobs.Submit(env.ID, start)
+		case protocol.TypeJobAck:
+			// JobAck has no fields, so the payload is not inspected.
+			a.jobs.Ack(env.ID)
+		default:
+			a.log.Warn("unknown_message_type", "type", env.Type)
+		}
 	}
-	if err := wsjson.Write(ctx, c, env); err != nil {
-		return fmt.Errorf("send %s: %w", typ, err)
-	}
-	return nil
 }
@@ -19,6 +19,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"

+	"github.com/shcizo/package-updater/internal/agent/updater"
 	"github.com/shcizo/package-updater/internal/pin"
 	"github.com/shcizo/package-updater/internal/protocol"
 )
@@ -113,12 +114,111 @@ var (
 	svcB = protocol.Service{Kind: protocol.KindContainer, ContainerName: "b", Image: "b:1", Status: "running"}
 )

+// fakeJobs records what the connection tells it and returns scripted state.
+type fakeJobs struct {
+	mu       sync.Mutex
+	known    []string
+	pending  []updater.Finished
+	progress []string // output chunks since the last TakeProgress
+	submits  map[string]protocol.JobStart
+	acks     []string
+	changed  chan struct{}
+}
+
+func newFakeJobs() *fakeJobs {
+	return &fakeJobs{submits: map[string]protocol.JobStart{}, changed: make(chan struct{}, 1)}
+}
+
+func (f *fakeJobs) Submit(id string, s protocol.JobStart) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.submits[id] = s
+}
+
+func (f *fakeJobs) Ack(id string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.acks = append(f.acks, id)
+	kept := f.pending[:0:0]
+	for _, p := range f.pending {
+		if p.ID != id {
+			kept = append(kept, p)
+		}
+	}
+	f.pending = kept
+}
+
+func (f *fakeJobs) Known() []string {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.known
+}
+
+func (f *fakeJobs) TakeProgress() []updater.Progress {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	if len(f.progress) == 0 {
+		return nil
+	}
+	out := updater.Progress{ID: "1", Output: strings.Join(f.progress, "")}
+	f.progress = nil
+	return []updater.Progress{out}
+}
+
+func (f *fakeJobs) Pending() []updater.Finished {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return append([]updater.Finished(nil), f.pending...)
+}
+
+func (f *fakeJobs) Changed() <-chan struct{} { return f.changed }
+
+func (f *fakeJobs) signal() {
+	select {
+	case f.changed <- struct{}{}:
+	default:
+	}
+}
+
+func (f *fakeJobs) addProgress(s string) {
+	f.mu.Lock()
+	f.progress = append(f.progress, s)
+	f.mu.Unlock()
+	f.signal()
+}
+
+func (f *fakeJobs) addResult(id string, r protocol.JobResult) {
+	f.mu.Lock()
+	f.pending = append(f.pending, updater.Finished{ID: id, Result: r})
+	f.mu.Unlock()
+	f.signal()
+}
+
+func (f *fakeJobs) submitted(id string) (protocol.JobStart, bool) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	s, ok := f.submits[id]
+	return s, ok
+}
+
+func (f *fakeJobs) acked() []string {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return append([]string(nil), f.acks...)
+}
+
 func startAgent(t *testing.T, fm *fakeMaster, fingerprint string, src Source) *Agent {
+	t.Helper()
+	return startAgentWith(t, fm, fingerprint, src, newFakeJobs(), 0)
+}
+
+func startAgentWith(t *testing.T, fm *fakeMaster, fingerprint string, src Source, jobs Jobs, progressInterval time.Duration) *Agent {
 	t.Helper()
 	a := New(Config{
 		URL: fm.url(), Token: "pua_test", Fingerprint: fingerprint, Version: "v-test",
 		MinBackoff: 10 * time.Millisecond, MaxBackoff: 50 * time.Millisecond,
-	}, src, slog.New(slog.NewTextHandler(io.Discard, nil)))
+		ProgressInterval: progressInterval,
+	}, src, jobs, slog.New(slog.NewTextHandler(io.Discard, nil)))
 	ctx, cancel := context.WithCancel(context.Background())
 	done := make(chan struct{})
 	go func() { _ = a.Run(ctx); close(done) }()
@@ -137,7 +237,7 @@ func TestAgent_SendsHelloThenSnapshotWithToken(t *testing.T) {
 	require.Equal(t, protocol.TypeHello, hello.Type)
 	h, err := protocol.Decode[protocol.Hello](hello)
 	require.NoError(t, err)
-	assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test"}, h)
+	assert.Equal(t, protocol.Hello{Hostname: "node-2", AgentVersion: "v-test", ProtocolVersion: protocol.Version}, h)

 	snap := fm.next(t)
 	require.Equal(t, protocol.TypeSnapshot, snap.Type)
@@ -198,3 +298,212 @@ func TestAgent_RefusesMasterWithWrongKey(t *testing.T) {
 	time.Sleep(200 * time.Millisecond)
 	assert.Empty(t, fm.auth, "the token must never be sent to a server with the wrong key")
 }
+
+func sendEnv(t *testing.T, c *websocket.Conn, env protocol.Envelope) {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	require.NoError(t, wsjson.Write(ctx, c, env))
+}
+
+func TestAgent_HelloCarriesVersionAndKnownJobs(t *testing.T) {
+	fm := newFakeMaster(t)
+	jobs := newFakeJobs()
+	jobs.known = []string{"3", "7"}
+	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
+
+	h, err := protocol.Decode[protocol.Hello](fm.next(t))
+	require.NoError(t, err)
+	assert.Equal(t, protocol.Version, h.ProtocolVersion)
+	assert.Equal(t, []string{"3", "7"}, h.KnownJobs)
+}
+
+func TestAgent_JobStartSubmitsAndAckAcks(t *testing.T) {
+	fm := newFakeMaster(t)
+	jobs := newFakeJobs()
+	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
+	fm.next(t)
+	fm.next(t)
+	c := <-fm.conns
+
+	start := protocol.JobStart{Kind: protocol.KindSwarm, SwarmClusterID: "c1", SwarmService: "gitea_gitea", Ref: "gitea/gitea:1.27"}
+	env, err := protocol.EncodeJob(protocol.TypeJobStart, "42", start)
+	require.NoError(t, err)
+	sendEnv(t, c, env)
+	require.Eventually(t, func() bool { _, ok := jobs.submitted("42"); return ok }, 2*time.Second, 5*time.Millisecond)
+	got, _ := jobs.submitted("42")
+	assert.Equal(t, start, got)
+
+	ack, err := protocol.EncodeJob(protocol.TypeJobAck, "42", protocol.JobAck{})
+	require.NoError(t, err)
+	sendEnv(t, c, ack)
+	require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond)
+	assert.Equal(t, []string{"42"}, jobs.acked())
+}
+
+func TestAgent_SendsPendingResultsAfterSnapshot(t *testing.T) {
+	fm := newFakeMaster(t)
+	jobs := newFakeJobs()
+	jobs.pending = []updater.Finished{
+		{ID: "5", Result: protocol.JobResult{Status: protocol.ResultSucceeded, Ref: "a:1@sha256:aa", Output: "ok\n"}},
+		{ID: "6", Result: protocol.JobResult{Status: protocol.ResultFailed, Error: "boom"}},
+	}
+	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
+
+	assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
+	assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
+	r5 := fm.next(t)
+	assert.Equal(t, protocol.TypeJobResult, r5.Type)
+	assert.Equal(t, "5", r5.ID)
+	res, err := protocol.Decode[protocol.JobResult](r5)
+	require.NoError(t, err)
+	assert.Equal(t, "a:1@sha256:aa", res.Ref)
+	r6 := fm.next(t)
+	assert.Equal(t, "6", r6.ID)
+}
+
+func TestAgent_ResendsUnackedResultAfterReconnect(t *testing.T) {
+	fm := newFakeMaster(t)
+	jobs := newFakeJobs()
+	jobs.addResult("5", protocol.JobResult{Status: protocol.ResultSucceeded})
+	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
+
+	fm.next(t) // hello
+	fm.next(t) // snapshot
+	assert.Equal(t, "5", fm.next(t).ID)
+	(<-fm.conns).CloseNow() // dropped without ack
+
+	assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
+	assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
+	again := fm.next(t)
+	assert.Equal(t, protocol.TypeJobResult, again.Type)
+	assert.Equal(t, "5", again.ID)
+
+	c := <-fm.conns
+	ack, err := protocol.EncodeJob(protocol.TypeJobAck, "5", protocol.JobAck{})
+	require.NoError(t, err)
+	sendEnv(t, c, ack)
+	require.Eventually(t, func() bool { return len(jobs.acked()) == 1 }, 2*time.Second, 5*time.Millisecond)
+	c.CloseNow()
+
+	assert.Equal(t, protocol.TypeHello, fm.next(t).Type)
+	assert.Equal(t, protocol.TypeSnapshot, fm.next(t).Type)
+	select {
+	case e := <-fm.envs:
+		t.Fatalf("acked result was sent again: %+v", e)
+	case <-time.After(100 * time.Millisecond):
+	}
+}
+
+func TestAgent_ProgressIsThrottled(t *testing.T) {
+	fm := newFakeMaster(t)
+	jobs := newFakeJobs()
+	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, 100*time.Millisecond)
+	fm.next(t)
+	fm.next(t)
+
+	for range 20 {
+		jobs.addProgress("x")
+		time.Sleep(2500 * time.Microsecond)
+	}
+	var batches int
+	var total strings.Builder
+	collect := func(d time.Duration) {
+		deadline := time.After(d)
+		for {
+			select {
+			case e := <-fm.envs:
+				if e.Type != protocol.TypeJobProgress {
+					continue
+				}
+				p, err := protocol.Decode[protocol.JobProgress](e)
+				require.NoError(t, err)
+				batches++
+				total.WriteString(p.Output)
+			case <-deadline:
+				return
+			}
+		}
+	}
+	collect(100 * time.Millisecond)
+	assert.LessOrEqual(t, batches, 2)
+	collect(400 * time.Millisecond)
+	assert.Equal(t, strings.Repeat("x", 20), total.String(), "the last chunk must be delivered")
+}
+
+func TestAgent_IgnoresUnknownType(t *testing.T) {
+	fm := newFakeMaster(t)
+	jobs := newFakeJobs()
+	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, jobs, time.Millisecond)
+	fm.next(t)
+	fm.next(t)
+	c := <-fm.conns
+
+	sendEnv(t, c, protocol.Envelope{Type: "future.thing", Payload: []byte(`{"a":1}`)})
+	env, err := protocol.EncodeJob(protocol.TypeJobStart, "9", protocol.JobStart{Kind: protocol.KindCompose, Ref: "a:1"})
+	require.NoError(t, err)
+	sendEnv(t, c, env)
+	require.Eventually(t, func() bool { _, ok := jobs.submitted("9"); return ok }, 2*time.Second, 5*time.Millisecond)
+	assert.Zero(t, len(fm.conns), "no reconnect expected")
+}
+
+func TestAgent_ClosesOnMalformedJobStart(t *testing.T) {
+	fm := newFakeMaster(t)
+	startAgentWith(t, fm, fm.fingerprint, &fakeSource{}, newFakeJobs(), time.Millisecond)
+	fm.next(t)
+	fm.next(t)
+	c := <-fm.conns
+
+	env, err := protocol.Encode(protocol.TypeJobStart, protocol.JobStart{Ref: "a:1"}) // no ID
+	require.NoError(t, err)
+	sendEnv(t, c, env)
+
+	select {
+	case <-fm.conns:
+	case <-time.After(2 * time.Second):
+		t.Fatal("agent did not reconnect after a malformed job.start")
+	}
+}
+
+func TestAgent_InventoryAndProgressConcurrently(t *testing.T) {
+	fm := newFakeMaster(t)
+	src := &fakeSource{}
+	src.set(svcA)
+	jobs := newFakeJobs()
+	a := startAgentWith(t, fm, fm.fingerprint, src, jobs, time.Millisecond)
+	fm.next(t)
+	fm.next(t)
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for i := range 30 {
+			jobs.addProgress("y")
+			if i%2 == 0 {
+				src.set(svcB)
+			} else {
+				src.set(svcA)
+			}
+			a.Notify()
+			time.Sleep(time.Millisecond)
+		}
+	}()
+	<-done
+	var progress, inventory int
+	deadline := time.After(500 * time.Millisecond)
+	for {
+		select {
+		case e := <-fm.envs:
+			switch e.Type {
+			case protocol.TypeJobProgress:
+				progress++
+			case protocol.TypeUpsert, protocol.TypeRemove:
+				inventory++
+			}
+		case <-deadline:
+			assert.Positive(t, progress)
+			assert.Positive(t, inventory)
+			return
+		}
+	}
+}