Refactor service

internal/service/service.go

@@ -65,6 +65,7 @@ type Service struct {
 
 	ctx           context.Context
 	cancel        context.CancelFunc
+	wg            sync.WaitGroup
 	mu            sync.Mutex
 	conversations map[string]*conversation.Conversation
 }
@@ -331,6 +332,18 @@ func (svc *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 // Run starts the service and blocks until shutdown.
 // Shutdown is triggered by SIGINT or SIGTERM.
+//
+// The return value tells the caller when it is safe to close the
+// database:
+//
+//   - nil: all HTTP handlers returned, and all inference calls and
+//     scheduled jobs finished. The caller can close the database.
+//   - non-nil: shutdown did not finish within the configured
+//     ShutdownTimeout, or the HTTP server failed. Handlers,
+//     inference calls, or jobs can still use the database. Do not
+//     close the database. Exit the process instead. The database
+//     runs in WAL mode, so committed data stays safe when the
+//     process exits without a close.
 func (svc *Service) Run(ctx context.Context) error {
 	svc.log.Info(
 		"starting odidere",
@@ -463,13 +476,21 @@ func (svc *Service) Run(ctx context.Context) error {
 	if err := svc.server.Shutdown(shutdownCtx); err != nil {
 		shutdownErr = fmt.Errorf("HTTP shutdown: %w", err)
 	} else {
-		// Jobs can save tool results after cancellation. Wait until
-		// their callbacks return before the caller closes the database.
+		// Jobs and HTTP inference can save tool results after
+		// cancellation. Wait until both return before the caller
+		// closes the database.
+		done := make(chan struct{})
+		go func() {
+			<-cronCtx.Done()
+			svc.wg.Wait()
+			close(done)
+		}()
 		select {
-		case <-cronCtx.Done():
+		case <-done:
 		case <-shutdownCtx.Done():
 			shutdownErr = fmt.Errorf(
-				"wait for jobs: %w", shutdownCtx.Err(),
+				"wait for background work: %w",
+				shutdownCtx.Err(),
 			)
 		}
 	}
@@ -487,6 +508,107 @@ func (svc *Service) Run(ctx context.Context) error {
 	return runErr
 }
 
+// newConversation saves a conversation and its initial system message.
+// It registers the handle only after the transaction commits.
+func (svc *Service) newConversation(
+	ctx context.Context,
+	userId string,
+	title *string,
+	system string,
+) (*db.Conversation, *conversation.Conversation, error) {
+	msg := llm.Message{
+		Role: llm.RoleSystem,
+		ContentParts: []llm.ContentPart{{
+			Type: llm.ContentTypeText,
+			Text: system,
+		}},
+	}
+	if err := msg.Validate(); err != nil {
+		return nil, nil, fmt.Errorf("invalid system message: %w", err)
+	}
+	content, err := json.Marshal(msg)
+	if err != nil {
+		return nil, nil, fmt.Errorf("marshal system message: %w", err)
+	}
+
+	tx, err := svc.db.Writer.BeginTx(ctx, nil)
+	if err != nil {
+		return nil, nil, fmt.Errorf("begin tx: %w", err)
+	}
+	defer tx.Rollback()
+
+	qtx := svc.dbw.WithTx(tx)
+	saved, err := qtx.InsertConversation(
+		ctx, db.InsertConversationParams{
+			Title:  title,
+			UserId: userId,
+		},
+	)
+	if err != nil {
+		return nil, nil, fmt.Errorf("insert conversation: %w", err)
+	}
+	if _, err := qtx.InsertMessage(ctx, db.InsertMessageParams{
+		Content:        content,
+		ConversationId: saved.Id,
+		Role:           db.MessageRoleSystem,
+		Sequence:       0,
+	}); err != nil {
+		return nil, nil, fmt.Errorf("insert system message: %w", err)
+	}
+
+	c, err := conversation.New(
+		conversation.Config{Id: saved.Id},
+		svc.db,
+		svc.log,
+		svc.tools,
+	)
+	if err != nil {
+		return nil, nil, fmt.Errorf("create conversation handle: %w", err)
+	}
+	if err := tx.Commit(); err != nil {
+		return nil, nil, fmt.Errorf("commit conversation tx: %w", err)
+	}
+
+	svc.mu.Lock()
+	// A retrieval can register a handle after the commit.
+	if existing := svc.conversations[saved.Id]; existing != nil {
+		c = existing
+	} else {
+		svc.conversations[saved.Id] = c
+	}
+	svc.mu.Unlock()
+
+	return saved, c, nil
+}
+
+// getConversation checks saved state before it returns a cached handle.
+func (svc *Service) getConversation(
+	ctx context.Context,
+	id string,
+) (*conversation.Conversation, error) {
+	if _, err := svc.dbr.SelectConversationId(ctx, id); err != nil {
+		return nil, fmt.Errorf("select conversation: %w", err)
+	}
+
+	svc.mu.Lock()
+	defer svc.mu.Unlock()
+
+	if c := svc.conversations[id]; c != nil {
+		return c, nil
+	}
+	c, err := conversation.New(
+		conversation.Config{Id: id},
+		svc.db,
+		svc.log,
+		svc.tools,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("create conversation handle: %w", err)
+	}
+	svc.conversations[id] = c
+	return c, nil
+}
+
 // ListJobs returns all registered jobs.
 func (svc *Service) ListJobs() []*job.Job {
 	return svc.jobs