From 0eca987ba1279f4bb75e2a70e2ee5429046948fc Mon Sep 17 00:00:00 2001 From: dwrz Date: Fri, 30 Oct 2026 00:45:48 +0000 Subject: [PATCH] Refactor conversation pkg --- internal/conversation/conversation.go | 337 ++++++++++++--------- internal/conversation/conversation_test.go | 141 ++++++++- 2 files changed, 336 insertions(+), 142 deletions(-) diff --git a/internal/conversation/conversation.go b/internal/conversation/conversation.go index 120247a..a9e0b55 100644 --- a/internal/conversation/conversation.go +++ b/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, diff --git a/internal/conversation/conversation_test.go b/internal/conversation/conversation_test.go index 97813ac..3527ced 100644 --- a/internal/conversation/conversation_test.go +++ b/internal/conversation/conversation_test.go @@ -71,6 +71,20 @@ func userMessage(text string) *llm.Message { } } +// infer reserves c and runs one inference call to completion. +func infer( + ctx context.Context, + c *Conversation, + client *llm.Client, + req InferenceRequest, +) error { + inf, err := c.Reserve(ctx, client, req) + if err != nil { + return err + } + return inf.Run() +} + // newTestModelClient returns an llm client pointed at a server that // answers with handler. func newTestModelClient( @@ -1440,7 +1454,7 @@ func TestSubscribeInference(t *testing.T) { errs = make(chan error, 1) ) go func() { - errs <- c.Inference(ctx, client, InferenceRequest{ + errs <- infer(ctx, c, client, InferenceRequest{ Message: userMessage("hi"), Model: "test-model", Stream: true, @@ -1556,7 +1570,7 @@ func TestSubscribeInferenceError(t *testing.T) { errs = make(chan error, 1) ) go func() { - errs <- c.Inference(ctx, client, InferenceRequest{ + errs <- infer(ctx, c, client, InferenceRequest{ Message: userMessage("hi"), Model: "test-model", Stream: true, @@ -1646,7 +1660,7 @@ func TestSubscribeToolResult(t *testing.T) { errs = make(chan error, 1) ) go func() { - errs <- c.Inference(ctx, client, InferenceRequest{ + errs <- infer(ctx, c, client, InferenceRequest{ Message: userMessage("hi"), Model: "test-model", }) @@ -1741,7 +1755,7 @@ func TestSubscribeCancel(t *testing.T) { errs = make(chan error, 1) ) go func() { - errs <- c.Inference(ctx, client, InferenceRequest{ + errs <- infer(ctx, c, client, InferenceRequest{ Message: userMessage("hi"), Model: "test-model", Stream: true, @@ -1850,7 +1864,7 @@ func TestSubscribeStartupError(t *testing.T) { }) errs := make(chan error, 1) go func() { - errs <- c.Inference(ctx, client, InferenceRequest{ + errs <- infer(ctx, c, client, InferenceRequest{ Model: "test-model", }) }() @@ -2445,7 +2459,7 @@ func TestInferenceApprovalAndMaxIterations(t *testing.T) { // First call: saves the assistant message and pending execution, // returning ErrApprovalRequired. - err = c.Inference(ctx, client, InferenceRequest{ + err = infer(ctx, c, client, InferenceRequest{ Message: userMessage("run guarded"), Model: "test-model", }) @@ -2479,7 +2493,7 @@ func TestInferenceApprovalAndMaxIterations(t *testing.T) { t.Fatalf("UpdateExecutionStatus: %v", err) } - err = c.Inference(ctx, client, InferenceRequest{ + err = infer(ctx, c, client, InferenceRequest{ Model: "test-model", MaxIterations: 1, }) @@ -2487,3 +2501,116 @@ func TestInferenceApprovalAndMaxIterations(t *testing.T) { t.Fatalf("Inference error = %v, want ErrMaxIterations", err) } } + +// TestReserve checks that a reservation makes the conversation busy +// before Run starts, and that Run releases it. +func TestReserve(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + client = newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "unavailable", http.StatusServiceUnavailable) + }) + req = InferenceRequest{Model: "test-model"} + ) + + if _, err := c.Reserve(ctx, nil, req); err == nil { + t.Fatal("Reserve(nil client) error = nil") + } + if _, err := c.Reserve( + ctx, client, InferenceRequest{MaxIterations: -1}, + ); err == nil { + t.Fatal("Reserve(invalid request) error = nil") + } + + inf, err := c.Reserve(ctx, client, req) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + if _, err := c.Reserve(ctx, client, req); !errors.Is(err, ErrBusy) { + t.Fatalf("second Reserve error = %v, want ErrBusy", err) + } + if _, err := c.InsertMessage( + ctx, *userMessage("edit"), + ); !errors.Is(err, ErrBusy) { + t.Fatalf("InsertMessage error = %v, want ErrBusy", err) + } + + if err := inf.Run(); err == nil { + t.Fatal("Run error = nil, want the model error") + } + if err := inf.Run(); err == nil { + t.Fatal("second Run error = nil") + } + + next, err := c.Reserve(ctx, client, req) + if err != nil { + t.Fatalf("Reserve after Run: %v", err) + } + next.Run() +} + +// TestInferenceEventsFastFailure checks that a subscriber that holds +// the Inference receives Done even when Run fails before the subscriber +// reads its first event. +func TestInferenceEventsFastFailure(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + client = newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "unavailable", http.StatusServiceUnavailable) + }) + ) + inf, err := c.Reserve(ctx, client, InferenceRequest{Model: "test-model"}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + runErr := inf.Run() + if runErr == nil { + t.Fatal("Run error = nil, want the model error") + } + + got := collectEvents(t, inf.Events(ctx)) + if len(got) != 1 { + t.Fatalf("events = %v, want only Done", got) + } + done, ok := got[0].(Done) + if !ok { + t.Fatalf("event = %v, want Done", got[0]) + } + if done.Err == nil || done.Err.Error() != runErr.Error() { + t.Fatalf("Done.Err = %v, want %v", done.Err, runErr) + } +} + +// TestStopReserved checks that Stop cancels a reserved Inference +// before Run starts. +func TestStopReserved(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + called = make(chan struct{}, 1) + client = newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + called <- struct{}{} + }) + ) + inf, err := c.Reserve(ctx, client, InferenceRequest{Model: "test-model"}) + if err != nil { + t.Fatalf("Reserve: %v", err) + } + c.Stop() + if err := inf.Run(); !errors.Is(err, context.Canceled) { + t.Fatalf("Run error = %v, want context.Canceled", err) + } + select { + case <-called: + t.Fatal("model was called after Stop") + default: + } +}