diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin.go index 4b5822e17..3477bc07f 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin.go +++ b/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 }, diff --git a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go b/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go index 2a59eefcf..ca659bb53 100644 --- a/plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go +++ b/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") + } +}