Refactor conversation pkg

internal/conversation/conversation.go

@@ -2,9 +2,10 @@
 // execution, and live event streaming for a single conversation.
 //
 // A [Conversation] coordinates access to one conversation in SQLite:
-//   - [Conversation.Inference] runs the model and tool-call loop. Only one
-//     inference call or message edit may run at a time; concurrent calls
-//     return [ErrBusy].
+//   - [Conversation.Reserve] reserves the conversation for one
+//     [Inference], and [Inference.Run] runs the model and tool-call loop.
+//     Only one inference call or message edit may run at a time;
+//     concurrent calls return [ErrBusy].
 //   - [Conversation.Subscribe] synchronizes a client's cached messages with
 //     saved history and streams live updates from an active inference run.
 //   - Message and tool-call edit methods modify saved history when no
@@ -132,14 +133,14 @@ type Conversation struct {
 	log   *slog.Logger
 	tools *tool.Registry
 
-	// mu synchronizes access to cancel and run.
+	// mu synchronizes access to cancel and inf.
 	mu sync.Mutex
 	// cancel is used to stop inference.
-	// When set, it indicates inference is ongoing.
+	// When set, it indicates inference is reserved or ongoing.
 	cancel context.CancelFunc
-	// run shares the ongoing Inference's events with subscribers.
-	// When set, it indicates inference is ongoing.
-	run *run
+	// inf shares the reserved Inference's events with subscribers.
+	// When set, it indicates inference is reserved or ongoing.
+	inf *Inference
 }
 
 // MessageManifest identifies one saved message cached by a client.
@@ -196,12 +197,22 @@ type Done struct {
 
 func (Done) Type() EventType { return EventDone }
 
-// run records events from a [Conversation.Inference] call.
-// [Conversation.Subscribe] uses these events to send past and live updates to
-// subscribers.
-type run struct {
-	// mu serializes access to the other fields.
+// Inference is one reserved inference call and its events.
+// [Conversation.Reserve] creates it, and [Inference.Run] runs it.
+// [Inference.Events] and [Conversation.Subscribe] use its events to send
+// past and live updates to subscribers.
+type Inference struct {
+	c      *Conversation
+	client *llm.Client
+	req    InferenceRequest
+	// ctx is canceled by cancel, which Stop and Delete also call.
+	ctx    context.Context
+	cancel context.CancelFunc
+
+	// mu serializes access to the fields below.
 	mu sync.Mutex
+	// started reports whether Run has been called.
+	started bool
 	// changed notifies waiting subscribers of new events or completion.
 	// Each publish closes changed and replaces it with a new channel;
 	// finish closes it without replacement.
@@ -216,20 +227,90 @@ type run struct {
 	events []Event
 }
 
-// finish marks the run complete and records err.
-// It closes changed to wake waiting subscribers.
-// The channel stays closed because the run is complete.
-func (r *run) finish(err error) {
-	r.mu.Lock()
-	defer r.mu.Unlock()
+// publish appends e to the events and wakes waiting subscribers.
+// The caller must hold inf.mu.
+func (inf *Inference) publish(e Event) {
+	inf.events = append(inf.events, e)
+	close(inf.changed)
+	inf.changed = make(chan struct{})
+}
 
-	r.done = true
-	r.err = err
-	close(r.changed)
+// finish marks the call complete and records err.
+// It closes changed to wake waiting subscribers.
+// The channel stays closed because the call is complete.
+func (inf *Inference) finish(err error) {
+	inf.mu.Lock()
+	defer inf.mu.Unlock()
+
+	inf.done = true
+	inf.err = err
+	close(inf.changed)
+}
+
+// Events streams the call's events from the start, then ends with
+// [Done]. The channel closes after Done or when ctx is canceled.
+// Canceling ctx stops delivery; it does not stop the call.
+func (inf *Inference) Events(ctx context.Context) <-chan Event {
+	ch := make(chan Event, 64)
+	go func() {
+		defer close(ch)
+		inf.follow(ctx, ch, nil)
+	}()
+	return ch
+}
+
+// follow sends the call's events to ch from index 0, skipping
+// [MessageSaved] and [AssistantUpdate] events whose message ID is in
+// skip. It waits for new events until the call finishes, then sends
+// [Done]. It returns false if ctx is canceled first.
+func (inf *Inference) follow(
+	ctx context.Context,
+	ch chan Event,
+	skip map[string]struct{},
+) bool {
+	pos := 0
+	for {
+		// Copy the slice header for unread events and the current
+		// notification channel. Because the call only appends to
+		// events, reading batch after Unlock is safe.
+		inf.mu.Lock()
+		batch := inf.events[pos:]
+		pos = len(inf.events)
+		wait, done, err := inf.changed, inf.done, inf.err
+		inf.mu.Unlock()
+
+		for _, e := range batch {
+			var id string
+			switch e := e.(type) {
+			case MessageSaved:
+				id = e.Message.Id
+			case AssistantUpdate:
+				id = e.MessageId
+			}
+			if id != "" {
+				if _, ok := skip[id]; ok {
+					continue
+				}
+			}
+			if !send(ctx, ch, e) {
+				return false
+			}
+		}
+		if done {
+			return send(ctx, ch, Done{Err: err})
+		}
+		// Wait for the call to close wait (when it appends new
+		// events or finishes), or for the subscriber to cancel ctx.
+		select {
+		case <-wait:
+		case <-ctx.Done():
+			return false
+		}
+	}
 }
 
 // Subscribe streams events that synchronize a client's message cache with
-// the saved conversation and any ongoing [Conversation.Inference] call.
+// the saved conversation and any reserved or ongoing [Inference].
 //
 // manifest lists the message IDs and timestamps the caller already holds.
 // An empty manifest requests all saved messages.
@@ -238,34 +319,35 @@ func (r *run) finish(err error) {
 //  1. Cache diff: [MessageSaved] for each saved message missing from or
 //     newer than manifest, followed by [MessageDeleted] for each ID in
 //     manifest that is no longer saved.
-//  2. Active inference: if [Conversation.Inference] is running when
-//     Subscribe loads the database, Subscribe streams the run's
+//  2. Active inference: if an [Inference] is reserved or running when
+//     Subscribe loads the database, Subscribe streams its
 //     [MessageSaved] and [AssistantUpdate] events, then ends with [Done].
-//     If no inference is running, Subscribe sends no [Done] event.
+//     If no inference is reserved or running, Subscribe sends no [Done]
+//     event.
 //
 // The channel closes when delivery finishes or ctx is canceled. Canceling
 // ctx stops delivery to this subscriber; it does not stop an ongoing
-// [Conversation.Inference] call.
+// [Inference].
 func (c *Conversation) Subscribe(
 	ctx context.Context,
 	manifest []MessageManifest,
 ) (<-chan Event, error) {
-	// Check for an active Inference run.
+	// Check for a reserved or ongoing Inference.
 	c.mu.Lock()
-	run := c.run
+	inf := c.inf
 	c.mu.Unlock()
 
-	// Inference holds run.mu while it commits a message and appends the
-	// MessageSaved event to run.events. Holding run.mu across this read
-	// ensures that any message committed after this point will appear in
-	// run.events, and any message already in rows will be recorded in the
-	// diff's sent map.
-	if run != nil {
-		run.mu.Lock()
+	// Inference.Run holds inf.mu while it commits a message and
+	// appends the MessageSaved event to inf.events. Holding inf.mu
+	// across this read ensures that any message committed after this
+	// point will appear in inf.events, and any message already in rows
+	// will be recorded in the diff's sent map.
+	if inf != nil {
+		inf.mu.Lock()
 	}
 	rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id)
-	if run != nil {
-		run.mu.Unlock()
+	if inf != nil {
+		inf.mu.Unlock()
 	}
 	if err != nil {
 		return nil, fmt.Errorf("load messages: %w", err)
@@ -283,58 +365,14 @@ func (c *Conversation) Subscribe(
 				return
 			}
 		}
-		if run == nil {
+		if inf == nil {
 			return
 		}
-
-		// Replay and follow run.events from index 0.
-		pos := 0
-		for {
-			// Copy the slice header for unread events and the
-			// current notification channel. Because Inference
-			// only appends to run.events, reading batch after
-			// Unlock is safe.
-			run.mu.Lock()
-			batch := run.events[pos:]
-			pos = len(run.events)
-			wait, done, err := run.changed, run.done, run.err
-			run.mu.Unlock()
-
-			// Skip events for messages that the initial diff
-			// already sent. If the client's manifest already had
-			// a message from this run, it is not in sent and will
-			// be re-sent here; the client deduplicates by ID and
-			// UpdatedAt.
-			for _, e := range batch {
-				var id string
-				switch e := e.(type) {
-				case MessageSaved:
-					id = e.Message.Id
-				case AssistantUpdate:
-					id = e.MessageId
-				}
-				if id != "" {
-					if _, ok := sent[id]; ok {
-						continue
-					}
-				}
-				if !send(ctx, ch, e) {
-					return
-				}
-			}
-			if done {
-				send(ctx, ch, Done{Err: err})
-				return
-			}
-			// Wait for Inference to close wait (when it appends
-			// new events or finishes), or for the subscriber to
-			// cancel ctx.
-			select {
-			case <-wait:
-			case <-ctx.Done():
-				return
-			}
-		}
+		// Skip events for messages that the initial diff already
+		// sent. If the client's manifest already had a message from
+		// this call, it is not in sent and will be re-sent here; the
+		// client deduplicates by ID and UpdatedAt.
+		inf.follow(ctx, ch, sent)
 	}()
 	return ch, nil
 }
@@ -438,22 +476,28 @@ func (r InferenceRequest) Validate() error {
 	return nil
 }
 
-// Inference resolves outstanding tool calls in saved history, saves
-// optional new user input, and runs the model/tool loop, blocking until
-// the work and result-save attempts finish. Pending approvals return
-// ErrApprovalRequired after sibling calls finish their save attempts;
-// new user input is not saved while calls remain unresolved.
-func (c *Conversation) Inference(
+// Reserve validates req and reserves the conversation for one inference
+// call. It returns ErrBusy if an inference call is reserved or running.
+//
+// ctx bounds the call's work, not the call to Reserve: the call stops
+// when ctx is canceled, when req.Timeout expires, or on Stop or Delete.
+// Pass a context that outlives any single HTTP request, so a client
+// disconnect does not stop the call.
+//
+// The caller must call [Inference.Run] on the result. Until Run returns,
+// the conversation stays busy: other calls to Reserve and all edits
+// return ErrBusy, even if Run is never called.
+func (c *Conversation) Reserve(
 	ctx context.Context,
 	client *llm.Client,
 	req InferenceRequest,
-) error {
+) (*Inference, error) {
 	// Reject invalid input before accepting work.
 	if client == nil {
-		return fmt.Errorf("missing llm client")
+		return nil, fmt.Errorf("missing llm client")
 	}
 	if err := req.Validate(); err != nil {
-		return fmt.Errorf("invalid inference request: %w", err)
+		return nil, fmt.Errorf("invalid inference request: %w", err)
 	}
 
 	// Limit the work duration and allow Stop to cancel it.
@@ -464,31 +508,62 @@ func (c *Conversation) Inference(
 		ctx, cancel = context.WithCancel(ctx)
 	}
 
-	// Reserve the conversation until work and save attempts finish.
-	// The run shares this call's events with subscribers.
-	r := &run{changed: make(chan struct{})}
+	// Check busy and reserve in one critical section, so no other
+	// Reserve or edit can start between the check and the reservation.
+	inf := &Inference{
+		c:       c,
+		client:  client,
+		req:     req,
+		ctx:     ctx,
+		cancel:  cancel,
+		changed: make(chan struct{}),
+	}
 	c.mu.Lock()
+	defer c.mu.Unlock()
 	if c.cancel != nil {
-		c.mu.Unlock()
 		cancel()
-		return ErrBusy
+		return nil, ErrBusy
 	}
 	c.cancel = cancel
-	c.run = r
-	c.mu.Unlock()
-	// Every failing path below must assign this err, not shadow it:
-	// the deferred run.finish call reports it to subscribers.
-	var err error
+	c.inf = inf
+	return inf, nil
+}
+
+// Run resolves outstanding tool calls in saved history, saves optional
+// new user input, and runs the model/tool loop, blocking until the work
+// and result-save attempts finish. Pending approvals return
+// ErrApprovalRequired after sibling calls finish their save attempts;
+// new user input is not saved while calls remain unresolved.
+//
+// When Run returns, it sends [Done] to subscribers and releases the
+// reservation. Call Run once; a second call returns an error.
+func (inf *Inference) Run() (err error) {
+	inf.mu.Lock()
+	started := inf.started
+	inf.started = true
+	inf.mu.Unlock()
+	if started {
+		return fmt.Errorf("inference already run")
+	}
+
+	c := inf.c
 	defer func() {
-		// Every exit path, including errors, finishes the run.
-		r.finish(err)
-		cancel()
+		// Every exit path, including errors, finishes the call.
+		// Every failing path below must assign err, not shadow it.
+		inf.finish(err)
+		inf.cancel()
 		c.mu.Lock()
 		c.cancel = nil
-		c.run = nil
+		c.inf = nil
 		c.mu.Unlock()
 	}()
 
+	var (
+		ctx    = inf.ctx
+		client = inf.client
+		req    = inf.req
+	)
+
 	// Check that the conversation exists.
 	conv, err := c.dbr.SelectConversationId(ctx, c.id)
 	if err != nil {
@@ -555,16 +630,14 @@ func (c *Conversation) Inference(
 		log.InfoContext(ctx, "persisting new message")
 
 		var saved *db.Message
-		r.mu.Lock()
+		inf.mu.Lock()
 		saved, err = c.insertMessage(
 			ctx, nil, *req.Message, nil, nil,
 		)
 		if err == nil {
-			r.events = append(r.events, MessageSaved{Message: *saved})
-			close(r.changed)
-			r.changed = make(chan struct{})
+			inf.publish(MessageSaved{Message: *saved})
 		}
-		r.mu.Unlock()
+		inf.mu.Unlock()
 		if err != nil {
 			return fmt.Errorf("insert message: %w", err)
 		}
@@ -609,14 +682,12 @@ func (c *Conversation) Inference(
 				break
 			}
 			for stream.Next() {
-				r.mu.Lock()
-				r.events = append(r.events, AssistantUpdate{
+				inf.mu.Lock()
+				inf.publish(AssistantUpdate{
 					MessageId: aid.String(),
 					Chunk:     stream.Current(),
 				})
-				close(r.changed)
-				r.changed = make(chan struct{})
-				r.mu.Unlock()
+				inf.mu.Unlock()
 			}
 			err = errors.Join(
 				stream.Err(), stream.Close(),
@@ -646,16 +717,14 @@ func (c *Conversation) Inference(
 		}
 		// Persist the assistant message before acting on it, so an
 		// interruption leaves a committed record.
-		r.mu.Lock()
+		inf.mu.Lock()
 		row, err = c.insertMessage(
 			ctx, &aid, res.Message, &res.Metadata, nil,
 		)
 		if err == nil {
-			r.events = append(r.events, MessageSaved{Message: *row})
-			close(r.changed)
-			r.changed = make(chan struct{})
+			inf.publish(MessageSaved{Message: *row})
 		}
-		r.mu.Unlock()
+		inf.mu.Unlock()
 		if err != nil {
 			err = fmt.Errorf("insert message: %w", err)
 			break
@@ -947,22 +1016,20 @@ func (c *Conversation) callTool(
 			Status: db.ExecutionStatusCompleted,
 		}
 	}
-	// c.run is set for the whole of Inference, the only caller path to
-	// callTool.
+	// c.inf is set for the whole of Inference.Run, the only caller path
+	// to callTool.
 	c.mu.Lock()
-	r := c.run
+	inf := c.inf
 	c.mu.Unlock()
 
-	// Hold r.mu across the commit so a concurrent Subscribe's diff
+	// Hold inf.mu across the commit so a concurrent Subscribe's diff
 	// sees the row and its event together.
-	r.mu.Lock()
+	inf.mu.Lock()
 	saved, err := c.insertMessage(saveCtx, nil, msg, nil, update)
 	if err == nil {
-		r.events = append(r.events, MessageSaved{Message: *saved})
-		close(r.changed)
-		r.changed = make(chan struct{})
+		inf.publish(MessageSaved{Message: *saved})
 	}
-	r.mu.Unlock()
+	inf.mu.Unlock()
 	if err != nil {
 		log.ErrorContext(
 			ctx,