Repository: google/adk-go
File: server/adkrest/internal/services/debugtelemetry.go (Go)

Code before:
	s.mu.Lock()
	defer s.mu.Unlock()
	for _, span := range spans {
		attrs := convertAttrs(span.Attributes())
		spanID := span.SpanContext().SpanID().String()
		record, ok := s.recordsBySpanID[spanID]
		if !ok {
			record = &spanRecord{}
			s.recordsBySpanID[spanID] = record
		}

		record.Span = &inMemorySpan{
			Name:         span.Name(),
			StartTime:    span.StartTime(),
			EndTime:      span.EndTime(),
			Context:      span.SpanContext(),
			ParentSpanID: span.Parent().SpanID(),
			Attributes:   attrs,
		}

		s.updateSpanIndexes(record.Span, record)
	}
	return nil
}

func (s *spanStore) updateSpanIndexes(span *inMemorySpan, records *spanRecord) {
	// Update session id -> trace id mapping.
	sessionIDKey := string(semconv.GenAIConversationIDKey)
	if sessionID, ok := span.Attributes[sessionIDKey]; ok {
		traces, ok := s.traceIDsBySessionID[sessionID]
		if !ok {
			traces = make(map[string]struct{})
			s.traceIDsBySessionID[sessionID] = traces
		}
		traceID := span.Context.TraceID().String()
		traces[traceID] = struct{}{}
	}
	// Update event id -> span id mapping.
	if eventID, ok := span.Attributes["gcp.vertex.agent.event_id"]; ok {
		s.recordsByEventID[eventID] = append(s.recordsByEventID[eventID], records)
	}
	// Update trace id -> span id mapping.
	traceID := span.Context.TraceID().String()
	s.recordsByTraceID[traceID] = append(s.recordsByTraceID[traceID], records)
}

// ForceFlush implements sdklog.Exporter and sdktrace.SpanProcessor.
func (s *spanStore) ForceFlush(ctx context.Context) error {
	return nil
}


Review comment (suggestion):
```suggestion
func (s *spanStore) updateSpanIndexes(span *inMemorySpan, record *spanRecord) {
```

records -> record as it's only one object

Code after:
		record, ok := s.recordsBySpanID[spanID]
		if !ok {
			record = &spanRecord{}
			s.recordsBySpanID[spanID] = record
		}

		record.Span = &inMemorySpan{
			Name:         span.Name(),
			StartTime:    span.StartTime(),
			EndTime:      span.EndTime(),
			Context:      span.SpanContext(),
			ParentSpanID: span.Parent().SpanID(),
			Attributes:   attrs,
		}

		s.updateSpanIndexes(record.Span, record)
	}
	return nil
}

func (s *spanStore) updateSpanIndexes(span *inMemorySpan, record *spanRecord) {
	// Update session id -> trace id mapping.
	sessionIDKey := string(semconv.GenAIConversationIDKey)
	if sessionID, ok := span.Attributes[sessionIDKey]; ok {
		traces, ok := s.traceIDsBySessionID[sessionID]
		if !ok {
			traces = make(map[string]struct{})
			s.traceIDsBySessionID[sessionID] = traces
		}
		traceID := span.Context.TraceID().String()
		traces[traceID] = struct{}{}
	}
	// Update event id -> span id mapping.
	if eventID, ok := span.Attributes[eventIDKey]; ok {
		s.recordsByEventID[eventID] = append(s.recordsByEventID[eventID], record)
	}
	// Update trace id -> span id mapping.
	traceID := span.Context.TraceID().String()
	s.recordsByTraceID[traceID] = append(s.recordsByTraceID[traceID], record)
}

// ForceFlush implements sdklog.Exporter and sdktrace.SpanProcessor.
func (s *spanStore) ForceFlush(ctx context.Context) error {
	return nil
}

// Shutdown implements sdklog.Exporter and sdktrace.SpanProcessor.
func (s *spanStore) Shutdown(ctx context.Context) error {
	return nil
}