fix: archive a parked question before its reply

go/core/internal/database/client_agent_instance_test.go

@@ -208,6 +208,55 @@ func TestConcurrentAgentInstanceMessageReplay(t *testing.T) {
 	}
 }
 
+// TestAgentInstanceTaskMessageWithoutTaskJoinsHistory pins what the gateway relies
+// on when it archives a parked question before the reply that replaces it: a bare
+// message event becomes history, in insertion order, without touching the task.
+func TestAgentInstanceTaskMessageWithoutTaskJoinsHistory(t *testing.T) {
+	db := setupTestDB(t)
+	ctx := context.Background()
+	if _, err := db.Exec(ctx, `
+		INSERT INTO a2a_context (id, namespace, user_id)
+		VALUES ('instance-1', 'team-a', 'alice');
+		INSERT INTO agent_instance (id, namespace, user_id, request_id, context_id, state, data)
+		VALUES ('instance-1', 'team-a', 'alice', 'request-1', 'instance-1', 'READY', '\\x00')
+	`); err != nil {
+		t.Fatal(err)
+	}
+	client := NewClient(db)
+	asked := &a2a.Message{ID: "message-1", Role: a2a.MessageRoleUser, TaskID: "task-1", ContextID: "instance-1"}
+	parked := &a2a.Task{
+		ID: "task-1", ContextID: "instance-1", History: []*a2a.Message{asked},
+		Status: a2a.TaskStatus{State: a2a.TaskStateInputRequired},
+	}
+	if _, _, err := client.CreateAgentInstanceTask(ctx, "instance-1", []byte("request-1"), parked); err != nil {
+		t.Fatal(err)
+	}
+
+	question := &a2a.Message{ID: "question-1", Role: a2a.MessageRoleAgent, TaskID: "task-1", ContextID: "instance-1"}
+	if err := client.StoreAgentInstanceTaskEvent(ctx, "instance-1", nil, question, nil); err != nil {
+		t.Fatal(err)
+	}
+	answer := &a2a.Message{ID: "answer-1", Role: a2a.MessageRoleUser, TaskID: "task-1", ContextID: "instance-1"}
+	resumed := *parked
+	resumed.History = []*a2a.Message{asked, question, answer}
+	resumed.Status = a2a.TaskStatus{State: a2a.TaskStateSubmitted}
+	if err := client.StoreAgentInstanceTaskEvent(ctx, "instance-1", &resumed, answer, nil); err != nil {
+		t.Fatal(err)
+	}
+
+	got, err := client.GetAgentInstanceTask(ctx, "instance-1", "task-1")
+	if err != nil {
+		t.Fatal(err)
+	}
+	ids := make([]string, 0, len(got.History))
+	for _, message := range got.History {
+		ids = append(ids, message.ID)
+	}
+	if strings.Join(ids, ",") != "message-1,question-1,answer-1" {
+		t.Fatalf("history = %v, want the question between the message it answers and its own answer", ids)
+	}
+}
+
 func TestAgentInstanceCheckpointRetainsRecordedBoundary(t *testing.T) {
 	db := setupTestDB(t)
 	ctx := context.Background()

go/core/v2/a2agateway/gateway.go

@@ -517,7 +517,21 @@ func (g *Gateway) prepareReply(ctx context.Context, instance *apiv1alpha1.AgentI
 	}
 	message.ContextID = stored.ContextID
 	attempt := *stored
-	attempt.History = append(append([]*a2atype.Message{}, stored.History...), message)
+	attempt.History = append([]*a2atype.Message{}, stored.History...)
+	if question := stored.Status.Message; question != nil {
+		attempt.History = append(attempt.History, question)
+		// A parked task holds its question in the status this reply replaces, and the
+		// store writes down the messages an event names rather than the history of the
+		// task it is handed — so the question is stored in its own right, and before
+		// the reply, because history is ordered by insertion.
+		if question.ID != "" {
+			question.TaskID, question.ContextID = stored.ID, stored.ContextID
+			if err := g.store.StoreAgentInstanceTaskEvent(ctx, instance.GetId(), nil, question, nil); err != nil {
+				return nil, g.storeError(ctx, err)
+			}
+		}
+	}
+	attempt.History = append(attempt.History, message)
 	now := time.Now()
 	attempt.Status = a2atype.TaskStatus{State: a2atype.TaskStateSubmitted, Timestamp: &now}
 	if err := g.store.StoreAgentInstanceTaskEvent(ctx, instance.GetId(), &attempt, message, nil); err != nil {

go/core/v2/a2agateway/gateway_test.go

@@ -381,7 +381,8 @@ func TestGatewayContinuesInputRequiredTask(t *testing.T) {
 	if !ok || task.Status.State != a2atype.TaskStateCompleted || !runtime.sent {
 		t.Fatalf("reply result = %#v, runtime sent = %v", result, runtime.sent)
 	}
-	if authorizer.verb != auth.VerbUpdate || reply.ContextID != gatewayTestID || len(store.stored) != 2 {
+	// Three: the question the task was parked on, the reply, and the finished turn.
+	if authorizer.verb != auth.VerbUpdate || reply.ContextID != gatewayTestID || len(store.stored) != 3 {
 		t.Fatalf("reply authorization = %s, context = %q, stored events = %d", authorizer.verb, reply.ContextID, len(store.stored))
 	}
 	if runtime.privateTask == nil || runtime.privateTask.Status.State != a2atype.TaskStateInputRequired || runtime.privateTask.Status.Message == nil || runtime.privateTask.Status.Message.ID != status.ID {
@@ -392,6 +393,35 @@ func TestGatewayContinuesInputRequiredTask(t *testing.T) {
 	}
 }
 
+func TestGatewayArchivesInputRequiredMessageBeforeReply(t *testing.T) {
+	question := a2atype.NewMessage(a2atype.MessageRoleAgent, a2atype.NewTextPart("Which database?"))
+	waiting := &a2atype.Task{
+		ID: "task-1", ContextID: gatewayTestID,
+		Status: a2atype.TaskStatus{State: a2atype.TaskStateInputRequired, Message: question},
+	}
+	reply := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("PostgreSQL"))
+	reply.TaskID = waiting.ID
+	store := &gatewayTestStore{task: waiting}
+	gateway := &Gateway{store: store}
+
+	prepared, err := gateway.prepareReply(t.Context(), gatewayTestInstance(), &a2atype.SendMessageRequest{Message: reply})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(prepared.task.History) != 2 || prepared.task.History[0] != question || prepared.task.History[1] != reply {
+		t.Fatalf("history = %#v, want question followed by reply", prepared.task.History)
+	}
+	// Archived, which the history above does not prove: a question that only reaches
+	// `prepared.task.History` is never written down, and the status that held it is
+	// gone. First, because history is ordered by insertion.
+	if len(store.stored) != 2 || store.stored[0] != question || store.stored[1] != reply {
+		t.Fatalf("stored events = %#v, want the question archived before the reply", store.stored)
+	}
+	if question.TaskID != waiting.ID || question.ContextID != waiting.ContextID {
+		t.Fatalf("archived question = task %q context %q, want the task it was asked in", question.TaskID, question.ContextID)
+	}
+}
+
 func TestGatewayClosesRuntimeAfterStreaming(t *testing.T) {
 	instance := gatewayTestInstance()
 	runtime := &gatewayTestRuntime{}