From d09c64fcee20e9eabb13ca9c8de483747aafa930 Mon Sep 17 00:00:00 2001
From: Lixeer <1612655510@qq.com>
Date: Tue, 24 Feb 2026 22:33:04 +0800
Subject: [PATCH] fix: implement code review suggestions

Address all feedback from PR review:
- Lock granularity
- Empty response handling
- Shutdown race condition
- Interface naming
---
 cmd/picoclaw/cmd_gateway.go              |  4 ++--
 pkg/providers/github_copilot_provider.go | 21 +++++++++++++--------
 pkg/providers/types.go                   |  2 +-
 pkg/tools/registry_test.go               | 20 ++++++++++----------
 4 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/cmd/picoclaw/cmd_gateway.go b/cmd/picoclaw/cmd_gateway.go
index 30d61aec3..fea569381 100644
--- a/cmd/picoclaw/cmd_gateway.go
+++ b/cmd/picoclaw/cmd_gateway.go
@@ -211,10 +211,10 @@ func gatewayCmd() {
 	<-sigChan
 
 	fmt.Println("\nShutting down...")
-	cancel()
-	if cp, ok := provider.(providers.SessionProvider); ok {
+	if cp, ok := provider.(providers.StatefulProvider); ok {
 		cp.Close()
 	}
+	cancel()
 	healthServer.Stop(context.Background())
 	deviceService.Stop()
 	heartbeatService.Stop()
diff --git a/pkg/providers/github_copilot_provider.go b/pkg/providers/github_copilot_provider.go
index c69658b44..9210021e1 100644
--- a/pkg/providers/github_copilot_provider.go
+++ b/pkg/providers/github_copilot_provider.go
@@ -94,19 +94,24 @@ func (p *GitHubCopilotProvider) Chat(
 		return nil, fmt.Errorf("marshal messages: %w", err)
 	}
 	p.mu.Lock()
-	defer p.mu.Unlock()
+	session := p.session
+	p.mu.Unlock()
 
-	resp, err := p.session.SendAndWait(ctx, copilot.MessageOptions{
+	if session == nil {
+		return nil, fmt.Errorf("provider closed")
+	}
+
+	resp, err := session.SendAndWait(ctx, copilot.MessageOptions{
 		Prompt: string(fullcontent),
 	})
-	if err != nil {
-		return nil, err
-	}
 
-	var content string
-	if resp != nil && resp.Data.Content != nil {
-		content = *resp.Data.Content
+	if resp == nil {
+		return nil, fmt.Errorf("empty response from copilot")
 	}
+	if resp.Data.Content == nil {
+		return nil, fmt.Errorf("no content in copilot response")
+	}
+	content := *resp.Data.Content
 
 	return &LLMResponse{
 		FinishReason: "stop",
diff --git a/pkg/providers/types.go b/pkg/providers/types.go
index 40ff6f7c8..b2dda04a5 100644
--- a/pkg/providers/types.go
+++ b/pkg/providers/types.go
@@ -30,7 +30,7 @@ type LLMProvider interface {
 	GetDefaultModel() string
 }
 
-type SessionProvider interface {
+type StatefulProvider interface {
 	LLMProvider
 	Close()
 }
diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go
index 8ae13b20c..33978e543 100644
--- a/pkg/tools/registry_test.go
+++ b/pkg/tools/registry_test.go
@@ -14,14 +14,14 @@ import (
 type mockRegistryTool struct {
 	name   string
 	desc   string
-	params map[string]any
+	params map[string]interface{}
 	result *ToolResult
 }
 
-func (m *mockRegistryTool) Name() string               { return m.name }
-func (m *mockRegistryTool) Description() string        { return m.desc }
-func (m *mockRegistryTool) Parameters() map[string]any { return m.params }
-func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]any) *ToolResult {
+func (m *mockRegistryTool) Name() string                       { return m.name }
+func (m *mockRegistryTool) Description() string                { return m.desc }
+func (m *mockRegistryTool) Parameters() map[string]interface{} { return m.params }
+func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]interface{}) *ToolResult {
 	return m.result
 }
 
@@ -51,7 +51,7 @@ func newMockTool(name, desc string) *mockRegistryTool {
 	return &mockRegistryTool{
 		name:   name,
 		desc:   desc,
-		params: map[string]any{"type": "object"},
+		params: map[string]interface{}{"type": "object"},
 		result: SilentResult("ok"),
 	}
 }
@@ -109,7 +109,7 @@ func TestToolRegistry_Execute_Success(t *testing.T) {
 	r.Register(&mockRegistryTool{
 		name:   "greet",
 		desc:   "says hello",
-		params: map[string]any{},
+		params: map[string]interface{}{},
 		result: SilentResult("hello"),
 	})
 
@@ -203,7 +203,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) {
 	if defs[0]["type"] != "function" {
 		t.Errorf("expected type 'function', got %v", defs[0]["type"])
 	}
-	fn, ok := defs[0]["function"].(map[string]any)
+	fn, ok := defs[0]["function"].(map[string]interface{})
 	if !ok {
 		t.Fatal("expected 'function' key to be a map")
 	}
@@ -217,7 +217,7 @@ func TestToolRegistry_GetDefinitions(t *testing.T) {
 
 func TestToolRegistry_ToProviderDefs(t *testing.T) {
 	r := NewToolRegistry()
-	params := map[string]any{"type": "object", "properties": map[string]any{}}
+	params := map[string]interface{}{"type": "object", "properties": map[string]interface{}{}}
 	r.Register(&mockRegistryTool{
 		name:   "beta",
 		desc:   "tool B",
@@ -310,7 +310,7 @@ func TestToolToSchema(t *testing.T) {
 	if schema["type"] != "function" {
 		t.Errorf("expected type 'function', got %v", schema["type"])
 	}
-	fn, ok := schema["function"].(map[string]any)
+	fn, ok := schema["function"].(map[string]interface{})
 	if !ok {
 		t.Fatal("expected 'function' to be a map")
 	}