feat: Implement tool provenance tracking for tool execution logging

plugin/agentanalytics/bigquery_agent_analytics_plugin.go

@@ -18,6 +18,8 @@ package agentanalytics
 import (
 	"context"
 	"fmt"
+	"reflect"
+	"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,
 			},
@@ -268,18 +274,25 @@ func NewBigQueryAgentAnalyticsPluginWithClients(
 		},
 
 		BeforeToolCallback: func(ctx agent.Context, t tool.Tool, args map[string]any) (map[string]any, error) {
-			attrs := map[string]any{"tool_name": t.Name()}
+			attrs := map[string]any{
+				"tool_name":   t.Name(),
+				"tool_origin": getToolOrigin(t),
+			}
 			logEvent(ctx, "TOOL_START", args, attrs)
 			return nil, nil
 		},
 		AfterToolCallback: func(ctx agent.Context, t tool.Tool, args, res map[string]any, err error) (map[string]any, error) {
-			attrs := map[string]any{"tool_name": t.Name()}
+			attrs := map[string]any{
+				"tool_name":   t.Name(),
+				"tool_origin": getToolOrigin(t),
+			}
 			logEvent(ctx, "TOOL_END", res, attrs)
 			return nil, nil
 		},
 		OnToolErrorCallback: func(ctx agent.Context, t tool.Tool, args map[string]any, err error) (map[string]any, error) {
 			attrs := map[string]any{
 				"tool_name":     t.Name(),
+				"tool_origin":   getToolOrigin(t),
 				"error_message": err.Error(),
 			}
 			logEvent(ctx, "TOOL_ERROR", nil, attrs)
@@ -297,3 +310,33 @@ func NewBigQueryAgentAnalyticsPluginWithClients(
 
 	return baseplugin.New(cfg)
 }
+
+func getToolOrigin(t tool.Tool) string {
+	tType := reflect.TypeOf(t).String()
+	if strings.Contains(tType, "mcptoolset.mcpTool") {
+		return "MCP"
+	}
+	if strings.Contains(tType, "agenttool.agentTool") {
+		val := reflect.ValueOf(t)
+		if val.Kind() == reflect.Ptr {
+			val = val.Elem()
+		}
+		if val.Kind() == reflect.Struct {
+			agentField := val.FieldByName("agent")
+			if agentField.IsValid() {
+				agentType := agentField.Type().String()
+				if strings.Contains(agentType, "remoteagent.a2aAgent") || strings.Contains(agentType, "remoteagent/v2.a2aAgent") {
+					return "A2A"
+				}
+			}
+		}
+		return "SUB_AGENT"
+	}
+	if strings.Contains(tType, "functiontool.functionTool") {
+		return "LOCAL"
+	}
+	if strings.Contains(tType, "TransferToAgent") {
+		return "TRANSFER_AGENT"
+	}
+	return "UNKNOWN"
+}

plugin/agentanalytics/bigquery_agent_analytics_plugin_test.go

@@ -15,6 +15,7 @@
 package agentanalytics
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"io"
@@ -37,6 +38,8 @@ import (
 	"google.golang.org/adk/v2/model"
 	baseplugin "google.golang.org/adk/v2/plugin"
 	"google.golang.org/adk/v2/session"
+	"google.golang.org/adk/v2/tool/agenttool"
+	"google.golang.org/adk/v2/tool/functiontool"
 )
 
 type mockTransport struct {
@@ -342,3 +345,112 @@ 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 TestToolProvenance(t *testing.T) {
+	// 1. Test LOCAL tool origin
+	localTool, err := functiontool.New[map[string]any, map[string]any](functiontool.Config{
+		Name:        "local_test_tool",
+		Description: "Local tool description",
+	}, func(ctx agent.Context, args map[string]any) (map[string]any, error) {
+		return map[string]any{"status": "ok"}, nil
+	})
+	if err != nil {
+		t.Fatalf("Failed to create functiontool: %v", err)
+	}
+
+	origin := getToolOrigin(localTool)
+	if origin != "LOCAL" {
+		t.Errorf("Expected LOCAL tool origin, got: %s", origin)
+	}
+
+	// 2. Test SUB_AGENT tool origin
+	subAgentTool := agenttool.New(nil, nil)
+	origin = getToolOrigin(subAgentTool)
+	if origin != "SUB_AGENT" {
+		t.Errorf("Expected SUB_AGENT tool origin, got: %s", origin)
+	}
+}