mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: tag chat-originating agent logs with chat_id (#25019)
Workspace-agent logs emitted while serving chatd-driven requests were not correlated with the originating chat, making agent logs hard to attribute to the corresponding/originating chat. This adds agent-side chat context middleware that parses `Coder-Chat-Id` once, enriches agent access logs and structured handler/background logs, and adds a chatd bridge log when chat headers are attached to an agent connection. Closes CODAGT-324
This commit is contained in:
Ethan
@@ -0,0 +1,35 @@
|
||||
package agentchat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
)
|
||||
|
||||
// extractContext reads chat identity headers from the request.
|
||||
// Returns zero values if headers are absent (non-chat request).
|
||||
func extractContext(r *http.Request) (chatID uuid.UUID, ancestorIDs []uuid.UUID, ok bool) {
|
||||
raw := r.Header.Get(workspacesdk.CoderChatIDHeader)
|
||||
if raw == "" {
|
||||
return uuid.Nil, nil, false
|
||||
}
|
||||
chatID, err := uuid.Parse(raw)
|
||||
if err != nil {
|
||||
return uuid.Nil, nil, false
|
||||
}
|
||||
rawAncestors := r.Header.Get(workspacesdk.CoderAncestorChatIDsHeader)
|
||||
if rawAncestors != "" {
|
||||
var ids []string
|
||||
if err := json.Unmarshal([]byte(rawAncestors), &ids); err == nil {
|
||||
for _, s := range ids {
|
||||
if id, err := uuid.Parse(s); err == nil {
|
||||
ancestorIDs = append(ancestorIDs, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return chatID, ancestorIDs, true
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package agentchat_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/agent/agentchat"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
)
|
||||
|
||||
func TestExtractContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
validID := uuid.MustParse("aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee")
|
||||
ancestor1 := uuid.MustParse("11111111-2222-3333-4444-555555555555")
|
||||
ancestor2 := uuid.MustParse("66666666-7777-8888-9999-aaaaaaaaaaaa")
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
chatID string // empty means header not set
|
||||
setChatID bool // whether to set the chat ID header at all
|
||||
ancestors string // empty means header not set
|
||||
setAncestors bool // whether to set the ancestor header at all
|
||||
wantChatID uuid.UUID
|
||||
wantAncestorIDs []uuid.UUID
|
||||
wantOK bool
|
||||
}{
|
||||
{
|
||||
name: "NoHeadersPresent",
|
||||
setChatID: false,
|
||||
setAncestors: false,
|
||||
wantChatID: uuid.Nil,
|
||||
wantAncestorIDs: nil,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "ValidChatID_NoAncestors",
|
||||
chatID: validID.String(),
|
||||
setChatID: true,
|
||||
setAncestors: false,
|
||||
wantChatID: validID,
|
||||
wantAncestorIDs: []uuid.UUID{},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "ValidChatID_ValidAncestors",
|
||||
chatID: validID.String(),
|
||||
setChatID: true,
|
||||
ancestors: mustMarshalJSON(t, []string{
|
||||
ancestor1.String(),
|
||||
ancestor2.String(),
|
||||
}),
|
||||
setAncestors: true,
|
||||
wantChatID: validID,
|
||||
wantAncestorIDs: []uuid.UUID{ancestor1, ancestor2},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
name: "MalformedChatID",
|
||||
chatID: "not-a-uuid",
|
||||
setChatID: true,
|
||||
setAncestors: false,
|
||||
wantChatID: uuid.Nil,
|
||||
wantAncestorIDs: nil,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "ValidChatID_MalformedAncestorJSON",
|
||||
chatID: validID.String(),
|
||||
setChatID: true,
|
||||
ancestors: `{this is not json}`,
|
||||
setAncestors: true,
|
||||
wantChatID: validID,
|
||||
wantAncestorIDs: []uuid.UUID{},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
// Only valid UUIDs in the array are returned; invalid
|
||||
// entries are silently skipped.
|
||||
name: "ValidChatID_PartialValidAncestorUUIDs",
|
||||
chatID: validID.String(),
|
||||
setChatID: true,
|
||||
ancestors: mustMarshalJSON(t, []string{
|
||||
ancestor1.String(),
|
||||
"bad-uuid",
|
||||
ancestor2.String(),
|
||||
}),
|
||||
setAncestors: true,
|
||||
wantChatID: validID,
|
||||
wantAncestorIDs: []uuid.UUID{ancestor1, ancestor2},
|
||||
wantOK: true,
|
||||
},
|
||||
{
|
||||
// Header is explicitly set to an empty string, which
|
||||
// Header.Get returns as "".
|
||||
name: "EmptyChatIDHeader",
|
||||
chatID: "",
|
||||
setChatID: true,
|
||||
setAncestors: false,
|
||||
wantChatID: uuid.Nil,
|
||||
wantAncestorIDs: nil,
|
||||
wantOK: false,
|
||||
},
|
||||
{
|
||||
name: "ValidChatID_EmptyAncestorHeader",
|
||||
chatID: validID.String(),
|
||||
setChatID: true,
|
||||
ancestors: "",
|
||||
setAncestors: true,
|
||||
wantChatID: validID,
|
||||
wantAncestorIDs: []uuid.UUID{},
|
||||
wantOK: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
tt := tt
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
r := httptest.NewRequest("GET", "/", nil)
|
||||
if tt.setChatID {
|
||||
r.Header.Set(workspacesdk.CoderChatIDHeader, tt.chatID)
|
||||
}
|
||||
if tt.setAncestors {
|
||||
r.Header.Set(workspacesdk.CoderAncestorChatIDsHeader, tt.ancestors)
|
||||
}
|
||||
|
||||
chatID, ancestorIDs, ok := extractContextForTest(r)
|
||||
|
||||
require.Equal(t, tt.wantOK, ok, "ok mismatch")
|
||||
require.Equal(t, tt.wantChatID, chatID, "chatID mismatch")
|
||||
require.Equal(t, tt.wantAncestorIDs, ancestorIDs, "ancestorIDs mismatch")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func extractContextForTest(r *http.Request) (uuid.UUID, []uuid.UUID, bool) {
|
||||
var chatContext agentchat.Context
|
||||
var ok bool
|
||||
agentchat.Middleware(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
|
||||
chatContext, ok = agentchat.FromContext(r.Context())
|
||||
})).ServeHTTP(httptest.NewRecorder(), r)
|
||||
if !ok {
|
||||
return uuid.Nil, nil, false
|
||||
}
|
||||
return chatContext.ID, chatContext.AncestorIDs, true
|
||||
}
|
||||
|
||||
// mustMarshalJSON marshals v to a JSON string, failing the test on error.
|
||||
func mustMarshalJSON(t *testing.T, v any) string {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(v)
|
||||
require.NoError(t, err)
|
||||
return string(b)
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package agentchat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/httpmw/loggermw"
|
||||
)
|
||||
|
||||
type chatContextKey struct{}
|
||||
|
||||
// Context carries the chat identity associated with an agent request.
|
||||
type Context struct {
|
||||
ID uuid.UUID
|
||||
AncestorIDs []uuid.UUID
|
||||
}
|
||||
|
||||
// FromContext returns the chat identity stored on the context.
|
||||
func FromContext(ctx context.Context) (Context, bool) {
|
||||
chatCtx, ok := ctx.Value(chatContextKey{}).(Context)
|
||||
if !ok || chatCtx.ID == uuid.Nil {
|
||||
return Context{}, false
|
||||
}
|
||||
return chatCtx, true
|
||||
}
|
||||
|
||||
// WithContext stores chat identity on the context for downstream logs.
|
||||
func WithContext(ctx context.Context, chatID uuid.UUID, ancestorIDs []uuid.UUID) context.Context {
|
||||
if chatID == uuid.Nil {
|
||||
return ctx
|
||||
}
|
||||
ancestors := make([]uuid.UUID, len(ancestorIDs))
|
||||
copy(ancestors, ancestorIDs)
|
||||
return context.WithValue(ctx, chatContextKey{}, Context{
|
||||
ID: chatID,
|
||||
AncestorIDs: ancestors,
|
||||
})
|
||||
}
|
||||
|
||||
// Fields returns structured log fields for the chat identity on ctx.
|
||||
func Fields(ctx context.Context) []slog.Field {
|
||||
chatCtx, ok := FromContext(ctx)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
return chatFields(chatCtx.ID, chatCtx.AncestorIDs)
|
||||
}
|
||||
|
||||
// Middleware tags agent logs for requests that originate from
|
||||
// chatd. Agent log lines emitted while serving a request with Coder-Chat-Id,
|
||||
// or by background work started by such a request, should include chat_id.
|
||||
// Install after loggermw.Logger so access-log enrichment can run.
|
||||
func Middleware(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
|
||||
chatID, ancestorIDs, ok := extractContext(r)
|
||||
if !ok {
|
||||
next.ServeHTTP(rw, r)
|
||||
return
|
||||
}
|
||||
|
||||
fields := chatFields(chatID, ancestorIDs)
|
||||
if requestLogger := loggermw.RequestLoggerFromContext(r.Context()); requestLogger != nil {
|
||||
requestLogger.WithFields(fields...)
|
||||
}
|
||||
|
||||
ctx := WithContext(r.Context(), chatID, ancestorIDs)
|
||||
next.ServeHTTP(rw, r.WithContext(ctx))
|
||||
})
|
||||
}
|
||||
|
||||
func chatFields(chatID uuid.UUID, ancestorIDs []uuid.UUID) []slog.Field {
|
||||
fields := []slog.Field{slog.F("chat_id", chatID.String())}
|
||||
if len(ancestorIDs) == 0 {
|
||||
return fields
|
||||
}
|
||||
|
||||
ancestors := make([]string, 0, len(ancestorIDs))
|
||||
for _, id := range ancestorIDs {
|
||||
ancestors = append(ancestors, id.String())
|
||||
}
|
||||
return append(fields, slog.F("ancestor_chat_ids", ancestors))
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package agentchat_test
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/agent/agentchat"
|
||||
"github.com/coder/coder/v2/coderd/httpmw/loggermw"
|
||||
"github.com/coder/coder/v2/coderd/tracing"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
func TestMiddlewareAccessLog(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
chatID := uuid.New()
|
||||
ancestorID := uuid.New()
|
||||
sink := testutil.NewFakeSink(t)
|
||||
handler := tracing.StatusWriterMiddleware(loggermw.Logger(sink.Logger())(
|
||||
agentchat.Middleware(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
|
||||
rw.WriteHeader(http.StatusNoContent)
|
||||
})),
|
||||
))
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.Header.Set(workspacesdk.CoderChatIDHeader, chatID.String())
|
||||
req.Header.Set(workspacesdk.CoderAncestorChatIDsHeader, mustMarshalJSON(t, []string{ancestorID.String()}))
|
||||
rw := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rw, req)
|
||||
require.Equal(t, http.StatusNoContent, rw.Code)
|
||||
|
||||
entries := sink.Entries()
|
||||
require.Len(t, entries, 1)
|
||||
fields := fieldsByName(entries[0].Fields)
|
||||
require.Equal(t, chatID.String(), fields["chat_id"])
|
||||
require.Equal(t, []string{ancestorID.String()}, fields["ancestor_chat_ids"])
|
||||
}
|
||||
|
||||
func TestMiddlewareWithoutChatHeader(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
sink := testutil.NewFakeSink(t)
|
||||
handler := tracing.StatusWriterMiddleware(loggermw.Logger(sink.Logger())(
|
||||
agentchat.Middleware(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) {
|
||||
rw.WriteHeader(http.StatusNoContent)
|
||||
})),
|
||||
))
|
||||
|
||||
rw := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rw, httptest.NewRequest(http.MethodGet, "/test", nil))
|
||||
require.Equal(t, http.StatusNoContent, rw.Code)
|
||||
|
||||
entries := sink.Entries()
|
||||
require.Len(t, entries, 1)
|
||||
fields := fieldsByName(entries[0].Fields)
|
||||
require.NotContains(t, fields, "chat_id")
|
||||
require.NotContains(t, fields, "ancestor_chat_ids")
|
||||
}
|
||||
|
||||
func TestMiddlewareContextFields(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
chatID := uuid.New()
|
||||
sink := testutil.NewFakeSink(t)
|
||||
handler := tracing.StatusWriterMiddleware(loggermw.Logger(sink.Logger())(
|
||||
agentchat.Middleware(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
|
||||
sink.Logger().With(agentchat.Fields(r.Context())...).Info(r.Context(), "handler log")
|
||||
rw.WriteHeader(http.StatusNoContent)
|
||||
})),
|
||||
))
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/test", nil)
|
||||
req.Header.Set(workspacesdk.CoderChatIDHeader, chatID.String())
|
||||
rw := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rw, req)
|
||||
require.Equal(t, http.StatusNoContent, rw.Code)
|
||||
|
||||
entries := sink.Entries()
|
||||
require.Len(t, entries, 2)
|
||||
for _, entry := range entries {
|
||||
if entry.Message != "handler log" {
|
||||
continue
|
||||
}
|
||||
fields := fieldsByName(entry.Fields)
|
||||
require.Equal(t, chatID.String(), fields["chat_id"])
|
||||
return
|
||||
}
|
||||
t.Fatal("handler log entry not found")
|
||||
}
|
||||
|
||||
func fieldsByName(fields []slog.Field) map[string]any {
|
||||
byName := make(map[string]any, len(fields))
|
||||
for _, field := range fields {
|
||||
byName[field.Name] = field.Value
|
||||
}
|
||||
return byName
|
||||
}
|
||||
Reference in New Issue
Block a user