feat: Pass rich structured event types to OnEventCallback (StateDelta, HITL, A2A)

plugin/agentanalytics/bigquery_agent_analytics_plugin.go

@@ -17,7 +17,9 @@ package agentanalytics
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
+	"strings"
 	"time"
 
 	bq "cloud.google.com/go/bigquery"
@@ -102,6 +104,10 @@ func NewBigQueryAgentAnalyticsPluginWithClients(
 		}
 		err = tableRef.Create(ctx, &bq.TableMetadata{
 			Schema: EventsSchema(),
+			TimePartitioning: &bq.TimePartitioning{
+				Field: "timestamp",
+				Type:  bq.DayPartitioningType,
+			},
 			Clustering: &bq.Clustering{
 				Fields: config.ClusteringFields,
 			},
@@ -232,6 +238,74 @@ func NewBigQueryAgentAnalyticsPluginWithClients(
 		},
 		OnEventCallback: func(ctx agent.InvocationContext, ev *session.Event) (*session.Event, error) {
 			attrs := map[string]any{"event_author": ev.Author}
+
+			// 1. State Delta Logging
+			if len(ev.Actions.StateDelta) > 0 {
+				stateAttrs := map[string]any{
+					"state_delta": ev.Actions.StateDelta,
+				}
+				for k, v := range attrs {
+					stateAttrs[k] = v
+				}
+				logEvent(ctx, "STATE_DELTA", nil, stateAttrs)
+			}
+
+			// 2. HITL Logging (Injected synthetic tools calls & responses)
+			if ev.Content != nil {
+				for _, part := range ev.Content.Parts {
+					if part.FunctionCall != nil {
+						switch part.FunctionCall.Name {
+						case "adk_request_credential":
+							logEvent(ctx, "HITL_CREDENTIAL_REQUEST", map[string]any{"tool": part.FunctionCall.Name, "args": part.FunctionCall.Args}, attrs)
+						case "adk_request_confirmation":
+							logEvent(ctx, "HITL_CONFIRMATION_REQUEST", map[string]any{"tool": part.FunctionCall.Name, "args": part.FunctionCall.Args}, attrs)
+						case "adk_request_input":
+							logEvent(ctx, "HITL_INPUT_REQUEST", map[string]any{"tool": part.FunctionCall.Name, "args": part.FunctionCall.Args}, attrs)
+						}
+					}
+					if part.FunctionResponse != nil {
+						switch part.FunctionResponse.Name {
+						case "adk_request_credential":
+							logEvent(ctx, "HITL_CREDENTIAL_REQUEST_COMPLETED", map[string]any{"tool": part.FunctionResponse.Name, "result": part.FunctionResponse.Response}, attrs)
+						case "adk_request_confirmation":
+							logEvent(ctx, "HITL_CONFIRMATION_REQUEST_COMPLETED", map[string]any{"tool": part.FunctionResponse.Name, "result": part.FunctionResponse.Response}, attrs)
+						case "adk_request_input":
+							logEvent(ctx, "HITL_INPUT_REQUEST_COMPLETED", map[string]any{"tool": part.FunctionResponse.Name, "result": part.FunctionResponse.Response}, attrs)
+						}
+					}
+				}
+			}
+
+			// 3. A2A Interaction Logging
+			if ev.CustomMetadata != nil {
+				a2aKeys := make(map[string]any)
+				for k, v := range ev.CustomMetadata {
+					if strings.HasPrefix(k, "a2a:") {
+						a2aKeys[k] = v
+					}
+				}
+				if len(a2aKeys) > 0 {
+					a2aTruncatedRaw, _, _ := SmartTruncate(a2aKeys, config.MaxContentLen)
+					var a2aTruncated map[string]any
+					_ = json.Unmarshal(a2aTruncatedRaw, &a2aTruncated)
+
+					responsePayload := a2aKeys["a2a:response"]
+					var contentDict any
+					if responsePayload != nil {
+						contentTruncRaw, _, _ := SmartTruncate(responsePayload, config.MaxContentLen)
+						_ = json.Unmarshal(contentTruncRaw, &contentDict)
+					}
+
+					stateAttrs := map[string]any{
+						"a2a_metadata": a2aTruncated,
+					}
+					for k, v := range attrs {
+						stateAttrs[k] = v
+					}
+					logEvent(ctx, "A2A_INTERACTION", contentDict, stateAttrs)
+				}
+			}
+
 			logEvent(ctx, "EVENT", ev.Content, attrs)
 			return ev, nil
 		},

plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go

@@ -15,6 +15,7 @@
 package agentanalytics
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"io"
@@ -342,3 +343,198 @@ func TestLogEvent_ExtractsTraceInfo(t *testing.T) {
 		t.Error("Timed out waiting for request")
 	}
 }
+
+func TestNewBigQueryAgentAnalyticsPlugin_CreateTable_WithPartitioning(t *testing.T) {
+	ctx := context.Background()
+	config := DefaultConfig()
+	config.Enabled = true
+	config.ProjectID = "test-project"
+	config.DatasetID = "test-dataset"
+	config.TableName = "test-table"
+
+	createCalled := false
+	var requestBody string
+
+	mockTransport := &mockTransport{
+		roundTrip: func(r *http.Request) (*http.Response, error) {
+			// Table metadata request: returns 404 Not Found to trigger creation
+			if r.Method == "GET" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables/test-table") {
+				return &http.Response{
+					StatusCode: http.StatusNotFound,
+					Body:       io.NopCloser(strings.NewReader(`{"error":{"code":404,"message":"Not found"}}`)),
+				}, nil
+			}
+			// Table creation request
+			if r.Method == "POST" && strings.Contains(r.URL.Path, "/datasets/test-dataset/tables") {
+				createCalled = true
+				bodyBytes, _ := io.ReadAll(r.Body)
+				requestBody = string(bodyBytes)
+				r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
+
+				return &http.Response{
+					StatusCode: http.StatusOK,
+					Body:       io.NopCloser(strings.NewReader("{}")),
+				}, nil
+			}
+			return &http.Response{
+				StatusCode: http.StatusOK,
+				Body:       io.NopCloser(strings.NewReader("{}")),
+			}, nil
+		},
+	}
+	httpClient := &http.Client{Transport: mockTransport}
+	bqClient, err := bq.NewClient(ctx, config.ProjectID, option.WithHTTPClient(httpClient))
+	if err != nil {
+		t.Fatalf("Failed to create bigquery client: %v", err)
+	}
+
+	lis, err := net.Listen("tcp", "localhost:0")
+	if err != nil {
+		t.Fatalf("failed to listen: %v", err)
+	}
+	gSrv := grpc.NewServer()
+	storagepb.RegisterBigQueryWriteServer(gSrv, &fakeBigQueryWriteServer{})
+	go func() { _ = gSrv.Serve(lis) }()
+	t.Cleanup(gSrv.Stop)
+
+	conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatalf("failed to dial test server: %v", err)
+	}
+	t.Cleanup(func() { _ = conn.Close() })
+
+	writeClient, err := bqstorage.NewBigQueryWriteClient(ctx, option.WithGRPCConn(conn))
+	if err != nil {
+		t.Fatalf("Failed to create BigQuery write client: %v", err)
+	}
+
+	_, err = NewBigQueryAgentAnalyticsPluginWithClients(ctx, config, bqClient, writeClient)
+	if err != nil {
+		t.Fatalf("Plugin initialization error: %v", err)
+	}
+
+	if !createCalled {
+		t.Error("Expected table creation to be called")
+	}
+
+	if !strings.Contains(requestBody, "timePartitioning") {
+		t.Errorf("Expected request body to contain 'timePartitioning', got: %s", requestBody)
+	}
+	if !strings.Contains(requestBody, "DAY") {
+		t.Errorf("Expected partitioning type to be 'DAY', got request body: %s", requestBody)
+	}
+	if !strings.Contains(requestBody, "timestamp") {
+		t.Errorf("Expected partitioning field to be 'timestamp', got request body: %s", requestBody)
+	}
+}
+
+func TestOnEventCallback_StateDelta(t *testing.T) {
+	p, requestsChan, mCtx := setupTestPlugin(t)
+
+	eventCb := p.OnEventCallback()
+	if eventCb == nil {
+		t.Fatal("OnEventCallback is nil")
+	}
+
+	_, err := eventCb(mCtx, &session.Event{
+		Author: "test-author",
+		Actions: session.EventActions{
+			StateDelta: map[string]any{
+				"key1": "value1",
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("OnEventCallback error: %v", err)
+	}
+
+	p.AfterRunCallback()(mCtx)
+
+	select {
+	case req := <-requestsChan:
+		if req == nil {
+			t.Error("Received nil request")
+		}
+	case <-time.After(5 * time.Second):
+		t.Error("Timed out waiting for request")
+	}
+}
+
+func TestOnEventCallback_HITL(t *testing.T) {
+	p, requestsChan, mCtx := setupTestPlugin(t)
+
+	eventCb := p.OnEventCallback()
+	if eventCb == nil {
+		t.Fatal("OnEventCallback is nil")
+	}
+
+	_, err := eventCb(mCtx, &session.Event{
+		Author: "test-author",
+		LLMResponse: model.LLMResponse{
+			Content: &genai.Content{
+				Parts: []*genai.Part{
+					{
+						FunctionCall: &genai.FunctionCall{
+							Name: "adk_request_credential",
+							Args: map[string]any{"prompt": "need password"},
+						},
+					},
+					{
+						FunctionResponse: &genai.FunctionResponse{
+							Name:     "adk_request_credential",
+							Response: map[string]any{"password": "123"},
+						},
+					},
+				},
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("OnEventCallback error: %v", err)
+	}
+
+	p.AfterRunCallback()(mCtx)
+
+	select {
+	case req := <-requestsChan:
+		if req == nil {
+			t.Error("Received nil request")
+		}
+	case <-time.After(5 * time.Second):
+		t.Error("Timed out waiting for request")
+	}
+}
+
+func TestOnEventCallback_A2A(t *testing.T) {
+	p, requestsChan, mCtx := setupTestPlugin(t)
+
+	eventCb := p.OnEventCallback()
+	if eventCb == nil {
+		t.Fatal("OnEventCallback is nil")
+	}
+
+	_, err := eventCb(mCtx, &session.Event{
+		Author: "test-author",
+		LLMResponse: model.LLMResponse{
+			CustomMetadata: map[string]any{
+				"a2a:task_id":    "task-123",
+				"a2a:context_id": "ctx-456",
+				"a2a:response":   `{"result": "success"}`,
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("OnEventCallback error: %v", err)
+	}
+
+	p.AfterRunCallback()(mCtx)
+
+	select {
+	case req := <-requestsChan:
+		if req == nil {
+			t.Error("Received nil request")
+		}
+	case <-time.After(5 * time.Second):
+		t.Error("Timed out waiting for request")
+	}
+}