fix unit test

Author: hyprh

client/client_test.go

@@ -206,7 +206,7 @@ func TestA2AClient_SetPushNotification(t *testing.T) {
 		TaskID: taskID,
 		PushNotificationConfig: jsonprotocol.PushNotificationConfig{
 			URL: "https://example.com/webhook",
-			Authentication: &jsonprotocol.AuthenticationInfo{
+			Authentication: &jsonprotocol.PushNotificationAuthenticationInfo{
 				Schemes: []string{"bearer"},
 			},
 		},
@@ -314,7 +314,7 @@ func TestA2AClient_GetPushNotification(t *testing.T) {
 			TaskID: taskID,
 			PushNotificationConfig: jsonprotocol.PushNotificationConfig{
 				URL: "https://example.com/webhook",
-				Authentication: &jsonprotocol.AuthenticationInfo{
+				Authentication: &jsonprotocol.PushNotificationAuthenticationInfo{
 					Schemes: []string{"bearer"},
 				},
 			},

protocol/jsonprotocol/types.go

@@ -659,25 +659,6 @@ func NewTaskArtifactUpdateEvent(taskID, contextID string, artifact Artifact, las
 	}
 }
 
-// SendTaskParams defines the parameters for the tasks_send and tasks_sendSubscribe RPC methods.
-// See A2A Spec section on RPC Methods.
-type SendTaskParams struct {
-	// RPCID is the ID of json-rpc.
-	RPCID string `json:"-"`
-	// ID is the ID of the task.
-	ID string `json:"id"`
-	// SessionID is the optional session ID.
-	SessionID *string `json:"sessionId,omitempty"`
-	// Message is the user's message initiating the task.
-	Message Message `json:"message"`
-	// PushNotification contains optional push notification settings.
-	PushNotification *PushNotificationConfig `json:"pushNotification,omitempty"`
-	// HistoryLength is the requested history length in response.
-	HistoryLength *int `json:"historyLength,omitempty"`
-	// Metadata is the optional metadata.
-	Metadata map[string]interface{} `json:"metadata,omitempty"`
-}
-
 // TaskQueryParams defines the parameters for the tasks_get RPC method.
 // See A2A Spec section on RPC Methods.
 type TaskQueryParams struct {

server/server.go

@@ -180,7 +180,7 @@ func (s *A2AServer) handleAgentCard(w http.ResponseWriter, r *http.Request) {
 }
 
 // handleJSONRPC is the main handler for all JSON-RPC 2.0 requests.
-// Routes methods like tasks/send, tasks/get, etc., as defined in A2A Spec.
+// Routes methods like message/send, tasks/get, tasks/cancel, etc., as defined in A2A Spec.
 func (s *A2AServer) handleJSONRPC(w http.ResponseWriter, r *http.Request) {
 	// --- CORS Handling ---
 	if s.corsEnabled {
@@ -555,7 +555,7 @@ func (s *A2AServer) handleTasksResubscribe(ctx context.Context, w http.ResponseW
 	}
 
 	// Use the helper function to handle the SSE stream
-	handleSSEStream(ctx, s.corsEnabled, w, flusher, eventsChan, request.ID.(string), true)
+	handleSSEStream(ctx, s.corsEnabled, w, flusher, eventsChan, request.ID, true)
 }
 
 // handleMessageSend handles the message_send method.
@@ -613,7 +613,7 @@ func (s *A2AServer) handleMessageStream(ctx context.Context, w http.ResponseWrit
 	}
 
 	// Use the helper function to handle the SSE stream
-	handleSSEStream(ctx, s.corsEnabled, w, flusher, eventsChan, request.ID.(string), false)
+	handleSSEStream(ctx, s.corsEnabled, w, flusher, eventsChan, request.ID, false)
 }
 
 // handleSSEStream handles an SSE stream for a task, including setup and event forwarding.
@@ -624,7 +624,7 @@ func handleSSEStream(
 	w http.ResponseWriter,
 	flusher http.Flusher,
 	eventsChan <-chan jsonprotocol.StreamingMessageEvent,
-	rpcID string,
+	rpcID interface{},
 	isResubscribe bool,
 ) {
 	// Set headers for SSE.
@@ -641,9 +641,9 @@ func handleSSEStream(
 
 	// Log appropriate message based on whether this is a new subscription or resubscribe
 	if isResubscribe {
-		log.Infof("SSE stream reopened for request ID: %v)", rpcID)
+		log.Debugf("SSE stream reopened for request ID: %v)", rpcID)
 	} else {
-		log.Infof("SSE stream opened for request ID: %v)", rpcID)
+		log.Debugf("SSE stream opened for request ID: %v)", rpcID)
 	}
 
 	// Use request context to detect client disconnection.
@@ -669,13 +669,13 @@ func handleSSEStream(
 			flusher.Flush()
 		case <-clientClosed:
 			// Client disconnected (request context canceled).
-			log.Infof("SSE client disconnected for request ID: %s. Closing stream.", rpcID)
+			log.Debugf("SSE client disconnected for request ID: %s. Closing stream.", rpcID)
 			return // Exit the handler.
 		}
 	}
 }
 
-func sendSSEEvent(w http.ResponseWriter, rpcID string, event interface{}) error {
+func sendSSEEvent(w http.ResponseWriter, rpcID interface{}, event interface{}) error {
 	// Determine event type string for SSE.
 	var eventType string
 	var actualEvent jsonprotocol.Event

server/server_handlers_test.go

@@ -13,7 +13,7 @@ import (
 
 	"github.com/a2aproject/a2a-go/auth"
 	"github.com/a2aproject/a2a-go/internal/jsonrpc"
-	jsonrpc1 "github.com/a2aproject/a2a-go/protocol/jsonprotocol"
+	"github.com/a2aproject/a2a-go/protocol/jsonprotocol"
 	"github.com/a2aproject/a2a-go/taskmanager"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -126,7 +126,7 @@ func TestA2AServer_HandlerErrors(t *testing.T) {
 	// Test wrong content type
 	t.Run("Wrong Content Type", func(t *testing.T) {
 		reqBody := bytes.NewBufferString(
-			`{"jsonrpc":"2.0","method":"tasks/send","params":{},"id":"test-id"}`)
+			`{"jsonrpc":"2.0","method":"tasks/get","params":{},"id":"test-id"}`)
 		testJSONRPCErrorResponse(t, testServer, http.MethodPost, reqBody, "text/plain",
 			jsonrpc.CodeInvalidRequest, "Content-Type")
 	})
@@ -141,7 +141,7 @@ func TestA2AServer_HandlerErrors(t *testing.T) {
 	// Test wrong JSONRPC version
 	t.Run("Wrong JSONRPC Version", func(t *testing.T) {
 		reqBody := bytes.NewBufferString(
-			`{"jsonrpc":"1.0","method":"tasks/send","params":{},"id":"test-id"}`)
+			`{"jsonrpc":"1.0","method":"tasks/get","params":{},"id":"test-id"}`)
 		testJSONRPCErrorResponse(t, testServer, http.MethodPost, reqBody, "application/json",
 			jsonrpc.CodeInvalidRequest, "jsonrpc field must be '2.0'")
 	})
@@ -156,10 +156,10 @@ func TestA2AServer_HandlerErrors(t *testing.T) {
 
 	// Test invalid parameters
 	t.Run("Invalid Parameters", func(t *testing.T) {
-		// Missing required fields in params (Message with no parts)
+		// Message with empty parts array for message/stream
 		reqBody := bytes.NewBufferString(
-			`{"jsonrpc":"2.0","method":"tasks/send",
-			"params":{"id":"test-task","message":{"role":"user","parts":[]}},"id":"test-id"}`)
+			`{"jsonrpc":"2.0","method":"message/stream",
+			"params":{"message":{"role":"user","parts":[],"messageId":"msg-test","kind":"message"}},"id":"test-id"}`)
 		testJSONRPCErrorResponse(t, testServer, http.MethodPost, reqBody, "application/json",
 			jsonrpc.CodeInvalidParams, "")
 	})
@@ -186,15 +186,15 @@ func TestA2AServer_AuthMiddleware(t *testing.T) {
 
 	t.Run("Auth Success", func(t *testing.T) {
 		// Configure mock task manager to succeed
-		mockTM.GetResponse = &jsonrpc1.Task{
+		mockTM.GetResponse = &jsonprotocol.Task{
 			ID:     "test-task-auth",
-			Status: jsonrpc1.TaskStatus{State: jsonrpc1.TaskStateCompleted},
+			Status: jsonprotocol.TaskStatus{State: jsonprotocol.TaskStateCompleted},
 		}
 		mockTM.GetError = nil
 
 		// Create valid request with auth header
-		req, _ := createJSONRPCRequest(t, jsonrpc1.MethodTasksGet,
-			jsonrpc1.TaskQueryParams{ID: "test-task-auth"}, "req-auth-1")
+		req, _ := createJSONRPCRequest(t, jsonprotocol.MethodTasksGet,
+			jsonprotocol.TaskQueryParams{ID: "test-task-auth"}, "req-auth-1")
 		req.Header.Set("X-API-Key", "test-api-key") // Valid API key
 
 		resp := executeRequest(t, testServer, req, testServer.URL+"/")
@@ -208,8 +208,8 @@ func TestA2AServer_AuthMiddleware(t *testing.T) {
 
 	t.Run("Auth Failure", func(t *testing.T) {
 		// Create request with invalid auth
-		req, _ := createJSONRPCRequest(t, jsonrpc1.MethodTasksGet,
-			jsonrpc1.TaskQueryParams{ID: "test-task-auth"}, "req-auth-2")
+		req, _ := createJSONRPCRequest(t, jsonprotocol.MethodTasksGet,
+			jsonprotocol.TaskQueryParams{ID: "test-task-auth"}, "req-auth-2")
 		req.Header.Set("X-API-Key", "invalid-key") // Invalid API key
 
 		resp := executeRequest(t, testServer, req, testServer.URL+"/")
@@ -220,7 +220,7 @@ func TestA2AServer_AuthMiddleware(t *testing.T) {
 
 	// Test that agent card endpoint is still accessible without auth
 	t.Run("AgentCard_NoAuth", func(t *testing.T) {
-		req, err := http.NewRequest(http.MethodGet, testServer.URL+jsonrpc1.AgentCardPath, nil)
+		req, err := http.NewRequest(http.MethodGet, testServer.URL+jsonprotocol.AgentCardPath, nil)
 		require.NoError(t, err)
 
 		resp, err := testServer.Client().Do(req)
@@ -244,44 +244,44 @@ func TestA2AServer_Resubscribe(t *testing.T) {
 
 	t.Run("Resubscribe_Success", func(t *testing.T) {
 		// Configure mock events
-		workingEvent := jsonrpc1.TaskStatusUpdateEvent{
+		workingEvent := jsonprotocol.TaskStatusUpdateEvent{
 			TaskID:    "resubscribe-task",
 			ContextID: "test-context",
-			Kind:      jsonrpc1.KindTaskStatusUpdate,
-			Status:    jsonrpc1.TaskStatus{State: jsonrpc1.TaskStateWorking},
+			Kind:      jsonprotocol.KindTaskStatusUpdate,
+			Status:    jsonprotocol.TaskStatus{State: jsonprotocol.TaskStateWorking},
 		}
 		finalPtr := true
-		completedEvent := jsonrpc1.TaskStatusUpdateEvent{
+		completedEvent := jsonprotocol.TaskStatusUpdateEvent{
 			TaskID:    "resubscribe-task",
 			ContextID: "test-context",
-			Kind:      jsonrpc1.KindTaskStatusUpdate,
-			Status:    jsonrpc1.TaskStatus{State: jsonrpc1.TaskStateCompleted},
+			Kind:      jsonprotocol.KindTaskStatusUpdate,
+			Status:    jsonprotocol.TaskStatus{State: jsonprotocol.TaskStateCompleted},
 			Final:     finalPtr,
 		}
-		mockTM.SubscribeEvents = []jsonrpc1.StreamingMessageEvent{
+		mockTM.SubscribeEvents = []jsonprotocol.StreamingMessageEvent{
 			{Result: &workingEvent},
 			{Result: &completedEvent},
 		}
 		mockTM.SubscribeError = nil
 
 		// Add task to mock task manager to ensure it exists
-		mockTM.tasks["resubscribe-task"] = &jsonrpc1.Task{
+		mockTM.tasks["resubscribe-task"] = &jsonprotocol.Task{
 			ID: "resubscribe-task",
-			Status: jsonrpc1.TaskStatus{
-				State:     jsonrpc1.TaskStateWorking,
+			Status: jsonprotocol.TaskStatus{
+				State:     jsonprotocol.TaskStateWorking,
 				Timestamp: getCurrentTimestamp(),
 			},
 		}
 
 		// Create request - resubscribe expects SSE response
-		params := jsonrpc1.TaskIDParams{
+		params := jsonprotocol.TaskIDParams{
 			ID: "resubscribe-task",
 		}
 		paramsBytes, _ := json.Marshal(params)
 
 		reqBody := jsonrpc.Request{
 			Message: jsonrpc.Message{JSONRPC: "2.0", ID: "req-resub-1"},
-			Method:  jsonrpc1.MethodTasksResubscribe,
+			Method:  jsonprotocol.MethodTasksResubscribe,
 			Params:  json.RawMessage(paramsBytes),
 		}
 		reqBytes, _ := json.Marshal(reqBody)
@@ -317,14 +317,14 @@ func TestA2AServer_Resubscribe(t *testing.T) {
 		mockTM.SubscribeError = taskmanager.ErrTaskNotFound("nonexistent-task")
 
 		// Create request
-		params := jsonrpc1.TaskIDParams{
+		params := jsonprotocol.TaskIDParams{
 			ID: "nonexistent-task",
 		}
 
 		resp := performJSONRPCRequest(
 			t,
 			testServer,
-			jsonrpc1.MethodTasksResubscribe,
+			jsonprotocol.MethodTasksResubscribe,
 			params,
 			"req-resub-err",
 		)