Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions api/pkg/org/infrastructure/runtime/helix/sessions.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,12 +31,17 @@ type SessionClient interface {
// - GetOutput: polled by the Spawner's pollUntilDone loop until the
// session reports a terminal status.
// - StopExternalAgent: used by the chat bridge's NewHandler.
// - SessionOwner: resolves the owning user ID so the transcript
// bridge can subscribe to the correct per-session pubsub topic.
// Helix publishes session updates to GetSessionQueue(owner, id), so
// the bridge must subscribe with the owner — not an empty string.
//
// Production impl is the in-process inProcHelixClient adapter.
type SpawnerClient interface {
SessionClient
GetOutput(ctx context.Context, sessionID string) (types.SessionOutputResponse, error)
StopExternalAgent(ctx context.Context, sessionID string) error
SessionOwner(ctx context.Context, sessionID string) (string, error)
}

// sendToSession pushes a message to an existing session via
Expand Down
20 changes: 19 additions & 1 deletion api/pkg/org/infrastructure/runtime/helix/spawner.go
Original file line number Diff line number Diff line change
Expand Up @@ -545,9 +545,27 @@ func transcriptSegmentFromEvent(e Event) (activation.TranscriptSegment, bool) {
}

func (b *bridge) run(ctx context.Context, cfg SpawnerConfig, sessionID string) {
// Resolve the session owner once. Helix publishes every session
// update (assistant text, tool_use, tool_result) to
// GetSessionQueue(session.Owner, sessionID); subscribing with an
// empty owner lands us on the wrong topic and the bridge receives
// zero frames — leaving only the spawner's own lifecycle markers on
// the activation stream. Owner never changes, so resolve it before
// the reconnect loop rather than on every reconnect.
ownerID := ""
if cfg.Client != nil {
if owner, err := cfg.Client.SessionOwner(ctx, sessionID); err != nil {
if cfg.Logger != nil {
cfg.Logger.Warn("helix transcript bridge: resolve session owner", "session", sessionID, "err", err)
}
} else {
ownerID = owner
}
}

delay := time.Second
for {
ch, err := SubscribeSessionUpdates(ctx, cfg.PubSub, cfg.Snapshotter, "", sessionID)
ch, err := SubscribeSessionUpdates(ctx, cfg.PubSub, cfg.Snapshotter, ownerID, sessionID)
if err != nil {
if cfg.Logger != nil {
cfg.Logger.Warn("helix subscribe", "session", sessionID, "err", err)
Expand Down
15 changes: 14 additions & 1 deletion api/pkg/org/infrastructure/runtime/helix/spawner_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ type fakeHelixClient struct {
outputCalls int32
subscribeCalls int32
startSessionID string
sessionOwner string // returned by SessionOwner; the transcript bridge subscribes to this owner's pubsub topic
startErr error
sendErr error
outputs []types.SessionOutputResponse
Expand Down Expand Up @@ -58,6 +59,9 @@ func (f *fakeHelixClient) GetOutput(_ context.Context, _ string) (types.SessionO
}

func (f *fakeHelixClient) StopExternalAgent(_ context.Context, _ string) error { return nil }
func (f *fakeHelixClient) SessionOwner(_ context.Context, _ string) (string, error) {
return f.sessionOwner, nil
}
func (f *fakeHelixClient) ServerStatus(_ context.Context) (ServerStatus, error) {
return ServerStatus{MaxConcurrentDesktops: 0, ActiveConcurrentDesktops: 0}, nil
}
Expand Down Expand Up @@ -422,6 +426,10 @@ func (c *concurrencyClient) StopExternalAgent(ctx context.Context, sid string) e
return c.inner.StopExternalAgent(ctx, sid)
}

func (c *concurrencyClient) SessionOwner(ctx context.Context, sid string) (string, error) {
return c.inner.SessionOwner(ctx, sid)
}

// TestSpawnerPublishesTranscriptViaEntryStream verifies the bridge
// subscribes via pubsub.GetSessionQueue, feeds frames through
// EntryStream, and republishes settled events as activation Stream
Expand All @@ -433,6 +441,11 @@ func TestSpawnerPublishesTranscriptViaEntryStream(t *testing.T) {
ps := newFakePubSub()
fc := &fakeHelixClient{
startSessionID: "ses_y",
// The session is owned by a real user; helix publishes its
// updates to that owner's pubsub topic. The bridge must resolve
// the owner (via SessionOwner) and subscribe there — subscribing
// with an empty owner is the regression this test guards.
sessionOwner: "u-owner",
// Several waiting outputs so the bridge has time to consume
// the pubsub frames before pollUntilDone terminates.
outputs: []types.SessionOutputResponse{
Expand All @@ -457,7 +470,7 @@ func TestSpawnerPublishesTranscriptViaEntryStream(t *testing.T) {

// Wait for the bridge to subscribe (handlers map populated).
deadline := time.Now().Add(2 * time.Second)
topic := pubsub.GetSessionQueue("", "ses_y")
topic := pubsub.GetSessionQueue("u-owner", "ses_y")
for time.Now().Before(deadline) {
if ps.handlerCount(topic) > 0 {
break
Expand Down
18 changes: 18 additions & 0 deletions api/pkg/server/helix_org_inproc.go
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,24 @@ func (c *inProcHelixClient) GetOutput(ctx context.Context, sessionID string) (ty
return *resp, nil
}

// SessionOwner returns the user ID that owns the session. The
// transcript bridge needs it to subscribe to the correct per-session
// pubsub topic: helix publishes session updates to
// GetSessionQueue(session.Owner, …), so subscribing with an empty owner
// (or the wrong user) silently yields zero frames — only the spawner's
// own lifecycle markers reach the activation stream. Mirrors the owner
// lookup in websocket_server_user.go.
func (c *inProcHelixClient) SessionOwner(ctx context.Context, sessionID string) (string, error) {
session, err := c.server.Store.GetSession(ctx, sessionID)
if err != nil {
return "", fmt.Errorf("get session %s: %w", sessionID, err)
}
if session == nil {
return "", fmt.Errorf("get session %s: not found", sessionID)
}
return session.Owner, nil
}

// StopExternalAgent stops a session's external Zed agent.
func (c *inProcHelixClient) StopExternalAgent(ctx context.Context, sessionID string) error {
r, err := c.newRequest(ctx, http.MethodDelete, "/api/v1/sessions/"+sessionID+"/stop-external-agent", nil, map[string]string{"id": sessionID})
Expand Down