feat(chat): make the thread part of reply affinity, not just the channel

pkg/hub/handlers_agent_messaging.go

@@ -169,16 +169,22 @@ func (s *Server) handleAgentOutboundMessage(w http.ResponseWriter, r *http.Reque
 
 	// Reply affinity (Phase 6, AC22): when the agent sends an untagged reply
 	// (no explicit channel), check webchat_conversation_context for the
-	// (recipient, project, agent) triple. If a row exists, route to the
-	// channel the user last spoke from. If no row exists, leave channel
-	// empty so the message fans out to all spokes (today's default behavior).
+	// (recipient, project, agent) triple. If a row exists, route to where the
+	// user last spoke from. If no row exists, leave the route empty so the
+	// message fans out to all spokes (today's default behavior).
+
 	if req.Channel == "" && recipientID != "" && s.webChatStore != nil && s.GetMessageBrokerProxy() != nil {
-		if lastCh, err := s.webChatStore.GetLastChannel(ctx, recipientID, agent.ProjectID, agent.ID); err != nil {
+		lastCh, lastThread, err := s.webChatStore.GetLastRoute(ctx, recipientID, agent.ProjectID, agent.ID)
+		switch {
+		case err != nil:
 			s.messageLog.Error("Failed to look up reply affinity",
 				"recipient_id", recipientID, "agent_id", agent.ID, "error", err)
 			// Non-fatal: fall through to fan-out-to-all behavior.
-		} else if lastCh != "" {
+		case lastCh != "":
 			req.Channel = lastCh
+			if req.ThreadID == "" {
+				req.ThreadID = lastThread
+			}
 		}
 	}
 

pkg/hub/handlers_agent_messaging_test.go

@@ -19,11 +19,15 @@ package hub
 import (
 	"bytes"
 	"context"
+	"database/sql"
 	"encoding/json"
+	"github.com/GoogleCloudPlatform/scion/pkg/eventbus"
+	"log/slog"
 	"net/http"
 	"net/http/httptest"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -264,3 +268,123 @@ func TestOutboundMessage_TranscriptMirrorDoesNotStarveAgentMessages(t *testing.T
 			rr.Code, rr.Body.String())
 	}
 }
+
+// recordingBus is a fan-out spoke that keeps what the handler published, so a
+// test can assert the route the handler chose without wiring persistence.
+type recordingBus struct {
+	mu   sync.Mutex
+	sent []*messages.StructuredMessage
+}
+
+func (b *recordingBus) Publish(_ context.Context, _ string, msg *messages.StructuredMessage) error {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	b.sent = append(b.sent, msg)
+	return nil
+}
+
+func (b *recordingBus) Subscribe(string, eventbus.EventHandler) (eventbus.Subscription, error) {
+	return nil, nil
+}
+
+func (b *recordingBus) Close() error { return nil }
+
+func (b *recordingBus) last() *messages.StructuredMessage {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	if len(b.sent) == 0 {
+		return nil
+	}
+	return b.sent[len(b.sent)-1]
+}
+
+// TestOutboundMessage_ReplyAffinityRestoresRoute exercises the affinity block
+// end to end: with a webchat store and a registered channel present, an agent
+// reply that names no route is put back where the user last spoke.
+func TestOutboundMessage_ReplyAffinityRestoresRoute(t *testing.T) {
+	srv, s := testServer(t)
+	ctx := context.Background()
+
+	db, err := sql.Open("sqlite3", ":memory:")
+	if err != nil {
+		t.Fatalf("open sqlite: %v", err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+	wcs := NewWebChatStore(db, "sqlite3")
+	if err := wcs.Init(); err != nil {
+		t.Fatalf("Init: %v", err)
+	}
+	srv.SetWebChatStore(wcs)
+
+	// A fan-out bus with a "web" spoke, so the channel the affinity row names
+	// passes the registered-channel check below it.
+	rec := &recordingBus{}
+	bus := eventbus.NewFanOutEventBus([]eventbus.NamedEventBus{
+		{Name: "web", ChannelID: "web", Bus: rec, Observer: true},
+	}, slog.Default())
+	srv.SetMessageBrokerProxy(NewMessageBrokerProxy(bus, s, srv.events, func() AgentDispatcher { return nil }, slog.Default()))
+
+	project := &store.Project{
+		ID: api.NewUUID(), Name: "aff", Slug: "aff", Visibility: store.VisibilityPrivate,
+	}
+	if err := s.CreateProject(ctx, project); err != nil {
+		t.Fatalf("CreateProject: %v", err)
+	}
+	human := &store.User{ID: api.NewUUID(), Email: "human@example.com", DisplayName: "Human", Status: store.UserStatusActive}
+	if err := s.CreateUser(ctx, human); err != nil {
+		t.Fatalf("CreateUser: %v", err)
+	}
+	agent := &store.Agent{
+		ID: api.NewUUID(), Name: "a", Slug: "a", ProjectID: project.ID,
+		Phase: "running", Visibility: store.VisibilityPrivate,
+	}
+	if err := s.CreateAgent(ctx, agent); err != nil {
+		t.Fatalf("CreateAgent: %v", err)
+	}
+
+	// The user last spoke in a web thread.
+	if err := wcs.RecordChannel(ctx, human.ID, project.ID, agent.ID, "web", "topic-42", time.Now()); err != nil {
+		t.Fatalf("RecordChannel: %v", err)
+	}
+
+	send := func(t *testing.T, threadID string) *messages.StructuredMessage {
+		t.Helper()
+		body, _ := json.Marshal(OutboundMessageRequest{
+			Recipient: "user:human@example.com",
+			Msg:       "reply",
+			ThreadID:  threadID,
+		})
+		req := httptest.NewRequest(http.MethodPost, "/api/v1/agents/"+agent.ID+"/outbound-message", bytes.NewReader(body))
+		req.Header.Set("Content-Type", "application/json")
+		req = req.WithContext(contextWithIdentity(req.Context(), &agentIdentityWrapper{&AgentTokenClaims{
+			Claims: jwt.Claims{Subject: agent.ID}, ProjectID: project.ID,
+		}}))
+		rr := httptest.NewRecorder()
+		srv.handleAgentOutboundMessage(rr, req, agent.ID)
+		if rr.Code != http.StatusOK {
+			t.Fatalf("expected 200, got %d: %s", rr.Code, rr.Body.String())
+		}
+		m := rec.last()
+		if m == nil {
+			t.Fatal("nothing was published to the spoke")
+		}
+		return m
+	}
+
+	t.Run("an untagged reply is routed back to the thread", func(t *testing.T) {
+		m := send(t, "")
+		if m.Channel != "web" {
+			t.Errorf("channel = %q, want web", m.Channel)
+		}
+		if m.ThreadID != "topic-42" {
+			t.Errorf("thread = %q, want topic-42; an untagged reply lands beside the conversation without it", m.ThreadID)
+		}
+	})
+
+	t.Run("a thread named by the caller is not overwritten", func(t *testing.T) {
+		m := send(t, "topic-99")
+		if m.ThreadID != "topic-99" {
+			t.Errorf("thread = %q, want topic-99", m.ThreadID)
+		}
+	})
+}

pkg/hub/handlers_broker_inbound.go

@@ -265,7 +265,7 @@ func (s *Server) handleBrokerInbound(w http.ResponseWriter, r *http.Request) {
 	// Only record for user-identity senders with a known channel.
 	if s.webChatStore != nil && req.Message.Channel != "" && strings.HasPrefix(req.Message.Sender, "user:") {
 		if senderUserID != "" {
-			if err := s.webChatStore.RecordChannel(r.Context(), senderUserID, agent.ProjectID, agent.ID, req.Message.Channel, now); err != nil {
+			if err := s.webChatStore.RecordChannel(r.Context(), senderUserID, agent.ProjectID, agent.ID, req.Message.Channel, req.Message.ThreadID, now); err != nil {
 				log.Error("Failed to record conversation context for broker inbound",
 					"user_id", senderUserID, "agent_id", agent.ID, "channel", req.Message.Channel, "error", err)
 			}

pkg/hub/webchannel.go

@@ -136,7 +136,7 @@ func (b *webChannelBus) Publish(ctx context.Context, topic string, msg *messages
 	}
 
 	// Reply affinity — still needed for cross-channel reply routing.
-	if err := b.store.RecordChannel(ctx, userID, projectID, agentID, "web", time.Now().UTC()); err != nil {
+	if err := b.store.RecordChannel(ctx, userID, projectID, agentID, "web", msg.ThreadID, time.Now().UTC()); err != nil {
 		b.log.Error("Failed to record conversation context",
 			"user_id", userID, "project_id", projectID, "agent_id", agentID, "error", err)
 		return err

pkg/hub/webchannel_store.go

@@ -50,16 +50,16 @@ type WebChatStore interface {
 	// indexed read instead of an aggregate query.
 	TouchThread(ctx context.Context, userID, projectID, agentID, messageID string, activityAt time.Time) error
 
-	// RecordChannel upserts reply-affinity context for (user, project, agent).
-	// Records the last channel a message was seen on, so the hub can route
-	// untagged replies back to the channel the user last spoke from.
-	RecordChannel(ctx context.Context, userID, projectID, agentID, channel string, messageAt time.Time) error
+	// RecordChannel upserts reply-affinity context for (user, project, agent):
+	// the channel and thread the user's message arrived on, so the hub can
+	// route an untagged reply back to where they were speaking.
+	RecordChannel(ctx context.Context, userID, projectID, agentID, channel, threadID string, messageAt time.Time) error
 
-	// GetLastChannel returns the last channel recorded for (user, project, agent),
-	// or "" if no row exists. Used for reply affinity: when an agent sends an
-	// untagged reply, the hub checks here to route to the channel the user last
-	// spoke from.
-	GetLastChannel(ctx context.Context, userID, projectID, agentID string) (string, error)
+	// GetLastRoute returns the last channel and thread recorded for (user, project, agent),
+	// or empty strings if no row exists. Used for reply affinity: when an agent
+	// sends an untagged reply, the hub routes it back to where the user last
+	// spoke.
+	GetLastRoute(ctx context.Context, userID, projectID, agentID string) (channel, threadID string, err error)
 
 	// GetThreadPrefs returns the display preferences for a (user, project, agent) thread.
 	// Returns default prefs (visibility_mode = "conversation") if no row exists.
@@ -361,6 +361,7 @@ CREATE TABLE IF NOT EXISTS webchat_conversation_context (
     project_id TEXT NOT NULL,
     agent_id TEXT NOT NULL,
     last_channel TEXT,
+    last_thread_id TEXT,
     last_message_at TEXT,
     PRIMARY KEY (user_id, project_id, agent_id)
 );
@@ -492,6 +493,17 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_webchat_topic_project_name
 	return nil
 }
 
+// addConversationContextThreadID adds last_thread_id to an existing
+// webchat_conversation_context. SQLite has no ADD COLUMN IF NOT EXISTS, so a
+// duplicate-column error is the success case on a database that already has it.
+func (s *sqliteWebChatStore) addConversationContextThreadID() error {
+	_, err := s.db.Exec(`ALTER TABLE webchat_conversation_context ADD COLUMN last_thread_id TEXT`)
+	if err != nil && !strings.Contains(strings.ToLower(err.Error()), "duplicate column") {
+		return err
+	}
+	return nil
+}
+
 // TouchThread upserts the thread watermark for the given (user, project, agent) triple.
 func (s *sqliteWebChatStore) TouchThread(ctx context.Context, userID, projectID, agentID, messageID string, activityAt time.Time) error {
 	const query = `
@@ -510,34 +522,37 @@ DO UPDATE SET
 }
 
 // RecordChannel upserts the reply-affinity context for the given (user, project, agent) triple.
-func (s *sqliteWebChatStore) RecordChannel(ctx context.Context, userID, projectID, agentID, channel string, messageAt time.Time) error {
+func (s *sqliteWebChatStore) RecordChannel(ctx context.Context, userID, projectID, agentID, channel, threadID string, messageAt time.Time) error {
 	const query = `
-INSERT INTO webchat_conversation_context (user_id, project_id, agent_id, last_channel, last_message_at)
-VALUES (?, ?, ?, ?, ?)
+INSERT INTO webchat_conversation_context (user_id, project_id, agent_id, last_channel, last_thread_id, last_message_at)
+VALUES (?, ?, ?, ?, ?, ?)
 ON CONFLICT (user_id, project_id, agent_id)
 DO UPDATE SET
     last_channel = excluded.last_channel,
+    last_thread_id = excluded.last_thread_id,
     last_message_at = excluded.last_message_at
 `
-	_, err := s.db.ExecContext(ctx, query, userID, projectID, agentID, channel, messageAt)
+	_, err := s.db.ExecContext(ctx, query, userID, projectID, agentID, channel, threadID, messageAt)
 	if err != nil {
 		return fmt.Errorf("webchat store: record channel: %w", err)
 	}
 	return nil
 }
 
-// GetLastChannel returns the last channel for (user, project, agent), or "" if no row exists.
-func (s *sqliteWebChatStore) GetLastChannel(ctx context.Context, userID, projectID, agentID string) (string, error) {
-	const query = `SELECT last_channel FROM webchat_conversation_context WHERE user_id = ? AND project_id = ? AND agent_id = ?`
-	var channel sql.NullString
-	err := s.db.QueryRowContext(ctx, query, userID, projectID, agentID).Scan(&channel)
+// GetLastRoute returns where this user last spoke to this agent. The thread is
+// empty for channels that do not carry threads, and both are empty when no row
+// exists.
+func (s *sqliteWebChatStore) GetLastRoute(ctx context.Context, userID, projectID, agentID string) (string, string, error) {
+	const query = `SELECT last_channel, last_thread_id FROM webchat_conversation_context WHERE user_id = ? AND project_id = ? AND agent_id = ?`
+	var channel, threadID sql.NullString
+	err := s.db.QueryRowContext(ctx, query, userID, projectID, agentID).Scan(&channel, &threadID)
 	if err != nil {
 		if err == sql.ErrNoRows {
-			return "", nil
+			return "", "", nil
 		}
-		return "", fmt.Errorf("webchat store: get last channel: %w", err)
+		return "", "", fmt.Errorf("webchat store: get last route: %w", err)
 	}
-	return channel.String, nil
+	return channel.String, threadID.String, nil
 }
 
 // GetThreadPrefs returns the display preferences for the given (user, project, agent) triple.
@@ -1215,6 +1230,9 @@ SELECT id, project_id, COALESCE(thread_id, ''), sender, msg, created
 
 // runMigrations executes idempotent data migrations.
 func (s *sqliteWebChatStore) runMigrations() error {
+	if err := s.addConversationContextThreadID(); err != nil {
+		return fmt.Errorf("conversation context thread_id: %w", err)
+	}
 	if err := s.migrateThreadIDs(DefaultMigrationBatchSize); err != nil {
 		return fmt.Errorf("thread_id backfill: %w", err)
 	}

pkg/hub/webchannel_store_postgres.go

@@ -52,6 +52,7 @@ CREATE TABLE IF NOT EXISTS webchat_conversation_context (
     project_id TEXT NOT NULL,
     agent_id TEXT NOT NULL,
     last_channel TEXT,
+    last_thread_id TEXT,
     last_message_at TIMESTAMPTZ,
     PRIMARY KEY (user_id, project_id, agent_id)
 );
@@ -173,6 +174,14 @@ CREATE UNIQUE INDEX IF NOT EXISTS idx_webchat_topic_project_name
 		return fmt.Errorf("webchat store: create name uniqueness index: %w", err)
 	}
 
+	// Adds last_thread_id to databases created before it was in the schema.
+	const ctxThreadCol = `
+ALTER TABLE webchat_conversation_context ADD COLUMN IF NOT EXISTS last_thread_id TEXT;
+`
+	if _, err := s.db.Exec(ctxThreadCol); err != nil {
+		return fmt.Errorf("webchat store: add conversation context thread_id: %w", err)
+	}
+
 	// Run idempotent migrations.
 	if err := s.runMigrations(); err != nil {
 		return fmt.Errorf("webchat store: migrations: %w", err)
@@ -199,34 +208,37 @@ DO UPDATE SET
 }
 
 // RecordChannel upserts the reply-affinity context for the given (user, project, agent) triple.
-func (s *pgWebChatStore) RecordChannel(ctx context.Context, userID, projectID, agentID, channel string, messageAt time.Time) error {
+func (s *pgWebChatStore) RecordChannel(ctx context.Context, userID, projectID, agentID, channel, threadID string, messageAt time.Time) error {
 	const query = `
-INSERT INTO webchat_conversation_context (user_id, project_id, agent_id, last_channel, last_message_at)
-VALUES ($1, $2, $3, $4, $5)
+INSERT INTO webchat_conversation_context (user_id, project_id, agent_id, last_channel, last_thread_id, last_message_at)
+VALUES ($1, $2, $3, $4, $5, $6)
 ON CONFLICT (user_id, project_id, agent_id)
 DO UPDATE SET
     last_channel = EXCLUDED.last_channel,
+    last_thread_id = EXCLUDED.last_thread_id,
     last_message_at = EXCLUDED.last_message_at
 `
-	_, err := s.db.ExecContext(ctx, query, userID, projectID, agentID, channel, messageAt)
+	_, err := s.db.ExecContext(ctx, query, userID, projectID, agentID, channel, threadID, messageAt)
 	if err != nil {
 		return fmt.Errorf("webchat store: record channel: %w", err)
 	}
 	return nil
 }
 
-// GetLastChannel returns the last channel for (user, project, agent), or "" if no row exists.
-func (s *pgWebChatStore) GetLastChannel(ctx context.Context, userID, projectID, agentID string) (string, error) {
-	const query = `SELECT last_channel FROM webchat_conversation_context WHERE user_id = $1 AND project_id = $2 AND agent_id = $3`
-	var channel sql.NullString
-	err := s.db.QueryRowContext(ctx, query, userID, projectID, agentID).Scan(&channel)
+// GetLastRoute returns where this user last spoke to this agent. The thread is
+// empty for channels that do not carry threads, and both are empty when no row
+// exists.
+func (s *pgWebChatStore) GetLastRoute(ctx context.Context, userID, projectID, agentID string) (string, string, error) {
+	const query = `SELECT last_channel, last_thread_id FROM webchat_conversation_context WHERE user_id = $1 AND project_id = $2 AND agent_id = $3`
+	var channel, threadID sql.NullString
+	err := s.db.QueryRowContext(ctx, query, userID, projectID, agentID).Scan(&channel, &threadID)
 	if err != nil {
 		if err == sql.ErrNoRows {
-			return "", nil
+			return "", "", nil
 		}
-		return "", fmt.Errorf("webchat store: get last channel: %w", err)
+		return "", "", fmt.Errorf("webchat store: get last route: %w", err)
 	}
-	return channel.String, nil
+	return channel.String, threadID.String, nil
 }
 
 // GetThreadPrefs returns the display preferences for the given (user, project, agent) triple.