From 420c3ccfa16f1dd1ec550b7f5db8281621374763 Mon Sep 17 00:00:00 2001 From: Nicholas Bucher Date: Fri, 28 Aug 2026 10:44:54 -0400 Subject: [PATCH 1/2] fix: archive a parked question before its reply MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An agent that stops to ask something parks the task with the question in its `status`, not in its history. Answering replaces that status, and the store writes down the messages each event names rather than the history of the task it is handed — so the question was written nowhere and vanished from the transcript on the next read. The copy still on screen was then appended below the answer it preceded, because a merge keeps what the server did not send. The question is now stored in its own right, and before the reply, because history is ordered by insertion. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Nicholas Bucher --- .../database/client_agent_instance_test.go | 49 +++++++++++++++++++ go/core/v2/a2agateway/gateway.go | 16 +++++- go/core/v2/a2agateway/gateway_test.go | 32 +++++++++++- 3 files changed, 95 insertions(+), 2 deletions(-) diff --git a/go/core/internal/database/client_agent_instance_test.go b/go/core/internal/database/client_agent_instance_test.go index 4650f2009..fb2a21038 100644 --- a/go/core/internal/database/client_agent_instance_test.go +++ b/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() diff --git a/go/core/v2/a2agateway/gateway.go b/go/core/v2/a2agateway/gateway.go index 1ff615a64..3b325b0ba 100644 --- a/go/core/v2/a2agateway/gateway.go +++ b/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 { diff --git a/go/core/v2/a2agateway/gateway_test.go b/go/core/v2/a2agateway/gateway_test.go index 991a13a71..1e6d4bf29 100644 --- a/go/core/v2/a2agateway/gateway_test.go +++ b/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{} From 11cc3832cdef90a043add64c505f6ad2976ea434 Mon Sep 17 00:00:00 2001 From: Eitan Yarmush Date: Fri, 28 Aug 2026 15:14:50 +0000 Subject: [PATCH 2/2] fix: archive reply history atomically Signed-off-by: Eitan Yarmush --- .../database/client_agent_instance_test.go | 12 ++------ go/core/internal/database/client_postgres.go | 10 +++++++ go/core/v2/a2agateway/gateway.go | 14 +++------ go/core/v2/a2agateway/gateway_test.go | 30 ++++++++++++++----- 4 files changed, 39 insertions(+), 27 deletions(-) diff --git a/go/core/internal/database/client_agent_instance_test.go b/go/core/internal/database/client_agent_instance_test.go index fb2a21038..f03938ad4 100644 --- a/go/core/internal/database/client_agent_instance_test.go +++ b/go/core/internal/database/client_agent_instance_test.go @@ -208,10 +208,7 @@ 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) { +func TestAgentInstanceReplyArchivesStatusMessageAtomically(t *testing.T) { db := setupTestDB(t) ctx := context.Background() if _, err := db.Exec(ctx, ` @@ -224,18 +221,15 @@ func TestAgentInstanceTaskMessageWithoutTaskJoinsHistory(t *testing.T) { } client := NewClient(db) asked := &a2a.Message{ID: "message-1", Role: a2a.MessageRoleUser, TaskID: "task-1", ContextID: "instance-1"} + question := &a2a.Message{ID: "question-1", Role: a2a.MessageRoleAgent, 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}, + Status: a2a.TaskStatus{State: a2a.TaskStateInputRequired, Message: question}, } 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} diff --git a/go/core/internal/database/client_postgres.go b/go/core/internal/database/client_postgres.go index 31240b144..2903d8e07 100644 --- a/go/core/internal/database/client_postgres.go +++ b/go/core/internal/database/client_postgres.go @@ -854,12 +854,19 @@ func (c *postgresClient) InterruptActiveAgentInstanceTask(ctx context.Context, i func (c *postgresClient) StoreAgentInstanceTaskEvent(ctx context.Context, instanceID string, task *a2a.Task, event a2a.Event, snapshot *dbpkg.AgentInstanceTaskSnapshot) error { err := c.withTx(ctx, func(q *dbgen.Queries) error { var sequence int64 + var replacedStatusMessage *a2a.Message if task != nil { if row, err := q.GetAgentInstanceTask(ctx, dbgen.GetAgentInstanceTaskParams{ContextID: instanceID, ID: string(task.ID)}); err == nil { previous, err := unmarshalAgentInstanceTask(row.Data) if err != nil { return err } + // A reply replaces the current status message, so archive both atomically. + if _, ok := event.(*a2a.Message); ok && previous.Status.Message != nil { + message := *previous.Status.Message + message.TaskID, message.ContextID = task.ID, task.ContextID + replacedStatusMessage = &message + } if len(previous.History) > 0 { sequence, err = storeAgentInstanceTaskMessages(ctx, q, instanceID, string(task.ID), previous.History) if err != nil { @@ -884,6 +891,9 @@ func (c *postgresClient) StoreAgentInstanceTaskEvent(ctx context.Context, instan } } messages := agentInstanceTaskEventMessages(task, event) + if replacedStatusMessage != nil { + messages = append([]*a2a.Message{replacedStatusMessage}, messages...) + } if len(messages) > 0 { var err error sequence, err = storeAgentInstanceTaskMessages(ctx, q, instanceID, string(event.TaskInfo().TaskID), messages) diff --git a/go/core/v2/a2agateway/gateway.go b/go/core/v2/a2agateway/gateway.go index 3b325b0ba..c6a02610f 100644 --- a/go/core/v2/a2agateway/gateway.go +++ b/go/core/v2/a2agateway/gateway.go @@ -519,17 +519,11 @@ func (g *Gateway) prepareReply(ctx context.Context, instance *apiv1alpha1.AgentI attempt := *stored 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) - } + if question.ID == "" { + return nil, a2atype.NewError(a2atype.ErrInternalError, "stored task status message has no ID") } + question.TaskID, question.ContextID = stored.ID, stored.ContextID + attempt.History = append(attempt.History, question) } attempt.History = append(attempt.History, message) now := time.Now() diff --git a/go/core/v2/a2agateway/gateway_test.go b/go/core/v2/a2agateway/gateway_test.go index 1e6d4bf29..9a9de157e 100644 --- a/go/core/v2/a2agateway/gateway_test.go +++ b/go/core/v2/a2agateway/gateway_test.go @@ -381,8 +381,7 @@ 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) } - // 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 { + if authorizer.verb != auth.VerbUpdate || reply.ContextID != gatewayTestID || len(store.stored) != 2 { 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 { @@ -393,7 +392,7 @@ func TestGatewayContinuesInputRequiredTask(t *testing.T) { } } -func TestGatewayArchivesInputRequiredMessageBeforeReply(t *testing.T) { +func TestGatewayMovesInputRequiredMessageBeforeReply(t *testing.T) { question := a2atype.NewMessage(a2atype.MessageRoleAgent, a2atype.NewTextPart("Which database?")) waiting := &a2atype.Task{ ID: "task-1", ContextID: gatewayTestID, @@ -411,17 +410,32 @@ func TestGatewayArchivesInputRequiredMessageBeforeReply(t *testing.T) { 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 len(store.stored) != 1 || store.stored[0] != reply { + t.Fatalf("stored events = %#v, want one atomic reply update", 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 TestGatewayRejectsInputRequiredMessageWithoutID(t *testing.T) { + waiting := &a2atype.Task{ + ID: "task-1", ContextID: gatewayTestID, + Status: a2atype.TaskStatus{State: a2atype.TaskStateInputRequired, Message: &a2atype.Message{}}, + } + store := &gatewayTestStore{task: waiting} + gateway := &Gateway{store: store} + reply := a2atype.NewMessage(a2atype.MessageRoleUser, a2atype.NewTextPart("PostgreSQL")) + reply.TaskID = waiting.ID + + if _, err := gateway.prepareReply(t.Context(), gatewayTestInstance(), &a2atype.SendMessageRequest{Message: reply}); err == nil { + t.Fatal("prepareReply() succeeded with an unidentifiable status message") + } + if len(store.stored) != 0 { + t.Fatalf("stored events = %#v, want no partial write", store.stored) + } +} + func TestGatewayClosesRuntimeAfterStreaming(t *testing.T) { instance := gatewayTestInstance() runtime := &gatewayTestRuntime{}