Refactor conversation pkg

internal/conversation/conversation_test.go

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