Merge pull request #1945 from dgageot/board/docker-agent-issues-cb91df7b

Fix listener resource leaks in serve commands

Author: dgageot

cmd/root/a2a.go

@@ -42,10 +42,11 @@ func (f *a2aFlags) runA2ACommand(cmd *cobra.Command, args []string) error {
 	out := cli.NewPrinter(cmd.OutOrStdout())
 	agentFilename := args[0]
 
-	ln, err := listenAndCloseOnCancel(ctx, f.listenAddr)
+	ln, cleanup, err := newListener(ctx, f.listenAddr)
 	if err != nil {
 		return err
 	}
+	defer cleanup()
 
 	out.Println("Listening on", ln.Addr().String())
 	return a2a.Run(ctx, agentFilename, f.agentName, &f.runConfig, ln)

cmd/root/api.go

@@ -86,10 +86,11 @@ func (f *apiFlags) runAPICommand(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("--pull-interval flag can only be used with OCI references, not local files")
 	}
 
-	ln, err := listenAndCloseOnCancel(ctx, f.listenAddr)
+	ln, lnCleanup, err := newListener(ctx, f.listenAddr)
 	if err != nil {
 		return err
 	}
+	defer lnCleanup()
 
 	out.Println("Listening on", ln.Addr().String())
 

cmd/root/flags.go

@@ -129,16 +129,21 @@ func parseModelShorthand(s string) *latest.ModelConfig {
 	return nil
 }
 
-// listenAndCloseOnCancel starts a listener and spawns a goroutine
-// that closes it when the context is cancelled.
-func listenAndCloseOnCancel(ctx context.Context, addr string) (net.Listener, error) {
+// newListener creates a TCP listener and returns a cleanup function that
+// must be deferred by the caller. The cleanup function closes the listener.
+// The listener is also closed if the context is cancelled, which unblocks
+// any in-progress Serve call.
+func newListener(ctx context.Context, addr string) (net.Listener, func(), error) {
 	ln, err := server.Listen(ctx, addr)
 	if err != nil {
-		return nil, fmt.Errorf("failed to listen on %s: %w", addr, err)
+		return nil, nil, fmt.Errorf("failed to listen on %s: %w", addr, err)
 	}
-	go func() {
-		<-ctx.Done()
+	stop := context.AfterFunc(ctx, func() {
 		_ = ln.Close()
-	}()
-	return ln, nil
+	})
+	cleanup := func() {
+		stop()
+		_ = ln.Close()
+	}
+	return ln, cleanup, nil
 }

cmd/root/mcp.go

@@ -48,10 +48,11 @@ func (f *mcpFlags) runMCPCommand(cmd *cobra.Command, args []string) error {
 		return mcp.StartMCPServer(ctx, agentFilename, f.agentName, &f.runConfig)
 	}
 
-	ln, err := listenAndCloseOnCancel(ctx, f.listenAddr)
+	ln, cleanup, err := newListener(ctx, f.listenAddr)
 	if err != nil {
 		return err
 	}
+	defer cleanup()
 
 	return mcp.StartHTTPServer(ctx, agentFilename, f.agentName, &f.runConfig, ln)
 }