From 2426509531933026cf585c53d75be1703d6cc136 Mon Sep 17 00:00:00 2001 From: dwrz Date: Fri, 23 Oct 2026 23:41:35 +0000 Subject: [PATCH] Add conversation pkg --- internal/conversation/conversation.go | 1708 ++++++++++++++ internal/conversation/conversation_test.go | 2489 ++++++++++++++++++++ internal/db/models.go | 4 +- internal/db/queries.go | 366 ++- internal/db/queries_test.go | 174 +- internal/db/schema/1.init.sql | 15 +- internal/job/job.go | 4 +- internal/job/job_test.go | 2 +- internal/llm/llm.go | 31 + 9 files changed, 4751 insertions(+), 42 deletions(-) create mode 100644 internal/conversation/conversation.go create mode 100644 internal/conversation/conversation_test.go diff --git a/internal/conversation/conversation.go b/internal/conversation/conversation.go new file mode 100644 index 0000000..120247a --- /dev/null +++ b/internal/conversation/conversation.go @@ -0,0 +1,1708 @@ +// Package conversation manages the saved history, model inference, tool +// 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.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 +// inference is running, using optimistic timestamp checks +// ([ErrRevisionConflict]) to prevent lost updates. +package conversation + +import ( + "context" + "database/sql" + "encoding/json/jsontext" + "encoding/json/v2" + "errors" + "fmt" + "log/slog" + "runtime" + "slices" + "sync" + "time" + "uuid" + + "code.chimeric.al/chimerical/odidere/internal/db" + "code.chimeric.al/chimerical/odidere/internal/llm" + "code.chimeric.al/chimerical/odidere/internal/tool" +) + +const ( + // interruptedWarning indicates a tool call started but did not finish + // normally, so its effects are unknown. + interruptedWarning = `▸ WARNING: the tool started, but its outcome is unknown. +It may have had effects. Do not assume it is safe to repeat.` + + // resultSaveTimeout is the time limit for saving one tool result. + resultSaveTimeout = time.Second +) + +var ( + // ErrApprovalRequired indicates that saved tool calls await user + // approval. + ErrApprovalRequired = errors.New("tool approval required") + + // ErrBusy indicates that inference is ongoing. + ErrBusy = errors.New("conversation busy") + + // ErrMaxIterations indicates that the model-request limit was reached + // without a final assistant message. + ErrMaxIterations = errors.New("max iterations exceeded") + + // ErrRevisionConflict indicates that the message was changed before + // the caller's revision. + ErrRevisionConflict = errors.New("message revision conflict") +) + +// DecodeMessages converts saved rows into llm.Messages, validating each. +// It stops at the first row that fails to unmarshal or validate; the +// returned slice is nil on error. +func DecodeMessages(rows []db.Message) ([]llm.Message, error) { + msgs := make([]llm.Message, len(rows)) + for i, row := range rows { + if err := json.Unmarshal( + row.Content, &msgs[i], + ); err != nil { + return nil, fmt.Errorf( + "unmarshal message %q: %w", row.Id, err, + ) + } + if err := msgs[i].Validate(); err != nil { + return nil, fmt.Errorf( + "invalid message %q: %w", row.Id, err, + ) + } + } + return msgs, nil +} + +// New constructs an in-memory conversation manager. +func New( + cfg Config, + sdb *db.SQLite, + log *slog.Logger, + tools *tool.Registry, +) (*Conversation, error) { + if err := cfg.Validate(); err != nil { + return nil, fmt.Errorf("invalid config: %w", err) + } + if sdb == nil { + return nil, fmt.Errorf("missing database") + } + if log == nil { + return nil, fmt.Errorf("missing logger") + } + if tools == nil { + return nil, fmt.Errorf("missing tool registry") + } + + return &Conversation{ + id: cfg.Id, + db: sdb, + dbr: db.New(sdb.Reader), + dbw: db.New(sdb.Writer), + log: log, + tools: tools, + }, nil +} + +// Config identifies the saved conversation to manage. +type Config struct { + Id string +} + +// Validate checks the configuration without accessing the database. +func (cfg Config) Validate() error { + if cfg.Id == "" { + return fmt.Errorf("missing conversation id") + } + return nil +} + +// Conversation is an in-memory handle for one saved conversation. +type Conversation struct { + id string + db *db.SQLite + dbr *db.Queries + dbw *db.Queries + log *slog.Logger + tools *tool.Registry + + // mu synchronizes access to cancel and run. + mu sync.Mutex + // cancel is used to stop inference. + // When set, it indicates inference is ongoing. + cancel context.CancelFunc + // run shares the ongoing Inference's events with subscribers. + // When set, it indicates inference is ongoing. + run *run +} + +// MessageManifest identifies one saved message cached by a client. +type MessageManifest struct { + Id string + UpdatedAt time.Time +} + +// EventType distinguishes the events delivered by Subscribe. +type EventType string + +const ( + EventAssistantUpdate EventType = "assistant_update" + EventDone EventType = "done" + EventMessageDeleted EventType = "message_deleted" + EventMessageSaved EventType = "message_saved" +) + +// Event is a conversation update delivered to subscribers. +type Event interface { + // Type reports the kind of the event. + Type() EventType +} + +// MessageSaved reports a message was saved. +type MessageSaved struct { + Message db.Message `json:"message"` +} + +func (MessageSaved) Type() EventType { return EventMessageSaved } + +// MessageDeleted reports a message ID the client's manifest holds but +// the saved history no longer contains. +type MessageDeleted struct { + MessageId string `json:"message_id"` +} + +func (MessageDeleted) Type() EventType { return EventMessageDeleted } + +// AssistantUpdate reports the next streamed fragment of one assistant +// message. +type AssistantUpdate struct { + MessageId string `json:"message_id"` + Chunk llm.StreamChunk `json:"chunk"` +} + +func (AssistantUpdate) Type() EventType { return EventAssistantUpdate } + +// Done reports that the Inference the subscriber joined has returned. +// Err is nil on success. +type Done struct { + Err error `json:"err"` +} + +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. + mu sync.Mutex + // 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. + changed chan struct{} + // done reports whether the inference work has finished. + done bool + // err reports why inference failed, or nil on success. + err error + // events holds all published events in order. + // Events are added, never removed. + // Each subscriber keeps its own read position in events. + 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() + + r.done = true + r.err = err + close(r.changed) +} + +// Subscribe streams events that synchronize a client's message cache with +// the saved conversation and any ongoing [Conversation.Inference] call. +// +// manifest lists the message IDs and timestamps the caller already holds. +// An empty manifest requests all saved messages. +// +// Events arrive on the returned channel in two phases: +// 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 +// [MessageSaved] and [AssistantUpdate] events, then ends with [Done]. +// If no inference is 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. +func (c *Conversation) Subscribe( + ctx context.Context, + manifest []MessageManifest, +) (<-chan Event, error) { + // Check for an active Inference run. + c.mu.Lock() + run := c.run + 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() + } + rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id) + if run != nil { + run.mu.Unlock() + } + if err != nil { + return nil, fmt.Errorf("load messages: %w", err) + } + events, sent := diffMessages(rows, manifest) + + // Buffer events so bursts (such as the initial diff or fast token + // streaming) do not cause a context switch on every event. + ch := make(chan Event, 64) + go func() { + defer close(ch) + + for _, e := range events { + if !send(ctx, ch, e) { + return + } + } + if run == 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 + } + } + }() + return ch, nil +} + +// send writes e to ch and returns true, or returns false if ctx is +// canceled before ch can accept e. This prevents a slow or disconnected +// subscriber from blocking the producer goroutine when ch is full. +func send(ctx context.Context, ch chan Event, e Event) bool { + select { + case ch <- e: + return true + case <-ctx.Done(): + return false + } +} + +// diffMessages compares saved rows against manifest. +// +// It returns: +// - events: [MessageSaved] for each row missing from or newer than +// manifest (in saved order), followed by [MessageDeleted] for each +// ID in manifest that is not in rows. +// - sent: the set of message IDs included in the [MessageSaved] events. +func diffMessages( + rows []db.Message, + manifest []MessageManifest, +) (events []Event, sent map[string]struct{}) { + var ( + // known maps each client-cached message id to its timestamp. + known = make(map[string]time.Time, len(manifest)) + // saved holds ids of all messages currently in the database. + saved = make(map[string]struct{}, len(rows)) + ) + for _, m := range manifest { + known[m.Id] = m.UpdatedAt + } + + // Emit MessageSaved for new or updated rows, in saved order. + sent = map[string]struct{}{} + for _, row := range rows { + saved[row.Id] = struct{}{} + if at, ok := known[row.Id]; ok && at.Equal(row.UpdatedAt) { + continue + } + events = append(events, MessageSaved{Message: row}) + sent[row.Id] = struct{}{} + } + + // Emit MessageDeleted for client-cached ids no longer in the database. + for id := range known { + if _, ok := saved[id]; !ok { + events = append(events, MessageDeleted{MessageId: id}) + } + } + + return events, sent +} + +// InferenceRequest supplies input and settings for one inference/tool loop. +// Runtime dependencies are passed separately to Inference. +type InferenceRequest struct { + // Message supplies new user input, including text and attachments. + // nil continues from saved history without adding a message. + Message *llm.Message + // Model is the LLM model ID. If empty, the default model is used. + Model string + // MaxIterations caps the agent loop iterations. Zero means unlimited. + MaxIterations int + // ThinkingBudgetTokens is the llama.cpp reasoning budget. + // (-1 = unlimited, 0 = off). + ThinkingBudgetTokens *int + // Timeout limits the inference time (e.g., "10m"). + Timeout time.Duration + // Stream selects provider streaming, not HTTP response delivery. + // Failed streams never save partial assistant messages. + // False requests a completed provider response. + Stream bool +} + +// Validate checks inference limits and optional user input. +func (r InferenceRequest) Validate() error { + if r.MaxIterations < 0 { + return fmt.Errorf("negative max_iterations") + } + if r.ThinkingBudgetTokens != nil && *r.ThinkingBudgetTokens < -1 { + return fmt.Errorf("invalid thinking_budget_tokens") + } + if r.Timeout < 0 { + return fmt.Errorf("negative timeout") + } + if m := r.Message; m != nil { + if m.Role != llm.RoleUser || len(m.ToolCalls) != 0 || + m.ToolCallID != "" || m.Reasoning != "" || + m.ReasoningContent != "" || m.Refusal != "" { + return fmt.Errorf("invalid message") + } + if err := m.Validate(); err != nil { + return fmt.Errorf("invalid message: %w", err) + } + } + 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( + ctx context.Context, + client *llm.Client, + req InferenceRequest, +) error { + // Reject invalid input before accepting work. + if client == nil { + return fmt.Errorf("missing llm client") + } + if err := req.Validate(); err != nil { + return fmt.Errorf("invalid inference request: %w", err) + } + + // Limit the work duration and allow Stop to cancel it. + var cancel context.CancelFunc + if req.Timeout > 0 { + ctx, cancel = context.WithTimeout(ctx, req.Timeout) + } else { + 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{})} + c.mu.Lock() + if c.cancel != nil { + c.mu.Unlock() + cancel() + return 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 + defer func() { + // Every exit path, including errors, finishes the run. + r.finish(err) + cancel() + c.mu.Lock() + c.cancel = nil + c.run = nil + c.mu.Unlock() + }() + + // Check that the conversation exists. + conv, err := c.dbr.SelectConversationId(ctx, c.id) + if err != nil { + return fmt.Errorf("load conversation: %w", err) + } + log := c.log.With( + slog.String("conversation_id", conv.Id), + ) + + // Load saved rows in conversation order. + rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil { + return fmt.Errorf("load history: %w", err) + } + base := len(rows) + + // Convert db.Message into llm.Message; validate. + history, err := DecodeMessages(rows) + if err != nil { + return fmt.Errorf("decode history: %w", err) + } + + // Find the last message that is not a tool result. + // For an assistant role with tool calls, run calls without a result. + var ( + aId string + remaining []llm.ToolCall + answered = map[string]struct{}{} + ) + for i, m := range slices.Backward(history) { + if m.Role == llm.RoleTool { + answered[m.ToolCallID] = struct{}{} + continue + } + if len(m.ToolCalls) != 0 { + aId = rows[i].Id + for _, call := range m.ToolCalls { + if _, ok := answered[call.ID]; !ok { + remaining = append(remaining, call) + } + } + } + // Done: we've found the last non-tool message. + break + } + + // Resolve tool calls before adding input or requesting inference. + if len(remaining) != 0 { + log.InfoContext( + ctx, + "executing remaining tool calls", + slog.Int("count", len(remaining)), + ) + var tres []llm.Message + tres, err = c.callTools(ctx, log, aId, remaining) + if err != nil { + return fmt.Errorf("call tools: %w", err) + } + history = append(history, tres...) + } + + // Save new user input before requesting a completion. + if req.Message != nil { + log.InfoContext(ctx, "persisting new message") + + var saved *db.Message + r.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{}) + } + r.mu.Unlock() + if err != nil { + return fmt.Errorf("insert message: %w", err) + } + log.InfoContext( + ctx, + "persisted new user message", + slog.String("id", saved.Id), + ) + + // Continue from the updated history. + history = append(history, *req.Message) + } + + // Loop for inference. + var ( + completed bool + start time.Time + tools = c.tools.ToAPI() + ) + log.InfoContext(ctx, "starting inference") + start = time.Now() + for i := 0; req.MaxIterations == 0 || i < req.MaxIterations; i++ { + var ( + cp = llm.CompletionsParams{ + Messages: history, + Model: req.Model, + ThinkingBudgetTokens: req.ThinkingBudgetTokens, + Tools: tools, + } + // aid identifies the assistant message this iteration + // produces, created before the request so the + // assistant_update events and the saved row share it. + aid = uuid.NewV4() + res *llm.CompletionResult + row *db.Message + stream *llm.CompletionStream + ) + if req.Stream { + stream, err = client.CompletionsStream(ctx, cp) + if err != nil { + err = fmt.Errorf("completions stream: %w", err) + break + } + for stream.Next() { + r.mu.Lock() + r.events = append(r.events, AssistantUpdate{ + MessageId: aid.String(), + Chunk: stream.Current(), + }) + close(r.changed) + r.changed = make(chan struct{}) + r.mu.Unlock() + } + err = errors.Join( + stream.Err(), stream.Close(), + ) + if err != nil { + break + } + sr := stream.Result() + res = &sr + } else { + res, err = client.Completions(ctx, cp) + if err != nil { + err = fmt.Errorf("completions: %w", err) + break + } + } + + // Reject invalid assistant output before saving it. + if res.Message.Role != llm.RoleAssistant { + err = fmt.Errorf( + "invalid message role: %q", res.Message.Role) + break + } + if err = res.Message.Validate(); err != nil { + err = fmt.Errorf("invalid assistant message: %w", err) + break + } + // Persist the assistant message before acting on it, so an + // interruption leaves a committed record. + r.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{}) + } + r.mu.Unlock() + if err != nil { + err = fmt.Errorf("insert message: %w", err) + break + } + history = append(history, res.Message) + + // If there are no more tool calls, we're done. + if len(res.Message.ToolCalls) == 0 { + completed = true + break + } + + // Resolve the reply's tool calls before the next request. + tres, toolErr := c.callTools( + ctx, log, row.Id, res.Message.ToolCalls, + ) + if toolErr != nil { + err = errors.Join(toolErr, ctx.Err()) + break + } + history = append(history, tres...) + } + duration := time.Since(start) + if err == nil && !completed { + err = ErrMaxIterations + } + if err != nil { + log.ErrorContext( + ctx, + "inference failed", + slog.Any("error", err), + slog.Duration("duration", duration), + ) + return err + } + log.InfoContext( + ctx, + "inference complete", + slog.Duration("duration", duration), + slog.Int("messages", len(history)-base), + ) + return nil +} + +// callTools runs the tool calls of the assistant message aId, up to +// twice the CPU count concurrently, and collects each persisted result. +// Each call saves its own result, so the returned order may differ from +// the saved order; results are keyed by tool call ID. The error joins +// every failure; if non-nil, at least one call has no result. +func (c *Conversation) callTools( + ctx context.Context, + log *slog.Logger, + aId string, + calls []llm.ToolCall, +) ([]llm.Message, error) { + // Load each call's execution status. + executions, err := c.dbr.SelectExecutionsMessageId( + ctx, aId, + ) + if err != nil { + return nil, fmt.Errorf("load executions: %w", err) + } + callExecutions := make(map[string]*db.Execution, len(executions)) + for i := range executions { + callExecutions[executions[i].ToolCallId] = &executions[i] + } + + // Run calls concurrently, up to twice the CPU count at once. + type outcome struct { + err error + msg llm.Message + } + var ( + completed = make(chan outcome, len(calls)) + sem = make(chan struct{}, runtime.NumCPU()*2) + ) + for _, call := range calls { + go func() { + select { + case sem <- struct{}{}: + defer func() { <-sem }() + case <-ctx.Done(): + completed <- outcome{err: ctx.Err()} + return + } + msg, err := c.callTool( + ctx, log, aId, call, callExecutions[call.ID], + ) + completed <- outcome{msg: msg, err: err} + }() + } + + // Wait for every worker, even if another worker fails. + var ( + errs []error + results []llm.Message + ) + for range calls { + result := <-completed + if result.err != nil { + log.ErrorContext( + ctx, + "tool call failed", + slog.Any("error", result.err), + ) + errs = append(errs, result.err) + continue + } + results = append(results, result.msg) + } + return results, errors.Join(errs...) +} + +// callTool runs one tool call and saves the tool message that reports +// the outcome to the model. +// +// The record is a row in the executions table, saved before the tool +// starts. The row shows whether the tool may have run: +// - No row: the tool never started. It is safe to run again. +// - Row with status 'pending': the tool requires approval and never +// started; it runs once approved. +// - Row with status 'running': the tool may have run and had effects. +// It is never retried. +// - Row with status 'completed': the tool finished or never started, +// and the tool message is saved. +// +// A nil execution is a new call: recorded running, or pending if the +// tool requires approval, in which case ErrApprovalRequired is returned +// before running. A recorded execution decides the rest: pending returns +// ErrApprovalRequired, approved is marked running and run, any other +// status is refused. +// +// Results: +// - A message and a nil error: the message is saved. The message +// reports the tool's outcome: exit code and output, the error that +// kept the tool from starting, or the error that interrupted it. +// - An error: the outcome was not recorded. The tool may have run and +// had effects. +// +// The message contains: +// - If the tool finished: its output. +// - If it never started: the error. +// - If it started but was stopped, or its output was cut off: the +// error, a warning that its effects are unknown, and any output. +// +// Limitations: +// - callTool relies on [tool.Registry.Execute] returning a nil result +// only when the process never started. +func (c *Conversation) callTool( + ctx context.Context, + log *slog.Logger, + aId string, + call llm.ToolCall, + execution *db.Execution, +) (llm.Message, error) { + log = log.With( + slog.String("tool_call_id", call.ID), + slog.String("name", call.Function.Name), + ) + + // Decide whether and how the call may run. A case that ends without + // returning has cleared the call to run and continues to the tail. + switch { + case execution == nil: + // Record a new call before running it or requesting approval. + status := db.ExecutionStatusRunning + if definition, ok := c.tools.Get(call.Function.Name); ok && + definition.RequiresApproval { + status = db.ExecutionStatusPending + } + var err error + execution, err = c.dbw.InsertExecution( + ctx, + db.InsertExecutionParams{ + MessageId: aId, + ToolCallId: call.ID, + Status: status, + }, + ) + if err != nil { + return llm.Message{}, fmt.Errorf( + "insert execution for call %q: %w", call.ID, err, + ) + } + if status == db.ExecutionStatusPending { + return llm.Message{}, fmt.Errorf( + "call %q: %w", call.ID, ErrApprovalRequired, + ) + } + case execution.Status == db.ExecutionStatusPending: + return llm.Message{}, fmt.Errorf( + "call %q: %w", call.ID, ErrApprovalRequired, + ) + case execution.Status == db.ExecutionStatusApproved: + if err := c.dbw.UpdateExecutionStatus( + ctx, + db.UpdateExecutionStatusParams{ + Id: execution.Id, + Status: db.ExecutionStatusRunning, + }, + ); err != nil { + return llm.Message{}, fmt.Errorf( + "start approved call %q: %w", call.ID, err, + ) + } + default: + // An existing row that may have run is never retried. + return llm.Message{}, fmt.Errorf( + "call %q has status %q but no result; will not retry", + call.ID, execution.Status, + ) + } + + log.InfoContext( + ctx, + "calling tool", + slog.String("args", call.Function.Arguments), + ) + // Run the tool. + tres, toolErr := c.tools.Execute( + ctx, call.Function.Name, call.Function.Arguments, + ) + var content string + switch { + case toolErr == nil: + log.InfoContext( + ctx, + "called tool", + slog.Duration( + "duration", + tres.FinishedAt.Sub(tres.StartedAt), + ), + ) + content = tres.FormatResult() + case tres == nil: + log.ErrorContext( + ctx, + "tool did not start", + slog.Any("error", toolErr), + ) + content = toolErr.Error() + default: + log.ErrorContext( + ctx, + "tool interrupted", + slog.Any("error", toolErr), + slog.Duration( + "duration", + tres.FinishedAt.Sub(tres.StartedAt), + ), + ) + // Retain output produced before interruption. + content = toolErr.Error() + "\n" + + interruptedWarning + "\n" + + tres.FormatResult() + } + + // Build the tool result msg. + msg := llm.Message{ + Role: llm.RoleTool, + Name: call.Function.Name, + ToolCallID: call.ID, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: content, + }}, + } + if err := msg.Validate(); err != nil { + return llm.Message{}, fmt.Errorf( + "invalid message %q: %w", call.ID, err, + ) + } + // Attempt to save the result, even in case of cancellation. + // Tool calls are expensive and may have side-effects; avoid throwing + // away completed work. The status update, if any, commits with the + // result. + saveCtx, cancel := context.WithTimeout( + context.WithoutCancel(ctx), resultSaveTimeout, + ) + defer cancel() + + // Mark finished tools and failed launches as completed. + // Keep interrupted calls as ExecutionStatusRunning; + // they may have had effects. + var update *db.UpdateExecutionStatusParams + if toolErr == nil || tres == nil { + update = &db.UpdateExecutionStatusParams{ + Id: execution.Id, + Status: db.ExecutionStatusCompleted, + } + } + // c.run is set for the whole of Inference, the only caller path to + // callTool. + c.mu.Lock() + r := c.run + c.mu.Unlock() + + // Hold r.mu across the commit so a concurrent Subscribe's diff + // sees the row and its event together. + r.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{}) + } + r.mu.Unlock() + if err != nil { + log.ErrorContext( + ctx, + "failed to save tool result", + slog.Any("error", err), + ) + return llm.Message{}, fmt.Errorf( + "save tool result %q: %w", call.ID, err, + ) + } + return msg, nil +} + +// Stop cancels ongoing inference. It returns without waiting. +// Inference may still be finishing after Stop is called. +// Stop copies cancel to prevent canceling a new inference request. +func (c *Conversation) Stop() { + c.mu.Lock() + cancel := c.cancel + c.mu.Unlock() + + if cancel != nil { + cancel() + } +} + +// insertMessage stores msg under id, or a new version 4 UUID if id is +// nil, with its optional inference metadata at the next sequence in +// the conversation; nil metadata stores NULL. If executionUpdate is +// non-nil, its status change commits in the same transaction. The +// caller chooses the execution and a valid status transition. +func (c *Conversation) insertMessage( + ctx context.Context, + id *uuid.UUID, + msg llm.Message, + metadata *llm.Metadata, + executionUpdate *db.UpdateExecutionStatusParams, +) (*db.Message, error) { + content, err := json.Marshal(msg) + if err != nil { + return nil, fmt.Errorf("marshal content: %w", err) + } + var md jsontext.Value + if metadata != nil { + encoded, err := json.Marshal( + metadata, json.FormatNilSliceAsNull(true), + ) + if err != nil { + return nil, fmt.Errorf("marshal metadata: %w", err) + } + md = jsontext.Value(encoded) + } + + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + sequence, err := qtx.NextSequence(ctx, c.id) + if err != nil { + return nil, fmt.Errorf("get next message sequence: %w", err) + } + saved, err := qtx.InsertMessage(ctx, db.InsertMessageParams{ + Content: content, + ConversationId: c.id, + Id: id, + Metadata: md, + Role: db.MessageRole(msg.Role), + Sequence: sequence, + }) + if err != nil { + return nil, fmt.Errorf("insert message: %w", err) + } + if executionUpdate != nil { + if err := qtx.UpdateExecutionStatus( + ctx, *executionUpdate, + ); err != nil { + return nil, fmt.Errorf( + "update execution status: %w", err, + ) + } + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit message tx: %w", err) + } + + return saved, nil +} + +// InsertMessage saves a new user or assistant message at the next +// sequence, validating it first. It returns ErrBusy while Inference is +// ongoing. +func (c *Conversation) InsertMessage( + ctx context.Context, + msg llm.Message, +) (*db.Message, error) { + if msg.Role != llm.RoleAssistant && msg.Role != llm.RoleUser { + return nil, fmt.Errorf( + "invalid message role: %q", msg.Role, + ) + } + if err := msg.Validate(); err != nil { + return nil, fmt.Errorf("invalid message: %w", err) + } + + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return nil, ErrBusy + } + + return c.insertMessage(ctx, nil, msg, nil, nil) +} + +// UpdateMessage replaces a message's model-visible content; its role may +// not change. A mismatch between updatedAt and the saved updated_at +// returns ErrRevisionConflict. Inference metadata is preserved, and may +// no longer reflect initially generated output. It returns ErrBusy while +// Inference is ongoing. +func (c *Conversation) UpdateMessage( + ctx context.Context, + mId string, + updatedAt time.Time, + message llm.Message, +) (*db.Message, error) { + if err := message.Validate(); err != nil { + return nil, fmt.Errorf("invalid message: %w", err) + } + + content, err := json.Marshal(message) + if err != nil { + return nil, fmt.Errorf("marshal message: %w", err) + } + + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return nil, ErrBusy + } + + row, err := c.dbr.SelectConversationMessageId(ctx, c.id, mId) + if err != nil { + return nil, fmt.Errorf("load message %q: %w", mId, err) + } + // Refuse to update a message that has been changed. + if !row.UpdatedAt.Equal(updatedAt) { + return nil, ErrRevisionConflict + } + // Validate the role update. + if llm.Role(row.Role) != message.Role { + return nil, fmt.Errorf("invalid role: %q", message.Role) + } + + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + updated, err := qtx.UpdateMessageContent( + ctx, + db.UpdateMessageContentParams{ + Id: row.Id, + Content: content, + }, + ) + if err != nil { + return nil, fmt.Errorf("update message %q: %w", row.Id, err) + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit message tx: %w", err) + } + return updated, nil +} + +// DeleteMessage removes one saved message, checking for conflicting +// changes. A mismatch between updatedAt and the saved updated_at returns +// ErrRevisionConflict. It returns ErrBusy while Inference is ongoing. +func (c *Conversation) DeleteMessage( + ctx context.Context, + mId string, + updatedAt time.Time, +) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return ErrBusy + } + + row, err := c.dbr.SelectConversationMessageId(ctx, c.id, mId) + if err != nil { + return fmt.Errorf("load message %q: %w", mId, err) + } + + // Refuse to delete a message if it has been updated. + if !row.UpdatedAt.Equal(updatedAt) { + return ErrRevisionConflict + } + + // Delete the row and close the sequence gap it leaves, in one + // transaction. + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + if err := qtx.DeleteMessageId(ctx, row.Id); err != nil { + return fmt.Errorf("delete message %q: %w", row.Id, err) + } + if err := qtx.UpdateMessagesSequence(ctx, c.id, row.Sequence); err != nil { + return fmt.Errorf( + "shift sequences after %d: %w", row.Sequence, err, + ) + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit message tx: %w", err) + } + return nil +} + +// AddToolCall appends a tool call to a saved assistant message. +// mId identifies that assistant; updatedAt is the message's +// updated_at from the caller's last read; a mismatch returns +// ErrRevisionConflict. This edit does not execute the tool: it creates +// no execution row and no result. It returns ErrBusy while Inference is +// ongoing. +func (c *Conversation) AddToolCall( + ctx context.Context, + mId string, + updatedAt time.Time, + call llm.ToolCall, +) (*db.Message, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return nil, ErrBusy + } + + row, err := c.dbr.SelectConversationMessageId(ctx, c.id, mId) + if err != nil { + return nil, fmt.Errorf("load message %q: %w", mId, err) + } + + // Refuse to update a message that has been changed. + if !row.UpdatedAt.Equal(updatedAt) { + return nil, ErrRevisionConflict + } + if row.Role != db.MessageRoleAssistant { + return nil, fmt.Errorf("invalid role: %q", row.Role) + } + + var message llm.Message + if err := json.Unmarshal( + row.Content, &message, + ); err != nil { + return nil, fmt.Errorf( + "unmarshal message %q: %w", row.Id, err, + ) + } + + if err := message.AddToolCall(call); err != nil { + return nil, fmt.Errorf("message %q: %w", row.Id, err) + } + + content, err := json.Marshal(message) + if err != nil { + return nil, fmt.Errorf("marshal updated message: %w", err) + } + + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + updated, err := qtx.UpdateMessageContent( + ctx, + db.UpdateMessageContentParams{ + Id: row.Id, + Content: content, + }, + ) + if err != nil { + return nil, fmt.Errorf("update message %q: %w", row.Id, err) + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit message tx: %w", err) + } + return updated, nil +} + +// UpdateToolCall replaces a call in a saved assistant message. mId +// identifies the message; callId identifies the call within it. The +// replacement must keep the existing call ID: the call's execution +// record and any saved result are keyed by it. updatedAt is the +// assistant message's updated_at from the caller's last read; a mismatch +// returns ErrRevisionConflict. It returns ErrBusy while Inference is +// ongoing. +func (c *Conversation) UpdateToolCall( + ctx context.Context, + mId string, + callId string, + updatedAt time.Time, + call llm.ToolCall, +) (*db.Message, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return nil, ErrBusy + } + + row, err := c.dbr.SelectConversationMessageId(ctx, c.id, mId) + if err != nil { + return nil, fmt.Errorf("load message %q: %w", mId, err) + } + + // Refuse to update a message that has been changed. + if !row.UpdatedAt.Equal(updatedAt) { + return nil, ErrRevisionConflict + } + // Validate the role update. + if row.Role != db.MessageRoleAssistant { + return nil, fmt.Errorf("invalid role: %q", row.Role) + } + + var message llm.Message + if err := json.Unmarshal( + row.Content, &message, + ); err != nil { + return nil, fmt.Errorf( + "unmarshal message %q: %w", row.Id, err, + ) + } + + if err := message.UpdateToolCall(call); err != nil { + return nil, fmt.Errorf("message %q: %w", row.Id, err) + } + + content, err := json.Marshal(message) + if err != nil { + return nil, fmt.Errorf("marshal updated message: %w", err) + } + + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + updated, err := qtx.UpdateMessageContent( + ctx, + db.UpdateMessageContentParams{ + Id: row.Id, + Content: content, + }, + ) + if err != nil { + return nil, fmt.Errorf("update message %q: %w", row.Id, err) + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit message tx: %w", err) + } + return updated, nil +} + +// DeleteToolCall removes a call from a saved assistant message, deleting +// the call's saved result and execution record with it. mId identifies +// the parent assistant. updatedAt is the assistant's updated_at from the +// caller's last read; a mismatch returns ErrRevisionConflict. It returns +// ErrBusy while Inference is ongoing. +func (c *Conversation) DeleteToolCall( + ctx context.Context, + mId string, + callId string, + updatedAt time.Time, +) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return ErrBusy + } + + row, err := c.dbr.SelectConversationMessageId(ctx, c.id, mId) + if err != nil { + return fmt.Errorf("load assistant %q: %w", mId, err) + } + + // Refuse to delete if the assistant has been changed. + if !row.UpdatedAt.Equal(updatedAt) { + return ErrRevisionConflict + } + if row.Role != db.MessageRoleAssistant { + return fmt.Errorf("message %q is not an assistant", mId) + } + var message llm.Message + if err := json.Unmarshal( + row.Content, &message, + ); err != nil { + return fmt.Errorf( + "unmarshal assistant %q: %w", row.Id, err, + ) + } + var hasCall bool + for _, call := range message.ToolCalls { + if call.ID == callId { + hasCall = true + break + } + } + if !hasCall { + return fmt.Errorf( + "assistant %q has no call %q", row.Id, callId, + ) + } + + // Remove the call from the assistant's content. + calls := make([]llm.ToolCall, 0, len(message.ToolCalls)-1) + for _, call := range message.ToolCalls { + if call.ID != callId { + calls = append(calls, call) + } + } + message.ToolCalls = calls + content, err := json.Marshal(message) + if err != nil { + return fmt.Errorf("marshal assistant: %w", err) + } + + // Delete the call, its result, and its execution record together. + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + sequence, err := qtx.SelectMessageSequenceToolResult( + ctx, c.id, row.Sequence, + ) + if err != nil { + return fmt.Errorf("find tool result sequence: %w", err) + } + result, err := qtx.SelectConversationToolResult( + ctx, c.id, row.Sequence, sequence, callId, + ) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("load tool result %q: %w", callId, err) + } + + if _, err := qtx.UpdateMessageContent(ctx, db.UpdateMessageContentParams{ + Id: row.Id, + Content: content, + }); err != nil { + return fmt.Errorf("update assistant %q: %w", row.Id, err) + } + if result != nil { + if err := qtx.DeleteMessageId(ctx, result.Id); err != nil { + return fmt.Errorf("delete result %q: %w", result.Id, err) + } + if err := qtx.UpdateMessagesSequence( + ctx, c.id, result.Sequence, + ); err != nil { + return fmt.Errorf( + "shift sequences after %d: %w", + result.Sequence, err, + ) + } + } + if err := qtx.DeleteExecutionMessageIdToolCallId( + ctx, row.Id, callId, + ); err != nil { + return fmt.Errorf( + "delete execution for call %q: %w", callId, err, + ) + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit tool call tx: %w", err) + } + return nil +} + +// AddToolResult saves textual output for a call in the assistant message +// mId, right after that assistant's last saved tool result; later +// messages shift down by one. It does not execute the tool, and it fails +// if the call already has a result. updatedAt is the assistant message's +// updated_at from the caller's last read; a mismatch returns +// ErrRevisionConflict. The call's execution is marked completed, or +// created as completed if missing. It returns ErrBusy while Inference is +// ongoing. +func (c *Conversation) AddToolResult( + ctx context.Context, + mId string, + callId string, + updatedAt time.Time, + content string, +) (*db.Message, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return nil, ErrBusy + } + + row, err := c.dbr.SelectConversationMessageId(ctx, c.id, mId) + if err != nil { + return nil, fmt.Errorf("load assistant %q: %w", mId, err) + } + + // Refuse to add a result if the assistant has been changed. + if !row.UpdatedAt.Equal(updatedAt) { + return nil, ErrRevisionConflict + } + if row.Role != db.MessageRoleAssistant { + return nil, fmt.Errorf("message %q is not an assistant", mId) + } + var message llm.Message + if err := json.Unmarshal( + row.Content, &message, + ); err != nil { + return nil, fmt.Errorf( + "unmarshal assistant %q: %w", row.Id, err, + ) + } + var call *llm.ToolCall + for i := range message.ToolCalls { + if message.ToolCalls[i].ID == callId { + call = &message.ToolCalls[i] + break + } + } + if call == nil { + return nil, fmt.Errorf( + "assistant %q has no call %q", row.Id, callId, + ) + } + + // Build the tool result message. + result := llm.Message{ + Role: llm.RoleTool, + Name: call.Function.Name, + ToolCallID: call.ID, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: content, + }}, + } + if err := result.Validate(); err != nil { + return nil, fmt.Errorf("invalid tool result: %w", err) + } + encoded, err := json.Marshal(result) + if err != nil { + return nil, fmt.Errorf( + "marshal tool result %q: %w", callId, err, + ) + } + + // Find the position, free it, insert the result, and complete the + // execution together, so the position and duplicate check stay valid. + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + sequence, err := qtx.SelectMessageSequenceToolResult( + ctx, c.id, row.Sequence, + ) + if err != nil { + return nil, fmt.Errorf("find tool result sequence: %w", err) + } + if _, err := qtx.SelectConversationToolResult( + ctx, c.id, row.Sequence, sequence, callId, + ); err == nil { + return nil, fmt.Errorf("call %q already has a result", callId) + } else if !errors.Is(err, sql.ErrNoRows) { + return nil, fmt.Errorf("check tool result %q: %w", callId, err) + } + if err := qtx.UpdateMessagesSequenceIncrement( + ctx, c.id, sequence, + ); err != nil { + return nil, fmt.Errorf( + "shift sequences from %d: %w", sequence, err, + ) + } + msg, err := qtx.InsertMessage(ctx, db.InsertMessageParams{ + Content: encoded, + ConversationId: c.id, + Role: db.MessageRoleTool, + Sequence: sequence, + }) + if err != nil { + return nil, fmt.Errorf("save tool result %q: %w", callId, err) + } + if err := qtx.InsertExecutionCompleted( + ctx, row.Id, callId, + ); err != nil { + return nil, fmt.Errorf("complete execution %q: %w", callId, err) + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit tool result tx: %w", err) + } + return msg, nil +} + +// UpdateToolResult replaces the saved textual output for a call in the +// assistant message mId. Unlike AddToolResult, updatedAt is the tool +// result message's updated_at (not the assistant message's); a mismatch +// returns ErrRevisionConflict. It does not execute the tool, and fails +// if the call has no result. The execution is marked completed, or +// created as completed if missing, settling an interrupted call's +// running row. It returns ErrBusy while Inference is ongoing. +func (c *Conversation) UpdateToolResult( + ctx context.Context, + mId string, + callId string, + updatedAt time.Time, + content string, +) (*db.Message, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + return nil, ErrBusy + } + + row, err := c.dbr.SelectConversationMessageId(ctx, c.id, mId) + if err != nil { + return nil, fmt.Errorf("load assistant %q: %w", mId, err) + } + if row.Role != db.MessageRoleAssistant { + return nil, fmt.Errorf("message %q is not an assistant", mId) + } + var message llm.Message + if err := json.Unmarshal( + row.Content, &message, + ); err != nil { + return nil, fmt.Errorf( + "unmarshal assistant %q: %w", row.Id, err, + ) + } + var call *llm.ToolCall + for i := range message.ToolCalls { + if message.ToolCalls[i].ID == callId { + call = &message.ToolCalls[i] + break + } + } + if call == nil { + return nil, fmt.Errorf( + "assistant %q has no call %q", row.Id, callId, + ) + } + + // Build the replacement tool result message. + result := llm.Message{ + Role: llm.RoleTool, + Name: call.Function.Name, + ToolCallID: call.ID, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: content, + }}, + } + if err := result.Validate(); err != nil { + return nil, fmt.Errorf("invalid tool result: %w", err) + } + encoded, err := json.Marshal(result) + if err != nil { + return nil, fmt.Errorf( + "marshal tool result %q: %w", callId, err, + ) + } + + // Load and update the result, and complete the execution + // together, so the revision check stays valid. + tx, err := c.db.Writer.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + qtx := c.dbw.WithTx(tx) + sequence, err := qtx.SelectMessageSequenceToolResult( + ctx, c.id, row.Sequence, + ) + if err != nil { + return nil, fmt.Errorf("find tool result sequence: %w", err) + } + saved, err := qtx.SelectConversationToolResult( + ctx, c.id, row.Sequence, sequence, callId, + ) + if errors.Is(err, sql.ErrNoRows) { + return nil, fmt.Errorf("call %q has no result", callId) + } + if err != nil { + return nil, fmt.Errorf("load tool result %q: %w", callId, err) + } + // Refuse to update a result that has been changed. + if !saved.UpdatedAt.Equal(updatedAt) { + return nil, ErrRevisionConflict + } + updated, err := qtx.UpdateMessageContent( + ctx, + db.UpdateMessageContentParams{ + Id: saved.Id, + Content: encoded, + }, + ) + if err != nil { + return nil, fmt.Errorf("update tool result %q: %w", callId, err) + } + if err := qtx.InsertExecutionCompleted( + ctx, row.Id, callId, + ); err != nil { + return nil, fmt.Errorf("complete execution %q: %w", callId, err) + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit tool result tx: %w", err) + } + return updated, nil +} + +// Delete removes the saved conversation and all its messages and +// execution records. +// If inference is ongoing, it is canceled. +func (c *Conversation) Delete(ctx context.Context) error { + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + c.cancel() + } + err := c.dbw.DeleteConversationId(ctx, c.id) + if err != nil { + return fmt.Errorf("delete conversation: %w", err) + } + return nil +} diff --git a/internal/conversation/conversation_test.go b/internal/conversation/conversation_test.go new file mode 100644 index 0000000..97813ac --- /dev/null +++ b/internal/conversation/conversation_test.go @@ -0,0 +1,2489 @@ +package conversation + +import ( + "context" + "encoding/json/v2" + "errors" + "fmt" + "log/slog" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "reflect" + "sync" + "testing" + "time" + + "code.chimeric.al/chimerical/odidere/internal/config" + "code.chimeric.al/chimerical/odidere/internal/db" + "code.chimeric.al/chimerical/odidere/internal/llm" + "code.chimeric.al/chimerical/odidere/internal/tool" +) + +// newTestConversation returns a Conversation over a fresh saved +// conversation, using tools (an empty registry when nil). +func newTestConversation(t *testing.T, tools *tool.Registry) *Conversation { + t.Helper() + var ( + ctx = t.Context() + log = slog.New(slog.DiscardHandler) + ) + database, err := db.Connect(ctx, db.Config{ + Path: filepath.Join( + t.TempDir(), "conversation.db", + ), + }) + if err != nil { + t.Fatalf("Connect: %v", err) + } + t.Cleanup(func() { database.Close() }) + if err := database.ApplySchema(ctx, log); err != nil { + t.Fatalf("ApplySchema: %v", err) + } + saved, err := db.New(database.Writer).InsertConversation( + ctx, db.InsertConversationParams{UserId: "default"}, + ) + if err != nil { + t.Fatalf("InsertConversation: %v", err) + } + if tools == nil { + tools, err = tool.NewRegistry(nil) + if err != nil { + t.Fatalf("NewRegistry: %v", err) + } + } + c, err := New(Config{Id: saved.Id}, database, log, tools) + if err != nil { + t.Fatalf("New: %v", err) + } + return c +} + +// userMessage returns a minimal user message holding text. +func userMessage(text string) *llm.Message { + return &llm.Message{ + Role: llm.RoleUser, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: text, + }}, + } +} + +// newTestModelClient returns an llm client pointed at a server that +// answers with handler. +func newTestModelClient( + t *testing.T, + handler http.HandlerFunc, +) *llm.Client { + t.Helper() + srv := httptest.NewServer(handler) + t.Cleanup(srv.Close) + client, err := llm.NewClient( + llm.Config{Provider: "test", URL: srv.URL}, + slog.New(slog.DiscardHandler), + ) + if err != nil { + t.Fatalf("NewClient: %v", err) + } + return client +} + +// collectEvents drains ch until it closes and returns the events. +// It fails the test if the stream does not close in time. +func collectEvents(t *testing.T, ch <-chan Event) []Event { + t.Helper() + var events []Event + for { + select { + case e, ok := <-ch: + if !ok { + return events + } + events = append(events, e) + case <-time.After(5 * time.Second): + t.Fatalf( + "timed out waiting for the stream to close; got %v so far", + events, + ) + return nil + } + } +} + +// TestNew checks dependency validation in New. +func TestNew(t *testing.T) { + var ( + ctx = t.Context() + log = slog.New(slog.DiscardHandler) + ) + database, err := db.Connect(ctx, db.Config{ + Path: filepath.Join(t.TempDir(), "conversation.db"), + }) + if err != nil { + t.Fatalf("Connect: %v", err) + } + t.Cleanup(func() { database.Close() }) + tools, err := tool.NewRegistry(nil) + if err != nil { + t.Fatalf("NewRegistry: %v", err) + } + validCfg := Config{Id: "conv_1"} + + tests := []struct { + name string + cfg Config + sdb *db.SQLite + log *slog.Logger + tools *tool.Registry + wantErr bool + }{ + { + name: "valid", + cfg: validCfg, + sdb: database, + log: log, + tools: tools, + wantErr: false, + }, + { + name: "invalid config", + cfg: Config{}, + sdb: database, + log: log, + tools: tools, + wantErr: true, + }, + { + name: "missing database", + cfg: validCfg, + sdb: nil, + log: log, + tools: tools, + wantErr: true, + }, + { + name: "missing logger", + cfg: validCfg, + sdb: database, + log: nil, + tools: tools, + wantErr: true, + }, + { + name: "missing tools", + cfg: validCfg, + sdb: database, + log: log, + tools: nil, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c, err := New(tt.cfg, tt.sdb, tt.log, tt.tools) + if (err != nil) != tt.wantErr { + t.Fatalf( + "New() error = %v, wantErr %v", + err, tt.wantErr, + ) + } + if !tt.wantErr && c == nil { + t.Fatal("New() returned nil Conversation") + } + }) + } +} + +// TestDecodeMessages checks decoding and validating saved message rows. +func TestDecodeMessages(t *testing.T) { + validContent, err := json.Marshal(userMessage("hello")) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + invalidContent, err := json.Marshal(llm.Message{ + Role: llm.RoleUser, + ToolCallID: "call_1", + }) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + + tests := []struct { + name string + rows []db.Message + wantLen int + wantErr bool + }{ + { + name: "empty", + rows: nil, + wantLen: 0, + wantErr: false, + }, + { + name: "valid rows", + rows: []db.Message{ + {Id: "m1", Content: validContent}, + {Id: "m2", Content: validContent}, + }, + wantLen: 2, + wantErr: false, + }, + { + name: "malformed json", + rows: []db.Message{ + {Id: "m1", Content: []byte("{not json")}, + }, + wantErr: true, + }, + { + name: "invalid message", + rows: []db.Message{ + {Id: "m1", Content: invalidContent}, + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + msgs, err := DecodeMessages(tt.rows) + if (err != nil) != tt.wantErr { + t.Fatalf( + "DecodeMessages() error = %v, wantErr %v", + err, tt.wantErr, + ) + } + if tt.wantErr && msgs != nil { + t.Fatalf("msgs = %v, want nil on error", msgs) + } + if !tt.wantErr && len(msgs) != tt.wantLen { + t.Fatalf( + "len(msgs) = %d, want %d", + len(msgs), tt.wantLen, + ) + } + }) + } +} + +// TestInferenceRequestValidate checks InferenceRequest validation rules. +func TestInferenceRequestValidate(t *testing.T) { + negTwo := -2 + negOne := -1 + zero := 0 + + tests := []struct { + name string + req InferenceRequest + wantErr bool + }{ + { + name: "empty request", + req: InferenceRequest{}, + wantErr: false, + }, + { + name: "valid user message and limits", + req: InferenceRequest{ + Message: userMessage("hi"), + MaxIterations: 5, + ThinkingBudgetTokens: &negOne, + Timeout: time.Minute, + }, + wantErr: false, + }, + { + name: "zero thinking budget", + req: InferenceRequest{ + ThinkingBudgetTokens: &zero, + }, + wantErr: false, + }, + { + name: "negative max iterations", + req: InferenceRequest{ + MaxIterations: -1, + }, + wantErr: true, + }, + { + name: "thinking budget below -1", + req: InferenceRequest{ + ThinkingBudgetTokens: &negTwo, + }, + wantErr: true, + }, + { + name: "negative timeout", + req: InferenceRequest{ + Timeout: -time.Second, + }, + wantErr: true, + }, + { + name: "non-user message role", + req: InferenceRequest{ + Message: &llm.Message{ + Role: llm.RoleAssistant, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "assistant", + }}, + }, + }, + wantErr: true, + }, + { + name: "user message with tool calls", + req: InferenceRequest{ + Message: &llm.Message{ + Role: llm.RoleUser, + ToolCalls: []llm.ToolCall{{ + ID: "call_1", + }}, + }, + }, + wantErr: true, + }, + { + name: "user message with refusal", + req: InferenceRequest{ + Message: &llm.Message{ + Role: llm.RoleUser, + Refusal: "nope", + }, + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.req.Validate() + if (err != nil) != tt.wantErr { + t.Fatalf( + "Validate() error = %v, wantErr %v", + err, tt.wantErr, + ) + } + }) + } +} + +// TestInsertMessage checks saved rows and concurrent appends. +func TestInsertMessage(t *testing.T) { + tests := []struct { + name string + role llm.Role + count int + }{ + { + name: "user message", + role: llm.RoleUser, + count: 1, + }, + { + name: "assistant message", + role: llm.RoleAssistant, + count: 1, + }, + { + name: "concurrent appends", + role: llm.RoleUser, + count: 8, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + msg = llm.Message{ + Role: tt.role, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "manually added", + }}, + } + ) + + // Start callers together; each must commit once. + // Their ordering is not guaranteed. + var ( + wg sync.WaitGroup + start = make(chan struct{}) + results = make(chan *db.Message, tt.count) + ) + for range tt.count { + wg.Go(func() { + <-start + saved, err := c.InsertMessage(ctx, msg) + if err != nil { + t.Errorf( + "InsertMessage: %v", + err, + ) + return + } + results <- saved + }) + } + close(start) + wg.Wait() + close(results) + + rows, err := c.dbr.SelectMessagesConversationId( + ctx, c.id, + ) + if err != nil { + t.Fatalf( + "SelectMessagesConversationId: %v", err, + ) + } + if len(rows) != tt.count { + t.Fatalf( + "saved rows = %d, want %d", + len(rows), tt.count, + ) + } + content, err := json.Marshal(msg) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + for i, row := range rows { + if row.Id == "" { + t.Errorf("row %d has no Id", i) + } + if row.ConversationId != c.id { + t.Errorf( + "row %d conversation = %q, want %q", + i, row.ConversationId, c.id, + ) + } + if row.Sequence != i+1 { + t.Errorf( + "row %d sequence = %d, want %d", + i, row.Sequence, i+1, + ) + } + if row.Role != db.MessageRole(tt.role) { + t.Errorf( + "row %d role = %q, want %q", + i, row.Role, tt.role, + ) + } + if string(row.Content) != string(content) { + t.Errorf( + "row %d content = %q, want %s", + i, row.Content, content, + ) + } + if row.Metadata != nil { + t.Errorf( + "row %d metadata = %s, want nil", + i, row.Metadata, + ) + } + if row.CreatedAt.IsZero() { + t.Errorf("row %d has no created_at", i) + } + if row.UpdatedAt.IsZero() { + t.Errorf("row %d has no updated_at", i) + } + } + var returned int + for saved := range results { + returned++ + if saved == nil { + t.Error("append returned nil") + continue + } + index := saved.Sequence - 1 + if index < 0 || index >= len(rows) { + t.Errorf( + "sequence %d outside history", + saved.Sequence, + ) + continue + } + if !reflect.DeepEqual(rows[index], *saved) { + t.Errorf( + "row %q differs from saved row", + saved.Id, + ) + } + } + if returned != tt.count { + t.Errorf( + "returned rows = %d, want %d", + returned, tt.count, + ) + } + }) + } +} + +// TestInsertMessageErrors checks rejected appends leave no saved message. +func TestInsertMessageErrors(t *testing.T) { + tests := []struct { + name string + role llm.Role + failure string + wantErr error + }{ + { + name: "inference running", + role: llm.RoleUser, + failure: "busy", + wantErr: ErrBusy, + }, + { + name: "inference stopped but not finished", + role: llm.RoleUser, + failure: "stopped", + wantErr: ErrBusy, + }, + { + name: "invalid role", + role: llm.Role("invalid"), + failure: "role", + }, + + { + name: "user with tool calls", + role: llm.RoleUser, + failure: "calls", + }, + { + name: "tool result in empty history", + role: llm.RoleTool, + failure: "tool", + }, + { + name: "missing conversation", + role: llm.RoleUser, + failure: "missing", + }, + { + name: "canceled context", + role: llm.RoleUser, + failure: "canceled", + wantErr: context.Canceled, + }, + { + name: "closed database", + role: llm.RoleUser, + failure: "closed", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + msg = llm.Message{ + Role: tt.role, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "manually added", + }}, + } + ) + + switch tt.failure { + case "calls": + msg.ToolCalls = []llm.ToolCall{{ID: "call"}} + case "busy", "stopped": + // Register work without calling Infer's stub. + _, cancel := context.WithCancel(ctx) + defer cancel() + c.cancel = cancel + if tt.failure == "stopped" { + // Stop must leave the call registered. + c.Stop() + } + case "missing": + c.id = "missing" + case "canceled": + canceled, cancel := context.WithCancel(ctx) + cancel() + ctx = canceled + case "closed": + if err := c.db.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + } + + saved, err := c.InsertMessage(ctx, msg) + if saved != nil { + t.Fatalf("returned row = %+v, want nil", saved) + } + if err == nil { + t.Fatal("append succeeded, want an error") + } + if tt.wantErr != nil && !errors.Is(err, tt.wantErr) { + t.Fatalf("error = %v, want %v", err, tt.wantErr) + } + if tt.failure == "closed" { + return + } + + rows, err := c.dbr.SelectMessagesConversationId( + t.Context(), c.id, + ) + if err != nil { + t.Fatalf( + "SelectMessagesConversationId: %v", err, + ) + } + if len(rows) != 0 { + t.Fatalf("saved rows = %d, want 0", len(rows)) + } + if tt.failure != "role" { + return + } + + // Retry after failure: the transaction must release the + // writer without consuming a sequence number. + msg.Role = llm.RoleUser + saved, err = c.InsertMessage(ctx, msg) + if err != nil { + t.Fatalf("append after rollback: %v", err) + } + if saved.Sequence != 1 { + t.Fatalf( + "sequence = %d, want 1", saved.Sequence, + ) + } + }) + } +} + +// TestDeleteMessage checks row removal, sequence compaction, and +// execution-row cascade. +func TestDeleteMessage(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + + // Seed: user, assistant with one tool call, the call's result, + // user. + user := llm.Message{ + Role: llm.RoleUser, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "hello", + }}, + } + seed := []struct { + msg llm.Message + role db.MessageRole + }{ + {user, db.MessageRoleUser}, + { + llm.Message{ + Role: llm.RoleAssistant, + ToolCalls: []llm.ToolCall{{ + ID: "call_1", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "search", + Arguments: `{}`, + }, + }}, + }, + db.MessageRoleAssistant, + }, + { + llm.Message{ + Role: llm.RoleTool, + Name: "search", + ToolCallID: "call_1", + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "results", + }}, + }, + db.MessageRoleTool, + }, + {user, db.MessageRoleUser}, + } + var rows []db.Message + for i, s := range seed { + content, err := json.Marshal(s.msg) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + row, err := c.dbw.InsertMessage(ctx, db.InsertMessageParams{ + Content: content, + ConversationId: c.id, + Role: s.role, + Sequence: i + 1, + }) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + rows = append(rows, *row) + } + if _, err := c.dbw.InsertExecution(ctx, db.InsertExecutionParams{ + MessageId: rows[1].Id, + ToolCallId: "call_1", + Status: db.ExecutionStatusCompleted, + }); err != nil { + t.Fatalf("InsertExecution: %v", err) + } + + // Deleting the assistant cascades its execution row. + if err := c.DeleteMessage( + ctx, rows[1].Id, rows[1].UpdatedAt, + ); err != nil { + t.Fatalf("DeleteMessage: %v", err) + } + remaining, err := c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(remaining) != 3 { + t.Fatalf("saved rows = %d, want 3", len(remaining)) + } + for i, row := range remaining { + if row.Sequence != i+1 { + t.Fatalf( + "row %d sequence = %d, want %d", + i, row.Sequence, i+1, + ) + } + } + if executions, err := c.dbr.SelectExecutionsMessageId( + ctx, rows[1].Id, + ); err != nil { + t.Fatalf("SelectExecutionsMessageId: %v", err) + } else if len(executions) != 0 { + t.Fatalf("execution rows survived: %d", len(executions)) + } + + // The tool result deletes too. The first delete's shift gave it a + // new updated_at, so the revision is the re-read one. + if err := c.DeleteMessage( + ctx, rows[2].Id, remaining[1].UpdatedAt, + ); err != nil { + t.Fatalf("DeleteMessage: %v", err) + } + remaining, err = c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(remaining) != 2 { + t.Fatalf("saved rows = %d, want 2", len(remaining)) + } + for i, row := range remaining { + if row.Sequence != i+1 { + t.Fatalf( + "row %d sequence = %d, want %d", + i, row.Sequence, i+1, + ) + } + } +} + +// TestDeleteMessageErrors checks revision and busy rejection. +func TestDeleteMessageErrors(t *testing.T) { + tests := []struct { + name string + failure string + wantErr error + }{ + { + name: "inference running", + failure: "busy", + wantErr: ErrBusy, + }, + { + name: "inference stopped but not finished", + failure: "stopped", + wantErr: ErrBusy, + }, + { + name: "missing message", + failure: "missing", + }, + { + name: "missing conversation", + failure: "conversation", + }, + { + name: "stale revision", + failure: "stale", + wantErr: ErrRevisionConflict, + }, + { + name: "canceled context", + failure: "canceled", + wantErr: context.Canceled, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + conversationId := c.id + row, err := c.InsertMessage(ctx, *userMessage("hello")) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + + messageId := row.Id + revision := row.UpdatedAt + switch tt.failure { + case "busy", "stopped": + // Register work without calling Infer's stub. + _, cancel := context.WithCancel(ctx) + defer cancel() + c.cancel = cancel + if tt.failure == "stopped" { + // Stop must leave the call registered. + c.Stop() + } + case "missing": + messageId = "missing" + case "conversation": + c.id = "missing" + case "stale": + revision = row.UpdatedAt.Add(time.Hour) + case "canceled": + canceled, cancel := context.WithCancel(ctx) + cancel() + ctx = canceled + } + + err = c.DeleteMessage(ctx, messageId, revision) + if err == nil { + t.Fatal("delete succeeded, want an error") + } + if tt.wantErr != nil && !errors.Is(err, tt.wantErr) { + t.Fatalf("error = %v, want %v", err, tt.wantErr) + } + + rows, err := c.dbr.SelectMessagesConversationId( + t.Context(), conversationId, + ) + if err != nil { + t.Fatalf( + "SelectMessagesConversationId: %v", err, + ) + } + if len(rows) != 1 { + t.Fatalf("saved rows = %d, want 1", len(rows)) + } + }) + } +} + +// TestUpdateToolResult replaces a saved result's output and completes +// the call's execution, which an interrupted call leaves running. +func TestUpdateToolResult(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + + // Seed: assistant with one tool call, the call's result, user. + seed := []struct { + msg llm.Message + role db.MessageRole + }{ + { + llm.Message{ + Role: llm.RoleAssistant, + ToolCalls: []llm.ToolCall{{ + ID: "call_1", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "search", + Arguments: `{}`, + }, + }}, + }, + db.MessageRoleAssistant, + }, + { + llm.Message{ + Role: llm.RoleTool, + Name: "search", + ToolCallID: "call_1", + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "results", + }}, + }, + db.MessageRoleTool, + }, + { + llm.Message{ + Role: llm.RoleUser, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "hello", + }}, + }, + db.MessageRoleUser, + }, + } + var rows []db.Message + for i, s := range seed { + content, err := json.Marshal(s.msg) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + row, err := c.dbw.InsertMessage(ctx, db.InsertMessageParams{ + Content: content, + ConversationId: c.id, + Role: s.role, + Sequence: i + 1, + }) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + rows = append(rows, *row) + } + if _, err := c.dbw.InsertExecution(ctx, db.InsertExecutionParams{ + MessageId: rows[0].Id, + ToolCallId: "call_1", + Status: db.ExecutionStatusRunning, + }); err != nil { + t.Fatalf("InsertExecution: %v", err) + } + original := append([]db.Message(nil), rows...) + + want := llm.Message{ + Role: llm.RoleTool, + Name: "search", + ToolCallID: "call_1", + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "edited", + }}, + } + updated, err := c.UpdateToolResult( + ctx, rows[0].Id, "call_1", rows[1].UpdatedAt, "edited", + ) + if err != nil { + t.Fatalf("UpdateToolResult: %v", err) + } + + // The replacement keeps the result row's identity and position. + if updated.Id != rows[1].Id { + t.Errorf("returned row %q, want %q", updated.Id, rows[1].Id) + } + if updated.Sequence != rows[1].Sequence { + t.Errorf( + "sequence = %d, want %d", + updated.Sequence, rows[1].Sequence, + ) + } + if updated.Role != db.MessageRoleTool { + t.Errorf("role = %q, want tool", updated.Role) + } + wantContent, err := json.Marshal(want) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + if string(updated.Content) != string(wantContent) { + t.Errorf("content = %q, want %s", updated.Content, wantContent) + } + + // The saved row matches; the other rows are untouched. + rows, err = c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(rows) != 3 { + t.Fatalf("saved rows = %d, want 3", len(rows)) + } + if !reflect.DeepEqual(rows[1], *updated) { + t.Errorf("saved result row differs from returned row") + } + for i := range original { + if i == 1 { + continue + } + if !reflect.DeepEqual(rows[i], original[i]) { + t.Errorf("row %d changed", i) + } + } + + // The running execution is completed. + executions, err := c.dbr.SelectExecutionsMessageId( + ctx, rows[0].Id, + ) + if err != nil { + t.Fatalf("SelectExecutionsMessageId: %v", err) + } + if len(executions) != 1 { + t.Fatalf("execution rows = %d, want 1", len(executions)) + } + if executions[0].Status != db.ExecutionStatusCompleted { + t.Errorf( + "status = %q, want %q", + executions[0].Status, db.ExecutionStatusCompleted, + ) + } +} + +// TestUpdateToolResultErrors checks busy, missing-result, stale-revision, +// and wrong-role rejection. +func TestUpdateToolResultErrors(t *testing.T) { + tests := []struct { + name string + failure string + wantErr error + }{ + { + name: "inference running", + failure: "busy", + wantErr: ErrBusy, + }, + { + name: "no saved result", + failure: "missing", + }, + { + name: "stale revision", + failure: "stale", + wantErr: ErrRevisionConflict, + }, + { + name: "not an assistant", + failure: "role", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + + // Seed: assistant with two calls, only the first + // answered, then a user message. + call1 := llm.ToolCall{ + ID: "call_1", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "search", + Arguments: `{}`, + }, + } + call2 := llm.ToolCall{ + ID: "call_2", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "search", + Arguments: `{}`, + }, + } + seed := []struct { + msg llm.Message + role db.MessageRole + }{ + { + llm.Message{ + Role: llm.RoleAssistant, + ToolCalls: []llm.ToolCall{call1, call2}, + }, + db.MessageRoleAssistant, + }, + { + llm.Message{ + Role: llm.RoleTool, + Name: "search", + ToolCallID: "call_1", + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "results", + }}, + }, + db.MessageRoleTool, + }, + { + *userMessage("hello"), + db.MessageRoleUser, + }, + } + var rows []db.Message + for i, s := range seed { + content, err := json.Marshal(s.msg) + if err != nil { + t.Fatalf("Marshal: %v", err) + } + row, err := c.dbw.InsertMessage( + ctx, + db.InsertMessageParams{ + Content: content, + ConversationId: c.id, + Role: s.role, + Sequence: i + 1, + }, + ) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + rows = append(rows, *row) + } + + mId := rows[0].Id + callId := "call_1" + revision := rows[1].UpdatedAt + switch tt.failure { + case "busy": + // Register work without calling Infer's stub. + _, cancel := context.WithCancel(ctx) + defer cancel() + c.cancel = cancel + case "missing": + callId = "call_2" + case "stale": + revision = rows[1].UpdatedAt.Add(time.Hour) + case "role": + mId = rows[2].Id + } + + _, err := c.UpdateToolResult( + ctx, mId, callId, revision, "edited", + ) + if err == nil { + t.Fatal("update succeeded, want an error") + } + if tt.wantErr != nil && !errors.Is(err, tt.wantErr) { + t.Fatalf("error = %v, want %v", err, tt.wantErr) + } + + // The failed update changes nothing. + saved, err := c.dbr.SelectMessagesConversationId( + t.Context(), c.id, + ) + if err != nil { + t.Fatalf( + "SelectMessagesConversationId: %v", err, + ) + } + if len(saved) != 3 { + t.Fatalf("saved rows = %d, want 3", len(saved)) + } + for i := range rows { + if string(saved[i].Content) != + string(rows[i].Content) { + t.Fatalf("row %d changed", i) + } + } + }) + } +} + +// TestStopIdle checks Stop does nothing when nothing is running. +func TestStopIdle(t *testing.T) { + c := &Conversation{} + c.Stop() + c.Stop() + + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel != nil { + t.Fatal("idle Stop registered work") + } +} + +// TestStopCancellation checks Stop cancels work without marking it finished. +func TestStopCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + c := &Conversation{cancel: cancel} + + c.Stop() + c.Stop() + + if ctx.Err() != context.Canceled { + t.Fatalf("context error = %v, want context.Canceled", ctx.Err()) + } + c.mu.Lock() + defer c.mu.Unlock() + if c.cancel == nil { + t.Fatal("Stop released the inference reservation") + } +} + +// TestConfigValidate checks Config validation rules. +func TestConfigValidate(t *testing.T) { + tests := []struct { + name string + config Config + wantErr bool + }{ + { + name: "missing conversation id", + config: Config{}, + wantErr: true, + }, + { + name: "valid conversation id", + config: Config{ + Id: "35afb18e-3ee1-4053-9519-1efba4d2932c", + }, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := tt.config.Validate() + if (err != nil) != tt.wantErr { + t.Errorf( + "Validate() error = %v, wantErr %v", + err, tt.wantErr, + ) + } + }) + } +} + +// TestSubscribeIdle checks the diff with no ongoing Inference: +// missing, changed, and deleted messages, and no Done. +func TestSubscribeIdle(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + var saved []*db.Message + for _, text := range []string{"one", "two", "three"} { + row, err := c.InsertMessage(ctx, *userMessage(text)) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + saved = append(saved, row) + } + manifest := func(rows ...*db.Message) []MessageManifest { + m := make([]MessageManifest, len(rows)) + for i, row := range rows { + m[i] = MessageManifest{ + Id: row.Id, + UpdatedAt: row.UpdatedAt, + } + } + return m + } + + tests := []struct { + name string + prepare func() []MessageManifest + want []EventType + wantText []string + }{ + { + name: "empty manifest", + prepare: func() []MessageManifest { return nil }, + want: []EventType{ + EventMessageSaved, + EventMessageSaved, + EventMessageSaved, + }, + wantText: []string{"one", "two", "three"}, + }, + { + name: "up to date", + prepare: func() []MessageManifest { + return manifest(saved...) + }, + want: nil, + }, + { + name: "missing middle", + prepare: func() []MessageManifest { + return manifest(saved[0], saved[2]) + }, + want: []EventType{EventMessageSaved}, + wantText: []string{"two"}, + }, + { + name: "stale entry", + prepare: func() []MessageManifest { + m := manifest(saved...) + m[0].UpdatedAt = time.Time{} + return m + }, + want: []EventType{EventMessageSaved}, + wantText: []string{"one"}, + }, + { + name: "deleted", + prepare: func() []MessageManifest { + // Sleep so the shift's updated_at is a later + // millisecond than the inserts' updated_at. + time.Sleep(10 * time.Millisecond) + if err := c.DeleteMessage( + ctx, saved[1].Id, saved[1].UpdatedAt, + ); err != nil { + t.Fatalf("DeleteMessage: %v", err) + } + return manifest(saved...) + }, + // Deleting "two" moves "three" into its slot. The moved + // row gets a new updated_at, so the diff reports it as + // saved before the deletion. + want: []EventType{ + EventMessageSaved, + EventMessageDeleted, + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ch, err := c.Subscribe(ctx, tt.prepare()) + if err != nil { + t.Fatalf("Subscribe: %v", err) + } + got := collectEvents(t, ch) + if len(got) != len(tt.want) { + t.Fatalf( + "got %d events, want %d: %v", + len(got), len(tt.want), got, + ) + } + for i, e := range got { + if e.Type() != tt.want[i] { + t.Errorf( + "event %d type = %v, want %v", + i, e.Type(), tt.want[i], + ) + } + } + for i, wantText := range tt.wantText { + s, ok := got[i].(MessageSaved) + if !ok { + t.Fatalf( + "event %d = %v, want MessageSaved", + i, got[i], + ) + } + var m llm.Message + if err := json.Unmarshal( + s.Message.Content, &m, + ); err != nil { + t.Fatalf( + "unmarshal message %q: %v", + s.Message.Id, err, + ) + } + if m.Text() != wantText { + t.Errorf( + "event %d text = %q, want %q", + i, m.Text(), wantText, + ) + } + } + }) + } +} + +// TestSubscribeInference checks a client joining an ongoing streaming +// Inference: the diff, the chunks, the saved assistant message, and +// Done with no error. +func TestSubscribeInference(t *testing.T) { + var ( + ctx = t.Context() + firstChunk = make(chan struct{}) + release = make(chan struct{}) + ) + client := newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + flusher := w.(http.Flusher) + fmt.Fprint( + w, + "data: {\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"}}]}\n\n", + ) + fmt.Fprint( + w, + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n", + ) + flusher.Flush() + close(firstChunk) + // End the handler if the test aborts, so server cleanup + // cannot hang on this request. + select { + case <-release: + case <-r.Context().Done(): + } + fmt.Fprint( + w, + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\" world\"}}]}\n\n", + ) + fmt.Fprint( + w, + "data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + ) + fmt.Fprint(w, "data: [DONE]\n\n") + }) + + var ( + c = newTestConversation(t, nil) + errs = make(chan error, 1) + ) + go func() { + errs <- c.Inference(ctx, client, InferenceRequest{ + Message: userMessage("hi"), + Model: "test-model", + Stream: true, + }) + }() + select { + case <-firstChunk: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the first chunk") + } + + // Join while the Inference is mid-stream. + ch, err := c.Subscribe(ctx, nil) + if err != nil { + t.Fatalf("Subscribe: %v", err) + } + close(release) + got := collectEvents(t, ch) + + // The diff delivers the user message first. + saved, ok := got[0].(MessageSaved) + if !ok || saved.Message.Role != db.MessageRoleUser { + t.Fatalf("first event = %v, want the user MessageSaved", got[0]) + } + var ( + chunkTexts []string + chunkIDs []string + asst *MessageSaved + ) + for _, e := range got[1:] { + switch e := e.(type) { + case AssistantUpdate: + if asst != nil { + t.Fatalf("AssistantUpdate after the assistant MessageSaved: %v", e) + } + chunkIDs = append(chunkIDs, e.MessageId) + chunkTexts = append( + chunkTexts, e.Chunk.Choices[0].Delta.Content, + ) + case MessageSaved: + if e.Message.Role != db.MessageRoleAssistant { + t.Fatalf("unexpected saved role %q", e.Message.Role) + } + asst = &e + case Done: + if e.Err != nil { + t.Errorf("Done.Err = %v, want nil", e.Err) + } + default: + t.Fatalf("unexpected event %v", e) + } + } + if asst == nil { + t.Fatalf("no assistant MessageSaved in %v", got) + } + var asstMsg llm.Message + if err := json.Unmarshal( + asst.Message.Content, &asstMsg, + ); err != nil { + t.Fatalf("unmarshal assistant message: %v", err) + } + if asstMsg.Text() != "hello world" { + t.Errorf("assistant text = %q, want %q", + asstMsg.Text(), "hello world") + } + if len(chunkTexts) == 0 { + t.Fatal("no AssistantUpdate events") + } + for i, id := range chunkIDs { + if id != asst.Message.Id { + t.Errorf("chunk %d message id = %q, want the assistant id %q", + i, id, asst.Message.Id) + } + } + if last := got[len(got)-1]; last.Type() != EventDone { + t.Fatalf("last event type = %v, want done", last.Type()) + } + if err := <-errs; err != nil { + t.Fatalf("Inference: %v", err) + } +} + +// TestSubscribeInferenceError checks that a dropped stream ends the +// subscriber with Done carrying the error, and saves no assistant +// message. +func TestSubscribeInferenceError(t *testing.T) { + var ( + ctx = t.Context() + firstChunk = make(chan struct{}) + release = make(chan struct{}) + ) + client := newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint( + w, + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n", + ) + w.(http.Flusher).Flush() + close(firstChunk) + // End the handler if the test aborts, so server cleanup + // cannot hang on this request. + select { + case <-release: + case <-r.Context().Done(): + } + // Returning without [DONE] drops the connection. + }) + + var ( + c = newTestConversation(t, nil) + errs = make(chan error, 1) + ) + go func() { + errs <- c.Inference(ctx, client, InferenceRequest{ + Message: userMessage("hi"), + Model: "test-model", + Stream: true, + }) + }() + select { + case <-firstChunk: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the first chunk") + } + + ch, err := c.Subscribe(ctx, nil) + if err != nil { + t.Fatalf("Subscribe: %v", err) + } + close(release) + got := collectEvents(t, ch) + + done, ok := got[len(got)-1].(Done) + if !ok { + t.Fatalf("last event = %v, want Done", got[len(got)-1]) + } + if done.Err == nil { + t.Fatal("Done.Err = nil, want the stream error") + } + rows, err := db.New(c.db.Reader). + SelectMessagesConversationId(ctx, c.id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + for _, row := range rows { + if row.Role != db.MessageRoleUser { + t.Fatalf("unexpected saved role %q", row.Role) + } + } + if err := <-errs; err == nil { + t.Fatal("Inference error = nil, want the stream error") + } +} + +// TestSubscribeToolResult checks that tool results reach subscribers. +func TestSubscribeToolResult(t *testing.T) { + var ( + ctx = t.Context() + requested = make(chan struct{}) + subscribed = make(chan struct{}) + ) + registry, err := tool.NewRegistry([]config.ToolConfig{{ + Name: "echo", + Description: "echoes its arguments", + Command: "/bin/echo", + Arguments: []string{"hi"}, + }}) + if err != nil { + t.Fatalf("NewRegistry: %v", err) + } + + var request int + client := newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + request++ + w.Header().Set("Content-Type", "application/json") + switch request { + case 1: + close(requested) + // End the handler if the test aborts, so server + // cleanup cannot hang on this request. + select { + case <-subscribed: + case <-r.Context().Done(): + } + fmt.Fprint( + w, + `{"choices":[{"message":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"echo","arguments":"{}"}}]}}]}`, + ) + case 2: + fmt.Fprint( + w, + `{"choices":[{"message":{"role":"assistant","content":"done"}}]}`, + ) + } + }) + + var ( + c = newTestConversation(t, registry) + errs = make(chan error, 1) + ) + go func() { + errs <- c.Inference(ctx, client, InferenceRequest{ + Message: userMessage("hi"), + Model: "test-model", + }) + }() + select { + case <-requested: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the first request") + } + + // Join while the Inference awaits the tool-call answer. + ch, err := c.Subscribe(ctx, nil) + if err != nil { + t.Fatalf("Subscribe: %v", err) + } + close(subscribed) + got := collectEvents(t, ch) + + var ( + aId string + sawTool bool + ) + for i, e := range got { + switch e := e.(type) { + case MessageSaved: + switch e.Message.Role { + case db.MessageRoleAssistant: + if aId == "" { + aId = e.Message.Id + } + case db.MessageRoleTool: + sawTool = true + var m llm.Message + if err := json.Unmarshal( + e.Message.Content, &m, + ); err != nil { + t.Fatalf( + "unmarshal message %q: %v", e.Message.Id, err, + ) + } + if m.Text() == "" { + t.Errorf("event %d: empty tool result", i) + } + default: + if i != 0 { + t.Errorf("event %d role = %q, want the user first", i, e.Message.Role) + } + } + case Done: + if e.Err != nil { + t.Errorf("Done.Err = %v, want nil", e.Err) + } + default: + t.Errorf("event %d = %v, want MessageSaved or Done", i, e) + } + } + if aId == "" { + t.Fatalf("no assistant MessageSaved in %v", got) + } + if !sawTool { + t.Fatalf("no tool MessageSaved in %v", got) + } + if err := <-errs; err != nil { + t.Fatalf("Inference: %v", err) + } +} + +// TestSubscribeCancel checks that canceling the subscriber's context +// closes the channel without Done, while Inference keeps running. +func TestSubscribeCancel(t *testing.T) { + var ( + ctx = t.Context() + firstChunk = make(chan struct{}) + ) + client := newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + fmt.Fprint( + w, + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n", + ) + w.(http.Flusher).Flush() + close(firstChunk) + // Hold the stream open until the request is canceled. + <-r.Context().Done() + }, + ) + + var ( + c = newTestConversation(t, nil) + errs = make(chan error, 1) + ) + go func() { + errs <- c.Inference(ctx, client, InferenceRequest{ + Message: userMessage("hi"), + Model: "test-model", + Stream: true, + }) + }() + select { + case <-firstChunk: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the first chunk") + } + + subCtx, cancel := context.WithCancel(ctx) + ch, err := c.Subscribe(subCtx, nil) + if err != nil { + t.Fatalf("Subscribe: %v", err) + } + cancel() + + // Delivery stops and ch closes without Done. + for _, e := range collectEvents(t, ch) { + if e.Type() == EventDone { + t.Fatalf("received Done after cancel: %v", e) + } + } + + // Inference kept running; Stop ends it. + c.Stop() + if err := <-errs; err == nil { + t.Fatal("Inference error = nil, want the cancellation") + } +} + +// TestSubscribeStartupError checks that a failure resolving the saved +// history's tool calls ends subscribers with Done carrying the error, +// not a success. +func TestSubscribeStartupError(t *testing.T) { + var ( + ctx = t.Context() + dir = t.TempDir() + started = filepath.Join(dir, "started") + release = filepath.Join(dir, "release") + script = filepath.Join(dir, "slow.sh") + ) + // The tool holds the Inference open until the test releases it, + // so the subscriber can join before the failure is reported. + if err := os.WriteFile(script, []byte( + "#!/bin/sh\ntouch "+started+"\nwhile [ ! -f "+release+ + " ]; do sleep 0.1; done\n", + ), 0o755); err != nil { + t.Fatalf("WriteFile: %v", err) + } + registry, err := tool.NewRegistry([]config.ToolConfig{{ + Name: "slow", + Description: "blocks until released", + Command: script, + }}) + if err != nil { + t.Fatalf("NewRegistry: %v", err) + } + + var ( + c = newTestConversation(t, registry) + msg = llm.Message{ + Role: llm.RoleAssistant, + ToolCalls: []llm.ToolCall{ + { + ID: "call_slow", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "slow", + Arguments: "{}", + }, + }, + { + ID: "call_boom", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "boom", + Arguments: "{}", + }, + }, + }, + } + ) + row, err := c.InsertMessage(ctx, msg) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + // call_boom "finished" earlier without a saved result, so it + // must be refused; call_slow keeps the resolution open. + if _, err := c.dbw.InsertExecution( + ctx, + db.InsertExecutionParams{ + MessageId: row.Id, + ToolCallId: "call_boom", + Status: db.ExecutionStatusCompleted, + }, + ); err != nil { + t.Fatalf("InsertExecution: %v", err) + } + + // The client is never used: the failure precedes any request. + client := newTestModelClient(t, func( + http.ResponseWriter, *http.Request, + ) { + }) + errs := make(chan error, 1) + go func() { + errs <- c.Inference(ctx, client, InferenceRequest{ + Model: "test-model", + }) + }() + + deadline := time.Now().Add(5 * time.Second) + for { + if _, err := os.Stat(started); err == nil { + break + } + if time.Now().After(deadline) { + t.Fatal("timed out waiting for the slow tool to start") + } + time.Sleep(10 * time.Millisecond) + } + + ch, err := c.Subscribe(ctx, nil) + if err != nil { + t.Fatalf("Subscribe: %v", err) + } + if err := os.WriteFile(release, nil, 0o644); err != nil { + t.Fatalf("WriteFile: %v", err) + } + got := collectEvents(t, ch) + + first, ok := got[0].(MessageSaved) + if !ok || first.Message.Role != db.MessageRoleAssistant { + t.Fatalf("first event = %v, want the assistant MessageSaved", got[0]) + } + var sawTool bool + for _, e := range got { + switch e := e.(type) { + case MessageSaved: + if e.Message.Role != db.MessageRoleTool { + continue + } + sawTool = true + case Done: + if e.Err == nil { + t.Error("Done.Err = nil, want the startup failure") + } + } + } + if !sawTool { + t.Errorf("no tool MessageSaved in %v", got) + } + if last := got[len(got)-1]; last.Type() != EventDone { + t.Fatalf("last event type = %v, want done", last.Type()) + } + if err := <-errs; err == nil { + t.Fatal("Inference error = nil, want the startup failure") + } +} + +// TestUpdateMessage checks replacing a saved message's content while +// preserving its identity, sequence, and inference metadata. +func TestUpdateMessage(t *testing.T) { + model := "test-model" + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + msg = llm.Message{ + Role: llm.RoleAssistant, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "initial", + }}, + } + meta = llm.Metadata{Model: &model} + ) + saved, err := c.insertMessage(ctx, nil, msg, &meta, nil) + if err != nil { + t.Fatalf("insertMessage: %v", err) + } + + replacement := llm.Message{ + Role: llm.RoleAssistant, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "updated", + }}, + } + updated, err := c.UpdateMessage( + ctx, saved.Id, saved.UpdatedAt, replacement, + ) + if err != nil { + t.Fatalf("UpdateMessage: %v", err) + } + if updated.Id != saved.Id { + t.Errorf("Id = %q, want %q", updated.Id, saved.Id) + } + if updated.Sequence != saved.Sequence { + t.Errorf( + "Sequence = %d, want %d", + updated.Sequence, saved.Sequence, + ) + } + if string(updated.Metadata) != string(saved.Metadata) { + t.Errorf( + "Metadata = %s, want %s", + updated.Metadata, saved.Metadata, + ) + } + var decoded llm.Message + if err := json.Unmarshal(updated.Content, &decoded); err != nil { + t.Fatalf("Unmarshal: %v", err) + } + if decoded.Text() != "updated" { + t.Errorf("Text() = %q, want %q", decoded.Text(), "updated") + } +} + +// TestUpdateMessageErrors checks busy, revision, role-change, and +// validation rejection in UpdateMessage. +func TestUpdateMessageErrors(t *testing.T) { + tests := []struct { + name string + failure string + wantErr error + }{ + { + name: "inference running", + failure: "busy", + wantErr: ErrBusy, + }, + { + name: "stale revision", + failure: "stale", + wantErr: ErrRevisionConflict, + }, + { + name: "role change", + failure: "role", + }, + { + name: "invalid message", + failure: "invalid", + }, + { + name: "missing message", + failure: "missing", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + saved, err := c.InsertMessage( + ctx, *userMessage("original"), + ) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + + mId := saved.Id + revision := saved.UpdatedAt + replacement := *userMessage("replacement") + + switch tt.failure { + case "busy": + _, cancel := context.WithCancel(ctx) + defer cancel() + c.cancel = cancel + case "stale": + revision = saved.UpdatedAt.Add(time.Hour) + case "role": + replacement.Role = llm.RoleAssistant + case "invalid": + replacement.ToolCallID = "call_1" + case "missing": + mId = "missing" + } + + _, err = c.UpdateMessage( + ctx, mId, revision, replacement, + ) + if err == nil { + t.Fatal("UpdateMessage succeeded, want error") + } + if tt.wantErr != nil && !errors.Is(err, tt.wantErr) { + t.Fatalf("error = %v, want %v", err, tt.wantErr) + } + }) + } +} + +// TestToolCallLifecycle checks AddToolCall, UpdateToolCall, and +// DeleteToolCall on a saved assistant message, including cascading +// removal of the call's tool result and execution row. +func TestToolCallLifecycle(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + asst, err := c.InsertMessage(ctx, llm.Message{ + Role: llm.RoleAssistant, + ContentParts: []llm.ContentPart{{ + Type: llm.ContentTypeText, + Text: "checking", + }}, + }) + if err != nil { + t.Fatalf("InsertMessage: %v", err) + } + + // Add a tool call. + call := llm.ToolCall{ + ID: "call_1", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "search", + Arguments: `{"q":"initial"}`, + }, + } + asst, err = c.AddToolCall(ctx, asst.Id, asst.UpdatedAt, call) + if err != nil { + t.Fatalf("AddToolCall: %v", err) + } + + // Update the tool call's arguments. + call.Function.Arguments = `{"q":"updated"}` + asst, err = c.UpdateToolCall( + ctx, asst.Id, "call_1", asst.UpdatedAt, call, + ) + if err != nil { + t.Fatalf("UpdateToolCall: %v", err) + } + var decoded llm.Message + if err := json.Unmarshal(asst.Content, &decoded); err != nil { + t.Fatalf("Unmarshal: %v", err) + } + if len(decoded.ToolCalls) != 1 || + decoded.ToolCalls[0].Function.Arguments != `{"q":"updated"}` { + t.Fatalf("ToolCalls = %+v", decoded.ToolCalls) + } + + // Add a tool result and a trailing user message so DeleteToolCall + // must remove the result and compact sequences. + if _, err := c.AddToolResult( + ctx, asst.Id, "call_1", asst.UpdatedAt, "search output", + ); err != nil { + t.Fatalf("AddToolResult: %v", err) + } + if _, err := c.InsertMessage( + ctx, *userMessage("follow-up"), + ); err != nil { + t.Fatalf("InsertMessage: %v", err) + } + asstRow, err := c.dbr.SelectConversationMessageId( + ctx, c.id, asst.Id, + ) + if err != nil { + t.Fatalf("SelectConversationMessageId: %v", err) + } + + // Delete the tool call: its result and execution record must be + // removed, and the trailing user message shifts to sequence 2. + if err := c.DeleteToolCall( + ctx, asst.Id, "call_1", asstRow.UpdatedAt, + ); err != nil { + t.Fatalf("DeleteToolCall: %v", err) + } + rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(rows) != 2 { + t.Fatalf("saved rows = %d, want 2", len(rows)) + } + if rows[1].Role != db.MessageRoleUser || rows[1].Sequence != 2 { + t.Errorf( + "row 1 = (role %q, seq %d), want (user, 2)", + rows[1].Role, rows[1].Sequence, + ) + } + executions, err := c.dbr.SelectExecutionsMessageId(ctx, asst.Id) + if err != nil { + t.Fatalf("SelectExecutionsMessageId: %v", err) + } + if len(executions) != 0 { + t.Errorf("executions = %d, want 0", len(executions)) + } +} + +// TestToolCallErrors checks error cases for AddToolCall, UpdateToolCall, +// and DeleteToolCall. +func TestToolCallErrors(t *testing.T) { + validCall := llm.ToolCall{ + ID: "call_1", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "search", + Arguments: `{}`, + }, + } + + type runFunc func( + ctx context.Context, + c *Conversation, + asst, user *db.Message, + ) error + + tests := []struct { + name string + run runFunc + wantErr error + }{ + { + name: "AddToolCall busy", + run: func( + ctx context.Context, + c *Conversation, + asst, _ *db.Message, + ) error { + _, cancel := context.WithCancel(ctx) + defer cancel() + c.cancel = cancel + _, err := c.AddToolCall( + ctx, asst.Id, asst.UpdatedAt, validCall, + ) + return err + }, + wantErr: ErrBusy, + }, + { + name: "AddToolCall stale revision", + run: func( + ctx context.Context, + c *Conversation, + asst, _ *db.Message, + ) error { + _, err := c.AddToolCall( + ctx, + asst.Id, + asst.UpdatedAt.Add(time.Hour), + validCall, + ) + return err + }, + wantErr: ErrRevisionConflict, + }, + { + name: "AddToolCall non-assistant role", + run: func( + ctx context.Context, + c *Conversation, + _, user *db.Message, + ) error { + _, err := c.AddToolCall( + ctx, user.Id, user.UpdatedAt, validCall, + ) + return err + }, + }, + { + name: "AddToolCall duplicate call id", + run: func( + ctx context.Context, + c *Conversation, + asst, _ *db.Message, + ) error { + _, err := c.AddToolCall( + ctx, asst.Id, asst.UpdatedAt, validCall, + ) + return err + }, + }, + { + name: "UpdateToolCall stale revision", + run: func( + ctx context.Context, + c *Conversation, + asst, _ *db.Message, + ) error { + _, err := c.UpdateToolCall( + ctx, + asst.Id, + "call_1", + asst.UpdatedAt.Add(time.Hour), + validCall, + ) + return err + }, + wantErr: ErrRevisionConflict, + }, + { + name: "DeleteToolCall missing call", + run: func( + ctx context.Context, + c *Conversation, + asst, _ *db.Message, + ) error { + return c.DeleteToolCall( + ctx, asst.Id, "missing", asst.UpdatedAt, + ) + }, + }, + { + name: "DeleteToolCall stale revision", + run: func( + ctx context.Context, + c *Conversation, + asst, _ *db.Message, + ) error { + return c.DeleteToolCall( + ctx, + asst.Id, + "call_1", + asst.UpdatedAt.Add(time.Hour), + ) + }, + wantErr: ErrRevisionConflict, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + asst, err := c.InsertMessage(ctx, llm.Message{ + Role: llm.RoleAssistant, + ToolCalls: []llm.ToolCall{validCall}, + }) + if err != nil { + t.Fatalf("InsertMessage assistant: %v", err) + } + user, err := c.InsertMessage(ctx, *userMessage("hi")) + if err != nil { + t.Fatalf("InsertMessage user: %v", err) + } + + err = tt.run(ctx, c, asst, user) + if err == nil { + t.Fatal("operation succeeded, want error") + } + if tt.wantErr != nil && !errors.Is(err, tt.wantErr) { + t.Fatalf("error = %v, want %v", err, tt.wantErr) + } + }) + } +} + +// TestAddToolResult checks inserting tool results right after an +// assistant message, shifting later messages down, and rejecting +// duplicate results or stale revisions. +func TestAddToolResult(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + ) + asst, err := c.InsertMessage(ctx, llm.Message{ + Role: llm.RoleAssistant, + ToolCalls: []llm.ToolCall{{ + ID: "call_1", + Type: llm.ToolTypeFunction, + Function: llm.Function{ + Name: "search", + Arguments: `{}`, + }, + }}, + }) + if err != nil { + t.Fatalf("InsertMessage assistant: %v", err) + } + user, err := c.InsertMessage(ctx, *userMessage("later")) + if err != nil { + t.Fatalf("InsertMessage user: %v", err) + } + + res, err := c.AddToolResult( + ctx, asst.Id, "call_1", asst.UpdatedAt, "found it", + ) + if err != nil { + t.Fatalf("AddToolResult: %v", err) + } + if res.Sequence != 2 { + t.Errorf("result sequence = %d, want 2", res.Sequence) + } + + // A second result for the same call is rejected. + if _, err := c.AddToolResult( + ctx, asst.Id, "call_1", asst.UpdatedAt, "duplicate", + ); err == nil { + t.Fatal("duplicate AddToolResult succeeded, want error") + } + + // Stale revision is rejected. + if _, err := c.AddToolResult( + ctx, + asst.Id, + "call_1", + asst.UpdatedAt.Add(time.Hour), + "stale", + ); !errors.Is(err, ErrRevisionConflict) { + t.Fatalf("error = %v, want ErrRevisionConflict", err) + } + + rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(rows) != 3 { + t.Fatalf("saved rows = %d, want 3", len(rows)) + } + if rows[2].Id != user.Id || rows[2].Sequence != 3 { + t.Errorf( + "trailing user = (id %q, seq %d), want (%q, 3)", + rows[2].Id, rows[2].Sequence, user.Id, + ) + } +} + +// TestDeleteConversation checks that Delete cancels active inference and +// removes the conversation and its messages. +func TestDeleteConversation(t *testing.T) { + var ( + ctx = t.Context() + c = newTestConversation(t, nil) + subCtx context.Context + cancel context.CancelFunc + ) + if _, err := c.InsertMessage(ctx, *userMessage("hello")); err != nil { + t.Fatalf("InsertMessage: %v", err) + } + subCtx, cancel = context.WithCancel(ctx) + defer cancel() + c.cancel = cancel + + if err := c.Delete(ctx); err != nil { + t.Fatalf("Delete: %v", err) + } + if subCtx.Err() != context.Canceled { + t.Errorf( + "active inference context = %v, want Canceled", + subCtx.Err(), + ) + } + rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(rows) != 0 { + t.Errorf("remaining messages = %d, want 0", len(rows)) + } +} + +// TestInferenceApprovalAndMaxIterations checks that tools requiring +// approval pause with ErrApprovalRequired until approved, and that +// MaxIterations caps the model loop. +func TestInferenceApprovalAndMaxIterations(t *testing.T) { + var ( + ctx = t.Context() + ) + registry, err := tool.NewRegistry([]config.ToolConfig{{ + Name: "guarded", + Description: "requires approval", + Command: "/bin/echo", + Arguments: []string{"approved-output"}, + RequiresApproval: true, + }}) + if err != nil { + t.Fatalf("NewRegistry: %v", err) + } + c := newTestConversation(t, registry) + + var request int + client := newTestModelClient( + t, + func(w http.ResponseWriter, r *http.Request) { + request++ + w.Header().Set("Content-Type", "application/json") + if request == 1 { + fmt.Fprint( + w, + `{"choices":[{"message":{"role":"assistant","tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"guarded","arguments":"{}"}}]}}]}`, + ) + return + } + // Second request returns an unknown tool (which does + // not require approval) so the 1-iteration loop hits + // ErrMaxIterations. + fmt.Fprint( + w, + `{"choices":[{"message":{"role":"assistant","tool_calls":[{"index":0,"id":"call_2","type":"function","function":{"name":"other","arguments":"{}"}}]}}]}`, + ) + }, + ) + + // First call: saves the assistant message and pending execution, + // returning ErrApprovalRequired. + err = c.Inference(ctx, client, InferenceRequest{ + Message: userMessage("run guarded"), + Model: "test-model", + }) + if !errors.Is(err, ErrApprovalRequired) { + t.Fatalf("Inference error = %v, want ErrApprovalRequired", err) + } + + // Approve the execution in the database, then run with + // MaxIterations = 1 so the model's second tool call hits + // ErrMaxIterations. + rows, err := c.dbr.SelectMessagesConversationId(ctx, c.id) + if err != nil || len(rows) != 2 { + t.Fatalf("rows = %d (err %v), want 2", len(rows), err) + } + executions, err := c.dbr.SelectExecutionsMessageId( + ctx, rows[1].Id, + ) + if err != nil || len(executions) != 1 { + t.Fatalf( + "executions = %d (err %v), want 1", + len(executions), err, + ) + } + if err := c.dbw.UpdateExecutionStatus( + ctx, + db.UpdateExecutionStatusParams{ + Id: executions[0].Id, + Status: db.ExecutionStatusApproved, + }, + ); err != nil { + t.Fatalf("UpdateExecutionStatus: %v", err) + } + + err = c.Inference(ctx, client, InferenceRequest{ + Model: "test-model", + MaxIterations: 1, + }) + if !errors.Is(err, ErrMaxIterations) { + t.Fatalf("Inference error = %v, want ErrMaxIterations", err) + } +} diff --git a/internal/db/models.go b/internal/db/models.go index b6ef9d1..03a15de 100644 --- a/internal/db/models.go +++ b/internal/db/models.go @@ -90,8 +90,8 @@ const ( // appears in the OpenAI Chat Completions messages array. type Message struct { // Content is the model-visible message JSON, without inference - // metadata. - Content string `json:"content"` + // metadata. The schema's check keeps it valid JSON. + Content jsontext.Value `json:"content"` // ConversationId identifies the conversation that owns this message // row. ConversationId string `json:"conversation_id"` diff --git a/internal/db/queries.go b/internal/db/queries.go index acdd1f8..57cbc4d 100644 --- a/internal/db/queries.go +++ b/internal/db/queries.go @@ -310,7 +310,8 @@ RETURNING content, conversation_id, created_at, id, json(metadata), role, sequen // InsertMessageParams holds arguments for InsertMessage. type InsertMessageParams struct { // Content is the model-visible message JSON saved in the new row. - Content string + // Malformed JSON causes insertion to fail. + Content jsontext.Value // ConversationId identifies the conversation that owns the new row. ConversationId string // Id is the new row's UUID. Nil generates a version 4 UUID. @@ -340,7 +341,7 @@ func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (* ctx, insertMessage, id.String(), - arg.Content, + string(arg.Content), arg.ConversationId, metadataArg, arg.Role, @@ -349,13 +350,14 @@ func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (* var ( m Message + content []byte metadata []byte createdAt string updatedAt string err error ) if err := row.Scan( - &m.Content, + &content, &m.ConversationId, &createdAt, &m.Id, @@ -366,6 +368,7 @@ func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (* ); err != nil { return nil, fmt.Errorf("InsertMessage: %w", err) } + m.Content = jsontext.Value(content) m.Metadata = jsontext.Value(metadata) if m.CreatedAt, err = time.Parse(time.RFC3339, createdAt); err != nil { return nil, fmt.Errorf("InsertMessage: parse created_at: %w", err) @@ -376,6 +379,58 @@ func (q *Queries) InsertMessage(ctx context.Context, arg InsertMessageParams) (* return &m, nil } +const selectConversationMessageId = ` +SELECT content, conversation_id, created_at, id, json(metadata), role, sequence, updated_at +FROM messages +WHERE conversation_id = ? AND id = ? +` + +// SelectConversationMessageId returns a conversation's message by ID. +func (q *Queries) SelectConversationMessageId( + ctx context.Context, + conversationId string, + id string, +) (*Message, error) { + row := q.db.QueryRowContext( + ctx, selectConversationMessageId, conversationId, id, + ) + + var ( + m Message + content []byte + metadata []byte + createdAt string + updatedAt string + ) + err := row.Scan( + &content, + &m.ConversationId, + &createdAt, + &m.Id, + &metadata, + &m.Role, + &m.Sequence, + &updatedAt, + ) + if err != nil { + return nil, fmt.Errorf("SelectConversationMessageId: %w", err) + } + m.Content = jsontext.Value(content) + m.Metadata = jsontext.Value(metadata) + if m.CreatedAt, err = time.Parse(time.RFC3339, createdAt); err != nil { + return nil, fmt.Errorf( + "SelectConversationMessageId: parse created_at: %w", err, + ) + } + if m.UpdatedAt, err = time.Parse(time.RFC3339, updatedAt); err != nil { + return nil, fmt.Errorf( + "SelectConversationMessageId: parse updated_at: %w", err, + ) + } + + return &m, nil +} + const selectMessagesConversationId = ` SELECT content, conversation_id, created_at, id, json(metadata), role, sequence, updated_at FROM messages @@ -396,12 +451,13 @@ func (q *Queries) SelectMessagesConversationId(ctx context.Context, conversation for rows.Next() { var ( m Message + content []byte metadata []byte createdAt string updatedAt string ) if err := rows.Scan( - &m.Content, + &content, &m.ConversationId, &createdAt, &m.Id, @@ -412,6 +468,7 @@ func (q *Queries) SelectMessagesConversationId(ctx context.Context, conversation ); err != nil { return nil, fmt.Errorf("SelectMessagesConversationId: %w", err) } + m.Content = jsontext.Value(content) m.Metadata = jsontext.Value(metadata) if m.CreatedAt, err = time.Parse(time.RFC3339, createdAt); err != nil { return nil, fmt.Errorf("SelectMessagesConversationId: parse created_at: %w", err) @@ -431,26 +488,54 @@ const updateMessageContent = ` UPDATE messages SET content = ?, updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') WHERE id = ? +RETURNING content, conversation_id, created_at, id, json(metadata), role, sequence, updated_at ` // UpdateMessageContentParams holds arguments for UpdateMessageContent. type UpdateMessageContentParams struct { - Content string + Content jsontext.Value Id string } -// UpdateMessageContent sets the content JSON for a message. -func (q *Queries) UpdateMessageContent(ctx context.Context, arg UpdateMessageContentParams) error { - _, err := q.db.ExecContext( +// UpdateMessageContent sets the content JSON for a message and returns +// the updated row. +func (q *Queries) UpdateMessageContent(ctx context.Context, arg UpdateMessageContentParams) (*Message, error) { + row := q.db.QueryRowContext( ctx, updateMessageContent, - arg.Content, + string(arg.Content), arg.Id, ) - if err != nil { - return fmt.Errorf("UpdateMessageContent: %w", err) + + var ( + m Message + content []byte + metadata []byte + createdAt string + updatedAt string + err error + ) + if err := row.Scan( + &content, + &m.ConversationId, + &createdAt, + &m.Id, + &metadata, + &m.Role, + &m.Sequence, + &updatedAt, + ); err != nil { + return nil, fmt.Errorf("UpdateMessageContent: %w", err) } - return nil + m.Content = jsontext.Value(content) + m.Metadata = jsontext.Value(metadata) + if m.CreatedAt, err = time.Parse(time.RFC3339, createdAt); err != nil { + return nil, fmt.Errorf("UpdateMessageContent: parse created_at: %w", err) + } + if m.UpdatedAt, err = time.Parse(time.RFC3339, updatedAt); err != nil { + return nil, fmt.Errorf("UpdateMessageContent: parse updated_at: %w", err) + } + return &m, nil } const deleteMessageId = ` @@ -467,6 +552,205 @@ func (q *Queries) DeleteMessageId(ctx context.Context, id string) error { return nil } +const updateMessagesSequence = ` +UPDATE messages +SET sequence = sequence - 1, updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE conversation_id = ? AND sequence > ? +` + +// UpdateMessagesSequence decrements the sequence of every message that +// follows the given sequence in the conversation. Every moved row gets a +// new updated_at because its position changed. +func (q *Queries) UpdateMessagesSequence( + ctx context.Context, + conversationId string, + sequence int, +) error { + _, err := q.db.ExecContext( + ctx, + updateMessagesSequence, + conversationId, + sequence, + ) + if err != nil { + return fmt.Errorf("UpdateMessagesSequence: %w", err) + } + return nil +} + +// The shift takes two statements because SQLite checks +// UNIQUE(conversation_id, sequence) row by row, not at the end of the +// statement. A single "sequence = sequence + 1" fails if it visits row n +// while row n+1 still exists. Step one negates the shifted rows to values +// no real row uses (sequences start at 1); step two flips them back. +const updateMessagesSequenceNegate = ` +UPDATE messages +SET sequence = -(sequence + 1) +WHERE conversation_id = ? AND sequence >= ? +` + +const updateMessagesSequenceRestore = ` +UPDATE messages +SET sequence = -sequence, updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +WHERE conversation_id = ? AND sequence < 0 +` + +// UpdateMessagesSequenceIncrement increments the sequence of every message +// at or after the given sequence in the conversation, freeing that sequence +// for an insert. Every moved row gets a new updated_at because its position +// changed. +// It runs two statements; call it inside a transaction so a failure +// between them cannot leave negative sequences behind. +func (q *Queries) UpdateMessagesSequenceIncrement( + ctx context.Context, + conversationId string, + sequence int, +) error { + if _, err := q.db.ExecContext( + ctx, + updateMessagesSequenceNegate, + conversationId, + sequence, + ); err != nil { + return fmt.Errorf("UpdateMessagesSequenceIncrement: %w", err) + } + if _, err := q.db.ExecContext( + ctx, + updateMessagesSequenceRestore, + conversationId, + ); err != nil { + return fmt.Errorf("UpdateMessagesSequenceIncrement: %w", err) + } + return nil +} + +// selectMessageSequenceToolResult returns the sequence where a new tool +// result belongs: right after the assistant's existing results. The caller +// frees that sequence with UpdateMessagesSequenceIncrement, then inserts. +// +// For example: +// +// 2 assistant (calls A, B, C) +// 3 tool (A) +// 4 tool (B) +// 5 user +// +// A result for C belongs at 5; the user message moves to 6. Without the +// user message, 5 would be the end of the conversation, and the result +// would still go at 5. +const selectMessageSequenceToolResult = ` +SELECT COALESCE( + -- The first message after the assistant that is not a tool result + -- (row 5 above) marks the end of the assistant's results. The new + -- result takes its place. If there is no such message, MIN returns + -- NULL and COALESCE uses the fallback below. + MIN(sequence), + -- Fallback: the results run to the end of the conversation, so + -- append. This is the same calculation as NextSequence. + ( + SELECT COALESCE(MAX(sequence), 0) + 1 + FROM messages + WHERE conversation_id = ? + ) +) +-- MIN always returns one row, even when nothing matches, so the +-- caller never gets sql.ErrNoRows. +FROM messages +WHERE conversation_id = ? + AND sequence > ? + AND role != 'tool' +` + +// SelectMessageSequenceToolResult returns the sequence at which to insert +// a new tool result for the assistant message at assistantSequence. +// See selectMessageSequenceToolResult. +func (q *Queries) SelectMessageSequenceToolResult( + ctx context.Context, + conversationId string, + assistantSequence int, +) (int, error) { + var sequence int + // The subquery's placeholder comes first in the query text. + err := q.db.QueryRowContext( + ctx, + selectMessageSequenceToolResult, + conversationId, + conversationId, + assistantSequence, + ).Scan(&sequence) + if err != nil { + return 0, fmt.Errorf("SelectMessageSequenceToolResult: %w", err) + } + return sequence, nil +} + +// selectConversationToolResult returns the tool result a call has +// among its assistant's results: after the assistant, before the +// sequence from selectMessageSequenceToolResult. The search stops there +// because tool call IDs are only unique within one assistant message; +// another assistant's results may reuse the ID. +const selectConversationToolResult = ` +SELECT content, conversation_id, created_at, id, json(metadata), role, sequence, updated_at +FROM messages +WHERE conversation_id = ? AND role = 'tool' + AND sequence > ? AND sequence < ? + AND json_extract(content, '$.tool_call_id') = ? +` + +// SelectConversationToolResult returns toolCallId's tool result with a +// sequence strictly between after and before. It returns +// sql.ErrNoRows if the call has no result in that range. +func (q *Queries) SelectConversationToolResult( + ctx context.Context, + conversationId string, + after int, + before int, + toolCallId string, +) (*Message, error) { + row := q.db.QueryRowContext( + ctx, + selectConversationToolResult, + conversationId, + after, + before, + toolCallId, + ) + + var ( + m Message + content []byte + metadata []byte + createdAt string + updatedAt string + ) + err := row.Scan( + &content, + &m.ConversationId, + &createdAt, + &m.Id, + &metadata, + &m.Role, + &m.Sequence, + &updatedAt, + ) + if err != nil { + return nil, fmt.Errorf("SelectConversationToolResult: %w", err) + } + m.Content = jsontext.Value(content) + m.Metadata = jsontext.Value(metadata) + if m.CreatedAt, err = time.Parse(time.RFC3339, createdAt); err != nil { + return nil, fmt.Errorf( + "SelectConversationToolResult: parse created_at: %w", err, + ) + } + if m.UpdatedAt, err = time.Parse(time.RFC3339, updatedAt); err != nil { + return nil, fmt.Errorf( + "SelectConversationToolResult: parse updated_at: %w", err, + ) + } + return &m, nil +} + const selectMaxSequence = ` SELECT COALESCE(MAX(sequence), 0) FROM messages @@ -691,3 +975,61 @@ func (q *Queries) UpdateExecutionStatus(ctx context.Context, arg UpdateExecution } return nil } + +const deleteExecutionMessageIdToolCallId = ` +DELETE FROM executions +WHERE message_id = ? AND tool_call_id = ? +` + +// DeleteExecutionMessageIdToolCallId deletes the execution record for one +// tool call within a message. It is a no-op if no record exists. +func (q *Queries) DeleteExecutionMessageIdToolCallId( + ctx context.Context, + messageId string, + toolCallId string, +) error { + _, err := q.db.ExecContext( + ctx, + deleteExecutionMessageIdToolCallId, + messageId, + toolCallId, + ) + if err != nil { + return fmt.Errorf( + "DeleteExecutionMessageIdToolCallId: %w", err, + ) + } + return nil +} + +// insertExecutionCompleted marks a call's execution completed, creating +// the row if the call never had one. An existing row is completed +// regardless of its status, so a pending or denied call is treated as +// answered and Inference will not wait on or run it. +const insertExecutionCompleted = ` +INSERT INTO executions (id, message_id, status, tool_call_id) +VALUES (?, ?, 'completed', ?) +ON CONFLICT (message_id, tool_call_id) DO UPDATE +SET status = 'completed', + updated_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') +` + +// InsertExecutionCompleted creates or completes the execution for one tool +// call within a message. See insertExecutionCompleted. +func (q *Queries) InsertExecutionCompleted( + ctx context.Context, + messageId string, + toolCallId string, +) error { + _, err := q.db.ExecContext( + ctx, + insertExecutionCompleted, + uuid.New().String(), + messageId, + toolCallId, + ) + if err != nil { + return fmt.Errorf("InsertExecutionCompleted: %w", err) + } + return nil +} diff --git a/internal/db/queries_test.go b/internal/db/queries_test.go index 279a3ae..a9991cd 100644 --- a/internal/db/queries_test.go +++ b/internal/db/queries_test.go @@ -8,6 +8,7 @@ import ( "log/slog" "reflect" "testing" + "time" "uuid" "github.com/mattn/go-sqlite3" @@ -33,6 +34,9 @@ func TestQueries(t *testing.T) { t.Run("Messaging", func(t *testing.T) { testMessaging(t, ctx, db) }) + t.Run("SequenceShift", func(t *testing.T) { + testSequenceShift(t, ctx, db) + }) t.Run("ConversationTimestamp", func(t *testing.T) { testConversationTimestamp(t, ctx, db) }) @@ -160,7 +164,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { ConversationId: c.Id, Sequence: seq, Role: MessageRoleSystem, - Content: `{"role":"system","content":"You are helpful"}`, + Content: jsontext.Value(`{"role":"system","content":"You are helpful"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) @@ -175,7 +179,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { ConversationId: c.Id, Sequence: seq, Role: MessageRoleUser, - Content: `{"role":"user","content":"Hello"}`, + Content: jsontext.Value(`{"role":"user","content":"Hello"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) @@ -192,7 +196,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { Sequence: seq, Role: MessageRoleAssistant, Metadata: jsontext.Value(metadata), - Content: `{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"NYC\"}"}}]}`, + Content: jsontext.Value(`{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"get_weather","arguments":"{\"location\":\"NYC\"}"}}]}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) @@ -207,7 +211,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { ConversationId: c.Id, Sequence: seq, Role: MessageRoleTool, - Content: `{"role":"tool","content":"{\"temperature\":72}","tool_call_id":"call_abc"}`, + Content: jsontext.Value(`{"role":"tool","content":"{\"temperature\":72}","tool_call_id":"call_abc"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) @@ -238,17 +242,17 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { t.Errorf("message %d metadata = %v, want nil", i, message.Metadata) } } - if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ + if _, err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ Id: messages[2].Id, - Content: `{"role":"assistant","content":"edited"}`, + Content: jsontext.Value(`{"role":"assistant","content":"edited"}`), }); err != nil { t.Fatalf("UpdateMessageContent: %v", err) } // Test UpdateMessageContent. - if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ + if _, err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ Id: toolMsg.Id, - Content: `{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}`, + Content: jsontext.Value(`{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}`), }); err != nil { t.Fatalf("UpdateMessageContent: %v", err) } @@ -260,7 +264,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { if updated[2].Metadata == nil || string(updated[2].Metadata) != metadata { t.Errorf("edited assistant metadata = %v", updated[2].Metadata) } - if updated[3].Content != `{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}` { + if string(updated[3].Content) != `{"role":"tool","content":"{\"temperature\":75}","tool_call_id":"call_abc"}` { t.Errorf("updated content mismatch: %s", updated[3].Content) } for _, tt := range []struct { @@ -283,7 +287,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { t.Fatal(err) } inserted, err := qw.InsertMessage(ctx, InsertMessageParams{ - Content: `{"role":"assistant","content":"answer"}`, + Content: jsontext.Value(`{"role":"assistant","content":"answer"}`), ConversationId: c.Id, Metadata: tt.metadata, Role: MessageRoleAssistant, @@ -361,7 +365,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { t.Fatal(err) } inserted, err := qw.InsertMessage(ctx, InsertMessageParams{ - Content: `{"role":"assistant","content":"answer"}`, + Content: jsontext.Value(`{"role":"assistant","content":"answer"}`), ConversationId: c.Id, Id: tt.id, Role: MessageRoleAssistant, @@ -397,7 +401,7 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { t.Fatal(err) } _, err = New(tx).InsertMessage(ctx, InsertMessageParams{ - Content: `{"role":"assistant","content":"rolled back"}`, + Content: jsontext.Value(`{"role":"assistant","content":"rolled back"}`), ConversationId: c.Id, Metadata: jsontext.Value(metadata), Role: MessageRoleAssistant, @@ -416,6 +420,140 @@ func testMessaging(t *testing.T, ctx context.Context, db *SQLite) { } } +func testSequenceShift(t *testing.T, ctx context.Context, db *SQLite) { + var ( + qr = New(db.Reader) + qw = New(db.Writer) + ) + + u, err := qw.InsertUser(ctx, InsertUserParams{ + Name: "sequence-shift", + Password: "test", + Role: UserRoleOwner, + }) + if err != nil { + t.Fatalf("InsertUser: %v", err) + } + + c, err := qw.InsertConversation(ctx, InsertConversationParams{ + UserId: u.Id, + }) + if err != nil { + t.Fatalf("InsertConversation: %v", err) + } + + insert := func(text string) { + t.Helper() + seq, err := qw.NextSequence(ctx, c.Id) + if err != nil { + t.Fatalf("NextSequence: %v", err) + } + if _, err := qw.InsertMessage(ctx, InsertMessageParams{ + ConversationId: c.Id, + Sequence: seq, + Role: MessageRoleUser, + Content: jsontext.Value(`{"role":"user","content":"` + text + `"}`), + }); err != nil { + t.Fatalf("InsertMessage: %v", err) + } + } + for _, text := range []string{"one", "two", "three", "four"} { + insert(text) + } + + // updated_at has millisecond precision, so sleep to put the + // shift in a later millisecond than the inserts. + time.Sleep(10 * time.Millisecond) + + before, err := qr.SelectMessagesConversationId(ctx, c.Id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + + // Deleting "two" and shifting closes the gap: "three" moves to + // 2 and "four" to 3, while "one" stays put. + if err := qw.DeleteMessageId(ctx, before[1].Id); err != nil { + t.Fatalf("DeleteMessageId: %v", err) + } + if err := qw.UpdateMessagesSequence(ctx, c.Id, before[1].Sequence); err != nil { + t.Fatalf("UpdateMessagesSequence: %v", err) + } + + moved, err := qr.SelectMessagesConversationId(ctx, c.Id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(moved) != 3 { + t.Fatalf("expected 3 messages, got %d", len(moved)) + } + if moved[0].Sequence != 1 { + t.Errorf("unshifted row sequence = %d, want 1", moved[0].Sequence) + } + if !moved[0].UpdatedAt.Equal(before[0].UpdatedAt) { + t.Errorf("unshifted row updated_at = %s, want %s", + moved[0].UpdatedAt, before[0].UpdatedAt) + } + for i, wantSeq := range []int{2, 3} { + row, prev := moved[i+1], before[i+2] + if row.Id != prev.Id { + t.Fatalf("row %d id = %s, want %s", i+1, row.Id, prev.Id) + } + if row.Sequence != wantSeq { + t.Errorf("row %d sequence = %d, want %d", + i+1, row.Sequence, wantSeq) + } + if !row.UpdatedAt.After(prev.UpdatedAt) { + t.Errorf("row %d updated_at = %s, want later than %s", + i+1, row.UpdatedAt, prev.UpdatedAt) + } + } + + // Free sequence 2 by incrementing rows 2 and 3, inside a + // transaction as the doc requires. + time.Sleep(10 * time.Millisecond) + + tx, err := db.Writer.BeginTx(ctx, nil) + if err != nil { + t.Fatalf("BeginTx: %v", err) + } + if err := New(tx).UpdateMessagesSequenceIncrement(ctx, c.Id, 2); err != nil { + tx.Rollback() + t.Fatalf("UpdateMessagesSequenceIncrement: %v", err) + } + if err := tx.Commit(); err != nil { + t.Fatalf("Commit: %v", err) + } + + after, err := qr.SelectMessagesConversationId(ctx, c.Id) + if err != nil { + t.Fatalf("SelectMessagesConversationId: %v", err) + } + if len(after) != 3 { + t.Fatalf("expected 3 messages, got %d", len(after)) + } + if after[0].Sequence != 1 { + t.Errorf("unshifted row sequence = %d, want 1", after[0].Sequence) + } + if !after[0].UpdatedAt.Equal(before[0].UpdatedAt) { + t.Errorf("unshifted row updated_at = %s, want %s", + after[0].UpdatedAt, before[0].UpdatedAt) + } + for i, wantSeq := range []int{3, 4} { + row, prev := after[i+1], moved[i+1] + if row.Id != prev.Id { + t.Fatalf("row %d id = %s, want %s", i+1, row.Id, prev.Id) + } + if row.Sequence != wantSeq { + t.Errorf("row %d sequence = %d, want %d", + i+1, row.Sequence, wantSeq) + } + if !row.UpdatedAt.After(prev.UpdatedAt) { + t.Errorf("row %d updated_at = %s, want later than %s", + i+1, row.UpdatedAt, prev.UpdatedAt) + } + } +} + func testConversationTimestamp(t *testing.T, ctx context.Context, database *SQLite) { var ( qr = New(database.Reader) @@ -448,7 +586,7 @@ func testConversationTimestamp(t *testing.T, ctx context.Context, database *SQLi ConversationId: c.Id, Sequence: 0, Role: MessageRoleUser, - Content: `{"role":"user","content":"Hello"}`, + Content: jsontext.Value(`{"role":"user","content":"Hello"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) @@ -469,9 +607,9 @@ func testConversationTimestamp(t *testing.T, ctx context.Context, database *SQLi ); err != nil { t.Fatalf("set conversation timestamp before update: %v", err) } - if err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ + if _, err := qw.UpdateMessageContent(ctx, UpdateMessageContentParams{ Id: msg.Id, - Content: `{"role":"user","content":"Updated"}`, + Content: jsontext.Value(`{"role":"user","content":"Updated"}`), }); err != nil { t.Fatalf("UpdateMessageContent: %v", err) } @@ -529,7 +667,7 @@ func testDelete(t *testing.T, ctx context.Context, db *SQLite) { ConversationId: c.Id, Sequence: 1, Role: MessageRoleUser, - Content: `{"role":"user","content":"Hello"}`, + Content: jsontext.Value(`{"role":"user","content":"Hello"}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) @@ -585,7 +723,7 @@ func testExecutions(t *testing.T, ctx context.Context, db *SQLite) { ConversationId: c.Id, Sequence: 1, Role: MessageRoleAssistant, - Content: `{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"host_run","arguments":"{\"command\":\"rm -rf /\"}"}}]}`, + Content: jsontext.Value(`{"role":"assistant","content":null,"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"host_run","arguments":"{\"command\":\"rm -rf /\"}"}}]}`), }) if err != nil { t.Fatalf("InsertMessage: %v", err) @@ -721,7 +859,7 @@ func testExecutions(t *testing.T, ctx context.Context, db *SQLite) { ConversationId: c.Id, Sequence: 2, Role: MessageRoleAssistant, - Content: `{"role":"assistant","content":"second"}`, + Content: jsontext.Value(`{"role":"assistant","content":"second"}`), }) if err != nil { t.Fatalf("InsertMessage second: %v", err) diff --git a/internal/db/schema/1.init.sql b/internal/db/schema/1.init.sql index 975c662..56479e5 100644 --- a/internal/db/schema/1.init.sql +++ b/internal/db/schema/1.init.sql @@ -53,15 +53,16 @@ CREATE TABLE IF NOT EXISTS conversations ( CREATE INDEX IF NOT EXISTS idx_conv_user ON conversations(user_id); -- Messages are ordered entries within a conversation. The content column --- stores the wire-format JSON for the message — the exact JSON object --- that appears in the OpenAI Chat Completions messages array. The role --- column is extracted for querying without JSON parsing. Nullable metadata --- stores inference JSONB separately from model-visible content. The sequence --- column defines ordering within the conversation. Deleting a message --- cascades to nothing further (all message data is in the row). +-- stores the wire-format JSON for the message — the exact JSON object that +-- appears in the OpenAI Chat Completions messages array; a CHECK constraint +-- keeps it valid JSON. The role column is extracted for querying without +-- JSON parsing. Nullable metadata stores inference JSONB separately from +-- model-visible content. The sequence column defines ordering within the +-- conversation. Deleting a message cascades to nothing further (all message +-- data is in the row). CREATE TABLE IF NOT EXISTS messages ( conversation_id TEXT NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, - content TEXT NOT NULL, + content TEXT NOT NULL CHECK (json_valid(content)), created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')), id TEXT PRIMARY KEY, metadata BLOB CHECK (metadata IS NULL OR json_valid(metadata, 8)), diff --git a/internal/job/job.go b/internal/job/job.go index a600eb9..fd13b69 100644 --- a/internal/job/job.go +++ b/internal/job/job.go @@ -352,7 +352,7 @@ func (j *Job) createConversation( return nil, fmt.Errorf("marshal initial message: %w", err) } if _, err := q.InsertMessage(ctx, db.InsertMessageParams{ - Content: string(content), + Content: content, ConversationId: c.Id, Role: db.MessageRole(message.Role), Sequence: sequence, @@ -411,7 +411,7 @@ func (j *Job) insertMessage( return nil, fmt.Errorf("get next message sequence: %w", err) } m, err := q.InsertMessage(ctx, db.InsertMessageParams{ - Content: string(content), + Content: content, ConversationId: cid, Metadata: metadataJSON, Role: db.MessageRole(message.Role), diff --git a/internal/job/job_test.go b/internal/job/job_test.go index fa873a8..ec577be 100644 --- a/internal/job/job_test.go +++ b/internal/job/job_test.go @@ -685,7 +685,7 @@ func TestCallTool(t *testing.T) { t.Errorf("tool metadata = %v, want nil", last.Metadata) } if last.Role != db.MessageRoleTool || - last.Content != string(want) { + string(last.Content) != string(want) { t.Errorf( "last stored message = role %q, %s; "+ "want %q, %s", diff --git a/internal/llm/llm.go b/internal/llm/llm.go index 582edd8..d29148f 100644 --- a/internal/llm/llm.go +++ b/internal/llm/llm.go @@ -548,6 +548,37 @@ func (m *Message) Validate() error { return nil } +// UpdateToolCall replaces the call identified by callId, keeping its +// position in ToolCalls. The replacement must have the same ID: the +// corresponding tool response carries it in Message.ToolCallID. +func (m *Message) UpdateToolCall(call ToolCall) error { + if err := call.Validate(); err != nil { + return fmt.Errorf("invalid tool call: %w", err) + } + for i := range m.ToolCalls { + if m.ToolCalls[i].ID == call.ID { + m.ToolCalls[i] = call + return nil + } + } + return fmt.Errorf("not found: call %q", call.ID) +} + +// AddToolCall appends call to the message's ToolCalls. The call must be +// valid and its ID must not already appear in the message. +func (m *Message) AddToolCall(call ToolCall) error { + if err := call.Validate(); err != nil { + return fmt.Errorf("invalid tool call: %w", err) + } + for _, existing := range m.ToolCalls { + if existing.ID == call.ID { + return fmt.Errorf("duplicate call id %q", call.ID) + } + } + m.ToolCalls = append(m.ToolCalls, call) + return nil +} + // UnmarshalJSON decodes a message while normalizing a string-valued content // field to one text ContentPart. An array-valued content field is decoded in // order. If reasoning_content is absent, reasoning is copied into