improve itk logging

Author: Krzysztof Dziedzic

.github/workflows/itk.yaml

@@ -38,5 +38,5 @@ jobs:
         run: bash run_itk.sh
         working-directory: itk
         env:
-          A2A_SAMPLES_REVISION: itk-v.0.11-alpha
+          A2A_SAMPLES_REVISION: itk-v.015-alpha
 

.gitignore

@@ -63,3 +63,4 @@ aiwt-*
 itk/instruction.proto
 itk/pb
 itk/a2a-samples
+itk/logs

itk/README.md

@@ -36,7 +36,7 @@ You must set the `A2A_SAMPLES_REVISION` environment variable to specify which re
 
 Example:
 ```bash
-export A2A_SAMPLES_REVISION=itk-v.0.11-alpha
+export A2A_SAMPLES_REVISION=itk-v.015-alpha
 ```
 
 ### 2. Execute Tests
@@ -52,3 +52,19 @@ The script will:
 - Checkout the specified revision.
 - Build the ITK service Docker image.
 - Run the tests and output results.
+
+## Debugging
+
+To enable detailed debug logging and capture logs from the agents:
+
+1. Set the `ITK_LOG_LEVEL` environment variable to `DEBUG`:
+   ```bash
+   export ITK_LOG_LEVEL=DEBUG
+   ```
+
+2. Run the tests as usual:
+   ```bash
+   ./run_itk.sh
+   ```
+
+When run with `DEBUG` level, the script will create a `logs` directory in this `itk` folder and mount it to the container. You can find detailed logs for each agent in the `logs/` directory.

itk/go.mod

@@ -10,7 +10,7 @@ require (
 )
 
 require (
-	github.com/a2aproject/a2a-go v0.3.12 // indirect
+	github.com/a2aproject/a2a-go v0.3.14-0.20260404065426-6a9878f8f8a8 // indirect
 	github.com/google/uuid v1.6.0 // indirect
 	golang.org/x/mod v0.33.0 // indirect
 	golang.org/x/net v0.48.0 // indirect

itk/main.go

@@ -1,12 +1,15 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"encoding/base64"
 	"encoding/json"
 	"flag"
 	"fmt"
+	"io"
 	"iter"
+	"log/slog"
 	"net"
 	"net/http"
 	"os"
@@ -276,6 +279,27 @@ func main() {
 func run() error {
 	flag.Parse()
 
+	logLevelStr := os.Getenv("ITK_LOG_LEVEL")
+	if logLevelStr == "" {
+		logLevelStr = "INFO"
+	}
+	var level slog.Level
+	switch strings.ToUpper(logLevelStr) {
+	case "DEBUG":
+		level = slog.LevelDebug
+	case "INFO":
+		level = slog.LevelInfo
+	case "WARN":
+		level = slog.LevelWarn
+	case "ERROR":
+		level = slog.LevelError
+	default:
+		level = slog.LevelInfo
+	}
+	
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: level}))
+	slog.SetDefault(logger)
+
 	jsonRPCV0Addr := fmt.Sprintf("http://127.0.0.1:%d", *httpPort)
 
 	agentCard := &a2a.AgentCard{
@@ -310,7 +334,10 @@ func run() error {
 	}
 
 	executor := &V10AgentExecutor{}
-	requestHandler := a2asrv.NewHandler(executor)
+	requestHandler := a2asrv.NewHandler(
+		executor,
+		a2asrv.WithCallInterceptors(a2asrv.NewLoggingInterceptor(&a2asrv.LoggingConfig{LogPayload: true})),
+	)
 
 	// Servers
 	mux := http.NewServeMux()
@@ -324,7 +351,7 @@ func run() error {
 
 	httpServer := &http.Server{
 		Addr:              fmt.Sprintf(":%d", *httpPort),
-		Handler:           mux,
+		Handler:           loggingMiddleware(logger, mux),
 		ReadHeaderTimeout: 3 * time.Second,
 	}
 
@@ -342,7 +369,10 @@ func run() error {
 		return nil
 	})
 
-	grpcServer := grpc.NewServer()
+	grpcServer := grpc.NewServer(
+		grpc.UnaryInterceptor(unaryLoggingInterceptor(logger)),
+		grpc.StreamInterceptor(streamLoggingInterceptor(logger)),
+	)
 	a2agrpc.NewHandler(requestHandler).RegisterWith(grpcServer)
 	g.Go(func() error {
 		lis, err := net.Listen("tcp", fmt.Sprintf(":%d", *grpcPort))
@@ -369,3 +399,34 @@ func run() error {
 
 	return g.Wait()
 }
+
+func loggingMiddleware(logger *slog.Logger, next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		var bodyBytes []byte
+		if r.Body != nil {
+			var err error
+			bodyBytes, err = io.ReadAll(r.Body)
+			if err != nil {
+				logger.Error("Failed to read request body", err)
+			} else {
+				r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
+			}
+		}
+		logger.Info("Incoming request", "method", r.Method, "path", r.URL.Path, "remote", r.RemoteAddr, "body", string(bodyBytes))
+		next.ServeHTTP(w, r)
+	})
+}
+
+func unaryLoggingInterceptor(logger *slog.Logger) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+		logger.Info("gRPC Unary Call", "method", info.FullMethod)
+		return handler(ctx, req)
+	}
+}
+
+func streamLoggingInterceptor(logger *slog.Logger) grpc.StreamServerInterceptor {
+	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		logger.Info("gRPC Stream Call", "method", info.FullMethod)
+		return handler(srv, ss)
+	}
+}

itk/run_itk.sh

@@ -1,6 +1,9 @@
 #!/bin/bash
 set -ex
 
+# Set default log level
+export ITK_LOG_LEVEL="${ITK_LOG_LEVEL:-INFO}"
+
 # Initialize default exit code
 RESULT=1
 
@@ -66,13 +69,25 @@ ITK_DIR=$(pwd)
 # Stop existing container if any
 docker rm -f itk-service || true
 
+# Create logs directory if debug
+if [ "${ITK_LOG_LEVEL^^}" = "DEBUG" ]; then
+  mkdir -p "$ITK_DIR/logs"
+fi
+
+DOCKER_MOUNT_LOGS=""
+if [ "${ITK_LOG_LEVEL^^}" = "DEBUG" ]; then
+  DOCKER_MOUNT_LOGS="-v $ITK_DIR/logs:/app/logs"
+fi
+
 # Run container with protobuf registration conflict set to 'warn'
 # This is necessary because the SDK v2 depends on its predecessor v0.x,
 # causing global proto registration conflicts.
 docker run -d --name itk-service \
   -e GOLANG_PROTOBUF_REGISTRATION_CONFLICT=warn \
+  -e ITK_LOG_LEVEL="$ITK_LOG_LEVEL" \
   -v "$A2A_GO_ROOT:/app/agents/repo" \
   -v "$ITK_DIR:/app/agents/repo/itk" \
+  $DOCKER_MOUNT_LOGS \
   -p 8000:8000 \
   itk_service