chimerical/odidere
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor conversation pkg

Browse Source
This commit is contained in:
dwrz
2026-10-30 00:45:48 +00:00
parent 2426509531
commit 0eca987ba1
2 changed files with 336 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/conversation/conversation.go
+202 -135
View File
@@ -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,
internal/conversation/conversation_test.go
+134 -7
View File
@@ -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:
}
}