fix issues

Author: noel2004

common/types/response.go

@@ -15,7 +15,16 @@ type Response struct {
 }
 
 func (resp *Response) DecodeData(out interface{}) error {
-	return mapstructure.Decode(resp.Data, out)
+	// Decode generically unmarshaled JSON (map[string]any, []any) into a typed struct
+	// honoring `json` tags and allowing weak type conversions.
+	dec, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{
+		TagName: "json",
+		Result:  out,
+	})
+	if err != nil {
+		return err
+	}
+	return dec.Decode(resp.Data)
 }
 
 // RenderJSON renders response with json

coordinator/internal/controller/proxy/controller.go

@@ -39,5 +39,5 @@ func InitController(cfg *config.ProxyConfig, reg prometheus.Registerer) {
 
 	Auth = NewAuthController(cfg, clients, proverManager)
 	GetTask = NewGetTaskController(cfg, clients, proverManager, priorityManager, reg)
-	SubmitProof = NewSubmitProofController(cfg, clients, proverManager, reg)
+	SubmitProof = NewSubmitProofController(cfg, clients, proverManager, priorityManager, reg)
 }

coordinator/internal/controller/proxy/get_task.go

@@ -146,6 +146,7 @@ func (ptc *GetTaskController) GetTasks(ctx *gin.Context) {
 			// TODO: log error
 		}
 	}
+	ptc.priorityUpstream.Delete(publicKey)
 
 	// Create a slice to hold the keys
 	keys := make([]string, 0, len(ptc.clients))

coordinator/internal/controller/proxy/prover_session.go

@@ -65,7 +65,7 @@ type proverSession struct {
 	completionCtx context.Context
 }
 
-func (c *proverSession) maintainLogin(ctx context.Context, cliMgr Client, up string, param *types.LoginParameter, phase uint) (*types.LoginSchema, error) {
+func (c *proverSession) maintainLogin(ctx context.Context, cliMgr Client, up string, param *types.LoginParameter, phase uint) (result *types.LoginSchema, nerr error) {
 	c.Lock()
 	curPhase := c.proverToken[up].phase
 	if c.completionCtx != nil {
@@ -89,6 +89,17 @@ func (c *proverSession) maintainLogin(ctx context.Context, cliMgr Client, up str
 	completeCtx, cf := context.WithCancel(ctx)
 	defer cf()
 	c.completionCtx = completeCtx
+	defer func() {
+		c.Lock()
+		c.completionCtx = nil
+		if result != nil {
+			c.proverToken[up] = loginToken{
+				LoginSchema: result,
+				phase:       curPhase + 1,
+			}
+		}
+		c.Unlock()
+	}()
 	c.Unlock()
 
 	cli := cliMgr.Client(ctx)
@@ -124,18 +135,9 @@ func (c *proverSession) maintainLogin(ctx context.Context, cliMgr Client, up str
 		return nil, err
 	}
 
-	c.Lock()
-	defer c.Unlock()
-
-	c.proverToken[up] = loginToken{
-		LoginSchema: &types.LoginSchema{
-			Token: loginResult.Token,
-		},
-		phase: curPhase + 1,
-	}
-	c.completionCtx = nil
-
-	return c.proverToken[up].LoginSchema, nil
+	return &types.LoginSchema{
+		Token: loginResult.Token,
+	}, nil
 }
 
 const expireTolerant = 10 * time.Minute

coordinator/internal/controller/proxy/submit_proof.go

@@ -14,24 +14,26 @@ import (
 
 // SubmitProofController the submit proof api controller
 type SubmitProofController struct {
-	proverMgr *ProverManager
-	clients   Clients
+	proverMgr        *ProverManager
+	clients          Clients
+	priorityUpstream *PriorityUpstreamManager
 }
 
 // NewSubmitProofController create the submit proof api controller instance
-func NewSubmitProofController(cfg *config.ProxyConfig, clients Clients, proverMgr *ProverManager, reg prometheus.Registerer) *SubmitProofController {
+func NewSubmitProofController(cfg *config.ProxyConfig, clients Clients, proverMgr *ProverManager, priorityMgr *PriorityUpstreamManager, reg prometheus.Registerer) *SubmitProofController {
 	return &SubmitProofController{
-		proverMgr: proverMgr,
-		clients:   clients,
+		proverMgr:        proverMgr,
+		clients:          clients,
+		priorityUpstream: priorityMgr,
 	}
 }
 
-func upstreamFromTaskName(taskID string) string {
-	parts, _, found := strings.Cut(taskID, ":")
+func upstreamFromTaskName(taskID string) (string, string) {
+	parts, rest, found := strings.Cut(taskID, ":")
 	if found {
-		return parts
+		return parts, rest
 	}
-	return ""
+	return "", parts
 }
 
 func formUpstreamWithTaskName(upstream string, taskID string) string {
@@ -40,6 +42,7 @@ func formUpstreamWithTaskName(upstream string, taskID string) string {
 
 // SubmitProof prover submit the proof to coordinator
 func (spc *SubmitProofController) SubmitProof(ctx *gin.Context) {
+
 	var submitParameter coordinatorType.SubmitProofParameter
 	if err := ctx.ShouldBind(&submitParameter); err != nil {
 		nerr := fmt.Errorf("prover submitProof parameter invalid, err:%w", err)
@@ -53,14 +56,15 @@ func (spc *SubmitProofController) SubmitProof(ctx *gin.Context) {
 	}
 
 	session := spc.proverMgr.Get(publicKey)
-	upstream := upstreamFromTaskName(submitParameter.TaskID)
+	upstream, realTaskID := upstreamFromTaskName(submitParameter.TaskID)
 	cli, existed := spc.clients[upstream]
 	if !existed {
 		// TODO: log error
 		nerr := fmt.Errorf("Invalid upstream name (%s) from taskID %s", upstream, submitParameter.TaskID)
 		types.RenderFailure(ctx, types.ErrCoordinatorParameterInvalidNo, nerr)
 		return
 	}
+	submitParameter.TaskID = realTaskID
 
 	resp, err := session.SubmitProof(ctx, &submitParameter, cli, upstream)
 	if err != nil {
@@ -71,6 +75,7 @@ func (spc *SubmitProofController) SubmitProof(ctx *gin.Context) {
 		types.RenderFailure(ctx, resp.ErrCode, fmt.Errorf("%s", resp.ErrMsg))
 		return
 	} else {
+		spc.priorityUpstream.Delete(upstream)
 		types.RenderSuccess(ctx, resp.Data)
 		return
 	}

coordinator/internal/types/response_test.go

@@ -0,0 +1,48 @@
+package types
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	"scroll-tech/common/types"
+)
+
+func TestResponseDecodeData_GetTaskSchema(t *testing.T) {
+	// Arrange: build a dummy payload and wrap it in Response
+	in := GetTaskSchema{
+		UUID:         "uuid-123",
+		TaskID:       "task-abc",
+		TaskType:     1,
+		UseSnark:     true,
+		TaskData:     "dummy-data",
+		HardForkName: "cancun",
+	}
+
+	resp := types.Response{
+		ErrCode: 0,
+		ErrMsg:  "",
+		Data:    in,
+	}
+
+	// Act: JSON round-trip the Response to simulate real HTTP encoding/decoding
+	b, err := json.Marshal(resp)
+	if err != nil {
+		t.Fatalf("marshal response: %v", err)
+	}
+
+	var decoded types.Response
+	if err := json.Unmarshal(b, &decoded); err != nil {
+		t.Fatalf("unmarshal response: %v", err)
+	}
+
+	var out GetTaskSchema
+	if err := decoded.DecodeData(&out); err != nil {
+		t.Fatalf("DecodeData error: %v", err)
+	}
+
+	// Assert: structs match after decode
+	if !reflect.DeepEqual(in, out) {
+		t.Fatalf("decoded struct mismatch:\nwant: %+v\n got: %+v", in, out)
+	}
+}