diff --git a/README.md b/README.md index 5a72516..296c38b 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,49 @@ from their usual local locations. configuration for trusted local or private-network deployments - Privacy controls for redaction, hashing, capture filters, and offline operation +## Annotated Trace Datasets + +Beacon annotations can mark a whole session, a transcript message, or a +specific event. The local JSON API exposes annotated traces directly for review, +evaluation, fine-tuning, and skill-development datasets. + +List annotated sessions and their annotated targets: + +```bash +curl 'http://localhost:4600/api/annotations/traces?label=dataset:eval&limit=25' +``` + +Export dataset-ready traces with session metadata, ordered event context, and +annotation records: + +```bash +curl 'http://localhost:4600/api/annotations/export?label=dataset:eval&event_limit=2000' +``` + +Discovery and export responses are paginated with `limit`, `offset`, and +`has_more`. Continue increasing `offset` until `has_more` is false when +collecting a complete dataset: + +```bash +offset=0 +while :; do + curl -fsS "http://localhost:4600/api/annotations/export?label=dataset:eval&event_limit=2000&limit=200&offset=${offset}" > "annotated-traces-${offset}.json" + jq -e '.has_more' "annotated-traces-${offset}.json" >/dev/null || break + offset=$((offset + 200)) +done +``` + +Each exported trace reports `event_truncated`, and the response includes +`warnings` when `event_limit` clipped ordered event context for a session. + +Both endpoints return versioned JSON schema markers: +`beacon.annotated_traces.index.v1` and +`beacon.annotated_traces.export.v1`. Supported filters include `session_id`, +`event_uid`, `target_type`, `label`, `author_type`, `source`, `category`, +`outcome`, `needs_followup`, `include_deleted`, and the usual Beacon scope +filters (`source_name`, `source_names`, `runtime`, `runtimes`, `project_key`, +`project_keys`). + Check the local setup: ```bash diff --git a/internal/store/annotations.go b/internal/store/annotations.go index c3ef245..15ceb52 100644 --- a/internal/store/annotations.go +++ b/internal/store/annotations.go @@ -22,12 +22,30 @@ type AnnotationFilter struct { AnnotationID string TargetType string SessionID string + SessionIDs []string EventUID string + AuthorType string + Source string + Category string + Outcome string + Label string + NeedsFollowup *bool IncludeDeleted bool Limit int Offset int } +type TraceAnnotationSessionSummary struct { + SessionID string + AnnotationCount int + SessionAnnotationCount int + MessageAnnotationCount int + EventAnnotationCount int + NeedsFollowupCount int + FirstAnnotationAt time.Time + LastAnnotationAt time.Time +} + type AnnotationUpdate struct { AuthorType *string AuthorID *string @@ -112,38 +130,9 @@ func ListTraceAnnotations(ctx context.Context, db *sql.DB, filter AnnotationFilt filter.TargetType = strings.TrimSpace(strings.ToLower(filter.TargetType)) filter.SessionID = strings.TrimSpace(filter.SessionID) filter.EventUID = strings.TrimSpace(filter.EventUID) - if filter.Limit <= 0 { - filter.Limit = defaultAnnotationLimit - } - if filter.Limit > maxAnnotationLimit { - filter.Limit = maxAnnotationLimit - } - if filter.Offset < 0 { - filter.Offset = 0 - } - - where := []string{"1 = 1"} - args := []any{} - if filter.AnnotationID != "" { - where = append(where, "annotation_id = ?") - args = append(args, filter.AnnotationID) - } - if filter.TargetType != "" { - where = append(where, "target_type = ?") - args = append(args, filter.TargetType) - } - if filter.SessionID != "" { - where = append(where, "session_id = ?") - args = append(args, filter.SessionID) - } - if filter.EventUID != "" { - where = append(where, "event_uid = ?") - args = append(args, filter.EventUID) - } - if !filter.IncludeDeleted { - where = append(where, "status != ?") - args = append(args, models.AnnotationStatusDeleted) - } + filter.Limit = normalizeAnnotationLimit(filter.Limit) + filter.Offset = normalizeAnnotationOffset(filter.Offset) + where, args := annotationFilterWhere(filter) args = append(args, filter.Limit, filter.Offset) rows, err := db.QueryContext(ctx, traceAnnotationSelectSQL("WHERE "+strings.Join(where, " AND ")+" ORDER BY created_at ASC, annotation_id ASC LIMIT ? OFFSET ?"), args...) @@ -166,6 +155,66 @@ func ListTraceAnnotations(ctx context.Context, db *sql.DB, filter AnnotationFilt return annotations, nil } +func ListTraceAnnotationSessionSummaries(ctx context.Context, db *sql.DB, filter AnnotationFilter) ([]TraceAnnotationSessionSummary, error) { + if db == nil { + return nil, fmt.Errorf("database is not configured") + } + filter.Limit = normalizeAnnotationLimit(filter.Limit) + filter.Offset = normalizeAnnotationOffset(filter.Offset) + where, args := annotationFilterWhere(filter) + args = append(args, filter.Limit, filter.Offset) + rows, err := db.QueryContext(ctx, `SELECT session_id, + count() AS annotation_count, + countIf(target_type = ?) AS session_annotation_count, + countIf(target_type = ?) AS message_annotation_count, + countIf(target_type = ?) AS event_annotation_count, + countIf(needs_followup != 0) AS needs_followup_count, + min(created_at) AS first_annotation_at, + max(updated_at) AS last_annotation_at + FROM trace_annotations FINAL + WHERE `+strings.Join(where, " AND ")+` + GROUP BY session_id + ORDER BY last_annotation_at DESC, session_id ASC + LIMIT ? OFFSET ?`, + append([]any{ + models.AnnotationTargetSession, + models.AnnotationTargetMessage, + models.AnnotationTargetEvent, + }, args...)...) + if err != nil { + return nil, err + } + defer rows.Close() + + summaries := make([]TraceAnnotationSessionSummary, 0) + for rows.Next() { + var summary TraceAnnotationSessionSummary + var annotationCount, sessionCount, messageCount, eventCount, followupCount uint64 + if err := rows.Scan( + &summary.SessionID, + &annotationCount, + &sessionCount, + &messageCount, + &eventCount, + &followupCount, + &summary.FirstAnnotationAt, + &summary.LastAnnotationAt, + ); err != nil { + return nil, err + } + summary.AnnotationCount = int(annotationCount) + summary.SessionAnnotationCount = int(sessionCount) + summary.MessageAnnotationCount = int(messageCount) + summary.EventAnnotationCount = int(eventCount) + summary.NeedsFollowupCount = int(followupCount) + summaries = append(summaries, summary) + } + if err := rows.Err(); err != nil { + return nil, err + } + return summaries, nil +} + func UpdateTraceAnnotation(ctx context.Context, db *sql.DB, annotationID string, update AnnotationUpdate) (models.TraceAnnotation, error) { current, err := GetTraceAnnotation(ctx, db, annotationID, false) if err != nil { @@ -417,6 +466,110 @@ func splitAnnotationLabels(value string) []string { return models.NormalizeAnnotationLabels(strings.Split(value, annotationLabelJoiner)) } +func annotationFilterWhere(filter AnnotationFilter) ([]string, []any) { + filter.AnnotationID = strings.TrimSpace(filter.AnnotationID) + filter.TargetType = strings.TrimSpace(strings.ToLower(filter.TargetType)) + filter.SessionID = strings.TrimSpace(filter.SessionID) + filter.EventUID = strings.TrimSpace(filter.EventUID) + filter.AuthorType = strings.TrimSpace(strings.ToLower(filter.AuthorType)) + filter.Source = strings.TrimSpace(strings.ToLower(filter.Source)) + filter.Category = strings.TrimSpace(strings.ToLower(filter.Category)) + filter.Outcome = strings.TrimSpace(strings.ToLower(filter.Outcome)) + filter.Label = strings.TrimSpace(strings.ToLower(filter.Label)) + filter.SessionIDs = compactAnnotationStrings(filter.SessionIDs) + + where := []string{"1 = 1"} + args := []any{} + if filter.AnnotationID != "" { + where = append(where, "annotation_id = ?") + args = append(args, filter.AnnotationID) + } + if filter.TargetType != "" { + where = append(where, "target_type = ?") + args = append(args, filter.TargetType) + } + if filter.SessionID != "" { + where = append(where, "session_id = ?") + args = append(args, filter.SessionID) + } + if len(filter.SessionIDs) > 0 { + where = append(where, "session_id IN ("+strings.TrimRight(strings.Repeat("?,", len(filter.SessionIDs)), ",")+")") + for _, sessionID := range filter.SessionIDs { + args = append(args, sessionID) + } + } + if filter.EventUID != "" { + where = append(where, "event_uid = ?") + args = append(args, filter.EventUID) + } + if filter.AuthorType != "" { + where = append(where, "author_type = ?") + args = append(args, filter.AuthorType) + } + if filter.Source != "" { + where = append(where, "source = ?") + args = append(args, filter.Source) + } + if filter.Category != "" { + where = append(where, "category = ?") + args = append(args, filter.Category) + } + if filter.Outcome != "" { + where = append(where, "outcome = ?") + args = append(args, filter.Outcome) + } + if filter.Label != "" { + where = append(where, "has(labels, ?)") + args = append(args, filter.Label) + } + if filter.NeedsFollowup != nil { + where = append(where, "needs_followup = ?") + args = append(args, boolToUInt8(*filter.NeedsFollowup)) + } + if !filter.IncludeDeleted { + where = append(where, "status != ?") + args = append(args, models.AnnotationStatusDeleted) + } + return where, args +} + +func normalizeAnnotationLimit(limit int) int { + if limit <= 0 { + return defaultAnnotationLimit + } + if limit > maxAnnotationLimit { + return maxAnnotationLimit + } + return limit +} + +func normalizeAnnotationOffset(offset int) int { + if offset < 0 { + return 0 + } + return offset +} + +func compactAnnotationStrings(values []string) []string { + if len(values) == 0 { + return nil + } + seen := make(map[string]struct{}, len(values)) + out := make([]string, 0, len(values)) + for _, value := range values { + value = strings.TrimSpace(value) + if value == "" { + continue + } + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + out = append(out, value) + } + return out +} + func boolToUInt8(value bool) uint8 { if value { return 1 diff --git a/internal/store/annotations_test.go b/internal/store/annotations_test.go index 062d964..09c104d 100644 --- a/internal/store/annotations_test.go +++ b/internal/store/annotations_test.go @@ -1,8 +1,11 @@ package store import ( + "reflect" "strings" "testing" + + "github.com/johnnygreco/beacon/internal/models" ) func TestTraceAnnotationSelectSQLTargetsFinalTable(t *testing.T) { @@ -27,6 +30,55 @@ func TestSplitAnnotationLabelsNormalizesValues(t *testing.T) { } } +func TestAnnotationFilterWhereSupportsDatasetFilters(t *testing.T) { + followup := true + where, args := annotationFilterWhere(AnnotationFilter{ + TargetType: models.AnnotationTargetMessage, + SessionIDs: []string{"session-2", "session-1", "session-2"}, + EventUID: "event-1", + AuthorType: models.AnnotationAuthorAgent, + Source: models.AnnotationSourceMCP, + Category: "quality", + Outcome: "useful", + Label: "dataset:eval", + NeedsFollowup: &followup, + IncludeDeleted: false, + }) + joined := strings.Join(where, " AND ") + for _, want := range []string{ + "target_type = ?", + "session_id IN (?,?)", + "event_uid = ?", + "author_type = ?", + "source = ?", + "category = ?", + "outcome = ?", + "has(labels, ?)", + "needs_followup = ?", + "status != ?", + } { + if !strings.Contains(joined, want) { + t.Fatalf("filter where missing %q:\n%s", want, joined) + } + } + wantArgs := []any{ + models.AnnotationTargetMessage, + "session-2", + "session-1", + "event-1", + models.AnnotationAuthorAgent, + models.AnnotationSourceMCP, + "quality", + "useful", + "dataset:eval", + uint8(1), + models.AnnotationStatusDeleted, + } + if !reflect.DeepEqual(args, wantArgs) { + t.Fatalf("filter args = %#v, want %#v", args, wantArgs) + } +} + func TestNewAnnotationIDUsesStablePrefix(t *testing.T) { id, err := newAnnotationID() if err != nil { diff --git a/internal/web/api_annotation_exports.go b/internal/web/api_annotation_exports.go new file mode 100644 index 0000000..71bf5dc --- /dev/null +++ b/internal/web/api_annotation_exports.go @@ -0,0 +1,437 @@ +package web + +import ( + "context" + "errors" + "fmt" + "net/http" + "sort" + "strings" + "time" + + "github.com/johnnygreco/beacon/internal/models" + "github.com/johnnygreco/beacon/internal/store" +) + +const ( + annotatedTraceIndexSchema = "beacon.annotated_traces.index.v1" + annotatedTraceExportSchema = "beacon.annotated_traces.export.v1" + annotatedTraceScanBatch = 100 +) + +type annotatedTraceGroup struct { + session APISessionSummary + counts APIAnnotationCounts + firstAnnotationAt time.Time + lastAnnotationAt time.Time + targets map[string]APIAnnotatedTargetSummary + annotations []APITraceAnnotation +} + +func (a *APIHandlers) ListAnnotatedTraces(w http.ResponseWriter, r *http.Request) { + req, err := parseAnnotatedTracesAPIRequest(r.URL.Query()) + if err != nil { + a.badRequest(w, err) + return + } + var scopeMetadata APIScopeMetadata + req.Scope, scopeMetadata = scopeForRequest(r.Context(), req.Scope) + if _, err := annotationFilterFromAnnotatedRequest(req, req.Limit, req.Offset); err != nil { + a.badRequest(w, err) + return + } + groups, hasMore, err := a.collectAnnotatedTraceGroups(r.Context(), r, req, false) + if err != nil { + a.annotationExportError(w, "failed to query annotated traces", err) + return + } + items := make([]APIAnnotatedTraceSummary, 0, len(groups)) + for _, group := range groups { + items = append(items, APIAnnotatedTraceSummary{ + Session: group.session, + Counts: group.counts, + FirstAnnotationAt: group.firstAnnotationAt, + LastAnnotationAt: group.lastAnnotationAt, + Targets: sortedAnnotatedTargets(group.targets), + }) + } + a.jsonResponse(w, APIAnnotatedTracesResponse{ + Schema: annotatedTraceIndexSchema, + Scope: scopeMetadata, + IncludeDeleted: req.IncludeDeleted, + Offset: req.Offset, + Limit: req.Limit, + HasMore: hasMore, + Items: items, + }) +} + +func (a *APIHandlers) ExportAnnotatedTraces(w http.ResponseWriter, r *http.Request) { + req, err := parseAnnotatedTracesAPIRequest(r.URL.Query()) + if err != nil { + a.badRequest(w, err) + return + } + var scopeMetadata APIScopeMetadata + req.Scope, scopeMetadata = scopeForRequest(r.Context(), req.Scope) + if _, err := annotationFilterFromAnnotatedRequest(req, req.Limit, req.Offset); err != nil { + a.badRequest(w, err) + return + } + groups, hasMore, err := a.collectAnnotatedTraceGroups(r.Context(), r, req, true) + if err != nil { + a.annotationExportError(w, "failed to export annotated traces", err) + return + } + traces := make([]APIAnnotatedTraceExport, 0, len(groups)) + warnings := []string{} + for _, group := range groups { + events, truncated, err := a.querySessionEventsForExport(r.Context(), group.session.ID, req.EventLimit, req.Scope) + if err != nil { + a.annotationExportError(w, "failed to export annotated traces", err) + return + } + if truncated { + warnings = append(warnings, fmt.Sprintf("session %s events truncated at event_limit=%d", group.session.ID, req.EventLimit)) + } + traces = append(traces, APIAnnotatedTraceExport{ + Session: group.session, + Counts: group.counts, + Annotations: group.annotations, + Events: events, + EventLimit: req.EventLimit, + EventTruncated: truncated, + }) + } + a.jsonResponse(w, APIAnnotatedTraceExportResponse{ + Schema: annotatedTraceExportSchema, + ExportedAt: time.Now().UTC(), + Scope: scopeMetadata, + IncludeDeleted: req.IncludeDeleted, + Offset: req.Offset, + Limit: req.Limit, + EventLimit: req.EventLimit, + HasMore: hasMore, + Traces: traces, + Warnings: warnings, + }) +} + +func (a *APIHandlers) collectAnnotatedTraceGroups(ctx context.Context, r *http.Request, req annotatedTracesAPIRequest, includeAnnotations bool) ([]annotatedTraceGroup, bool, error) { + collected := []annotatedTraceGroup{} + scanOffset := 0 + for { + filter, err := annotationFilterFromAnnotatedRequest(req, annotatedTraceScanBatch, scanOffset) + if err != nil { + return nil, false, err + } + candidates, err := store.ListTraceAnnotationSessionSummaries(ctx, a.db, filter) + if err != nil { + return nil, false, err + } + if len(candidates) == 0 { + break + } + scanOffset += len(candidates) + sessionIDs := annotationSessionIDs(candidates) + sessions, err := a.querySessionSummariesByID(ctx, sessionIDs, req.Scope) + if err != nil { + return nil, false, err + } + annotations, err := a.listTraceAnnotationsForSessionBatch(ctx, sessionIDs, req) + if err != nil { + return nil, false, err + } + visible := make(map[string][]models.TraceAnnotation, len(sessionIDs)) + for _, annotation := range annotations { + if err := a.ensureAnnotationTargetInScope(r, annotation, req.Scope); err != nil { + if errors.Is(err, store.ErrAnnotationTargetNotFound) { + continue + } + return nil, false, err + } + if _, ok := sessions[annotation.SessionID]; !ok { + continue + } + visible[annotation.SessionID] = append(visible[annotation.SessionID], annotation) + } + for _, candidate := range candidates { + session, ok := sessions[candidate.SessionID] + if !ok { + continue + } + sessionAnnotations := visible[candidate.SessionID] + if len(sessionAnnotations) == 0 { + continue + } + group := annotatedTraceGroup{session: session, targets: map[string]APIAnnotatedTargetSummary{}} + for _, annotation := range sessionAnnotations { + addAnnotationToTraceGroup(&group, annotation, includeAnnotations) + } + collected = append(collected, group) + } + if len(candidates) < annotatedTraceScanBatch { + break + } + } + sort.Slice(collected, func(i, j int) bool { + if collected[i].lastAnnotationAt.Equal(collected[j].lastAnnotationAt) { + return collected[i].session.ID < collected[j].session.ID + } + return collected[i].lastAnnotationAt.After(collected[j].lastAnnotationAt) + }) + hasMore := len(collected) > req.Offset+req.Limit + if req.Offset >= len(collected) { + return []annotatedTraceGroup{}, hasMore, nil + } + end := req.Offset + req.Limit + if end > len(collected) { + end = len(collected) + } + return collected[req.Offset:end], hasMore, nil +} + +func (a *APIHandlers) listTraceAnnotationsForSessionBatch(ctx context.Context, sessionIDs []string, req annotatedTracesAPIRequest) ([]models.TraceAnnotation, error) { + if len(sessionIDs) == 0 { + return nil, nil + } + annotations := []models.TraceAnnotation{} + for offset := 0; ; { + page, err := store.ListTraceAnnotations(ctx, a.db, store.AnnotationFilter{ + SessionIDs: sessionIDs, + TargetType: req.TargetType, + EventUID: req.EventUID, + AuthorType: req.AuthorType, + Source: req.Source, + Category: req.Category, + Outcome: req.Outcome, + Label: req.Label, + NeedsFollowup: req.NeedsFollowup, + IncludeDeleted: req.IncludeDeleted, + Limit: maxAnnotationsAPILimit, + Offset: offset, + }) + if err != nil { + return nil, err + } + annotations = append(annotations, page...) + if len(page) < maxAnnotationsAPILimit { + break + } + offset += len(page) + } + return annotations, nil +} + +func annotationFilterFromAnnotatedRequest(req annotatedTracesAPIRequest, limit, offset int) (store.AnnotationFilter, error) { + targetType := strings.TrimSpace(strings.ToLower(req.TargetType)) + switch targetType { + case "", models.AnnotationTargetSession, models.AnnotationTargetMessage, models.AnnotationTargetEvent: + default: + return store.AnnotationFilter{}, fmt.Errorf("target_type must be session, message, or event") + } + return store.AnnotationFilter{ + TargetType: targetType, + SessionID: strings.TrimSpace(req.SessionID), + EventUID: strings.TrimSpace(req.EventUID), + AuthorType: strings.TrimSpace(req.AuthorType), + Source: strings.TrimSpace(req.Source), + Category: strings.TrimSpace(req.Category), + Outcome: strings.TrimSpace(req.Outcome), + Label: strings.TrimSpace(req.Label), + NeedsFollowup: req.NeedsFollowup, + IncludeDeleted: req.IncludeDeleted, + Limit: limit, + Offset: offset, + }, nil +} + +func (a *APIHandlers) querySessionSummariesByID(ctx context.Context, ids []string, scope APIScopeFilters) (map[string]APISessionSummary, error) { + out := make(map[string]APISessionSummary, len(ids)) + ids = compactScopeValues(ids) + if len(ids) == 0 { + return out, nil + } + placeholders := make([]string, len(ids)) + for i := range ids { + placeholders[i] = "?" + } + now := time.Now() + cutoff := now.Add(-idleThreshold) + sessionSource, sourceArgs := sessionProjectionSubqueryForScope("", scope) + sessionScope := scope.withoutProjectKeys() + scopeClause, scopeArgs := sessionScope.sqlAndClause("") + args := reopenedFlagArgs(scope, cutoff) + args = append(args, sourceArgs...) + for _, id := range ids { + args = append(args, id) + } + args = append(args, scopeArgs...) + rows, err := a.db.QueryContext(ctx, `SELECT `+sessionSummaryColumnsWithReopenedFlagScoped(scope)+` + FROM `+sessionSource+` + WHERE session_id IN (`+strings.Join(placeholders, ",")+`)`+scopeClause, args...) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + session, err := scanSessionSummaryIncludingReopened(rows, now) + if err != nil { + a.logSkippedRow("annotated trace session summary", err) + continue + } + out[session.ID] = apiSessionSummaryFromView(session) + } + if err := rows.Err(); err != nil { + return nil, err + } + return out, nil +} + +func (a *APIHandlers) querySessionEventsForExport(ctx context.Context, sessionID string, eventLimit int, scope APIScopeFilters) ([]APISessionEvent, bool, error) { + sessionScope := scope.withoutProjectKeys() + sessionScopeClause := "" + sessionScopeArgs := []any{} + scopedSessionSQL := "SELECT ? AS session_id" + if len(compactScopeValues(scope.ProjectKeys)) == 0 { + sessionScopeClause, sessionScopeArgs = sessionScope.sqlAndClause("") + scopedSessionSQL = `SELECT session_id + FROM session_projection FINAL + WHERE session_id = ?` + sessionScopeClause + ` + LIMIT 1` + } + args := []any{sessionID} + args = append(args, sessionScopeArgs...) + eventScopeClause, eventScopeArgs := scope.eventAndSessionProjectSQLAndClause("e", "e.cwd", "s") + args = append(args, eventScopeArgs...) + args = append(args, eventLimit+1) + rows, err := a.db.QueryContext(ctx, + `WITH scoped_session AS ( + `+scopedSessionSQL+` + ), + session_events AS ( + SELECT e.event_uid, e.session_id, e.event_kind, e.payload_type, e.actor_role, + e.timestamp, e.text_preview, e.tool_name, e.tool_use_id, e.model, + e.input_tokens + e.output_tokens AS tokens, e.duration_ms + FROM `+latestActivityEventsSubquery("ae.session_id IN (SELECT session_id FROM scoped_session)")+` AS e + LEFT JOIN `+sessionProjectFallbackSubquery("ae.session_id IN (SELECT session_id FROM scoped_session)")+` AS s ON s.session_id = e.session_id + WHERE 1 = 1`+eventScopeClause+` + ORDER BY timestamp, event_uid + LIMIT ? + ), + payload_previews AS ( + SELECT event_uid, + argMax(input_preview, captured_at) AS input_preview, + argMax(output_preview, captured_at) AS output_preview + FROM tool_payloads + WHERE event_uid IN (SELECT event_uid FROM session_events) + GROUP BY event_uid + ) + SELECT e.event_uid, e.session_id, e.event_kind, e.payload_type, e.actor_role, + e.timestamp, e.text_preview, e.tool_name, e.tool_use_id, e.model, + e.tokens, e.duration_ms, + COALESCE(p.input_preview, ''), COALESCE(p.output_preview, '') + FROM session_events e + LEFT JOIN payload_previews p ON e.event_uid = p.event_uid + ORDER BY e.timestamp, e.event_uid`, args...) + if err != nil { + return nil, false, err + } + defer rows.Close() + events := make([]APISessionEvent, 0, eventLimit) + for rows.Next() { + var event APISessionEvent + if err := rows.Scan(&event.EventUID, &event.SessionID, &event.EventKind, &event.PayloadType, &event.ActorRole, + &event.Timestamp, &event.TextPreview, &event.ToolName, &event.ToolUseID, &event.Model, &event.Tokens, &event.DurationMs, + &event.InputPreview, &event.OutputPreview); err != nil { + a.logSkippedRow("annotated trace events", err) + continue + } + events = append(events, event) + } + if err := rows.Err(); err != nil { + return nil, false, err + } + truncated := len(events) > eventLimit + if truncated { + events = events[:eventLimit] + } + return events, truncated, nil +} + +func annotationSessionIDs(summaries []store.TraceAnnotationSessionSummary) []string { + ids := make([]string, 0, len(summaries)) + for _, summary := range summaries { + ids = append(ids, summary.SessionID) + } + return ids +} + +func addAnnotationToTraceGroup(group *annotatedTraceGroup, annotation models.TraceAnnotation, includeAnnotation bool) { + apiAnnotation := apiTraceAnnotationFromModel(annotation) + group.counts.AnnotationCount++ + switch annotation.TargetType { + case models.AnnotationTargetSession: + group.counts.SessionAnnotationCount++ + case models.AnnotationTargetMessage: + group.counts.MessageAnnotationCount++ + case models.AnnotationTargetEvent: + group.counts.EventAnnotationCount++ + } + if annotation.NeedsFollowup { + group.counts.NeedsFollowupCount++ + } + if group.firstAnnotationAt.IsZero() || annotation.CreatedAt.Before(group.firstAnnotationAt) { + group.firstAnnotationAt = annotation.CreatedAt + } + if group.lastAnnotationAt.IsZero() || annotation.UpdatedAt.After(group.lastAnnotationAt) { + group.lastAnnotationAt = annotation.UpdatedAt + } + targetKey := annotation.TargetType + "\x00" + annotation.EventUID + target := group.targets[targetKey] + if target.TargetType == "" { + target.TargetType = annotation.TargetType + target.EventUID = annotation.EventUID + target.FirstAnnotationAt = annotation.CreatedAt + target.LastAnnotationAt = annotation.UpdatedAt + } + target.AnnotationCount++ + if annotation.CreatedAt.Before(target.FirstAnnotationAt) { + target.FirstAnnotationAt = annotation.CreatedAt + } + if annotation.UpdatedAt.After(target.LastAnnotationAt) { + target.LastAnnotationAt = annotation.UpdatedAt + } + group.targets[targetKey] = target + if includeAnnotation { + group.annotations = append(group.annotations, apiAnnotation) + } +} + +func sortedAnnotatedTargets(targets map[string]APIAnnotatedTargetSummary) []APIAnnotatedTargetSummary { + out := make([]APIAnnotatedTargetSummary, 0, len(targets)) + for _, target := range targets { + out = append(out, target) + } + sort.Slice(out, func(i, j int) bool { + if out[i].TargetType != out[j].TargetType { + return out[i].TargetType < out[j].TargetType + } + return out[i].EventUID < out[j].EventUID + }) + return out +} + +func (a *APIHandlers) annotationExportError(w http.ResponseWriter, publicMessage string, err error) { + var validation *models.AnnotationValidationError + switch { + case err == nil: + return + case errors.As(err, &validation): + a.jsonError(w, validation.Message, http.StatusBadRequest) + case errors.Is(err, store.ErrAnnotationTargetNotFound): + a.jsonError(w, "annotation target not found", http.StatusNotFound) + default: + a.internalError(w, publicMessage, err) + } +} diff --git a/internal/web/api_annotations_test.go b/internal/web/api_annotations_test.go index 9000171..7bb9a07 100644 --- a/internal/web/api_annotations_test.go +++ b/internal/web/api_annotations_test.go @@ -6,6 +6,7 @@ import ( "database/sql/driver" "encoding/json" "errors" + "fmt" "io" "net/http" "net/http/httptest" @@ -284,6 +285,182 @@ func TestAnnotationAPIListByAnnotationIDAppliesScope(t *testing.T) { } } +func TestAnnotationAPIAnnotatedTracesListsSessionsAndTargets(t *testing.T) { + now := time.Date(2026, 6, 16, 10, 0, 0, 0, time.UTC) + fake := newAnnotationAPIFake() + fake.sessions["session-1"] = annotationAPISession{sourceName: "source-a", runtime: "runtime-a", startedAt: now.Add(-time.Hour), endedAt: now} + fake.events["message-1"] = annotationAPIEvent{sessionID: "session-1", eventKind: "message", sourceName: "source-a", runtime: "runtime-a", timestamp: now.Add(-time.Minute), textPreview: "user asks for help"} + fake.annotations["ann-session"] = testTraceAnnotation("ann-session", models.AnnotationTargetSession, "session-1", "", now, "session note", []string{"dataset:eval"}) + fake.annotations["ann-message"] = testTraceAnnotation("ann-message", models.AnnotationTargetMessage, "session-1", "message-1", now.Add(time.Second), "message note", []string{"dataset:eval"}) + hidden := testTraceAnnotation("ann-hidden", models.AnnotationTargetMessage, "session-1", "hidden-message", now.Add(2*time.Second), "hidden note", []string{"dataset:eval"}) + fake.annotations[hidden.AnnotationID] = hidden + fake.events["hidden-message"] = annotationAPIEvent{sessionID: "session-1", eventKind: "message", sourceName: "source-b", runtime: "runtime-a", timestamp: now} + handlers := &APIHandlers{db: newAnnotationAPIDB(t, fake), logger: testLogger()} + + req := annotationAPIRequest(http.MethodGet, "/api/annotations/traces?label=dataset:eval&source_name=source-a", "") + w := httptest.NewRecorder() + handlers.ListAnnotatedTraces(w, req) + if w.Code != http.StatusOK { + t.Fatalf("annotated traces status = %d body=%s", w.Code, w.Body.String()) + } + var got APIAnnotatedTracesResponse + if err := json.NewDecoder(w.Body).Decode(&got); err != nil { + t.Fatalf("decode annotated traces: %v", err) + } + if got.Schema != annotatedTraceIndexSchema || len(got.Items) != 1 { + t.Fatalf("annotated traces response = %#v", got) + } + item := got.Items[0] + if item.Session.ID != "session-1" || item.Counts.AnnotationCount != 2 || item.Counts.SessionAnnotationCount != 1 || item.Counts.MessageAnnotationCount != 1 { + t.Fatalf("annotated trace item = %#v", item) + } + if len(item.Targets) != 2 { + t.Fatalf("targets = %#v, want session and visible message targets", item.Targets) + } + for _, target := range item.Targets { + if target.EventUID == "hidden-message" { + t.Fatalf("out-of-scope target leaked: %#v", item.Targets) + } + } +} + +func TestAnnotationAPIAnnotatedTracesOrdersByVisibleScopedAnnotations(t *testing.T) { + now := time.Date(2026, 6, 16, 10, 0, 0, 0, time.UTC) + fake := newAnnotationAPIFake() + fake.sessions["session-a"] = annotationAPISession{sourceName: "source-a", runtime: "runtime-a", startedAt: now.Add(-2 * time.Hour), endedAt: now} + fake.sessions["session-b"] = annotationAPISession{sourceName: "source-a", runtime: "runtime-a", startedAt: now.Add(-time.Hour), endedAt: now} + fake.events["hidden-message"] = annotationAPIEvent{sessionID: "session-a", eventKind: "message", sourceName: "source-b", runtime: "runtime-a", timestamp: now} + fake.annotations["ann-a-visible"] = testTraceAnnotation("ann-a-visible", models.AnnotationTargetSession, "session-a", "", now.Add(time.Second), "older visible note", []string{"dataset:eval"}) + fake.annotations["ann-b-visible"] = testTraceAnnotation("ann-b-visible", models.AnnotationTargetSession, "session-b", "", now.Add(5*time.Second), "newer visible note", []string{"dataset:eval"}) + fake.annotations["ann-a-hidden"] = testTraceAnnotation("ann-a-hidden", models.AnnotationTargetMessage, "session-a", "hidden-message", now.Add(10*time.Second), "newest hidden note", []string{"dataset:eval"}) + handlers := &APIHandlers{db: newAnnotationAPIDB(t, fake), logger: testLogger()} + + w := httptest.NewRecorder() + handlers.ListAnnotatedTraces(w, annotationAPIRequest(http.MethodGet, "/api/annotations/traces?source_name=source-a&label=dataset:eval&limit=1", "")) + if w.Code != http.StatusOK { + t.Fatalf("annotated traces status = %d body=%s", w.Code, w.Body.String()) + } + var firstPage APIAnnotatedTracesResponse + if err := json.NewDecoder(w.Body).Decode(&firstPage); err != nil { + t.Fatalf("decode first page: %v", err) + } + if len(firstPage.Items) != 1 || firstPage.Items[0].Session.ID != "session-b" || !firstPage.HasMore { + t.Fatalf("first page = %#v, want session-b first with more results", firstPage) + } + + w = httptest.NewRecorder() + handlers.ListAnnotatedTraces(w, annotationAPIRequest(http.MethodGet, "/api/annotations/traces?source_name=source-a&label=dataset:eval&limit=1&offset=1", "")) + if w.Code != http.StatusOK { + t.Fatalf("annotated traces second page status = %d body=%s", w.Code, w.Body.String()) + } + var secondPage APIAnnotatedTracesResponse + if err := json.NewDecoder(w.Body).Decode(&secondPage); err != nil { + t.Fatalf("decode second page: %v", err) + } + if len(secondPage.Items) != 1 || secondPage.Items[0].Session.ID != "session-a" || secondPage.HasMore { + t.Fatalf("second page = %#v, want session-a without hidden target influence", secondPage) + } +} + +func TestAnnotationAPIExportAnnotatedTracesIncludesContextAndDeleted(t *testing.T) { + now := time.Date(2026, 6, 16, 10, 0, 0, 0, time.UTC) + fake := newAnnotationAPIFake() + fake.sessions["session-1"] = annotationAPISession{sourceName: "source-a", runtime: "runtime-a", startedAt: now.Add(-time.Hour), endedAt: now} + fake.events["message-1"] = annotationAPIEvent{sessionID: "session-1", eventKind: "message", sourceName: "source-a", runtime: "runtime-a", timestamp: now.Add(-2 * time.Minute), textPreview: "first message", actorRole: "user", model: "gpt-test", tokens: 12} + fake.events["event-1"] = annotationAPIEvent{sessionID: "session-1", eventKind: "tool_call", sourceName: "source-a", runtime: "runtime-a", timestamp: now.Add(-time.Minute), textPreview: "tool call", toolName: "shell", tokens: 3} + active := testTraceAnnotation("ann-session", models.AnnotationTargetSession, "session-1", "", now, "session note", []string{"dataset:eval"}) + deleted := testTraceAnnotation("ann-event", models.AnnotationTargetEvent, "session-1", "event-1", now.Add(time.Second), "deleted event note", []string{"dataset:eval"}) + deleted.Status = models.AnnotationStatusDeleted + deletedAt := now.Add(2 * time.Second) + deleted.DeletedAt = &deletedAt + fake.annotations[active.AnnotationID] = active + fake.annotations[deleted.AnnotationID] = deleted + handlers := &APIHandlers{db: newAnnotationAPIDB(t, fake), logger: testLogger()} + + w := httptest.NewRecorder() + handlers.ExportAnnotatedTraces(w, annotationAPIRequest(http.MethodGet, "/api/annotations/export?session_id=session-1&include_deleted=1&event_limit=1", "")) + if w.Code != http.StatusOK { + t.Fatalf("export status = %d body=%s", w.Code, w.Body.String()) + } + var got APIAnnotatedTraceExportResponse + if err := json.NewDecoder(w.Body).Decode(&got); err != nil { + t.Fatalf("decode export: %v", err) + } + if got.Schema != annotatedTraceExportSchema || len(got.Traces) != 1 { + t.Fatalf("export response = %#v", got) + } + trace := got.Traces[0] + if trace.Session.ID != "session-1" || trace.Counts.AnnotationCount != 2 || trace.Counts.EventAnnotationCount != 1 { + t.Fatalf("export trace counts = %#v", trace) + } + if len(trace.Annotations) != 2 || trace.Annotations[1].Status != models.AnnotationStatusDeleted { + t.Fatalf("export annotations = %#v", trace.Annotations) + } + if len(trace.Events) != 1 || trace.Events[0].EventUID != "message-1" || !trace.EventTruncated || len(got.Warnings) != 1 { + t.Fatalf("export events/truncation = events:%#v truncated:%v warnings:%#v", trace.Events, trace.EventTruncated, got.Warnings) + } +} + +func TestAnnotationAPIExportAnnotatedTracesPaginatesAnnotations(t *testing.T) { + now := time.Date(2026, 6, 16, 10, 0, 0, 0, time.UTC) + fake := newAnnotationAPIFake() + fake.sessions["session-1"] = annotationAPISession{sourceName: "source-a", runtime: "runtime-a", startedAt: now.Add(-time.Hour), endedAt: now} + for i := 0; i < maxAnnotationsAPILimit+5; i++ { + id := fmt.Sprintf("ann-%03d", i) + fake.annotations[id] = testTraceAnnotation(id, models.AnnotationTargetSession, "session-1", "", now.Add(time.Duration(i)*time.Second), "session note", []string{"dataset:eval"}) + } + handlers := &APIHandlers{db: newAnnotationAPIDB(t, fake), logger: testLogger()} + + w := httptest.NewRecorder() + handlers.ExportAnnotatedTraces(w, annotationAPIRequest(http.MethodGet, "/api/annotations/export?session_id=session-1&label=dataset:eval", "")) + if w.Code != http.StatusOK { + t.Fatalf("export status = %d body=%s", w.Code, w.Body.String()) + } + var got APIAnnotatedTraceExportResponse + if err := json.NewDecoder(w.Body).Decode(&got); err != nil { + t.Fatalf("decode export: %v", err) + } + if len(got.Traces) != 1 { + t.Fatalf("trace count = %d, want 1", len(got.Traces)) + } + trace := got.Traces[0] + if trace.Counts.AnnotationCount != maxAnnotationsAPILimit+5 || len(trace.Annotations) != maxAnnotationsAPILimit+5 { + t.Fatalf("paginated annotations = count:%d len:%d, want %d", trace.Counts.AnnotationCount, len(trace.Annotations), maxAnnotationsAPILimit+5) + } +} + +func TestAnnotationAPIAnnotatedTracesEmptyResults(t *testing.T) { + handlers := &APIHandlers{db: newAnnotationAPIDB(t, newAnnotationAPIFake()), logger: testLogger()} + w := httptest.NewRecorder() + handlers.ListAnnotatedTraces(w, annotationAPIRequest(http.MethodGet, "/api/annotations/traces", "")) + if w.Code != http.StatusOK { + t.Fatalf("empty traces status = %d body=%s", w.Code, w.Body.String()) + } + var got APIAnnotatedTracesResponse + if err := json.NewDecoder(w.Body).Decode(&got); err != nil { + t.Fatalf("decode empty traces: %v", err) + } + if got.Schema != annotatedTraceIndexSchema || len(got.Items) != 0 { + t.Fatalf("empty traces = %#v", got) + } +} + +func TestAnnotationAPIAnnotatedTracesCapsOffset(t *testing.T) { + handlers := &APIHandlers{db: newAnnotationAPIDB(t, newAnnotationAPIFake()), logger: testLogger()} + w := httptest.NewRecorder() + handlers.ListAnnotatedTraces(w, annotationAPIRequest(http.MethodGet, "/api/annotations/traces?offset=999999999", "")) + if w.Code != http.StatusOK { + t.Fatalf("large offset status = %d body=%s", w.Code, w.Body.String()) + } + var got APIAnnotatedTracesResponse + if err := json.NewDecoder(w.Body).Decode(&got); err != nil { + t.Fatalf("decode large offset: %v", err) + } + if got.Offset != maxAnnotatedTracesOffset || len(got.Items) != 0 || got.HasMore { + t.Fatalf("large offset response = %#v", got) + } +} + func annotationAPIRequest(method, target, body string, routeParams ...string) *http.Request { req := httptest.NewRequest(method, target, strings.NewReader(body)) if len(routeParams) == 0 { @@ -296,6 +473,24 @@ func annotationAPIRequest(method, target, body string, routeParams ...string) *h return req.WithContext(context.WithValue(req.Context(), chi.RouteCtxKey, rctx)) } +func testTraceAnnotation(id, targetType, sessionID, eventUID string, at time.Time, note string, labels []string) models.TraceAnnotation { + return models.NormalizeTraceAnnotation(models.TraceAnnotation{ + AnnotationID: id, + Revision: 1, + TargetType: targetType, + SessionID: sessionID, + EventUID: eventUID, + AuthorType: models.AnnotationAuthorAgent, + Source: models.AnnotationSourceMCP, + Labels: labels, + Note: note, + Status: models.AnnotationStatusActive, + SchemaVersion: models.AnnotationSchemaVersion, + CreatedAt: at, + UpdatedAt: at, + }) +} + type annotationAPIFake struct { sessions map[string]annotationAPISession events map[string]annotationAPIEvent @@ -305,13 +500,27 @@ type annotationAPIFake struct { type annotationAPISession struct { sourceName string runtime string + provider string + startedAt time.Time + endedAt time.Time } type annotationAPIEvent struct { - sessionID string - eventKind string - sourceName string - runtime string + sessionID string + eventKind string + payloadType string + actorRole string + sourceName string + runtime string + timestamp time.Time + textPreview string + toolName string + toolUseID string + model string + tokens int64 + durationMs int64 + inputPreview string + outputPreview string } func newAnnotationAPIFake() *annotationAPIFake { @@ -411,8 +620,14 @@ func (c annotationAPIConn) ExecContext(_ context.Context, query string, args []d func (c annotationAPIConn) QueryContext(_ context.Context, query string, args []driver.NamedValue) (driver.Rows, error) { values := namedValues(args) switch { + case strings.Contains(query, "trace_annotations FINAL") && strings.Contains(query, "GROUP BY session_id"): + return c.annotationSessionSummaryRows(query, values), nil case strings.Contains(query, "trace_annotations FINAL"): return c.annotationRows(query, values), nil + case strings.Contains(query, "session_events AS"): + return c.sessionEventRows(query, values), nil + case strings.Contains(query, "COALESCE(source_name") && strings.Contains(query, "session_projection"): + return c.sessionSummaryRows(query, values), nil case strings.Contains(query, "session_projection"): sessionID := values.stringAt(0) session, ok := c.fake.sessions[sessionID] @@ -432,28 +647,138 @@ func (c annotationAPIConn) QueryContext(_ context.Context, query string, args [] } } +func (c annotationAPIConn) annotationSessionSummaryRows(query string, values annotationNamedValues) driver.Rows { + annotations := c.filteredAnnotations(query, values, 3) + type summary struct { + sessionID string + count uint64 + sessionCount uint64 + messageCount uint64 + eventCount uint64 + followupCount uint64 + first time.Time + last time.Time + } + bySession := map[string]*summary{} + for _, annotation := range annotations { + s := bySession[annotation.SessionID] + if s == nil { + s = &summary{sessionID: annotation.SessionID} + bySession[annotation.SessionID] = s + } + s.count++ + switch annotation.TargetType { + case models.AnnotationTargetSession: + s.sessionCount++ + case models.AnnotationTargetMessage: + s.messageCount++ + case models.AnnotationTargetEvent: + s.eventCount++ + } + if annotation.NeedsFollowup { + s.followupCount++ + } + if s.first.IsZero() || annotation.CreatedAt.Before(s.first) { + s.first = annotation.CreatedAt + } + if s.last.IsZero() || annotation.UpdatedAt.After(s.last) { + s.last = annotation.UpdatedAt + } + } + summaries := make([]*summary, 0, len(bySession)) + for _, s := range bySession { + summaries = append(summaries, s) + } + sort.Slice(summaries, func(i, j int) bool { + if summaries[i].last.Equal(summaries[j].last) { + return summaries[i].sessionID < summaries[j].sessionID + } + return summaries[i].last.After(summaries[j].last) + }) + rows := make([][]driver.Value, 0, len(summaries)) + for _, s := range summaries { + rows = append(rows, []driver.Value{s.sessionID, s.count, s.sessionCount, s.messageCount, s.eventCount, s.followupCount, s.first, s.last}) + } + rows = paginateDriverRows(rows, values) + return &annotationRows{ + columns: []string{"session_id", "annotation_count", "session_annotation_count", "message_annotation_count", "event_annotation_count", "needs_followup_count", "first_annotation_at", "last_annotation_at"}, + rows: rows, + } +} + func (c annotationAPIConn) annotationRows(query string, values annotationNamedValues) driver.Rows { - idx := 0 + annotations := c.filteredAnnotations(query, values, 0) + annotations = paginateTraceAnnotations(annotations, values) + rows := make([][]driver.Value, 0, len(annotations)) + for _, annotation := range annotations { + rows = append(rows, annotationRowValues(annotation)) + } + return &annotationRows{ + columns: []string{"annotation_id", "revision", "target_type", "session_id", "event_uid", "author_type", "author_id", "author_name", "source", "category", "outcome", "quality_score", "confidence", "needs_followup", "labels", "note", "metadata_json", "status", "schema_version", "created_at", "updated_at", "deleted_at"}, + rows: rows, + } +} + +func (c annotationAPIConn) filteredAnnotations(query string, values annotationNamedValues, idx int) []models.TraceAnnotation { var annotationID, targetType, sessionID, eventUID, excludedStatus string - if strings.Contains(query, "annotation_id = ?") { + var sessionIDs []string + var authorType, source, category, outcome, label string + var needsFollowup *bool + if hasAnnotationFilter(query, "annotation_id = ?") { annotationID = values.stringAt(idx) idx++ } - if strings.Contains(query, "target_type = ?") { + if hasAnnotationFilter(query, "target_type = ?") { targetType = values.stringAt(idx) idx++ } - if strings.Contains(query, "session_id = ?") { + if hasAnnotationFilter(query, "session_id = ?") { sessionID = values.stringAt(idx) idx++ } - if strings.Contains(query, "event_uid = ?") { + if hasAnnotationFilter(query, "session_id IN (") { + count := placeholderCountInClause(query, "session_id IN (") + for i := 0; i < count; i++ { + sessionIDs = append(sessionIDs, values.stringAt(idx)) + idx++ + } + } + if hasAnnotationFilter(query, "event_uid = ?") { eventUID = values.stringAt(idx) idx++ } - if strings.Contains(query, "status != ?") { + if hasAnnotationFilter(query, "author_type = ?") { + authorType = values.stringAt(idx) + idx++ + } + if hasAnnotationFilter(query, "source = ?") { + source = values.stringAt(idx) + idx++ + } + if hasAnnotationFilter(query, "category = ?") { + category = values.stringAt(idx) + idx++ + } + if hasAnnotationFilter(query, "outcome = ?") { + outcome = values.stringAt(idx) + idx++ + } + if hasAnnotationFilter(query, "has(labels, ?)") { + label = values.stringAt(idx) + idx++ + } + if hasAnnotationFilter(query, "needs_followup = ?") { + value := values.intAt(idx) != 0 + needsFollowup = &value + idx++ + } + if hasAnnotationFilter(query, "status != ?") { excludedStatus = values.stringAt(idx) } + sessionIDSet := make(map[string]struct{}, len(sessionIDs)) + for _, id := range sessionIDs { + sessionIDSet[id] = struct{}{} + } var annotations []models.TraceAnnotation for _, annotation := range c.fake.annotations { @@ -466,9 +791,32 @@ func (c annotationAPIConn) annotationRows(query string, values annotationNamedVa if sessionID != "" && annotation.SessionID != sessionID { continue } + if len(sessionIDSet) > 0 { + if _, ok := sessionIDSet[annotation.SessionID]; !ok { + continue + } + } if eventUID != "" && annotation.EventUID != eventUID { continue } + if authorType != "" && annotation.AuthorType != authorType { + continue + } + if source != "" && annotation.Source != source { + continue + } + if category != "" && annotation.Category != category { + continue + } + if outcome != "" && annotation.Outcome != outcome { + continue + } + if label != "" && !annotationHasLabel(annotation, label) { + continue + } + if needsFollowup != nil && annotation.NeedsFollowup != *needsFollowup { + continue + } if excludedStatus != "" && annotation.Status == excludedStatus { continue } @@ -480,43 +828,229 @@ func (c annotationAPIConn) annotationRows(query string, values annotationNamedVa } return annotations[i].CreatedAt.Before(annotations[j].CreatedAt) }) - rows := make([][]driver.Value, 0, len(annotations)) - for _, annotation := range annotations { - deletedAt := time.Unix(0, 0).UTC() - if annotation.DeletedAt != nil { - deletedAt = *annotation.DeletedAt + return annotations +} + +func (c annotationAPIConn) sessionSummaryRows(query string, values annotationNamedValues) driver.Rows { + sessionIDs := map[string]struct{}{} + for _, value := range values { + if id, ok := value.(string); ok { + if _, exists := c.fake.sessions[id]; exists { + sessionIDs[id] = struct{}{} + } + } + } + rows := make([][]driver.Value, 0, len(sessionIDs)) + for id := range sessionIDs { + session := c.fake.sessions[id] + if !queryScopeMatches(query, values, firstSessionSummaryScopeArg(query, values), session.sourceName, session.runtime) { + continue } + rows = append(rows, annotationAPISessionSummaryRow(id, session)) + } + sort.Slice(rows, func(i, j int) bool { + return rows[i][0].(string) < rows[j][0].(string) + }) + return &annotationRows{ + columns: []string{"session_id", "source_name", "runtime", "provider", "format", "project_key", "project_path", "started_at", "ended_at", "turn_count", "total_tokens", "total_input_tokens", "total_output_tokens", "total_cache_read_tokens", "total_cache_create_tokens", "tool_call_count", "mcp_call_count", "error_count", "last_model", "working_dir", "parent_session_id", "has_session_end", "completion_state", "total_cost_usd", "cost_event_count", "cost_provenance", "attention_score", "attention_reasons", "archive_reason", "archived_at", "reopened"}, + rows: rows, + } +} + +func (c annotationAPIConn) sessionEventRows(_ string, values annotationNamedValues) driver.Rows { + sessionID := values.stringAt(0) + type keyedEvent struct { + uid string + event annotationAPIEvent + } + var events []keyedEvent + for eventUID, event := range c.fake.events { + if event.sessionID == sessionID { + events = append(events, keyedEvent{uid: eventUID, event: event}) + } + } + sort.Slice(events, func(i, j int) bool { + if events[i].event.timestamp.Equal(events[j].event.timestamp) { + return events[i].uid < events[j].uid + } + return events[i].event.timestamp.Before(events[j].event.timestamp) + }) + limit := values.intAt(len(values) - 1) + if limit > 0 && len(events) > limit { + events = events[:limit] + } + rows := make([][]driver.Value, 0, len(events)) + for _, item := range events { + event := item.event rows = append(rows, []driver.Value{ - annotation.AnnotationID, - int64(annotation.Revision), - annotation.TargetType, - annotation.SessionID, - annotation.EventUID, - annotation.AuthorType, - annotation.AuthorID, - annotation.AuthorName, - annotation.Source, - annotation.Category, - annotation.Outcome, - int64(annotation.QualityScore), - int64(annotation.Confidence), - boolAsInt64(annotation.NeedsFollowup), - strings.Join(annotation.Labels, ","), - annotation.Note, - annotation.MetadataJSON, - annotation.Status, - int64(annotation.SchemaVersion), - annotation.CreatedAt, - annotation.UpdatedAt, - deletedAt, + item.uid, + event.sessionID, + firstNonEmpty(event.eventKind, "message"), + event.payloadType, + event.actorRole, + event.timestamp, + event.textPreview, + event.toolName, + event.toolUseID, + event.model, + event.tokens, + event.durationMs, + event.inputPreview, + event.outputPreview, }) } return &annotationRows{ - columns: []string{"annotation_id", "revision", "target_type", "session_id", "event_uid", "author_type", "author_id", "author_name", "source", "category", "outcome", "quality_score", "confidence", "needs_followup", "labels", "note", "metadata_json", "status", "schema_version", "created_at", "updated_at", "deleted_at"}, + columns: []string{"event_uid", "session_id", "event_kind", "payload_type", "actor_role", "timestamp", "text_preview", "tool_name", "tool_use_id", "model", "tokens", "duration_ms", "input_preview", "output_preview"}, rows: rows, } } +func annotationAPISessionSummaryRow(id string, session annotationAPISession) []driver.Value { + startedAt := session.startedAt + if startedAt.IsZero() { + startedAt = time.Date(2026, 6, 16, 10, 0, 0, 0, time.UTC) + } + endedAt := session.endedAt + if endedAt.IsZero() { + endedAt = startedAt.Add(time.Minute) + } + return []driver.Value{ + id, + firstNonEmpty(session.sourceName, "source-a"), + firstNonEmpty(session.runtime, "runtime-a"), + firstNonEmpty(session.provider, "provider-a"), + "jsonl", + "beacon", + "/work/beacon", + startedAt, + endedAt, + int64(2), + int64(30), + int64(10), + int64(20), + int64(0), + int64(0), + int64(1), + int64(0), + int64(0), + "gpt-test", + "/work/beacon", + "", + int64(1), + "completed", + float64(0), + int64(0), + "none", + int64(0), + []string{}, + "", + time.Unix(0, 0).UTC(), + int64(0), + } +} + +func annotationRowValues(annotation models.TraceAnnotation) []driver.Value { + deletedAt := time.Unix(0, 0).UTC() + if annotation.DeletedAt != nil { + deletedAt = *annotation.DeletedAt + } + return []driver.Value{ + annotation.AnnotationID, + int64(annotation.Revision), + annotation.TargetType, + annotation.SessionID, + annotation.EventUID, + annotation.AuthorType, + annotation.AuthorID, + annotation.AuthorName, + annotation.Source, + annotation.Category, + annotation.Outcome, + int64(annotation.QualityScore), + int64(annotation.Confidence), + boolAsInt64(annotation.NeedsFollowup), + strings.Join(annotation.Labels, ","), + annotation.Note, + annotation.MetadataJSON, + annotation.Status, + int64(annotation.SchemaVersion), + annotation.CreatedAt, + annotation.UpdatedAt, + deletedAt, + } +} + +func placeholderCountInClause(query, marker string) int { + start := strings.Index(query, marker) + if start < 0 { + return 0 + } + start += len(marker) + end := strings.Index(query[start:], ")") + if end < 0 { + return 0 + } + return strings.Count(query[start:start+end], "?") +} + +func paginateTraceAnnotations(annotations []models.TraceAnnotation, values annotationNamedValues) []models.TraceAnnotation { + limit, offset := queryLimitOffset(values) + if offset >= len(annotations) { + return nil + } + end := len(annotations) + if limit > 0 && offset+limit < end { + end = offset + limit + } + return annotations[offset:end] +} + +func paginateDriverRows(rows [][]driver.Value, values annotationNamedValues) [][]driver.Value { + limit, offset := queryLimitOffset(values) + if offset >= len(rows) { + return nil + } + end := len(rows) + if limit > 0 && offset+limit < end { + end = offset + limit + } + return rows[offset:end] +} + +func queryLimitOffset(values annotationNamedValues) (int, int) { + if len(values) < 2 { + return 0, 0 + } + return values.intAt(len(values) - 2), values.intAt(len(values) - 1) +} + +func hasAnnotationFilter(query, clause string) bool { + return strings.Contains(query, "AND "+clause) || strings.Contains(query, "WHERE "+clause) +} + +func firstSessionSummaryScopeArg(query string, values annotationNamedValues) int { + scopeArgs := 0 + if strings.Contains(query, "source_name IN") { + scopeArgs++ + } + if strings.Contains(query, "runtime IN") { + scopeArgs++ + } + if scopeArgs == 0 { + return len(values) + } + return len(values) - scopeArgs +} + +func annotationHasLabel(annotation models.TraceAnnotation, label string) bool { + for _, candidate := range annotation.Labels { + if candidate == label { + return true + } + } + return false +} + type annotationNamedValues []driver.Value func namedValues(args []driver.NamedValue) annotationNamedValues { diff --git a/internal/web/api_contract_test.go b/internal/web/api_contract_test.go index fe34967..b6a8f41 100644 --- a/internal/web/api_contract_test.go +++ b/internal/web/api_contract_test.go @@ -20,24 +20,30 @@ type apiContract struct { func TestAPIContractsMatchGoStructTags(t *testing.T) { contracts := loadAPIContracts(t) types := map[string]reflect.Type{ - "APISessionSummary": reflect.TypeOf(APISessionSummary{}), - "APIDashboardSessionsResponse": reflect.TypeOf(APIDashboardSessionsResponse{}), - "APIDashboardSearchResult": reflect.TypeOf(APIDashboardSearchResult{}), - "APIDashboardSearchResponse": reflect.TypeOf(APIDashboardSearchResponse{}), - "APIScopeMetadata": reflect.TypeOf(APIScopeMetadata{}), - "APIScopeFilters": reflect.TypeOf(APIScopeFilters{}), - "APIActivityItem": reflect.TypeOf(APIActivityItem{}), - "APIDashboardCharts": reflect.TypeOf(APIDashboardCharts{}), - "ModelSeriesChart": reflect.TypeOf(views.ModelSeriesChart{}), - "ModelMetricChart": reflect.TypeOf(views.ModelMetricChart{}), - "ModelMetricSeries": reflect.TypeOf(views.ModelMetricSeries{}), - "ModelSeriesDataset": reflect.TypeOf(views.ModelSeriesDataset{}), - "ModelAnalyticsSummary": reflect.TypeOf(views.ModelAnalyticsSummary{}), - "APISessionDetail": reflect.TypeOf(APISessionDetail{}), - "APISessionEvent": reflect.TypeOf(APISessionEvent{}), - "APIToolPayload": reflect.TypeOf(APIToolPayload{}), - "APITraceAnnotation": reflect.TypeOf(APITraceAnnotation{}), - "APITraceAnnotationListResponse": reflect.TypeOf(APITraceAnnotationListResponse{}), + "APISessionSummary": reflect.TypeOf(APISessionSummary{}), + "APIDashboardSessionsResponse": reflect.TypeOf(APIDashboardSessionsResponse{}), + "APIDashboardSearchResult": reflect.TypeOf(APIDashboardSearchResult{}), + "APIDashboardSearchResponse": reflect.TypeOf(APIDashboardSearchResponse{}), + "APIScopeMetadata": reflect.TypeOf(APIScopeMetadata{}), + "APIScopeFilters": reflect.TypeOf(APIScopeFilters{}), + "APIActivityItem": reflect.TypeOf(APIActivityItem{}), + "APIDashboardCharts": reflect.TypeOf(APIDashboardCharts{}), + "ModelSeriesChart": reflect.TypeOf(views.ModelSeriesChart{}), + "ModelMetricChart": reflect.TypeOf(views.ModelMetricChart{}), + "ModelMetricSeries": reflect.TypeOf(views.ModelMetricSeries{}), + "ModelSeriesDataset": reflect.TypeOf(views.ModelSeriesDataset{}), + "ModelAnalyticsSummary": reflect.TypeOf(views.ModelAnalyticsSummary{}), + "APISessionDetail": reflect.TypeOf(APISessionDetail{}), + "APISessionEvent": reflect.TypeOf(APISessionEvent{}), + "APIToolPayload": reflect.TypeOf(APIToolPayload{}), + "APITraceAnnotation": reflect.TypeOf(APITraceAnnotation{}), + "APITraceAnnotationListResponse": reflect.TypeOf(APITraceAnnotationListResponse{}), + "APIAnnotationCounts": reflect.TypeOf(APIAnnotationCounts{}), + "APIAnnotatedTargetSummary": reflect.TypeOf(APIAnnotatedTargetSummary{}), + "APIAnnotatedTraceSummary": reflect.TypeOf(APIAnnotatedTraceSummary{}), + "APIAnnotatedTracesResponse": reflect.TypeOf(APIAnnotatedTracesResponse{}), + "APIAnnotatedTraceExport": reflect.TypeOf(APIAnnotatedTraceExport{}), + "APIAnnotatedTraceExportResponse": reflect.TypeOf(APIAnnotatedTraceExportResponse{}), } for name, contract := range contracts { diff --git a/internal/web/api_requests.go b/internal/web/api_requests.go index 458af2e..d9ab8a1 100644 --- a/internal/web/api_requests.go +++ b/internal/web/api_requests.go @@ -19,6 +19,11 @@ const ( maxDashboardSearchAPILimit = 240 defaultAnnotationsAPILimit = 200 maxAnnotationsAPILimit = 500 + defaultAnnotatedTracesLimit = 50 + maxAnnotatedTracesLimit = 200 + maxAnnotatedTracesOffset = 100000 + defaultAnnotatedEventsLimit = 1000 + maxAnnotatedEventsLimit = 5000 ) type apiIntParam struct { @@ -268,6 +273,23 @@ type annotationsAPIRequest struct { Scope APIScopeFilters } +type annotatedTracesAPIRequest struct { + TargetType string + SessionID string + EventUID string + AuthorType string + Source string + Category string + Outcome string + Label string + NeedsFollowup *bool + IncludeDeleted bool + Limit int + Offset int + EventLimit int + Scope APIScopeFilters +} + func parseAnnotationsAPIRequest(values url.Values) (annotationsAPIRequest, error) { limit, err := parseAPIIntParam(values, apiIntParam{ Name: "limit", @@ -295,3 +317,67 @@ func parseAnnotationsAPIRequest(values url.Values) (annotationsAPIRequest, error Scope: parseAPIScopeFilters(values), }, nil } + +func parseAnnotatedTracesAPIRequest(values url.Values) (annotatedTracesAPIRequest, error) { + limit, err := parseAPIIntParam(values, apiIntParam{ + Name: "limit", + Default: defaultAnnotatedTracesLimit, + Min: 1, + Max: maxAnnotatedTracesLimit, + }) + if err != nil { + return annotatedTracesAPIRequest{}, err + } + offset, err := parseAPIIntParam(values, apiIntParam{Name: "offset", Default: 0, Min: 0, Max: maxAnnotatedTracesOffset}) + if err != nil { + return annotatedTracesAPIRequest{}, err + } + eventLimit, err := parseAPIIntParam(values, apiIntParam{ + Name: "event_limit", + Default: defaultAnnotatedEventsLimit, + Min: 1, + Max: maxAnnotatedEventsLimit, + }) + if err != nil { + return annotatedTracesAPIRequest{}, err + } + includeDeletedValue := values.Get("include_deleted") + includeDeleted := includeDeletedValue == "1" || strings.EqualFold(includeDeletedValue, "true") + needsFollowup, err := parseOptionalAPIBool(values.Get("needs_followup"), "needs_followup") + if err != nil { + return annotatedTracesAPIRequest{}, err + } + return annotatedTracesAPIRequest{ + TargetType: strings.TrimSpace(values.Get("target_type")), + SessionID: strings.TrimSpace(values.Get("session_id")), + EventUID: strings.TrimSpace(values.Get("event_uid")), + AuthorType: strings.TrimSpace(values.Get("author_type")), + Source: strings.TrimSpace(values.Get("source")), + Category: strings.TrimSpace(values.Get("category")), + Outcome: strings.TrimSpace(values.Get("outcome")), + Label: strings.TrimSpace(values.Get("label")), + NeedsFollowup: needsFollowup, + IncludeDeleted: includeDeleted, + Limit: limit, + Offset: offset, + EventLimit: eventLimit, + Scope: parseAPIScopeFilters(values), + }, nil +} + +func parseOptionalAPIBool(raw, name string) (*bool, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, nil + } + switch strings.ToLower(raw) { + case "1", "true": + value := true + return &value, nil + case "0", "false": + value := false + return &value, nil + default: + return nil, fmt.Errorf("invalid %s", name) + } +} diff --git a/internal/web/router.go b/internal/web/router.go index f69e8a5..0d0015c 100644 --- a/internal/web/router.go +++ b/internal/web/router.go @@ -72,6 +72,8 @@ func NewRouter( r.Get("/events/{event_id}", apiHandlers.GetEvent) r.Get("/events/{event_id}/annotations", apiHandlers.GetEventAnnotations) r.Get("/tool-payloads/{event_id}", apiHandlers.GetToolPayload) + r.Get("/annotations/traces", apiHandlers.ListAnnotatedTraces) + r.Get("/annotations/export", apiHandlers.ExportAnnotatedTraces) r.Get("/annotations", apiHandlers.ListAnnotations) r.Post("/annotations", apiHandlers.CreateAnnotation) r.Get("/annotations/{annotation_id}", apiHandlers.GetAnnotation) diff --git a/internal/web/viewmodels.go b/internal/web/viewmodels.go index 2add03a..310e9e4 100644 --- a/internal/web/viewmodels.go +++ b/internal/web/viewmodels.go @@ -201,3 +201,59 @@ type APITraceAnnotation struct { type APITraceAnnotationListResponse struct { Items []APITraceAnnotation `json:"items"` } + +type APIAnnotationCounts struct { + AnnotationCount int `json:"annotation_count"` + SessionAnnotationCount int `json:"session_annotation_count"` + MessageAnnotationCount int `json:"message_annotation_count"` + EventAnnotationCount int `json:"event_annotation_count"` + NeedsFollowupCount int `json:"needs_followup_count"` +} + +type APIAnnotatedTargetSummary struct { + TargetType string `json:"target_type"` + EventUID string `json:"event_uid,omitempty"` + AnnotationCount int `json:"annotation_count"` + FirstAnnotationAt time.Time `json:"first_annotation_at"` + LastAnnotationAt time.Time `json:"last_annotation_at"` +} + +type APIAnnotatedTraceSummary struct { + Session APISessionSummary `json:"session"` + Counts APIAnnotationCounts `json:"counts"` + FirstAnnotationAt time.Time `json:"first_annotation_at"` + LastAnnotationAt time.Time `json:"last_annotation_at"` + Targets []APIAnnotatedTargetSummary `json:"targets"` +} + +type APIAnnotatedTracesResponse struct { + Schema string `json:"schema"` + Scope APIScopeMetadata `json:"scope"` + IncludeDeleted bool `json:"include_deleted"` + Offset int `json:"offset"` + Limit int `json:"limit"` + HasMore bool `json:"has_more"` + Items []APIAnnotatedTraceSummary `json:"items"` +} + +type APIAnnotatedTraceExport struct { + Session APISessionSummary `json:"session"` + Counts APIAnnotationCounts `json:"counts"` + Annotations []APITraceAnnotation `json:"annotations"` + Events []APISessionEvent `json:"events"` + EventLimit int `json:"event_limit"` + EventTruncated bool `json:"event_truncated"` +} + +type APIAnnotatedTraceExportResponse struct { + Schema string `json:"schema"` + ExportedAt time.Time `json:"exported_at"` + Scope APIScopeMetadata `json:"scope"` + IncludeDeleted bool `json:"include_deleted"` + Offset int `json:"offset"` + Limit int `json:"limit"` + EventLimit int `json:"event_limit"` + HasMore bool `json:"has_more"` + Traces []APIAnnotatedTraceExport `json:"traces"` + Warnings []string `json:"warnings"` +} diff --git a/tests/contracts/api-contracts.json b/tests/contracts/api-contracts.json index 3b52879..1399d8d 100644 --- a/tests/contracts/api-contracts.json +++ b/tests/contracts/api-contracts.json @@ -249,6 +249,75 @@ }, "optional": {} }, + "APIAnnotationCounts": { + "required": { + "annotation_count": "integer", + "session_annotation_count": "integer", + "message_annotation_count": "integer", + "event_annotation_count": "integer", + "needs_followup_count": "integer" + }, + "optional": {} + }, + "APIAnnotatedTargetSummary": { + "required": { + "target_type": "string", + "annotation_count": "integer", + "first_annotation_at": "string", + "last_annotation_at": "string" + }, + "optional": { + "event_uid": "string" + } + }, + "APIAnnotatedTraceSummary": { + "required": { + "session": "APISessionSummary", + "counts": "APIAnnotationCounts", + "first_annotation_at": "string", + "last_annotation_at": "string", + "targets": "APIAnnotatedTargetSummary[]" + }, + "optional": {} + }, + "APIAnnotatedTracesResponse": { + "required": { + "schema": "string", + "scope": "APIScopeMetadata", + "include_deleted": "boolean", + "offset": "integer", + "limit": "integer", + "has_more": "boolean", + "items": "APIAnnotatedTraceSummary[]" + }, + "optional": {} + }, + "APIAnnotatedTraceExport": { + "required": { + "session": "APISessionSummary", + "counts": "APIAnnotationCounts", + "annotations": "APITraceAnnotation[]", + "events": "APISessionEvent[]", + "event_limit": "integer", + "event_truncated": "boolean" + }, + "optional": {} + }, + "APIAnnotatedTraceExportResponse": { + "required": { + "schema": "string", + "exported_at": "string", + "scope": "APIScopeMetadata", + "include_deleted": "boolean", + "offset": "integer", + "limit": "integer", + "event_limit": "integer", + "has_more": "boolean", + "traces": "APIAnnotatedTraceExport[]", + "warnings": "string[]" + }, + "optional": {} + }, "APIToolPayload": { "required": { "event_uid": "string",